frame-alloc 0.1.0

//! White-box tests for *nested wrappers* — a per-CPU cache placed **inside**
//! each region (`RegionedAllocator<.., DepotAllocator<..>, ..>`).
//!
//! The plain occupancy contract (no aliasing, in range, conservation) is held
//! over the composition by the black-box suites; the tests here pin the behaviour
//! that is *unique* to the stack and invisible to that contract:
//!
//! * deallocation routes by physical address to the **owning** region's cache,
//!   not the freeing CPU's home region;
//! * each region's cache is private — region-`i` allocations never hand out a
//!   frame cached for region-`j`;
//! * the regioned layer still work-steals across regions when a home region
//!   (cache included) is exhausted, and nothing leaks through the two layers.

extern crate std;

use crate::tests::common::{N1, OwnedRegion, TestProvenance};
use crate::{
    AllocError, CpuId, DepotAllocator, NoInterruptControl, PageSize, PhysRange, PhysicalAllocator,
    RegionedAllocator, SummaryBuddyAllocator,
};
use core::cell::Cell;
use core::sync::atomic::AtomicUsize;
use core::sync::atomic::Ordering::Relaxed;
use std::collections::HashSet;
use std::sync::{Arc, Mutex};
use std::vec::Vec;

const BASE: PageSize = PageSize::from_log2(6);
const ORDERS: usize = 3;
const REGIONS: usize = 2;
const SLOTS: usize = 2;
const CAP: usize = 8;

fn max_block() -> usize {
    BASE.bytes() << (ORDERS - 1)
}

// A CPU selector whose "current CPU" is settable.
std::thread_local! {
    static CPU: Cell<usize> = const { Cell::new(0) };
}

fn set_cpu(v: usize) {
    CPU.set(v);
}

struct TestCpuId;
impl CpuId for TestCpuId {
    fn current_cpu() -> usize {
        CPU.with(Cell::get)
    }
}

pub type RegionedDepot<
    const REGIONS: usize,
    A,
    S,
    const SLOTS: usize,
    const CAP: usize = 128,
    const DEPOT_CAP: usize = 512,
    I = NoInterruptControl,
> = RegionedAllocator<REGIONS, DepotAllocator<A, S, SLOTS, CAP, DEPOT_CAP, I>, S>;

type Comp =
    RegionedDepot<REGIONS, SummaryBuddyAllocator<ORDERS, TestProvenance>, TestCpuId, SLOTS, CAP>;

/// Two-region allocator, each region a magazine-cached buddy over its own
/// disjoint span of `frames` frames. Returns the wrapper plus the two backing
/// regions (kept alive by the caller).
fn pool(frames: usize) -> (Comp, OwnedRegion, OwnedRegion) {
    let bytes = frames * BASE.bytes();
    let r0 = OwnedRegion::new(bytes, max_block());
    let r1 = OwnedRegion::new(bytes, max_block());
    let alloc: Comp = RegionedAllocator::new(
        BASE,
        [const {
            DepotAllocator::new(
                BASE,
                SummaryBuddyAllocator::<ORDERS, TestProvenance>::new(BASE),
            )
        }; REGIONS],
    );
    unsafe {
        alloc.init_at(
            0,
            r0.addr(),
            bytes,
            &[PhysRange {
                base: r0.addr(),
                len: bytes,
            }],
        );
        alloc.init_at(
            1,
            r1.addr(),
            bytes,
            &[PhysRange {
                base: r1.addr(),
                len: bytes,
            }],
        );
    }
    (alloc, r0, r1)
}

/// `[lo, hi)` byte bounds of a region's span.
fn bounds(r: &OwnedRegion, frames: usize) -> (usize, usize) {
    (r.addr(), r.addr() + frames * BASE.bytes())
}

// `new` must stay `const fn`.
static _STATIC_COMP: Comp = RegionedAllocator::new(
    BASE,
    [const {
        DepotAllocator::new(
            BASE,
            SummaryBuddyAllocator::<ORDERS, TestProvenance>::new(BASE),
        )
    }; REGIONS],
);

#[test]
fn fast_path_cache_hit_stays_in_home_region() {
    set_cpu(0);
    let (a, r0, _r1) = pool(64);
    let (lo0, hi0) = bounds(&r0, 64);

    let p = a.allocate_physical(BASE, N1).expect("alloc failed");
    assert!((lo0..hi0).contains(&p), "home alloc must land in region 0");

    // Same CPU: the free caches into region 0's magazine, the re-alloc pops it.
    unsafe { a.deallocate_physical(BASE, N1, p) };
    let p2 = a.allocate_physical(BASE, N1).expect("re-alloc failed");
    assert_eq!(
        p2, p,
        "fast path should return the just-freed frame (LIFO cache)"
    );
    assert!(
        a.regions().any(|m| m.cached_frames() > 0),
        "the hit must have come from a per-region magazine, not a backend round-trip"
    );
    unsafe { a.deallocate_physical(BASE, N1, p2) };
}

#[test]
fn dealloc_routes_to_owning_region_not_freeing_cpu() {
    let (a, r0, _r1) = pool(64);
    let (lo0, hi0) = bounds(&r0, 64);

    set_cpu(0);
    let p = a
        .alloc_in_region(0, BASE, N1)
        .expect("region-0 alloc failed");
    assert!((lo0..hi0).contains(&p), "frame not in region 0");

    // Free it while the current CPU's *home* is region 1. Routing is by address,
    // so the frame must go back into region 0's cache.
    set_cpu(1);
    unsafe { a.deallocate_physical(BASE, N1, p) };
    let q = a
        .alloc_in_region(0, BASE, N1)
        .expect("region-0 re-alloc failed");
    assert_eq!(
        q, p,
        "a frame freed on a foreign-home CPU must return to its owning region's cache"
    );
    unsafe { a.deallocate_physical(BASE, N1, q) };
}

#[test]
fn per_region_caches_are_private() {
    set_cpu(0);
    let (a, _r0, r1) = pool(64);
    let (lo1, hi1) = bounds(&r1, 64);

    // Warm region 0's cache: allocate then free a batch back into it.
    let held: Vec<usize> = (0..CAP)
        .map(|_| {
            a.alloc_in_region(0, BASE, N1)
                .expect("region-0 alloc failed")
        })
        .collect();
    let cached0: HashSet<usize> = held.iter().copied().collect();
    for &p in &held {
        unsafe { a.deallocate_physical(BASE, N1, p) };
    }

    // Draining region 1 must never surface a frame cached for region 0.
    while let Ok(q) = a.alloc_in_region(1, BASE, N1) {
        assert!(
            (lo1..hi1).contains(&q),
            "region-1 alloc strayed out of region 1: {q:#x}"
        );
        assert!(
            !cached0.contains(&q),
            "region-1 alloc handed out a frame cached for region 0 — a cache crossed a boundary"
        );
    }
}

#[test]
fn home_exhaustion_steals_from_other_region() {
    let (a, r0, r1) = pool(32);
    let (lo1, hi1) = bounds(&r1, 32);
    set_cpu(0);

    // Drain region 0 (home of CPU 0), cache included, via the pinned path.
    let mut held = Vec::new();
    while let Ok(p) = a.alloc_in_region(0, BASE, N1) {
        let (lo0, hi0) = bounds(&r0, 32);
        assert!((lo0..hi0).contains(&p), "pinned drain left region 0");
        held.push(p);
    }

    let stolen = a.allocate_physical(BASE, N1).expect("work-steal failed");
    assert!(
        (lo1..hi1).contains(&stolen),
        "expected a region-1 frame via work-stealing, got {stolen:#x}"
    );

    unsafe { a.deallocate_physical(BASE, N1, stolen) };
    for p in held {
        unsafe { a.deallocate_physical(BASE, N1, p) };
    }
}

#[test]
fn full_drain_through_both_layers_conserves() {
    // Single CPU throughout: one region start, one magazine slot per region, so
    // the drain can never strand frames in an idle sibling slot.
    set_cpu(0);
    let (a, r0, r1) = pool(32);
    let in_pool = {
        let (lo0, hi0) = bounds(&r0, 32);
        let (lo1, hi1) = bounds(&r1, 32);
        move |p: usize| (lo0..hi0).contains(&p) || (lo1..hi1).contains(&p)
    };

    let mut addrs = HashSet::new();
    while let Ok(p) = a.allocate_physical(BASE, N1) {
        assert!(in_pool(p), "frame {p:#x} outside both regions");
        assert!(addrs.insert(p), "frame {p:#x} handed out twice");
    }
    let total = addrs.len();
    assert!(total > 0, "pool drained nothing");

    for &p in &addrs {
        unsafe { a.deallocate_physical(BASE, N1, p) };
    }
    let mut recovered = 0;
    while a.allocate_physical(BASE, N1).is_ok() {
        recovered += 1;
    }
    assert_eq!(
        recovered, total,
        "frames leaked through the regioned-over-magazine stack"
    );
}

#[test]
fn pinned_region_never_steals_even_with_cache() {
    let (a, _r0, r1) = pool(32);
    let (lo1, hi1) = bounds(&r1, 32);
    set_cpu(0);

    let mut held = Vec::new();
    while let Ok(p) = a.alloc_in_region(0, BASE, N1) {
        held.push(p);
    }
    assert_eq!(
        a.alloc_in_region(0, BASE, N1),
        Err(AllocError::OutOfMemory),
        "an exhausted pinned region must not steal another region"
    );
    // Region 1 is still untouched and allocatable.
    let other = a
        .alloc_in_region(1, BASE, N1)
        .expect("region 1 should be free");
    assert!((lo1..hi1).contains(&other));

    unsafe { a.deallocate_physical(BASE, N1, other) };
    for p in held {
        unsafe { a.deallocate_physical(BASE, N1, p) };
    }
}

#[cfg_attr(audit, ignore = "audit builds are sequential-only")]
#[test]
fn concurrent_cross_cpu_no_double_alloc_or_leak() {
    const ITERS: usize = if cfg!(miri) { 16 } else { 5_000 };
    let (alloc, _r0, _r1) = pool(if cfg!(miri) { 64 } else { 1024 });
    let alloc = Arc::new(alloc);
    let checked_out: Arc<Mutex<HashSet<usize>>> = Arc::new(Mutex::new(HashSet::new()));
    let next_cpu = Arc::new(AtomicUsize::new(0));

    let handles: Vec<_> = (0..4)
        .map(|_| {
            let alloc = Arc::clone(&alloc);
            let checked_out = Arc::clone(&checked_out);
            let next_cpu = Arc::clone(&next_cpu);
            std::thread::spawn(move || {
                let cpu = next_cpu.fetch_add(1, Relaxed);
                let neighbour = cpu + 1;
                for _ in 0..ITERS {
                    set_cpu(cpu);
                    if let Ok(p) = alloc.allocate_physical(BASE, N1) {
                        assert!(
                            checked_out.lock().unwrap().insert(p),
                            "frame {p:#x} handed out twice"
                        );
                        assert!(checked_out.lock().unwrap().remove(&p), "frame vanished");
                        // Free from a different CPU so the dealloc must route by
                        // address across both layers, not by the freeing CPU.
                        set_cpu(neighbour);
                        unsafe { alloc.deallocate_physical(BASE, N1, p) };
                    }
                }
            })
        })
        .collect();
    for h in handles {
        h.join().unwrap();
    }

    assert!(
        checked_out.lock().unwrap().is_empty(),
        "frames still checked out after all workers finished"
    );

    set_cpu(0);
    let mut remaining = 0;
    while alloc.allocate_physical(BASE, N1).is_ok() {
        remaining += 1;
    }
    assert!(remaining > 0, "pool empty after churn — frames leaked");
}