frame-alloc 0.1.0

//! End-to-end tour of the default API surface.
//!
//! Walks the path a kernel actually takes: declare a `const`-constructed `static`
//! allocator, hand it a *holey* boot memory map (usable RAM punctured by the
//! kernel image and an MMIO hole), then allocate and free physical frames. The
//! second half wraps the same backend in a [`DepotAllocator`] to show that the
//! wrappers compose over any allocator through the shared traits.
//!
//! Run it:
//!
//! ```text
//! cargo run --example basic                  # the allocation walk
//! cargo run --example basic --features stats # …plus AllocatorStats readouts
//! ```
//!
//! The pool here is a heap allocation standing in for physical RAM; provenance is
//! recovered with [`with_exposed_provenance_mut`](core::ptr::with_exposed_provenance_mut)
//! via the `IdentityProv` harness. A real kernel would supply a higher-half
//! direct-map pointer carrying provenance established once with inline asm, as the
//! `Provenance` / `RegionInit` docs describe.

use frame_alloc::{
    CpuId, DepotAllocator, PageSize, PhysRange, PhysicalAllocator, Provenance, RegionInit,
    SummaryBuddyAllocator,
};
use std::alloc::{Layout, alloc, dealloc};
use std::num::NonZeroUsize;
use std::ptr;
use std::ptr::NonNull;

// Scaffolding: stand-ins for what a kernel would already have.

/// The "physical" pool is a heap allocation whose provenance is exposed in
/// `Region::addr`, so `with_exposed_provenance_mut` can recover it.
pub struct IdentityProv;

unsafe impl Provenance for IdentityProv {
    unsafe fn create(phys: usize) -> NonNull<u8> {
        unsafe { NonNull::new_unchecked(ptr::with_exposed_provenance_mut(phys)) }
    }
    unsafe fn destroy<T>(p: NonNull<T>) -> usize {
        p.addr().get()
    }
}

/// An aligned heap allocation standing in for a span of physical RAM.
pub struct Region {
    ptr: *mut u8,
    layout: Layout,
}

impl Region {
    pub fn new(bytes: usize, align: usize) -> Self {
        let layout = Layout::from_size_align(bytes, align).unwrap();
        let ptr = unsafe { alloc(layout) };
        assert!(!ptr.is_null(), "region allocation failed");
        Self { ptr, layout }
    }

    pub fn addr(&self) -> usize {
        self.ptr.expose_provenance()
    }
}

impl Drop for Region {
    fn drop(&mut self) {
        unsafe { dealloc(self.ptr, self.layout) }
    }
}

pub fn n(v: usize) -> NonZeroUsize {
    NonZeroUsize::new(v).expect("non-zero count")
}

// Configuration:

/// Base frame: 4 KiB.
const BASE: PageSize = PageSize::from_log2(12);
/// Order ceiling. Order `k` is a `BASE << k` block, so order 8 is a 1 MiB block.
const ORDERS: usize = 9;
const MAX_BLOCK: usize = BASE.bytes() << (ORDERS - 1);

/// Total span the allocator manages, in base frames. 256 × 4 KiB = 1 MiB.
const SPAN_FRAMES: usize = 256;

/// A `const`-constructed `static`: no runtime initialiser, nothing on the heap.
/// `init` happens once at boot; from then on it is shared, lock-free for the
/// fast path, across every CPU.
static PHYS: SummaryBuddyAllocator<ORDERS, IdentityProv> = SummaryBuddyAllocator::new(BASE);

fn main() {
    // A boot memory map with holes:
    //
    // The backing pool spans the whole window; we then declare only the *usable*
    // sub-ranges. Two holes stay reserved and are never handed out:
    //
    //   [  0 .. 100)  usable   <- the buddy carves its in-pool bitmap from here
    //   [100 .. 104)  RESERVED  (kernel image)
    //   [104 .. 200)  usable
    //   [200 .. 205)  RESERVED  (MMIO window)
    //   [205 .. 256)  usable
    let pool = Region::new(SPAN_FRAMES * BASE.bytes(), MAX_BLOCK);
    let phys_base = pool.addr(); // the "physical" origin

    let frame = BASE.bytes();
    let range = |lo: usize, hi: usize| PhysRange {
        base: phys_base + lo * frame,
        len: (hi - lo) * frame,
    };
    let usable = [range(0, 100), range(104, 200), range(205, 256)];

    // SAFETY: single-threaded, called once before any allocation. `phys_base` is
    // MAX_BLOCK-aligned (Region honours the requested alignment); the bitmap host
    // range is exclusively owned and reachable through `IdentityProv::create`; the
    // usable ranges are sorted, non-overlapping, base-frame-aligned, and within
    // the span.
    unsafe { PHYS.init(phys_base, SPAN_FRAMES * frame, &usable) };

    println!("Initialised SummaryBuddy over a 1 MiB span with two reserved holes.");
    print_stats("after init", &PHYS);

    // A few allocate / deallocate cycles:
    //
    // One single base frame, then a 4-frame (order-2) contiguous block. Every
    // address is physical and aligned to the request size.
    let single = PHYS
        .allocate_physical(BASE, n(1))
        .expect("single-frame alloc");
    let block = PHYS
        .allocate_physical(BASE, n(4))
        .expect("4-frame contiguous alloc");
    println!("\nallocate_physical(1 frame)  -> {single:#x}");
    println!("allocate_physical(4 frames) -> {block:#x}");
    print_stats("with 5 frames out", &PHYS);

    // SAFETY: each address came from `allocate_physical` with the same page size
    // and count and is not used afterwards.
    unsafe {
        PHYS.deallocate_physical(BASE, n(1), single);
        PHYS.deallocate_physical(BASE, n(4), block);
    }
    println!("\nfreed both — buddies merge back.");
    print_stats("after free", &PHYS);

    // Composition: a per-CPU magazine + shared depot over a fresh SummaryBuddy:
    compose_with_depot();
}

/// Wrap a `SummaryBuddyAllocator` in a [`DepotAllocator`].
fn compose_with_depot() {
    /// Uniprocessor selector: every CPU maps to slot 0. A real kernel returns an
    /// APIC id / `TPIDR_EL1` here.
    struct OneCpu;
    impl CpuId for OneCpu {
        fn current_cpu() -> usize {
            0
        }
    }
    const SLOTS: usize = 8; // magazines (>= CPUs you want disjoint)

    // Same const-new composability: the whole stack is one `static`-able value.
    let mag: DepotAllocator<SummaryBuddyAllocator<ORDERS, IdentityProv>, OneCpu, SLOTS> =
        DepotAllocator::new(BASE, SummaryBuddyAllocator::new(BASE));

    let pool = Region::new(SPAN_FRAMES * BASE.bytes(), MAX_BLOCK);
    // SAFETY: as above; here the whole region is usable, so the single-range
    // convenience applies.
    unsafe { mag.init_region(pool.addr(), SPAN_FRAMES * BASE.bytes()) };

    println!("\n── DepotAllocator<SummaryBuddyAllocator> ──");
    let a = mag.allocate_physical(BASE, n(1)).expect("mag alloc");
    // SAFETY: `a` came from this allocator; freed once, then not reused by us.
    unsafe { mag.deallocate_physical(BASE, n(1), a) };
    let b = mag.allocate_physical(BASE, n(1)).expect("mag re-alloc");
    println!("alloc {a:#x} -> free -> alloc {b:#x}");
    assert_eq!(a, b, "the magazine should return the just-freed frame");
    println!("re-alloc returned the cached frame (no backend round-trip).");
    // SAFETY: final free of a live frame.
    unsafe { mag.deallocate_physical(BASE, n(1), b) };
}

// Stats readout:

/// Print [`AllocatorStats`](frame_alloc::AllocatorStats) for `a` under `--features stats`.
#[cfg(feature = "stats")]
fn print_stats(label: &str, a: &impl frame_alloc::AllocatorStats) {
    let kib = |b: usize| b / 1024;
    println!(
        "  [{label}] total {} KiB · free {} KiB · largest run {} KiB",
        kib(a.total_bytes()),
        kib(a.free_bytes()),
        kib(a.largest_free_bytes()),
    );
}

#[cfg(not(feature = "stats"))]
fn print_stats<A>(_label: &str, _a: &A) {}