frame-alloc 0.1.0

//! White-box, depot allocator specific tests.
//!
//! Allocator-agnostic API conformance (alloc/dealloc round-trips, holes, stats
//! invariants, concurrency) and the randomised occupancy oracle live in the
//! external black-box suite under `tests/`.

extern crate std;

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

const BASE: PageSize = PageSize::from_log2(6);
const ORDERS: usize = 5;
const SLOTS: usize = 2;
const CAP: usize = 8;
const DEPOT_CAP: usize = 16;

// 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)
    }
}

type Dep = DepotAllocator<
    SummaryBuddyAllocator<ORDERS, TestProvenance>,
    TestCpuId,
    SLOTS,
    CAP,
    DEPOT_CAP,
    NoInterruptControl,
>;

fn pool(total_frames: usize) -> (Dep, OwnedRegion) {
    let total_bytes = total_frames * BASE.bytes();
    let max_block = BASE.bytes() << (ORDERS - 1);
    let region = OwnedRegion::new(total_bytes, max_block);
    let alloc: Dep = DepotAllocator::new(
        BASE,
        SummaryBuddyAllocator::<ORDERS, TestProvenance>::new(BASE),
    );
    unsafe {
        alloc.init_region(region.addr(), total_bytes);
    }
    (alloc, region)
}

/// Fill the depot to capacity by allocating `CAP + DEPOT_CAP` base frames on CPU 0
/// and freeing them all back: the magazine keeps `CAP`, the rest overflow into the
/// shared depot. Returns once `depot_len() > 0`.
fn fill_depot<A: PhysicalAllocator>(a: &A) {
    let held: Vec<usize> = (0..(CAP + DEPOT_CAP))
        .map(|_| a.allocate_physical(BASE, N1).unwrap())
        .collect();
    for &p in &held {
        unsafe { a.deallocate_physical(BASE, N1, p) };
    }
}

fn find_aligned_run(frames: &[usize], count: usize) -> Vec<usize> {
    let available: HashSet<usize> = frames.iter().copied().collect();
    for &base in frames {
        if base % (count * BASE.bytes()) == 0
            && (0..count).all(|i| available.contains(&(base + i * BASE.bytes())))
        {
            return (0..count).map(|i| base + i * BASE.bytes()).collect();
        }
    }
    panic!("exhausted pool must contain an aligned {count}-frame run");
}

fn assert_multiframe_recovery(request_frames: usize) {
    set_cpu(0);
    let (a, _region) = pool(128);
    let mut held = Vec::new();
    while let Ok(p) = a.allocate_physical(BASE, N1) {
        held.push(p);
    }
    held.sort_unstable();

    let run = find_aligned_run(&held, request_frames);
    let release_count = (SLOTS * CAP + CAP / 2)
        .max(request_frames)
        .min(SLOTS * CAP + DEPOT_CAP);
    let mut released = run.clone();
    for &p in &held {
        if released.len() == release_count {
            break;
        }
        if !released.contains(&p) {
            released.push(p);
        }
    }
    held.retain(|p| !released.contains(p));

    for (i, &p) in released.iter().enumerate() {
        set_cpu(i % SLOTS);
        unsafe { a.deallocate_physical(BASE, N1, p) };
    }
    assert!(a.depot_len() > 0, "setup must populate the shared depot");
    assert!(
        a.cached_frames() > a.depot_len(),
        "setup must retain frames in magazines as well"
    );

    set_cpu(0);
    let count = NonZeroUsize::new(request_frames).unwrap();
    let p = a
        .allocate_physical(BASE, count)
        .expect("wrapper caches must be reclaimed after backend OOM");
    let released_set: HashSet<usize> = released.iter().copied().collect();
    assert!(
        (0..request_frames).all(|i| released_set.contains(&(p + i * BASE.bytes()))),
        "allocation must come from frames that were sequestered in wrapper caches"
    );
    assert!(a.frames_flushed() >= request_frames);
    if request_frames == 2 {
        assert!(
            a.cached_frames() > 0,
            "progressive recovery should stop once the request is satisfiable"
        );
    }

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

// `new` must be `const fn`.
static _STATIC_DEPOT: Dep = DepotAllocator::new(
    BASE,
    SummaryBuddyAllocator::<ORDERS, TestProvenance>::new(BASE),
);

#[test]
fn alloc_dealloc_single_is_lifo_and_cached() {
    set_cpu(0);
    let (a, _region) = pool(64);

    let p = a.allocate_physical(BASE, N1).expect("alloc failed");
    assert_eq!(p % BASE.bytes(), 0, "address not frame-aligned");
    assert!(a.cached_frames() > 0, "refill should leave frames cached");

    unsafe { a.deallocate_physical(BASE, N1, p) };
    let p2 = a.allocate_physical(BASE, N1).expect("re-alloc failed");
    assert_eq!(p2, p, "expected the just-freed frame back (LIFO)");
    unsafe { a.deallocate_physical(BASE, N1, p2) };
}

#[test]
fn multiframe_requests_pass_through() {
    set_cpu(0);
    let (a, _region) = pool(64);

    let n4 = NonZeroUsize::new(4).unwrap();
    let p = a
        .allocate_physical(BASE, n4)
        .expect("contiguous alloc failed");
    assert_eq!(p % (4 * BASE.bytes()), 0, "order-2 block not aligned");
    assert_eq!(
        a.cached_frames(),
        0,
        "multi-frame traffic must not touch the magazines or depot"
    );
    unsafe { a.deallocate_physical(BASE, n4, p) };
    assert_eq!(
        a.cached_frames(),
        0,
        "multi-frame free must not touch the magazines or depot"
    );
}

#[test]
fn dealloc_caches_in_current_cpu_magazine() {
    let (a, _region) = pool(64);

    set_cpu(0);
    let p = a.allocate_physical(BASE, N1).expect("alloc failed");
    let cached_after_alloc = a.cached_frames();

    set_cpu(1);
    unsafe { a.deallocate_physical(BASE, N1, p) };
    assert_eq!(
        a.cached_frames(),
        cached_after_alloc + 1,
        "free must cache in the current CPU's magazine"
    );

    let p1 = a.allocate_physical(BASE, N1).expect("re-alloc failed");
    assert_eq!(p1, p, "slot 1 should return the frame it just cached");
    unsafe { a.deallocate_physical(BASE, N1, p1) };
}

#[test]
fn frees_overflow_into_depot_and_other_cpu_reuses() {
    set_cpu(0);
    let (a, _region) = pool(256);

    let held: Vec<usize> = (0..(CAP + DEPOT_CAP))
        .map(|_| a.allocate_physical(BASE, N1).unwrap())
        .collect();
    for &p in &held {
        unsafe { a.deallocate_physical(BASE, N1, p) };
    }
    assert!(
        a.depot_len() > 0,
        "magazine overflow should have populated the depot"
    );
    assert!(a.peak_depot_len() >= a.depot_len());

    set_cpu(1);
    let freed: HashSet<usize> = held.iter().copied().collect();
    let p = a
        .allocate_physical(BASE, N1)
        .expect("cross-CPU alloc failed");
    assert!(
        freed.contains(&p),
        "CPU 1 should reuse a frame CPU 0 freed via the depot"
    );
    unsafe { a.deallocate_physical(BASE, N1, p) };
}

#[test]
fn successful_backend_request_preserves_caches() {
    set_cpu(0);
    let (a, _region) = pool(256);

    let held: Vec<usize> = (0..(CAP + DEPOT_CAP))
        .map(|_| a.allocate_physical(BASE, N1).unwrap())
        .collect();
    for &p in &held {
        unsafe { a.deallocate_physical(BASE, N1, p) };
    }
    let cached_before = a.cached_frames();
    let depot_before = a.depot_len();
    assert!(depot_before > 0, "depot should be populated");

    let big = NonZeroUsize::new(16).unwrap();
    let p = a
        .allocate_physical(BASE, big)
        .expect("contiguous allocation failed");
    assert_eq!(
        a.depot_len(),
        depot_before,
        "recovery must not drain before the backend reports OOM"
    );
    assert_eq!(a.cached_frames(), cached_before);
    assert_eq!(a.frames_flushed(), 0);
    unsafe { a.deallocate_physical(BASE, big, p) };
}

#[test]
fn single_frame_oom_steals_from_sibling_magazine() {
    set_cpu(0);
    let (a, _region) = pool(64);
    let mut held = Vec::new();
    while let Ok(p) = a.allocate_physical(BASE, N1) {
        held.push(p);
    }

    let stranded = held.pop().expect("pool must contain a frame");
    set_cpu(1);
    unsafe { a.deallocate_physical(BASE, N1, stranded) };

    set_cpu(0);
    assert_eq!(
        a.allocate_physical(BASE, N1),
        Ok(stranded),
        "backend OOM must not hide a frame in a sibling magazine"
    );
}

#[test]
fn progressive_recovery_handles_requests_below_equal_and_above_cap() {
    for request_frames in [2, CAP, 2 * CAP] {
        assert_multiframe_recovery(request_frames);
    }
}

#[test]
fn explicit_flush_drains_depot_and_magazines() {
    set_cpu(0);
    let (a, _region) = pool(256);
    fill_depot(&a);
    let cached_before = a.cached_frames();
    assert!(a.depot_len() > 0);
    assert!(cached_before > a.depot_len());

    a.flush();

    assert_eq!(a.depot_len(), 0);
    assert_eq!(a.cached_frames(), 0);
    assert_eq!(a.frames_flushed(), cached_before);
}

#[test]
fn non_oom_error_preserves_all_caches() {
    set_cpu(0);
    let (a, _region) = pool(256);
    fill_depot(&a);
    let cached_before = a.cached_frames();
    let flushed_before = a.frames_flushed();

    let too_large = NonZeroUsize::new(1usize << ORDERS).unwrap();
    assert_eq!(
        a.allocate_physical(BASE, too_large),
        Err(crate::AllocError::RequestTooLarge)
    );
    assert_eq!(a.cached_frames(), cached_before);
    assert_eq!(a.frames_flushed(), flushed_before);
}

#[test]
fn small_multiframe_bypasses_and_keeps_depot() {
    set_cpu(0);
    let (a, _region) = pool(256);

    let held: Vec<usize> = (0..(CAP + DEPOT_CAP))
        .map(|_| a.allocate_physical(BASE, N1).unwrap())
        .collect();
    for &p in &held {
        unsafe { a.deallocate_physical(BASE, N1, p) };
    }
    let depot_before = a.depot_len();
    assert!(depot_before > 0, "depot should be populated");

    let n4 = NonZeroUsize::new(4).unwrap();
    let p = a
        .allocate_physical(BASE, n4)
        .expect("multi-frame alloc failed");
    assert_eq!(
        a.depot_len(),
        depot_before,
        "a small multi-frame request must not flush the depot"
    );
    unsafe { a.deallocate_physical(BASE, n4, p) };
}

#[test]
fn satisfiable_contiguous_request_preserves_depot() {
    set_cpu(0);
    let (a, _region) = pool(256);
    fill_depot(&a);

    let depot_before = a.depot_len();
    assert!(depot_before > 0, "depot should be populated");

    let big = NonZeroUsize::new(16).unwrap();
    let p = a
        .allocate_physical(BASE, big)
        .expect("contiguous alloc should succeed from the untouched backend");

    assert_eq!(
        a.depot_len(),
        depot_before,
        "recovery must not drain the depot when the request is already satisfiable"
    );
    assert_eq!(
        a.frames_flushed(),
        0,
        "no frames should have been returned to the backend"
    );
    unsafe { a.deallocate_physical(BASE, big, p) };
}

#[test]
fn unsatisfiable_request_drains_all_caches() {
    set_cpu(0);
    let (a, _region) = pool(16);
    let mut held = Vec::new();
    while let Ok(p) = a.allocate_physical(BASE, N1) {
        held.push(p);
    }
    for &p in &held {
        unsafe { a.deallocate_physical(BASE, N1, p) };
    }

    let depot_before = a.depot_len();
    let cached_before = a.cached_frames();
    assert!(depot_before > 0, "depot should be populated");

    let block = NonZeroUsize::new(16).unwrap(); // order-4 = pool's max block
    let r = a.allocate_physical(BASE, block);
    assert!(
        r.is_err(),
        "request must OOM — the cache cannot rebuild order-4"
    );

    assert_eq!(
        a.depot_len(),
        0,
        "the failing request should have drained the whole depot"
    );
    assert_eq!(
        a.cached_frames(),
        0,
        "the failing request should also have drained every magazine"
    );
    assert_eq!(
        a.frames_flushed(),
        cached_before,
        "all cached frames should have been returned"
    );
}

#[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, _region) = pool(if cfg!(miri) { 128 } 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) % SLOTS;
                for i 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");
                        set_cpu(neighbour);
                        unsafe { alloc.deallocate_physical(BASE, N1, p) };
                    }
                    let big = NonZeroUsize::new(16).unwrap();
                    if i % 512 == 0
                        && let Ok(p) = alloc.allocate_physical(BASE, big)
                    {
                        unsafe { alloc.deallocate_physical(BASE, big, 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");
}

#[test]
fn try_init_forwards_backend_error() {
    // The wrapper's `try_init` forwards the backend's `InitError` unchanged.
    let total_bytes = 16 * BASE.bytes();
    let max_block = BASE.bytes() << (ORDERS - 1);
    let region = OwnedRegion::new(total_bytes, max_block);
    let alloc: Dep = DepotAllocator::new(
        BASE,
        SummaryBuddyAllocator::<ORDERS, TestProvenance>::new(BASE),
    );
    // phys_base + 1 is not base-frame aligned — a backend `Misaligned`.
    let e = unsafe { alloc.try_init_region(region.addr() + 1, total_bytes) };
    assert_eq!(
        e,
        Err(InitError::Misaligned {
            required: BASE.bytes()
        })
    );
    // Untouched: a corrected retry succeeds and the allocator works.
    let ok = unsafe { alloc.try_init_region(region.addr(), total_bytes) };
    assert_eq!(ok, Ok(()));
    assert!(alloc.allocate_physical(BASE, N1).is_ok());
}

#[test]
#[should_panic(expected = "SLOTS must be > 0")]
fn zero_slots_panics() {
    let _ = DepotAllocator::<
        SummaryBuddyAllocator<ORDERS, TestProvenance>,
        TestCpuId,
        0,
        CAP,
        DEPOT_CAP,
    >::new(
        BASE,
        SummaryBuddyAllocator::<ORDERS, TestProvenance>::new(BASE),
    );
}

#[test]
#[should_panic(expected = "CAP must be >= 2")]
fn cap_below_two_panics() {
    let _ = DepotAllocator::<
        SummaryBuddyAllocator<ORDERS, TestProvenance>,
        TestCpuId,
        SLOTS,
        1,
        DEPOT_CAP,
    >::new(
        BASE,
        SummaryBuddyAllocator::<ORDERS, TestProvenance>::new(BASE),
    );
}

#[test]
#[should_panic(expected = "DEPOT_CAP must be >= 1")]
fn depot_cap_zero_panics() {
    let _ = DepotAllocator::<SummaryBuddyAllocator<ORDERS, TestProvenance>, TestCpuId, SLOTS, CAP, 0>::new(
        BASE,
        SummaryBuddyAllocator::<ORDERS, TestProvenance>::new(BASE),
    );
}