krypteia-memory 0.1.0

//! TLSF (Two-Level Segregated Fit) allocator.
//!
//! O(1) alloc and free, bounded fragmentation, deterministic
//! latency. Designed for real-time embedded systems (RTEMS, L4,
//! FreeRTOS). Reference: Masmano et al., 2004.
//!
//! The allocator is initialised once via [`krypteia_init`] with a
//! contiguous RAM block supplied by the caller. All subsequent
//! `alloc` / `dealloc` calls draw from that block.
//!
//! # Thread safety
//!
//! A minimal spin-lock serialises access. On single-core bare-metal
//! this compiles to a plain flag (the `AtomicBool` CAS is
//! implemented by the Rust runtime with interrupt-disable on
//! targets that lack hardware atomics).

use core::alloc::{GlobalAlloc, Layout};
use core::cell::UnsafeCell;
use core::ptr;
use core::sync::atomic::{AtomicBool, Ordering};

// ====================================================================
// Configuration
// ====================================================================

const WORD: usize = core::mem::size_of::<usize>();

/// Minimum block: header + next_free + prev_free + footer = 4 words.
/// 16 bytes on 32-bit, 32 bytes on 64-bit.
const MIN_BLOCK: usize = 4 * WORD;

const FREE_BIT: usize = 1;
const PREV_FREE_BIT: usize = 2;
const FLAGS: usize = FREE_BIT | PREV_FREE_BIT;

/// Second-level subdivision bits (4 sub-classes per FL level).
const SL_BITS: u32 = 2;
const SL_COUNT: usize = 1 << SL_BITS;

/// Number of first-level classes. With FL_MIN derived from
/// MIN_BLOCK, 8 FL levels cover up to 2^(FL_MIN+8) bytes.
/// On 64-bit: 2^(5+8) = 8 KB per single block; the last bin
/// holds everything larger. On 32-bit: 2^(4+8) = 4 KB.
const FL_COUNT: usize = 8;
const BINS: usize = FL_COUNT * SL_COUNT; // 32 — fits in a u32 bitmap

// ====================================================================
// Block helpers (all unsafe, operating on raw heap pointers)
// ====================================================================

const fn log2_floor(x: usize) -> u32 {
    (usize::BITS - 1) - (x.leading_zeros())
}

const FL_MIN: u32 = log2_floor(MIN_BLOCK);

#[inline(always)]
fn round_up(x: usize, align: usize) -> usize {
    (x + align - 1) & !(align - 1)
}

// ---- read / write a usize at an arbitrary byte pointer ----

#[inline(always)]
unsafe fn read_word(p: *const u8) -> usize {
    ptr::read_unaligned(p as *const usize)
}

#[inline(always)]
unsafe fn write_word(p: *mut u8, v: usize) {
    ptr::write_unaligned(p as *mut usize, v);
}

// ---- block header ----

unsafe fn block_size(b: *const u8) -> usize {
    read_word(b) & !FLAGS
}
unsafe fn block_is_free(b: *const u8) -> bool {
    read_word(b) & FREE_BIT != 0
}
unsafe fn block_is_prev_free(b: *const u8) -> bool {
    read_word(b) & PREV_FREE_BIT != 0
}

unsafe fn block_set_header(b: *mut u8, size: usize, flags: usize) {
    write_word(b, size | (flags & FLAGS));
}

/// Pointer to the usable body (right after the header word).
unsafe fn block_body(b: *mut u8) -> *mut u8 {
    b.add(WORD)
}

/// Recover block pointer from body pointer.
unsafe fn body_to_block(body: *mut u8) -> *mut u8 {
    body.sub(WORD)
}

/// Next physical block = current + current.size.
unsafe fn block_next(b: *const u8) -> *mut u8 {
    (b as *mut u8).add(block_size(b))
}

/// Previous physical block (via footer of the previous free block).
unsafe fn block_prev(b: *const u8) -> *mut u8 {
    let prev_size = read_word((b as *const u8).sub(WORD)) & !FLAGS;
    (b as *mut u8).sub(prev_size)
}

// ---- free-list pointers (stored in the body of free blocks) ----

unsafe fn free_next(b: *const u8) -> *mut u8 {
    read_word((b as *const u8).add(WORD)) as *mut u8
}
unsafe fn free_set_next(b: *mut u8, n: *mut u8) {
    write_word(b.add(WORD), n as usize);
}
unsafe fn free_prev(b: *const u8) -> *mut u8 {
    read_word((b as *const u8).add(2 * WORD)) as *mut u8
}
unsafe fn free_set_prev(b: *mut u8, p: *mut u8) {
    write_word(b.add(2 * WORD), p as usize);
}

/// Write a footer at the last word of the block body (copy of the
/// header). Used for backward merging.
unsafe fn block_write_footer(b: *mut u8) {
    let sz = block_size(b);
    write_word(b.add(sz - WORD), read_word(b));
}

// ---- bin mapping ----

fn bin_index(size: usize) -> usize {
    debug_assert!(size >= MIN_BLOCK);
    let fl = log2_floor(size);
    let sl = (size >> (fl.saturating_sub(SL_BITS) as usize)) & (SL_COUNT - 1);
    let fi = (fl - FL_MIN) as usize;
    (fi * SL_COUNT + sl).min(BINS - 1)
}

/// Return the bin index to start searching from for `size` bytes.
/// Blocks in this bin may or may not be large enough (the SL
/// granularity is coarse), so [`TlsfInner::find`] checks the
/// actual block size and falls back to the next bin if needed.
fn search_index(size: usize) -> usize {
    bin_index(size.max(MIN_BLOCK))
}

// ====================================================================
// TLSF inner state
// ====================================================================

struct TlsfInner {
    bitmap: u32,
    bins: [*mut u8; BINS],
    heap_end: *mut u8,
    used: usize,
    peak: usize,
    total: usize,
    ready: bool,
}

// *mut u8 is not Send, but we protect with SpinMutex.
unsafe impl Send for TlsfInner {}

impl TlsfInner {
    const fn new() -> Self {
        Self {
            bitmap: 0,
            bins: [ptr::null_mut(); BINS],
            heap_end: ptr::null_mut(),
            used: 0,
            peak: 0,
            total: 0,
            ready: false,
        }
    }

    unsafe fn init(&mut self, base: *mut u8, size: usize) {
        let start = round_up(base as usize, WORD) as *mut u8;
        let end = ((base as usize + size) & !(WORD - 1)) as *mut u8;
        let usable = end as usize - start as usize;

        assert!(
            usable >= MIN_BLOCK + WORD,
            "memory: heap too small (need >= {} bytes, got {})",
            MIN_BLOCK + WORD,
            usable,
        );

        self.heap_end = end;
        self.total = usable;
        self.bitmap = 0;
        self.bins = [ptr::null_mut(); BINS];
        self.used = 0;
        self.peak = 0;

        // One free block spanning the whole usable area minus
        // a WORD-sized sentinel at the end.
        let blk_size = usable - WORD;
        block_set_header(start, blk_size, FREE_BIT);
        block_write_footer(start);

        // Sentinel: size 0, not free, prev is free.
        let sentinel = start.add(blk_size);
        block_set_header(sentinel, 0, PREV_FREE_BIT);

        self.insert(start);
        self.ready = true;
    }

    // ---- free-list management ----

    unsafe fn insert(&mut self, block: *mut u8) {
        let idx = bin_index(block_size(block));
        let head = self.bins[idx];
        free_set_next(block, head);
        free_set_prev(block, ptr::null_mut());
        if !head.is_null() {
            free_set_prev(head, block);
        }
        self.bins[idx] = block;
        self.bitmap |= 1u32 << idx;
    }

    unsafe fn remove(&mut self, block: *mut u8) {
        let nx = free_next(block);
        let pv = free_prev(block);
        if !nx.is_null() {
            free_set_prev(nx, pv);
        }
        if !pv.is_null() {
            free_set_next(pv, nx);
        } else {
            let idx = bin_index(block_size(block));
            self.bins[idx] = nx;
            if nx.is_null() {
                self.bitmap &= !(1u32 << idx);
            }
        }
    }

    /// Find the first free block that can satisfy `size` bytes.
    ///
    /// Starts at the bin for `size` and checks the head block. If
    /// the head is too small (SL granularity can put a slightly
    /// smaller block in the same bin), falls back to the next
    /// non-empty bin whose blocks are guaranteed large enough.
    /// At most 2 bitmap scans → still O(1).
    unsafe fn find(&self, size: usize) -> *mut u8 {
        let idx = search_index(size);
        let mut mask = self.bitmap & !((1u32 << idx) - 1);
        while mask != 0 {
            let found = mask.trailing_zeros() as usize;
            let block = self.bins[found];
            if block_size(block) >= size {
                return block;
            }
            // Head of this bin is too small; clear it and try the
            // next non-empty bin.
            mask &= !(1u32 << found);
        }
        ptr::null_mut()
    }

    // ---- alloc / dealloc ----

    unsafe fn alloc(&mut self, layout: Layout) -> *mut u8 {
        if !self.ready {
            return ptr::null_mut();
        }
        let needed = Self::block_size_for(layout);
        let block = self.find(needed);
        if block.is_null() {
            return ptr::null_mut();
        }

        self.remove(block);
        let bsz = block_size(block);
        let prev_flag = read_word(block) & PREV_FREE_BIT;
        let remain = bsz - needed;

        if remain >= MIN_BLOCK {
            // Split: shrink current, create a free remainder.
            block_set_header(block, needed, prev_flag); // used

            let split = block.add(needed);
            block_set_header(split, remain, FREE_BIT); // free, prev is used
            block_write_footer(split);

            // The block after `split` keeps its PREV_FREE_BIT set
            // (split is free).
            let after = split.add(remain);
            if (after as usize) < self.heap_end as usize {
                let h = read_word(after);
                write_word(after, h | PREV_FREE_BIT);
            }
            self.insert(split);
        } else {
            // Use the whole block.
            block_set_header(block, bsz, prev_flag); // used

            // Next block: clear PREV_FREE (this block is now used).
            let next = block.add(bsz);
            if (next as usize) < self.heap_end as usize {
                let h = read_word(next);
                write_word(next, h & !PREV_FREE_BIT);
            }
        }

        let alloc_sz = block_size(block);
        self.used += alloc_sz;
        if self.used > self.peak {
            self.peak = self.used;
        }

        // Return user pointer (may need alignment adjustment).
        let body = block_body(block);
        if layout.align() <= WORD {
            body
        } else {
            let aligned = round_up(body as usize + WORD, layout.align()) as *mut u8;
            // Store block address just before the aligned pointer.
            write_word(aligned.sub(WORD), block as usize);
            aligned
        }
    }

    unsafe fn dealloc(&mut self, ptr: *mut u8, layout: Layout) {
        if ptr.is_null() {
            return;
        }
        let block = if layout.align() <= WORD {
            body_to_block(ptr)
        } else {
            read_word(ptr.sub(WORD)) as *mut u8
        };

        let mut current = block;
        let mut total = block_size(block);
        self.used -= total;

        // Merge with the next physical block if it is free.
        let next = block.add(total);
        if (next as usize) < self.heap_end as usize && block_is_free(next) {
            self.remove(next);
            total += block_size(next);
        }

        // Merge with the previous physical block if it is free.
        if block_is_prev_free(block) {
            let prev = block_prev(block);
            self.remove(prev);
            total += block_size(prev);
            current = prev;
        }

        // The merged block's PREV_FREE_BIT reflects its own
        // predecessor (unchanged by the merge).
        let pf = if current == block {
            0 // prev wasn't free (we didn't merge backward)
        } else {
            read_word(current) & PREV_FREE_BIT
        };

        block_set_header(current, total, FREE_BIT | pf);
        block_write_footer(current);

        // The block after the merged region now has a free predecessor.
        let after = current.add(total);
        if (after as usize) < self.heap_end as usize {
            let h = read_word(after);
            write_word(after, h | PREV_FREE_BIT);
        }

        self.insert(current);
    }

    /// Compute the block size (header included) needed for a given
    /// `Layout`, respecting the minimum block constraint.
    fn block_size_for(layout: Layout) -> usize {
        let align = layout.align().max(WORD);
        let extra = if align > WORD { WORD + align - 1 } else { 0 };
        let body = (layout.size() + extra).max(MIN_BLOCK - WORD);
        let body = round_up(body, WORD);
        body + WORD
    }
}

// ====================================================================
// Spin mutex (minimal, no_std)
// ====================================================================

struct SpinMutex<T> {
    locked: AtomicBool,
    data: UnsafeCell<T>,
}

unsafe impl<T: Send> Sync for SpinMutex<T> {}

impl<T> SpinMutex<T> {
    const fn new(data: T) -> Self {
        Self {
            locked: AtomicBool::new(false),
            data: UnsafeCell::new(data),
        }
    }

    fn lock(&self) -> SpinGuard<'_, T> {
        while self
            .locked
            .compare_exchange_weak(false, true, Ordering::Acquire, Ordering::Relaxed)
            .is_err()
        {
            core::hint::spin_loop();
        }
        SpinGuard { mutex: self }
    }
}

struct SpinGuard<'a, T> {
    mutex: &'a SpinMutex<T>,
}

impl<T> core::ops::Deref for SpinGuard<'_, T> {
    type Target = T;
    fn deref(&self) -> &T {
        unsafe { &*self.mutex.data.get() }
    }
}

impl<T> core::ops::DerefMut for SpinGuard<'_, T> {
    fn deref_mut(&mut self) -> &mut T {
        unsafe { &mut *self.mutex.data.get() }
    }
}

impl<T> Drop for SpinGuard<'_, T> {
    fn drop(&mut self) {
        self.mutex.locked.store(false, Ordering::Release);
    }
}

// ====================================================================
// Global allocator wrapper
// ====================================================================

/// The krypteia TLSF allocator, suitable as a `#[global_allocator]`.
pub struct KrypteiaAlloc {
    inner: SpinMutex<TlsfInner>,
}

impl KrypteiaAlloc {
    /// Create the allocator in an uninitialised state.
    pub const fn new() -> Self {
        Self {
            inner: SpinMutex::new(TlsfInner::new()),
        }
    }

    /// Initialise the allocator with a contiguous RAM block.
    ///
    /// # Safety
    ///
    /// `base` must point to a region of at least `size` bytes that
    /// remains valid for the lifetime of the program. The caller
    /// must not read or write this region after calling `init`.
    pub unsafe fn init(&self, base: *mut u8, size: usize) {
        self.inner.lock().init(base, size);
    }
}

unsafe impl GlobalAlloc for KrypteiaAlloc {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        self.inner.lock().alloc(layout)
    }

    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        self.inner.lock().dealloc(ptr, layout);
    }
}

// ====================================================================
// The singleton + public API
// ====================================================================

/// The singleton TLSF allocator instance.
///
/// When the `global-alloc` feature is active this is registered as
/// the Rust `#[global_allocator]`. Otherwise it is a plain static
/// that callers can interact with via [`krypteia_init`] and
/// [`memory_stats`].
#[cfg(feature = "global-alloc")]
#[global_allocator]
static ALLOCATOR: KrypteiaAlloc = KrypteiaAlloc::new();

#[cfg(not(feature = "global-alloc"))]
static ALLOCATOR: KrypteiaAlloc = KrypteiaAlloc::new();

/// Memory allocation statistics.
#[derive(Clone, Copy, Debug, Default)]
pub struct MemStats {
    /// Bytes currently allocated (including block headers).
    pub used: usize,
    /// Bytes still available in the heap.
    pub free: usize,
    /// High-water mark since init.
    pub peak: usize,
}

/// Initialise the library's internal TLSF heap.
///
/// Must be called **once** before any cryptographic operation.
///
/// # Safety
///
/// `base` must point to a region of `size` bytes that remains
/// valid for the lifetime of the program. The caller must not
/// access this region after init.
///
/// Returns `0` on success, `-1` if the region is too small.
pub unsafe fn krypteia_init(base: *mut u8, size: usize) -> i32 {
    if size < MIN_BLOCK + WORD {
        return -1;
    }
    ALLOCATOR.init(base, size);
    0
}

/// Return current heap statistics.
pub fn memory_stats() -> MemStats {
    let inner = ALLOCATOR.inner.lock();
    MemStats {
        used: inner.used,
        free: inner.total.saturating_sub(inner.used),
        peak: inner.peak,
    }
}

// ====================================================================
// Tests (run on the host with `cargo test -p memory
//         --no-default-features --features self-alloc`)
// ====================================================================

#[cfg(test)]
mod tests {
    extern crate std;
    use super::*;
    use core::alloc::Layout;
    use std::vec;
    use std::vec::Vec;

    /// Create a small test heap on the stack and run a closure
    /// against a fresh TlsfInner.
    fn with_heap<F: FnOnce(&mut TlsfInner)>(size: usize, f: F) {
        let mut buf = vec![0u8; size];
        let mut tlsf = TlsfInner::new();
        unsafe { tlsf.init(buf.as_mut_ptr(), size) };
        f(&mut tlsf);
    }

    #[test]
    fn init_creates_one_free_block() {
        with_heap(1024, |tlsf| {
            assert!(tlsf.bitmap != 0, "bitmap should have at least one bit set");
            assert_eq!(tlsf.used, 0);
        });
    }

    #[test]
    fn single_alloc_dealloc() {
        with_heap(1024, |tlsf| {
            let layout = Layout::from_size_align(64, 8).unwrap();
            let ptr = unsafe { tlsf.alloc(layout) };
            assert!(!ptr.is_null());
            assert!(tlsf.used > 0);

            unsafe { tlsf.dealloc(ptr, layout) };
            assert_eq!(tlsf.used, 0, "used should be 0 after freeing everything");
        });
    }

    #[test]
    fn multiple_alloc_dealloc() {
        with_heap(4096, |tlsf| {
            let layout = Layout::from_size_align(128, 8).unwrap();
            let mut ptrs = Vec::new();

            // Allocate 8 blocks.
            for _ in 0..8 {
                let p = unsafe { tlsf.alloc(layout) };
                assert!(!p.is_null(), "allocation should succeed");
                ptrs.push(p);
            }

            // Free in reverse order.
            for p in ptrs.into_iter().rev() {
                unsafe { tlsf.dealloc(p, layout) };
            }
            assert_eq!(tlsf.used, 0);
        });
    }

    #[test]
    fn interleaved_alloc_free() {
        with_heap(4096, |tlsf| {
            let la = Layout::from_size_align(32, 8).unwrap();
            let lb = Layout::from_size_align(256, 8).unwrap();

            let a = unsafe { tlsf.alloc(la) };
            let b = unsafe { tlsf.alloc(lb) };
            let c = unsafe { tlsf.alloc(la) };
            assert!(!a.is_null() && !b.is_null() && !c.is_null());

            // Free the middle one → should not merge with a or c.
            unsafe { tlsf.dealloc(b, lb) };

            // Re-allocate the same size → should reuse the freed block.
            let d = unsafe { tlsf.alloc(lb) };
            assert!(!d.is_null());

            unsafe {
                tlsf.dealloc(a, la);
                tlsf.dealloc(c, la);
                tlsf.dealloc(d, lb);
            }
            assert_eq!(tlsf.used, 0);
        });
    }

    #[test]
    fn coalescing_restores_full_heap() {
        with_heap(2048, |tlsf| {
            let la = Layout::from_size_align(64, 8).unwrap();
            let a = unsafe { tlsf.alloc(la) };
            let b = unsafe { tlsf.alloc(la) };
            let c = unsafe { tlsf.alloc(la) };
            assert!(!a.is_null() && !b.is_null() && !c.is_null());

            // Free in order: should coalesce into one big block.
            unsafe {
                tlsf.dealloc(a, la);
                tlsf.dealloc(b, la);
                tlsf.dealloc(c, la);
            }
            assert_eq!(tlsf.used, 0);

            // The heap should be able to allocate nearly all of it.
            let big = Layout::from_size_align(1800, 8).unwrap();
            let p = unsafe { tlsf.alloc(big) };
            assert!(!p.is_null(), "coalescing should restore the full heap");
            unsafe { tlsf.dealloc(p, big) };
        });
    }

    #[test]
    fn alloc_too_large_returns_null() {
        with_heap(256, |tlsf| {
            let big = Layout::from_size_align(1024, 8).unwrap();
            let p = unsafe { tlsf.alloc(big) };
            assert!(p.is_null());
        });
    }

    #[test]
    fn peak_tracking() {
        with_heap(4096, |tlsf| {
            let la = Layout::from_size_align(128, 8).unwrap();
            let a = unsafe { tlsf.alloc(la) };
            let b = unsafe { tlsf.alloc(la) };
            let peak_after_two = tlsf.peak;
            unsafe { tlsf.dealloc(b, la) };
            assert_eq!(tlsf.peak, peak_after_two, "peak should not decrease");
            unsafe { tlsf.dealloc(a, la) };
            assert_eq!(tlsf.peak, peak_after_two);
        });
    }

    #[test]
    fn write_then_read_back() {
        with_heap(4096, |tlsf| {
            let layout = Layout::from_size_align(256, 8).unwrap();
            let ptr = unsafe { tlsf.alloc(layout) };
            assert!(!ptr.is_null());

            // Write a known pattern.
            unsafe {
                for i in 0..256 {
                    *ptr.add(i) = (i & 0xFF) as u8;
                }
            }

            // Read it back.
            unsafe {
                for i in 0..256 {
                    assert_eq!(*ptr.add(i), (i & 0xFF) as u8);
                }
            }
            unsafe { tlsf.dealloc(ptr, layout) };
        });
    }

    #[test]
    fn stress_random_sizes() {
        with_heap(8192, |tlsf| {
            let mut ptrs: Vec<(*mut u8, Layout)> = Vec::new();
            let sizes = [16, 32, 48, 64, 100, 128, 200, 256, 512];

            // Allocate.
            for &sz in &sizes {
                let la = Layout::from_size_align(sz, 8).unwrap();
                let p = unsafe { tlsf.alloc(la) };
                assert!(!p.is_null(), "failed to alloc {} bytes", sz);
                ptrs.push((p, la));
            }

            // Free every other one.
            let mut kept = Vec::new();
            for (i, (p, la)) in ptrs.into_iter().enumerate() {
                if i % 2 == 0 {
                    unsafe { tlsf.dealloc(p, la) };
                } else {
                    kept.push((p, la));
                }
            }

            // Reallocate freed slots.
            for &sz in &[16, 48, 100, 200, 512] {
                let la = Layout::from_size_align(sz, 8).unwrap();
                let p = unsafe { tlsf.alloc(la) };
                assert!(!p.is_null(), "realloc of {} bytes failed", sz);
                kept.push((p, la));
            }

            // Free everything.
            for (p, la) in kept {
                unsafe { tlsf.dealloc(p, la) };
            }
            assert_eq!(tlsf.used, 0);
        });
    }

    // ---------- Alignment tests ----------

    #[test]
    fn alloc_respects_word_alignment() {
        with_heap(4096, |tlsf| {
            for &sz in &[1usize, 3, 7, 13, 31, 64] {
                let la = Layout::from_size_align(sz, WORD).unwrap();
                let p = unsafe { tlsf.alloc(la) };
                assert!(!p.is_null());
                assert_eq!(
                    p as usize % WORD,
                    0,
                    "pointer for size {} is not {}-byte aligned",
                    sz,
                    WORD,
                );
                unsafe { tlsf.dealloc(p, la) };
            }
        });
    }

    #[test]
    fn alloc_respects_large_alignment() {
        with_heap(4096, |tlsf| {
            // 16-byte and 32-byte alignments (common for SIMD)
            for &align in &[16usize, 32] {
                let la = Layout::from_size_align(64, align).unwrap();
                let p = unsafe { tlsf.alloc(la) };
                assert!(!p.is_null());
                assert_eq!(p as usize % align, 0, "pointer not {}-byte aligned", align,);
                // Write and read back to ensure the pointer is valid.
                unsafe {
                    for i in 0..64 {
                        *p.add(i) = 0xAA;
                    }
                    for i in 0..64 {
                        assert_eq!(*p.add(i), 0xAA);
                    }
                    tlsf.dealloc(p, la);
                }
            }
        });
    }

    // ---------- Coalescing directions ----------

    #[test]
    fn coalesce_forward_only() {
        // Free B then A (A is before B) → B merges forward into A.
        with_heap(2048, |tlsf| {
            let la = Layout::from_size_align(64, 8).unwrap();
            let a = unsafe { tlsf.alloc(la) };
            let b = unsafe { tlsf.alloc(la) };
            let c = unsafe { tlsf.alloc(la) };
            assert!(!a.is_null() && !b.is_null() && !c.is_null());

            unsafe { tlsf.dealloc(b, la) }; // free middle
            unsafe { tlsf.dealloc(a, la) }; // free first → should coalesce with B

            // A+B should now be one free block, big enough for ~140 bytes.
            let big = Layout::from_size_align(140, 8).unwrap();
            let p = unsafe { tlsf.alloc(big) };
            assert!(!p.is_null(), "forward coalescing should create a larger block");
            unsafe {
                tlsf.dealloc(p, big);
                tlsf.dealloc(c, la);
            }
            assert_eq!(tlsf.used, 0);
        });
    }

    #[test]
    fn coalesce_backward_only() {
        // Free A then B → B merges backward into A.
        with_heap(2048, |tlsf| {
            let la = Layout::from_size_align(64, 8).unwrap();
            let a = unsafe { tlsf.alloc(la) };
            let b = unsafe { tlsf.alloc(la) };
            let c = unsafe { tlsf.alloc(la) };
            assert!(!a.is_null() && !b.is_null() && !c.is_null());

            unsafe { tlsf.dealloc(a, la) }; // free first
            unsafe { tlsf.dealloc(b, la) }; // free second → should coalesce with A

            let big = Layout::from_size_align(140, 8).unwrap();
            let p = unsafe { tlsf.alloc(big) };
            assert!(!p.is_null(), "backward coalescing should create a larger block");
            unsafe {
                tlsf.dealloc(p, big);
                tlsf.dealloc(c, la);
            }
            assert_eq!(tlsf.used, 0);
        });
    }

    #[test]
    fn coalesce_both_directions() {
        // Free A, free C, then free B → B should coalesce with both A and C.
        with_heap(4096, |tlsf| {
            let la = Layout::from_size_align(64, 8).unwrap();
            let a = unsafe { tlsf.alloc(la) };
            let b = unsafe { tlsf.alloc(la) };
            let c = unsafe { tlsf.alloc(la) };
            let d = unsafe { tlsf.alloc(la) };
            assert!(!a.is_null() && !b.is_null() && !c.is_null() && !d.is_null());

            unsafe { tlsf.dealloc(a, la) };
            unsafe { tlsf.dealloc(c, la) };
            unsafe { tlsf.dealloc(b, la) }; // merges with A (backward) and C (forward)

            // A+B+C should be one big block now.
            let big = Layout::from_size_align(210, 8).unwrap();
            let p = unsafe { tlsf.alloc(big) };
            assert!(!p.is_null(), "both-direction coalescing should work");
            unsafe {
                tlsf.dealloc(p, big);
                tlsf.dealloc(d, la);
            }
            assert_eq!(tlsf.used, 0);
        });
    }

    // ---------- Edge cases ----------

    #[test]
    fn alloc_1_byte() {
        with_heap(1024, |tlsf| {
            let la = Layout::from_size_align(1, 1).unwrap();
            let p = unsafe { tlsf.alloc(la) };
            assert!(!p.is_null());
            unsafe {
                *p = 0x42;
                assert_eq!(*p, 0x42);
                tlsf.dealloc(p, la);
            }
            assert_eq!(tlsf.used, 0);
        });
    }

    #[test]
    fn alloc_exact_heap_minus_overhead() {
        // Try to allocate the maximum possible from a small heap.
        let heap_size = 512;
        with_heap(heap_size, |tlsf| {
            // The usable free block is heap - WORD (sentinel).
            // The block header takes WORD, so max user data ≈ heap - 2*WORD.
            let max_user = heap_size - 2 * WORD;
            let la = Layout::from_size_align(max_user, WORD).unwrap();
            let p = unsafe { tlsf.alloc(la) };
            assert!(!p.is_null(), "should be able to allocate max_user={}", max_user);
            unsafe { tlsf.dealloc(p, la) };
            assert_eq!(tlsf.used, 0);
        });
    }

    #[test]
    fn exhaust_then_free_then_reuse() {
        with_heap(512, |tlsf| {
            let la = Layout::from_size_align(64, 8).unwrap();
            let mut ptrs = Vec::new();

            // Allocate until failure.
            loop {
                let p = unsafe { tlsf.alloc(la) };
                if p.is_null() {
                    break;
                }
                ptrs.push(p);
            }
            let count = ptrs.len();
            assert!(count >= 3, "should fit at least 3 × 64-byte blocks in 512");

            // Free everything.
            for p in ptrs.drain(..) {
                unsafe { tlsf.dealloc(p, la) };
            }
            assert_eq!(tlsf.used, 0);

            // Reallocate the same number — all should succeed again.
            for _ in 0..count {
                let p = unsafe { tlsf.alloc(la) };
                assert!(!p.is_null(), "reallocation after full free should succeed");
                ptrs.push(p);
            }
            for p in ptrs {
                unsafe { tlsf.dealloc(p, la) };
            }
            assert_eq!(tlsf.used, 0);
        });
    }

    #[test]
    fn no_overlap_between_allocations() {
        with_heap(4096, |tlsf| {
            let sizes = [32usize, 64, 128, 48, 96, 200];
            let mut regions: Vec<(*mut u8, usize)> = Vec::new();

            for &sz in &sizes {
                let la = Layout::from_size_align(sz, 8).unwrap();
                let p = unsafe { tlsf.alloc(la) };
                assert!(!p.is_null());
                regions.push((p, sz));
            }

            // Verify no pair of allocations overlaps.
            for i in 0..regions.len() {
                let (pi, si) = regions[i];
                let ri = pi as usize..pi as usize + si;
                for j in (i + 1)..regions.len() {
                    let (pj, sj) = regions[j];
                    let rj = pj as usize..pj as usize + sj;
                    assert!(
                        ri.end <= rj.start || rj.end <= ri.start,
                        "allocations {} ([{:#x}..{:#x}]) and {} ([{:#x}..{:#x}]) overlap",
                        i,
                        ri.start,
                        ri.end,
                        j,
                        rj.start,
                        rj.end,
                    );
                }
            }

            for (p, sz) in regions {
                let la = Layout::from_size_align(sz, 8).unwrap();
                unsafe { tlsf.dealloc(p, la) };
            }
            assert_eq!(tlsf.used, 0);
        });
    }

    #[test]
    fn data_survives_adjacent_alloc_free() {
        // Allocate A, write a pattern, alloc B, free B, verify A
        // is undamaged. Guards against off-by-one in block headers.
        with_heap(4096, |tlsf| {
            let la = Layout::from_size_align(128, 8).unwrap();
            let a = unsafe { tlsf.alloc(la) };
            assert!(!a.is_null());

            // Fill A with 0xAA.
            unsafe {
                for i in 0..128 {
                    *a.add(i) = 0xAA;
                }
            }

            // Allocate and free B (neighbor of A in the heap).
            let b = unsafe { tlsf.alloc(la) };
            assert!(!b.is_null());
            unsafe {
                for i in 0..128 {
                    *b.add(i) = 0xBB;
                }
                tlsf.dealloc(b, la);
            }

            // A must be intact.
            unsafe {
                for i in 0..128 {
                    assert_eq!(*a.add(i), 0xAA, "A corrupted at byte {} after adjacent alloc/free", i,);
                }
                tlsf.dealloc(a, la);
            }
            assert_eq!(tlsf.used, 0);
        });
    }

    // ---------- Fragmentation / realistic crypto scenario ----------

    #[test]
    fn simulated_crypto_session() {
        // Simulate a TLS-like session: Cipher ctx + Mac ctx + temp
        // buffers, with interleaved alloc/free.
        with_heap(8192, |tlsf| {
            let cipher_layout = Layout::from_size_align(488, 8).unwrap();
            let mac_layout = Layout::from_size_align(752, 8).unwrap();
            let tmp_layout = Layout::from_size_align(256, 8).unwrap();

            // Open session: allocate cipher + mac.
            let cipher = unsafe { tlsf.alloc(cipher_layout) };
            let mac = unsafe { tlsf.alloc(mac_layout) };
            assert!(!cipher.is_null() && !mac.is_null());

            // Process 5 "records": each allocates a temp buffer.
            for _ in 0..5 {
                let tmp = unsafe { tlsf.alloc(tmp_layout) };
                assert!(!tmp.is_null(), "temp buffer alloc failed");
                unsafe {
                    for i in 0..256 {
                        *tmp.add(i) = 0xCC;
                    }
                    tlsf.dealloc(tmp, tmp_layout);
                }
            }

            // Close session.
            unsafe {
                tlsf.dealloc(mac, mac_layout);
                tlsf.dealloc(cipher, cipher_layout);
            }
            assert_eq!(tlsf.used, 0);

            // Another session should work fine (no permanent frag).
            let cipher2 = unsafe { tlsf.alloc(cipher_layout) };
            let mac2 = unsafe { tlsf.alloc(mac_layout) };
            assert!(!cipher2.is_null() && !mac2.is_null());
            unsafe {
                tlsf.dealloc(cipher2, cipher_layout);
                tlsf.dealloc(mac2, mac_layout);
            }
            assert_eq!(tlsf.used, 0);
        });
    }

    #[test]
    fn repeated_init_resets_state() {
        let mut buf = vec![0u8; 2048];
        let mut tlsf = TlsfInner::new();
        unsafe { tlsf.init(buf.as_mut_ptr(), 2048) };

        let la = Layout::from_size_align(128, 8).unwrap();
        let p = unsafe { tlsf.alloc(la) };
        assert!(!p.is_null());
        // Don't free — just re-init. Should reset everything.
        unsafe { tlsf.init(buf.as_mut_ptr(), 2048) };
        assert_eq!(tlsf.used, 0);
        assert_eq!(tlsf.peak, 0);

        // Should be able to allocate again.
        let p2 = unsafe { tlsf.alloc(la) };
        assert!(!p2.is_null());
        unsafe { tlsf.dealloc(p2, la) };
    }

    #[test]
    fn minimum_heap_size() {
        // The smallest heap that should work: MIN_BLOCK + WORD.
        let min = MIN_BLOCK + WORD;
        with_heap(min, |tlsf| {
            let la = Layout::from_size_align(1, 1).unwrap();
            let p = unsafe { tlsf.alloc(la) };
            assert!(!p.is_null(), "1-byte alloc on minimum heap should work");
            unsafe { tlsf.dealloc(p, la) };
            assert_eq!(tlsf.used, 0);
        });
    }

    #[test]
    fn dealloc_null_is_noop() {
        with_heap(1024, |tlsf| {
            let la = Layout::from_size_align(8, 8).unwrap();
            // Should not panic or corrupt state.
            unsafe { tlsf.dealloc(ptr::null_mut(), la) };
            assert_eq!(tlsf.used, 0);
        });
    }
}