frame-alloc 0.1.0

//! White-box, regioned 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::strategies::cpu_id::NoCpuId;
use crate::tests::common::{N1, OwnedRegion, TestProvenance};
use crate::{
    AllocError, AllocatorStats, CpuId, InitError, PageSize, PhysRange, PhysicalAllocator,
    RegionedAllocator, SummaryBuddyAllocator,
};
use core::cell::Cell;
use std::collections::HashSet;

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

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

type Regioned =
    RegionedAllocator<REGIONS, SummaryBuddyAllocator<ORDERS, TestProvenance>, TestCpuId>;

/// Build a 2-region allocator over two *disjoint* heap spans of `frames` frames
/// each, with the CPU selector `S`.
fn pool_with<S>(
    frames: usize,
) -> (
    RegionedAllocator<REGIONS, SummaryBuddyAllocator<ORDERS, TestProvenance>, S>,
    OwnedRegion,
    OwnedRegion,
) {
    let bytes = frames * BASE.bytes();
    let r0 = OwnedRegion::new(bytes, max_block());
    let r1 = OwnedRegion::new(bytes, max_block());
    let alloc: RegionedAllocator<REGIONS, SummaryBuddyAllocator<ORDERS, TestProvenance>, S> =
        RegionedAllocator::new(
            BASE,
            [const { 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)
}

fn pool(frames: usize) -> (Regioned, OwnedRegion, OwnedRegion) {
    pool_with::<TestCpuId>(frames)
}

// `new` must be `const fn`.
static _STATIC_REGIONED: Regioned = RegionedAllocator::new(
    BASE,
    [const { SummaryBuddyAllocator::<ORDERS, TestProvenance>::new(BASE) }; REGIONS],
);

#[test]
fn dealloc_routes_by_address_not_current_cpu() {
    let (alloc, r0, _r1) = pool(32);
    let r0_lo = r0.addr();
    let r0_hi = r0.addr() + 32 * BASE.bytes();

    set_cpu(0);
    let phys = alloc.allocate_physical(BASE, N1).expect("alloc failed");
    assert!((r0_lo..r0_hi).contains(&phys), "home alloc not in region 0");

    set_cpu(1);
    unsafe { alloc.deallocate_physical(BASE, N1, phys) };

    set_cpu(0);
    let phys2 = alloc.allocate_physical(BASE, N1).expect("re-alloc failed");
    assert_eq!(
        phys2, phys,
        "freed frame did not return to its owning region"
    );
    unsafe { alloc.deallocate_physical(BASE, N1, phys2) };
}

#[test]
fn local_oom_steals_from_other_region() {
    let (alloc, r0, r1) = pool(32);
    set_cpu(0);

    let reserved0 = alloc.regions().next().unwrap().reserved_frames();
    let allocatable0 = 32 - reserved0;

    let r0_lo = r0.addr();
    let r0_hi = r0.addr() + 32 * BASE.bytes();
    for _ in 0..allocatable0 {
        let a = alloc
            .allocate_physical(BASE, N1)
            .expect("region-0 drain failed");
        assert!((r0_lo..r0_hi).contains(&a), "drain strayed out of region 0");
    }

    let r1_lo = r1.addr();
    let r1_hi = r1.addr() + 32 * BASE.bytes();
    let stolen = alloc
        .allocate_physical(BASE, N1)
        .expect("work-steal failed");
    assert!(
        (r1_lo..r1_hi).contains(&stolen),
        "expected a region-1 frame via work-stealing, got {stolen:#x}"
    );
}

#[test]
fn alloc_in_region_pins_and_never_steals() {
    let (alloc, r0, r1) = pool(32);
    let r0_lo = r0.addr();
    let r0_hi = r0.addr() + 32 * BASE.bytes();
    let r1_lo = r1.addr();
    let r1_hi = r1.addr() + 32 * BASE.bytes();

    set_cpu(0);
    let p = alloc
        .alloc_in_region(1, BASE, N1)
        .expect("pinned alloc failed");
    assert!(
        (r1_lo..r1_hi).contains(&p),
        "alloc_in_region(1) strayed out of region 1"
    );
    unsafe { alloc.deallocate_physical(BASE, N1, p) };

    let reserved0 = alloc.regions().next().unwrap().reserved_frames();
    let mut held = std::vec::Vec::new();
    for _ in 0..(32 - reserved0) {
        let a = alloc
            .alloc_in_region(0, BASE, N1)
            .expect("region-0 drain failed");
        assert!((r0_lo..r0_hi).contains(&a), "pinned drain left region 0");
        held.push(a);
    }
    assert_eq!(
        alloc.alloc_in_region(0, BASE, N1),
        Err(AllocError::OutOfMemory),
        "exhausted pin must not steal another region"
    );
    let still = alloc
        .alloc_in_region(1, BASE, N1)
        .expect("region 1 should be free");
    assert!((r1_lo..r1_hi).contains(&still));
    unsafe { alloc.deallocate_physical(BASE, N1, still) };
    for a in held {
        unsafe { alloc.deallocate_physical(BASE, N1, a) };
    }
}

#[test]
fn alloc_in_chain_follows_order_and_falls_back() {
    let (alloc, r0, r1) = pool(32);
    let r0_lo = r0.addr();
    let r0_hi = r0.addr() + 32 * BASE.bytes();
    let r1_lo = r1.addr();
    let r1_hi = r1.addr() + 32 * BASE.bytes();

    set_cpu(0);
    let p = alloc
        .alloc_in_chain(&[1, 0], BASE, N1)
        .expect("chain alloc failed");
    assert!(
        (r1_lo..r1_hi).contains(&p),
        "chain ignored its first preference"
    );
    unsafe { alloc.deallocate_physical(BASE, N1, p) };

    let reserved1 = alloc.regions().nth(1).unwrap().reserved_frames();
    let mut held = std::vec::Vec::new();
    for _ in 0..(32 - reserved1) {
        held.push(
            alloc
                .alloc_in_region(1, BASE, N1)
                .expect("region-1 drain failed"),
        );
    }
    let spilled = alloc
        .alloc_in_chain(&[1, 0], BASE, N1)
        .expect("chain fallback failed");
    assert!(
        (r0_lo..r0_hi).contains(&spilled),
        "chain did not fall back to region 0"
    );
    unsafe { alloc.deallocate_physical(BASE, N1, spilled) };
    for a in held {
        unsafe { alloc.deallocate_physical(BASE, N1, a) };
    }

    let mut held1 = std::vec::Vec::new();
    for _ in 0..(32 - reserved1) {
        held1.push(
            alloc
                .alloc_in_chain(&[1], BASE, N1)
                .expect("region-1 chain drain failed"),
        );
    }
    assert_eq!(
        alloc.alloc_in_chain(&[1], BASE, N1),
        Err(AllocError::OutOfMemory),
        "single-region chain must enforce the constraint by omission"
    );
    for a in held1 {
        unsafe { alloc.deallocate_physical(BASE, N1, a) };
    }
}

#[test]
fn alloc_in_chain_empty_is_oom() {
    let (alloc, _r0, _r1) = pool(32);
    assert_eq!(
        alloc.alloc_in_chain(&[], BASE, N1),
        Err(AllocError::OutOfMemory),
        "empty chain must be OutOfMemory"
    );
}

#[test]
fn drains_both_regions_and_conserves() {
    let (alloc, r0, r1) = pool(32);
    set_cpu(0);

    let reserved: usize = alloc.regions().map(|a| a.reserved_frames()).sum();
    let expected = 2 * 32 - reserved;

    let bounds = [
        (r0.addr(), r0.addr() + 32 * BASE.bytes()),
        (r1.addr(), r1.addr() + 32 * BASE.bytes()),
    ];

    let mut addrs = HashSet::new();
    while let Ok(a) = alloc.allocate_physical(BASE, N1) {
        assert!(
            bounds.iter().any(|&(lo, hi)| (lo..hi).contains(&a)),
            "frame {a:#x} outside both regions"
        );
        assert!(addrs.insert(a), "frame {a:#x} handed out twice");
    }
    assert_eq!(addrs.len(), expected, "first drain count wrong");

    for &a in &addrs {
        unsafe { alloc.deallocate_physical(BASE, N1, a) };
    }
    let mut recovered = 0;
    while alloc.allocate_physical(BASE, N1).is_ok() {
        recovered += 1;
    }
    assert_eq!(recovered, expected, "frames leaked after a full free cycle");
}

#[test]
fn default_selector_alloc_dealloc() {
    let (alloc, r0, _r1) = pool_with::<NoCpuId>(32);
    let r0_lo = r0.addr();
    let r0_hi = r0.addr() + 32 * BASE.bytes();

    let p = alloc.allocate_physical(BASE, N1).expect("alloc failed");
    assert!(
        (r0_lo..r0_hi).contains(&p),
        "NoCpuId must start allocation in region 0, got {p:#x}"
    );
    unsafe { alloc.deallocate_physical(BASE, N1, p) };
}

#[test]
fn add_usable_routes_to_owning_region() {
    const HOLE: usize = 8;
    let frames = 32;
    let fb = BASE.bytes();
    let bytes = frames * fb;
    let usable_len = (frames - HOLE) * fb;

    let r0 = OwnedRegion::new(bytes, max_block());
    let r1 = OwnedRegion::new(bytes, max_block());
    let alloc: Regioned = RegionedAllocator::new(
        BASE,
        [const { 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: usable_len,
            }],
        );
    }

    let mut it = alloc.regions();
    let r0_before = it.next().unwrap().free_bytes();
    let r1_before = it.next().unwrap().free_bytes();

    // The freed range lies in region 1's span.
    unsafe { alloc.add_usable(r1.addr() + usable_len, HOLE * fb) };

    let mut it = alloc.regions();
    let r0_after = it.next().unwrap().free_bytes();
    let r1_after = it.next().unwrap().free_bytes();

    assert_eq!(
        r0_after, r0_before,
        "an address in region 1 must not touch region 0"
    );
    assert_eq!(
        r1_after,
        r1_before + HOLE * fb,
        "the freed hole must be added to region 1's pool"
    );
}

#[test]
fn pad_cells_moves_each_element_exactly_once() {
    use crate::implementations::wrappers::regioned::pad_cells;
    use core::sync::atomic::{AtomicUsize, Ordering};

    static DROPS: AtomicUsize = AtomicUsize::new(0);

    struct Boom(usize);
    impl Drop for Boom {
        fn drop(&mut self) {
            DROPS.fetch_add(1, Ordering::Relaxed);
        }
    }

    DROPS.store(0, Ordering::Relaxed);
    let padded = pad_cells([Boom(0), Boom(1), Boom(2)]);
    // Moving the elements into their CachePadded cells must not drop anything.
    assert_eq!(
        DROPS.load(Ordering::Relaxed),
        0,
        "pad_cells dropped an element during the move"
    );
    // Each value survived the move intact and stayed in order (field access
    // autoderefs through CachePadded).
    assert_eq!([padded[0].0, padded[1].0, padded[2].0], [0, 1, 2]);

    drop(padded);
    // Every element is dropped exactly once: no leak (source array forgotten),
    // no double-drop (MaybeUninit scratch never dropped).
    assert_eq!(
        DROPS.load(Ordering::Relaxed),
        3,
        "expected exactly N drops after dropping the padded array"
    );
}

#[test]
fn init_at_double_init_errors() {
    let bytes = 32 * BASE.bytes();
    let r = OwnedRegion::new(bytes, max_block());
    let alloc: Regioned = RegionedAllocator::new(
        BASE,
        [const { SummaryBuddyAllocator::<ORDERS, TestProvenance>::new(BASE) }; REGIONS],
    );
    let usable = [PhysRange {
        base: r.addr(),
        len: bytes,
    }];
    unsafe {
        alloc
            .try_init_at(0, r.addr(), bytes, &usable)
            .expect("first init");
        let e = alloc.try_init_at(0, r.addr(), bytes, &usable);
        assert_eq!(e, Err(InitError::AlreadyInitialized));
    }
}

#[test]
fn init_at_overlapping_spans_errors() {
    let bytes = 32 * BASE.bytes();
    let r = OwnedRegion::new(bytes, max_block());
    let alloc: Regioned = RegionedAllocator::new(
        BASE,
        [const { SummaryBuddyAllocator::<ORDERS, TestProvenance>::new(BASE) }; REGIONS],
    );
    unsafe {
        alloc
            .try_init_at(
                0,
                r.addr(),
                bytes,
                &[PhysRange {
                    base: r.addr(),
                    len: bytes,
                }],
            )
            .expect("first init");
        // Region 1's span overlaps region 0 exactly; the reported `other` is 0.
        let e = alloc.try_init_at(
            1,
            r.addr(),
            bytes,
            &[PhysRange {
                base: r.addr(),
                len: bytes,
            }],
        );
        assert_eq!(e, Err(InitError::OverlapsRegion { other: 0 }));
    }
}

#[test]
#[should_panic(expected = "REGIONS must be > 0")]
fn zero_regions_panics() {
    let _ = RegionedAllocator::<0, SummaryBuddyAllocator<ORDERS, TestProvenance>, TestCpuId>::new(
        BASE,
        [],
    );
}

#[test]
#[should_panic(expected = "out of range")]
fn init_at_out_of_range_panics() {
    let bytes = 32 * BASE.bytes();
    let r = OwnedRegion::new(bytes, max_block());
    let alloc: Regioned = RegionedAllocator::new(
        BASE,
        [const { SummaryBuddyAllocator::<ORDERS, TestProvenance>::new(BASE) }; REGIONS],
    );
    unsafe {
        alloc.init_at(
            REGIONS,
            r.addr(),
            bytes,
            &[PhysRange {
                base: r.addr(),
                len: bytes,
            }],
        )
    };
}

#[test]
fn init_at_zero_span_errors() {
    let r = OwnedRegion::new(32 * BASE.bytes(), max_block());
    let alloc: Regioned = RegionedAllocator::new(
        BASE,
        [const { SummaryBuddyAllocator::<ORDERS, TestProvenance>::new(BASE) }; REGIONS],
    );
    let e = unsafe { alloc.try_init_at(0, r.addr(), 0, &[]) };
    assert_eq!(e, Err(InitError::InvalidSpan));
    // The failed attempt left region 0 uninitialised and retryable.
    assert_eq!(alloc.region_bounds(0), (0, 0));
}

#[test]
fn init_at_unaligned_span_errors() {
    let r = OwnedRegion::new(32 * BASE.bytes(), max_block());
    let alloc: Regioned = RegionedAllocator::new(
        BASE,
        [const { SummaryBuddyAllocator::<ORDERS, TestProvenance>::new(BASE) }; REGIONS],
    );
    // One byte past a whole number of base frames.
    let e = unsafe { alloc.try_init_at(0, r.addr(), BASE.bytes() + 1, &[]) };
    assert_eq!(e, Err(InitError::InvalidSpan));
}

#[test]
fn init_at_unaligned_base_errors() {
    let r = OwnedRegion::new(32 * BASE.bytes(), max_block());
    let alloc: Regioned = RegionedAllocator::new(
        BASE,
        [const { SummaryBuddyAllocator::<ORDERS, TestProvenance>::new(BASE) }; REGIONS],
    );
    let e = unsafe {
        alloc.try_init_at(
            0,
            r.addr() + 1,
            BASE.bytes(),
            &[PhysRange {
                base: r.addr() + 1,
                len: BASE.bytes(),
            }],
        )
    };
    assert_eq!(
        e,
        Err(InitError::Misaligned {
            required: BASE.bytes()
        })
    );
}

#[test]
fn init_at_err_leaves_region_retryable() {
    // A failed try_init_at leaves the region uninitialised (bounds 0,0) and a
    // corrected retry succeeds.
    let bytes = 32 * BASE.bytes();
    let r = OwnedRegion::new(bytes, max_block());
    let alloc: Regioned = RegionedAllocator::new(
        BASE,
        [const { SummaryBuddyAllocator::<ORDERS, TestProvenance>::new(BASE) }; REGIONS],
    );
    let usable = [PhysRange {
        base: r.addr(),
        len: bytes,
    }];

    // Misaligned span_len rejects before the region records any bounds.
    let bad = unsafe { alloc.try_init_at(0, r.addr(), bytes + 1, &usable) };
    assert_eq!(bad, Err(InitError::InvalidSpan));
    assert_eq!(alloc.region_bounds(0), (0, 0));

    // Corrected retry over the same region index succeeds.
    let ok = unsafe { alloc.try_init_at(0, r.addr(), bytes, &usable) };
    assert_eq!(ok, Ok(()));
    assert_eq!(alloc.region_bounds(0), (r.addr(), r.addr() + bytes));
    assert!(alloc.alloc_in_region(0, BASE, N1).is_ok());
}

#[test]
#[should_panic(expected = "not owned by any region")]
fn dealloc_unowned_address_panics_in_debug() {
    let (alloc, r0, r1) = pool(32);
    let foreign = r0.addr().max(r1.addr()) + 1_000 * 32 * BASE.bytes();
    unsafe { alloc.deallocate_physical(BASE, N1, foreign) };
}

#[test]
#[should_panic(expected = "out of range")]
fn alloc_in_region_out_of_range_panics() {
    let (alloc, _r0, _r1) = pool(32);
    let _ = alloc.alloc_in_region(REGIONS, BASE, N1);
}

#[test]
#[should_panic(expected = "in chain out of range")]
fn alloc_in_chain_out_of_range_panics() {
    let (alloc, _r0, _r1) = pool(32);
    let _ = alloc.alloc_in_chain(&[REGIONS], BASE, N1);
}