#![feature(allocator_api, box_as_ptr)]

#[path = "shared/shared.rs"]
mod util;

unsafe extern "C" {
    unsafe fn bump_init();
    unsafe fn bump_alloc(size: usize, align: usize) -> *mut u8;
}

struct BumpAllocator;

unsafe impl std::alloc::Allocator for BumpAllocator {
    fn allocate(
        &self,
        layout: std::alloc::Layout,
    ) -> Result<std::ptr::NonNull<[u8]>, std::alloc::AllocError> {
        unsafe {
            let ptr = bump_alloc(layout.size(), layout.align());
            if ptr.is_null() {
                Err(std::alloc::AllocError)
            } else {
                Ok(std::ptr::NonNull::slice_from_raw_parts(
                    std::ptr::NonNull::new_unchecked(ptr),
                    layout.size(),
                ))
            }
        }
    }

    unsafe fn deallocate(&self, _ptr: std::ptr::NonNull<u8>, _layout: std::alloc::Layout) {
        // Bump allocator does not deallocate individual allocations
    }
}

fn main() {
    unsafe {
        bump_init();
    }

    let a = Box::new_in(42u32, BumpAllocator);
    let b = Box::new_in(42u32, BumpAllocator);
    let c = Box::new_in(52u32, BumpAllocator);
    eprintln!("a: {a}, b: {b}, c: {c}");
    assert_ne!(Box::as_ptr(&a), Box::as_ptr(&b));
    assert_eq!(*a, *b);

    struct BigType {
        data: [u8; 0x1010],
    }

    let mut big = Box::new_in(BigType { data: [0; 0x1010] }, BumpAllocator);
    assert_ne!(Box::as_ptr(&big) as *const (), Box::as_ptr(&a) as *const ());
    big.data[47] = 123;
    assert_eq!(big.data[47], 123);

    #[repr(align(256))]
    struct AlignedType {
        #[allow(dead_code)]
        data: [u8; 512],
    }
    let aligned = Box::new_in(AlignedType { data: [0; 512] }, BumpAllocator);
    assert_eq!(
        (Box::as_ptr(&aligned) as usize) % 256,
        0,
        "Aligned allocation should be aligned to 256 bytes"
    );
}