#[allow(unsafe_code)]
mod arena {
    use std::alloc;
    use std::marker::PhantomData;

    const PAGE_SIZE: usize = 4096;
    const MAX_PAGE_SIZE: usize = 512 * PAGE_SIZE;

    pub struct Arena<'arena> {
        offset: usize,
        block: Block<'arena>,
        phantom: PhantomData<&'arena ()>,
    }

    #[derive(Copy, Clone)]
    struct Block<'arena> {
        data: *mut u8,
        layout: alloc::Layout,
        prev: Option<&'arena Block<'arena>>,
    }

    impl<'arena> Arena<'arena> {
        pub fn new() -> Arena<'arena> {
            Arena { offset: 0, block: Block::alloc(PAGE_SIZE), phantom: PhantomData }
        }

        pub fn alloc<T: Copy>(&mut self, val: T) -> &'arena T {
            let offset = self.offset_raw::<T>(std::mem::size_of::<T>());
            unsafe {
                *offset = val;
                &*offset
            }
        }

        pub fn alloc_slice<T: Copy>(&mut self, val: &[T]) -> &'arena [T] {
            let offset = self.offset_raw::<T>(std::mem::size_of_val(val));
            unsafe {
                std::ptr::copy_nonoverlapping(val.as_ptr(), offset, val.len());
                std::slice::from_raw_parts(offset as *const T, val.len())
            }
        }

        pub fn alloc_slice_with_value<T: Copy>(&mut self, val: T, len: usize) -> &'arena [T] {
            let offset = self.offset_raw::<T>(len * std::mem::size_of::<T>());
            unsafe {
                let write_ptr = offset;
                for idx in 0..len {
                    std::ptr::write(write_ptr.add(idx), val);
                }
                std::slice::from_raw_parts(offset as *const T, len)
            }
        }

        pub fn alloc_str(&mut self, val: &str) -> &'arena str {
            let bytes = self.alloc_slice(val.as_bytes());
            unsafe { std::str::from_utf8_unchecked(bytes) }
        }

        fn offset_raw<T: Copy>(&mut self, size: usize) -> *mut T {
            let align = align_of::<T>();
            self.offset = (self.offset + align - 1) & !(align - 1);

            if self.offset + size > self.block.layout.size() {
                self.grow(size);
            }
            unsafe {
                let offset = self.block.data.add(self.offset) as *mut T;
                self.offset += size;
                offset
            }
        }

        fn grow(&mut self, size: usize) {
            let prev = self.block;
            let block_size = (self.block.layout.size() * 2).min(MAX_PAGE_SIZE).max(size);
            self.offset = 0;
            self.block = Block::alloc(block_size);
            self.block.prev = Some(self.alloc(prev));
        }

        pub fn mem_usage(&self) -> (usize, usize) {
            let mut total = 0;
            let mut current = Some(self.block);
            while let Some(block) = current {
                total += block.layout.size();
                current = block.prev.copied();
            }
            (total - self.block.layout.size() + self.offset, total)
        }
    }

    impl<'arena> Block<'arena> {
        fn alloc(size: usize) -> Block<'arena> {
            let layout = alloc::Layout::from_size_align(size, PAGE_SIZE).unwrap();
            let data = unsafe { alloc::alloc(layout) };
            Block { data, layout, prev: None }
        }
        fn dealloc(&self) {
            unsafe { alloc::dealloc(self.data, self.layout) }
        }
    }

    impl Drop for Arena<'_> {
        fn drop(&mut self) {
            let mut current = Some(self.block);
            while let Some(block) = current {
                current = block.prev.copied();
                block.dealloc();
            }
        }
    }
}