//! Bounded, single-CPU heap experiment over four already mapped 4 KiB pages.
//! This is not the global allocator. A generation-tagged handle lets the
//! isolated test distinguish stale/double release from a newly reused VA.

use core::ptr::{read_unaligned, write_unaligned};

pub const CAPACITY: usize = 4 * 4096;
const MAX_ALLOCS: usize = 64;
const MAX_FREE: usize = MAX_ALLOCS + 1;
const CANARY_BYTES: usize = 8;
const CANARY_KEY: u64 = 0x5455_524b_4845_4150;

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Error {
    Invalid,
    Overflow,
    OutOfMemory,
    MetadataFull,
    DoubleFree,
    CorruptedCanary,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Handle {
    pub address: usize,
    pub generation: u64,
}

#[derive(Clone, Copy, Default)]
struct Region {
    start: usize,
    end: usize,
}

#[derive(Clone, Copy, PartialEq, Eq)]
enum State {
    Empty,
    Active,
    Freed,
}

#[derive(Clone, Copy)]
struct Allocation {
    address: usize,
    size: usize,
    end: usize,
    generation: u64,
    state: State,
}

impl Allocation {
    const EMPTY: Self = Self {
        address: 0,
        size: 0,
        end: 0,
        generation: 0,
        state: State::Empty,
    };
}

pub struct Heap {
    base: usize,
    free: [Region; MAX_FREE],
    free_len: usize,
    allocations: [Allocation; MAX_ALLOCS],
    active: usize,
    next_generation: u64,
}

impl Heap {
    /// Initialize metadata directly in reserved storage, without first
    /// materializing two 3.5 KiB copies on the small early boot stack.
    ///
    /// # Safety
    /// `storage` must be valid, aligned, exclusively owned, and remain alive
    /// while any handle is outstanding. The backing pages must satisfy `new`.
    pub unsafe fn init_in_place(storage: *mut Self, base: usize) -> Result<(), Error> {
        if storage.is_null() || base == 0 || base.checked_add(CAPACITY).is_none() {
            return Err(Error::Invalid);
        }
        // SAFETY: caller guarantees exclusive valid storage; each field is
        // initialized before any reference to the complete Heap is created.
        unsafe {
            core::ptr::addr_of_mut!((*storage).base).write(base);
            let regions = core::ptr::addr_of_mut!((*storage).free).cast::<Region>();
            for index in 0..MAX_FREE {
                regions.add(index).write(Region::default());
            }
            regions.write(Region {
                start: 0,
                end: CAPACITY,
            });
            core::ptr::addr_of_mut!((*storage).free_len).write(1);
            let allocations = core::ptr::addr_of_mut!((*storage).allocations).cast::<Allocation>();
            for index in 0..MAX_ALLOCS {
                allocations.add(index).write(Allocation::EMPTY);
            }
            core::ptr::addr_of_mut!((*storage).active).write(0);
            core::ptr::addr_of_mut!((*storage).next_generation).write(1);
        }
        Ok(())
    }

    /// # Safety
    /// `base..base+CAPACITY` must be exclusively owned writable memory for
    /// this heap's entire lifetime. The caller must serialize all access.
    pub unsafe fn new(base: usize) -> Result<Self, Error> {
        if base == 0 || base.checked_add(CAPACITY).is_none() {
            return Err(Error::Invalid);
        }
        let mut free = [Region::default(); MAX_FREE];
        free[0] = Region {
            start: 0,
            end: CAPACITY,
        };
        Ok(Self {
            base,
            free,
            free_len: 1,
            allocations: [Allocation::EMPTY; MAX_ALLOCS],
            active: 0,
            next_generation: 1,
        })
    }

    pub fn active_count(&self) -> usize {
        self.active
    }

    pub fn free_bytes(&self) -> usize {
        self.free[..self.free_len]
            .iter()
            .map(|region| region.end - region.start)
            .sum()
    }

    pub fn used_bytes(&self) -> usize {
        CAPACITY - self.free_bytes()
    }

    pub fn largest_free(&self) -> usize {
        self.free[..self.free_len]
            .iter()
            .map(|region| region.end - region.start)
            .max()
            .unwrap_or(0)
    }

    pub fn allocate(&mut self, size: usize, align: usize) -> Result<Handle, Error> {
        if size == 0 || align == 0 || !align.is_power_of_two() || align > 4096 {
            return Err(Error::Invalid);
        }
        let _ = size.checked_add(CANARY_BYTES).ok_or(Error::Overflow)?;
        let slot = self
            .allocations
            .iter()
            .position(|entry| entry.state != State::Active)
            .ok_or(Error::MetadataFull)?;
        let generation = self.next_generation;
        let next_generation = generation.checked_add(1).ok_or(Error::Overflow)?;
        for index in 0..self.free_len {
            let region = self.free[index];
            let first = self.base.checked_add(region.start).ok_or(Error::Overflow)?;
            let Some(aligned) = first
                .checked_add(align - 1)
                .map(|address| address & !(align - 1))
            else {
                return Err(Error::Overflow);
            };
            let start = aligned - self.base;
            let Some(end) = start
                .checked_add(size)
                .and_then(|last| last.checked_add(CANARY_BYTES))
            else {
                return Err(Error::Overflow);
            };
            if end > region.end {
                continue;
            }
            let prefix = start > region.start;
            let suffix = end < region.end;
            if prefix && suffix && self.free_len == MAX_FREE {
                return Err(Error::MetadataFull);
            }
            match (prefix, suffix) {
                (true, true) => {
                    self.free.copy_within(index + 1..self.free_len, index + 2);
                    self.free[index] = Region {
                        start: region.start,
                        end: start,
                    };
                    self.free[index + 1] = Region {
                        start: end,
                        end: region.end,
                    };
                    self.free_len += 1;
                }
                (true, false) => self.free[index].end = start,
                (false, true) => self.free[index].start = end,
                (false, false) => {
                    self.free.copy_within(index + 1..self.free_len, index);
                    self.free_len -= 1;
                }
            }
            let handle = Handle {
                address: aligned,
                generation,
            };
            let canary = Self::canary(handle, size);
            // SAFETY: allocation reserved [start,end) in caller-owned mapped
            // memory; unaligned write touches only the 8-byte trailer.
            unsafe { write_unaligned((self.base + end - CANARY_BYTES) as *mut u64, canary) };
            self.allocations[slot] = Allocation {
                address: aligned,
                size,
                end,
                generation,
                state: State::Active,
            };
            self.active += 1;
            self.next_generation = next_generation;
            return Ok(handle);
        }
        Err(Error::OutOfMemory)
    }

    pub fn release(&mut self, handle: Handle) -> Result<(), Error> {
        let Some(index) = self.allocations.iter().position(|entry| {
            entry.generation == handle.generation && entry.address == handle.address
        }) else {
            return Err(Error::Invalid);
        };
        let allocation = self.allocations[index];
        if allocation.state == State::Freed {
            return Err(Error::DoubleFree);
        }
        if allocation.state != State::Active {
            return Err(Error::Invalid);
        }
        // SAFETY: the active allocation's trailer remains inside the mapped
        // heap. A failed canary check leaves ownership unchanged.
        let observed =
            unsafe { read_unaligned((self.base + allocation.end - CANARY_BYTES) as *const u64) };
        if observed != Self::canary(handle, allocation.size) {
            return Err(Error::CorruptedCanary);
        }
        let start = handle.address - self.base;
        let end = allocation.end;
        let insert = self.free[..self.free_len]
            .iter()
            .position(|region| region.start >= start)
            .unwrap_or(self.free_len);
        if (insert > 0 && self.free[insert - 1].end > start)
            || (insert < self.free_len && self.free[insert].start < end)
        {
            return Err(Error::Invalid);
        }
        let left = insert > 0 && self.free[insert - 1].end == start;
        let right = insert < self.free_len && self.free[insert].start == end;
        if left && right {
            self.free[insert - 1].end = self.free[insert].end;
            self.free.copy_within(insert + 1..self.free_len, insert);
            self.free_len -= 1;
        } else if left {
            self.free[insert - 1].end = end;
        } else if right {
            self.free[insert].start = start;
        } else {
            if self.free_len == MAX_FREE {
                return Err(Error::MetadataFull);
            }
            self.free.copy_within(insert..self.free_len, insert + 1);
            self.free[insert] = Region { start, end };
            self.free_len += 1;
        }
        self.allocations[index].state = State::Freed;
        self.active -= 1;
        Ok(())
    }

    fn canary(handle: Handle, size: usize) -> u64 {
        CANARY_KEY ^ handle.address as u64 ^ handle.generation ^ size as u64
    }
}

#[cfg(test)]
mod tests {
    extern crate std;

    use super::*;
    use std::boxed::Box;

    #[test]
    fn alignment_page_crossing_and_reuse() {
        let mut backing = Box::new([0u8; CAPACITY]);
        let mut heap = unsafe { Heap::new(backing.as_mut_ptr() as usize) }.unwrap();
        let small = heap.allocate(16, 16).unwrap();
        assert_eq!(small.address % 16, 0);
        let aligned = heap.allocate(32, 64).unwrap();
        assert_eq!(aligned.address % 64, 0);
        let page = heap.allocate(4096, 4096).unwrap();
        assert_eq!(page.address % 4096, 0);
        let crossing = heap.allocate(5000, 8).unwrap();
        assert!(crossing.address / 4096 != (crossing.address + 4999) / 4096);
        for handle in [crossing, page, aligned, small] {
            heap.release(handle).unwrap();
        }
        assert_eq!(heap.active_count(), 0);
        assert_eq!(heap.free_bytes(), CAPACITY);
        assert_eq!(heap.largest_free(), CAPACITY);
    }

    #[test]
    fn zero_overflow_double_free_foreign_and_canary_fail_closed() {
        let mut backing = Box::new([0u8; CAPACITY]);
        let mut heap = unsafe { Heap::new(backing.as_mut_ptr() as usize) }.unwrap();
        assert_eq!(heap.allocate(0, 8), Err(Error::Invalid));
        assert_eq!(heap.allocate(1, 3), Err(Error::Invalid));
        assert_eq!(heap.allocate(usize::MAX, 8), Err(Error::Overflow));
        let handle = heap.allocate(16, 8).unwrap();
        assert_eq!(
            heap.release(Handle {
                address: handle.address + 1,
                ..handle
            }),
            Err(Error::Invalid)
        );
        // SAFETY: this byte is the first byte of the allocation trailer.
        unsafe {
            let trailer = (handle.address as *mut u8).add(16);
            trailer.write(trailer.read() ^ 1);
        };
        assert_eq!(heap.release(handle), Err(Error::CorruptedCanary));
        assert_eq!(heap.active_count(), 1);
        let canary = Heap::canary(handle, 16);
        unsafe { write_unaligned((handle.address + 16) as *mut u64, canary) };
        heap.release(handle).unwrap();
        assert_eq!(heap.release(handle), Err(Error::DoubleFree));
        let reused = heap.allocate(16, 8).unwrap();
        assert_eq!(reused.address, handle.address);
        assert_ne!(reused.generation, handle.generation);
        assert_eq!(heap.release(handle), Err(Error::Invalid));
        heap.release(reused).unwrap();
    }

    #[test]
    fn many_small_allocations_reverse_release_and_oom_recovery() {
        let mut backing = Box::new([0u8; CAPACITY]);
        let mut heap = unsafe { Heap::new(backing.as_mut_ptr() as usize) }.unwrap();
        let mut handles: [Option<Handle>; 32] = [None; 32];
        for index in 0..handles.len() {
            let handle = heap.allocate(32, 16).unwrap();
            assert!(
                !handles[..index]
                    .iter()
                    .flatten()
                    .any(|prior| prior.address == handle.address)
            );
            handles[index] = Some(handle);
        }
        assert_eq!(heap.active_count(), 32);
        assert_eq!(heap.allocate(CAPACITY, 8), Err(Error::OutOfMemory));
        for handle in handles.into_iter().rev().flatten() {
            heap.release(handle).unwrap();
        }
        assert_eq!(heap.free_bytes(), CAPACITY);
        let again = heap.allocate(CAPACITY - CANARY_BYTES, 1).unwrap();
        assert_eq!(heap.allocate(1, 1), Err(Error::OutOfMemory));
        heap.release(again).unwrap();
        assert_eq!(heap.largest_free(), CAPACITY);
    }
}