frame-alloc 0.1.0

use crate::util::cache::CachePadded;
use crate::util::lifo::{Lifo, LifoGuard};
use crate::{
    AllocError, CpuId, InterruptControl, NoInterruptControl, PageSize, PhysicalAllocator,
    RegionInit,
};
use core::marker::PhantomData;
use core::num::NonZeroUsize;

const N1: NonZeroUsize = NonZeroUsize::MIN;

/// A per-CPU magazine cache with a shared **depot** tier, over a backing physical
/// allocator. This extends [`MagazineAllocator`](crate::MagazineAllocator) with a second,
/// shared cache tier between the per-CPU magazines and the backend.
///
/// `SLOTS` is the number of per-CPU magazines, `CAP` each magazine's depth, and
/// `DEPOT_CAP` the shared depot's capacity in frames. [`CpuId::current_cpu`] is taken
/// modulo `SLOTS`. `CAP` defaults to 128, `DEPOT_CAP` to 512.
/// `I` is the [`InterruptControl`] strategy for the magazine and depot locks; it
/// defaults to [`NoInterruptControl`].
/// Multi-frame requests first use the backend unchanged. On backend OOM, the
/// wrapper progressively returns cached frames and retries, preserving cache
/// contents once the request becomes satisfiable.
pub struct DepotAllocator<
    A,
    S,
    const SLOTS: usize,
    const CAP: usize = 128,
    const DEPOT_CAP: usize = 512,
    I: InterruptControl = NoInterruptControl,
> {
    backend: A,
    mags: [CachePadded<Lifo<CAP, I>>; SLOTS],
    depot: CachePadded<Lifo<DEPOT_CAP, I>>,
    base_frame: PageSize,
    /// Total cached base frames returned to the backend by explicit flushes or
    /// OOM recovery across the allocator's life. Diagnostic only.
    #[cfg(any(feature = "stats", test))]
    frames_flushed: core::sync::atomic::AtomicUsize,
    /// Maximum number of frames simultaneously held by the shared depot.
    #[cfg(any(feature = "stats", test))]
    peak_depot_len: core::sync::atomic::AtomicUsize,
    _selector: PhantomData<fn() -> S>,
}

// SAFETY: the magazine and depot cells are only touched through a `LifoGuard` (i.e.
// under their lock), and the backend is shared only if it is itself `Sync`. So the
// wrapper is `Sync`/`Send` exactly when the backend is.
unsafe impl<
    A: Sync,
    S,
    const SLOTS: usize,
    const CAP: usize,
    const DEPOT_CAP: usize,
    I: InterruptControl,
> Sync for DepotAllocator<A, S, SLOTS, CAP, DEPOT_CAP, I>
{
}
unsafe impl<
    A: Send,
    S,
    const SLOTS: usize,
    const CAP: usize,
    const DEPOT_CAP: usize,
    I: InterruptControl,
> Send for DepotAllocator<A, S, SLOTS, CAP, DEPOT_CAP, I>
{
}

impl<A, S, const SLOTS: usize, const CAP: usize, const DEPOT_CAP: usize, I: InterruptControl>
    DepotAllocator<A, S, SLOTS, CAP, DEPOT_CAP, I>
{
    /// Create a depot cache over `backend`. Call
    /// [`init_region`](RegionInit::init_region) before use.
    ///
    /// `base_frame` must match the backend's base frame size.
    pub const fn new(base_frame: PageSize, backend: A) -> Self {
        assert!(SLOTS > 0, "SLOTS must be > 0");
        assert!(CAP >= 2, "CAP must be >= 2");
        assert!(DEPOT_CAP >= 1, "DEPOT_CAP must be >= 1");
        Self {
            backend,
            mags: [const { CachePadded::new(Lifo::new()) }; SLOTS],
            depot: CachePadded::new(Lifo::new()),
            base_frame,
            #[cfg(any(feature = "stats", test))]
            frames_flushed: core::sync::atomic::AtomicUsize::new(0),
            #[cfg(any(feature = "stats", test))]
            peak_depot_len: core::sync::atomic::AtomicUsize::new(0),
            _selector: PhantomData,
        }
    }

    /// The backing allocator, for diagnostics.
    #[cfg(any(feature = "stats", test))]
    pub(crate) fn backend(&self) -> &A {
        &self.backend
    }

    /// Number of base frames currently cached across all Depots. These are
    /// frames the backend considers allocated but that are immediately available
    /// on the fast path. Non-linearizable under concurrent use.
    #[cfg(any(feature = "stats", test))]
    pub fn cached_frames(&self) -> usize {
        let mut total = 0;
        for mag in &self.mags {
            total += mag.lock().len();
        }
        total + self.depot_len()
    }

    /// Base frames currently held in the shared depot. Non-linearizable under
    /// concurrent use.
    #[cfg(any(feature = "stats", test))]
    pub fn depot_len(&self) -> usize {
        self.depot.lock().len()
    }

    /// High-water mark of base frames held in the shared depot.
    #[cfg(any(feature = "stats", test))]
    pub fn peak_depot_len(&self) -> usize {
        self.peak_depot_len
            .load(core::sync::atomic::Ordering::Relaxed)
    }

    /// Total cached base frames returned to the backend by [`flush`](Self::flush)
    /// or OOM recovery over this allocator's life. Includes depot and magazine
    /// frames. Monotonic and safe to read concurrently.
    #[cfg(any(feature = "stats", test))]
    pub fn frames_flushed(&self) -> usize {
        self.frames_flushed
            .load(core::sync::atomic::Ordering::Relaxed)
    }
}

impl<
    A: PhysicalAllocator,
    S: CpuId,
    const SLOTS: usize,
    const CAP: usize,
    const DEPOT_CAP: usize,
    I: InterruptControl,
> DepotAllocator<A, S, SLOTS, CAP, DEPOT_CAP, I>
{
    /// Frames moved per refill / drain. Half the capacity gives boundary
    /// hysteresis; `new` asserts `CAP >= 2`, so this is always >= 1.
    #[inline(always)]
    const fn batch() -> usize {
        CAP / 2
    }

    /// Move up to `want` frames from the shared depot into `mag`, returning the
    /// number moved. The caller holds `mag`'s lock; this briefly takes the depot
    /// lock (always magazine-before-depot, the same order everywhere).
    fn depot_to_mag(&self, mag: &mut LifoGuard<'_, CAP, I>, want: usize) -> usize {
        let mut depot = self.depot.lock();
        let mut moved = 0;
        while moved < want {
            match depot.pop() {
                Some(addr) => {
                    mag.push(addr);
                    moved += 1;
                }
                None => break,
            }
        }
        moved
    }

    /// Push as many of `src` as fit into the depot. Returns the number accepted.
    /// The caller holds the source magazine's lock.
    fn push_to_depot(&self, src: &[usize]) -> usize {
        let mut depot = self.depot.lock();
        let pushed = depot.push_slice(src);
        #[cfg(any(feature = "stats", test))]
        self.peak_depot_len
            .fetch_max(depot.len(), core::sync::atomic::Ordering::Relaxed);
        pushed
    }

    fn alloc_one(&self) -> Result<usize, AllocError> {
        let current = S::current_cpu() % SLOTS;
        {
            let mut mag = self.mags[current].lock();
            if let Some(addr) = mag.pop() {
                return Ok(addr);
            }

            // Empty — pull a batch from the shared depot, then hand out one.
            if self.depot_to_mag(&mut mag, Self::batch()) > 0
                && let Some(addr) = mag.pop()
            {
                return Ok(addr);
            }

            // Depot empty — refill a batch from the backend, then hand out one.
            let mut filled = 0usize;
            while filled < Self::batch() {
                match self.backend.allocate_physical(self.base_frame, N1) {
                    Ok(addr) => {
                        mag.push(addr);
                        filled += 1;
                    }
                    Err(AllocError::OutOfMemory) => break,
                    Err(e) => return Err(e),
                }
            }

            if let Some(addr) = mag.pop() {
                return Ok(addr);
            }
        }

        // A sibling magazine can still own a free frame after the shared tiers
        // are exhausted. Reuse it directly for an order-0 request.
        for offset in 1..SLOTS {
            let slot = (current + offset) % SLOTS;
            if let Some(addr) = self.mags[slot].lock().pop() {
                return Ok(addr);
            }
        }

        Err(AllocError::OutOfMemory)
    }

    /// # Safety
    ///
    /// `addr` must be a `(base, 1)` frame previously obtained from this allocator.
    unsafe fn free_one(&self, addr: usize) {
        let mut mag = self.mags[S::current_cpu() % SLOTS].lock();
        if mag.is_full() {
            // Move a batch off the top: into the depot first, backend for the rest.
            let overflow = mag.take_top(Self::batch());
            let pushed = self.push_to_depot(overflow);
            for &a in &overflow[pushed..] {
                // SAFETY: `a` was handed out by the backend as a (base, 1) frame
                // and is not referenced elsewhere once drained.
                unsafe { self.backend.deallocate_physical(self.base_frame, N1, a) };
            }
        }
        mag.push(addr);
    }

    /// Return up to `want` depot-cached frames to the backend. Drained in `CAP`-sized
    /// sub-batches. Returns the number drained.
    fn drain_chunk(&self, want: usize) -> usize {
        let mut moved = 0;
        let mut batch = [0usize; CAP];
        while moved < want {
            let take = (want - moved).min(CAP);
            let n = {
                let mut depot = self.depot.lock();
                let mut got = 0;
                while got < take {
                    match depot.pop() {
                        Some(addr) => {
                            batch[got] = addr;
                            got += 1;
                        }
                        None => break,
                    }
                }
                got
            };
            if n == 0 {
                break;
            }
            for &addr in &batch[..n] {
                // SAFETY: addr was handed out by the backend as a (base, 1) frame and
                // was owned by the depot until popped into `batch` above.
                unsafe { self.backend.deallocate_physical(self.base_frame, N1, addr) };
            }
            moved += n;
        }
        #[cfg(any(feature = "stats", test))]
        self.frames_flushed
            .fetch_add(moved, core::sync::atomic::Ordering::Relaxed);
        moved
    }

    /// Return up to `want` frames from one magazine to the backend. Backend calls
    /// happen only after releasing the magazine lock.
    fn drain_magazine(&self, slot: usize, want: usize) -> usize {
        let mut moved = 0;
        let mut batch = [0usize; CAP];
        while moved < want {
            let take = (want - moved).min(CAP);
            let n = {
                let mut mag = self.mags[slot].lock();
                let mut got = 0;
                while got < take {
                    match mag.pop() {
                        Some(addr) => {
                            batch[got] = addr;
                            got += 1;
                        }
                        None => break,
                    }
                }
                got
            };
            if n == 0 {
                break;
            }
            for &addr in &batch[..n] {
                // SAFETY: the frame was owned by this magazine until it was
                // popped into `batch` above.
                unsafe { self.backend.deallocate_physical(self.base_frame, N1, addr) };
            }
            moved += n;
        }
        #[cfg(any(feature = "stats", test))]
        self.frames_flushed
            .fetch_add(moved, core::sync::atomic::Ordering::Relaxed);
        moved
    }

    /// Return the current bounded snapshot of every cache tier to the backend.
    /// No cache lock is held while invoking the backend. Concurrent frees may
    /// repopulate a tier after its snapshot is taken.
    pub fn flush(&self) {
        let depot = self.depot.lock().len();
        self.drain_chunk(depot);
        for slot in 0..SLOTS {
            let magazine = self.mags[slot].lock().len();
            self.drain_magazine(slot, magazine);
        }
    }

    /// Continue a multi-frame allocation after its initial backend attempt OOMed,
    /// draining the shared depot first and then the per-CPU magazines.
    fn recover_after_oom(&self, ps: PageSize, count: NonZeroUsize) -> Result<usize, AllocError> {
        let mut chunk = Self::batch();
        loop {
            let moved = self.drain_chunk(chunk);
            if moved == 0 {
                break;
            }
            match self.backend.allocate_physical(ps, count) {
                Err(AllocError::OutOfMemory) => {}
                other => return other,
            }
            chunk = chunk.saturating_mul(2);
        }

        let current = S::current_cpu() % SLOTS;
        for offset in 0..SLOTS {
            let slot = (current + offset) % SLOTS;
            let mut remaining = self.mags[slot].lock().len();
            while remaining > 0 {
                let moved = self.drain_magazine(slot, remaining.min(Self::batch()));
                if moved == 0 {
                    break;
                }
                remaining -= moved;
                match self.backend.allocate_physical(ps, count) {
                    Err(AllocError::OutOfMemory) => {}
                    other => return other,
                }
            }
        }

        Err(AllocError::OutOfMemory)
    }

    fn alloc_multiframe(&self, ps: PageSize, count: NonZeroUsize) -> Result<usize, AllocError> {
        match self.backend.allocate_physical(ps, count) {
            Err(AllocError::OutOfMemory) => {}
            other => return other,
        }

        self.recover_after_oom(ps, count)
    }
}

unsafe impl<
    A: PhysicalAllocator,
    S: CpuId,
    const SLOTS: usize,
    const CAP: usize,
    const DEPOT_CAP: usize,
    I: InterruptControl,
> PhysicalAllocator for DepotAllocator<A, S, SLOTS, CAP, DEPOT_CAP, I>
{
    fn allocate_physical(&self, ps: PageSize, count: NonZeroUsize) -> Result<usize, AllocError> {
        if ps == self.base_frame && count == N1 {
            return self.alloc_one();
        }
        self.alloc_multiframe(ps, count)
    }

    unsafe fn deallocate_physical(&self, ps: PageSize, count: NonZeroUsize, phys: usize) {
        if ps == self.base_frame && count == N1 {
            // SAFETY: caller upholds the per-frame contract; (base, 1) frames are
            // fungible and cached.
            unsafe { self.free_one(phys) };
        } else {
            // SAFETY: forwarded unchanged to the backend.
            unsafe { self.backend.deallocate_physical(ps, count, phys) };
        }
    }
}

unsafe impl<
    A: RegionInit,
    S,
    const SLOTS: usize,
    const CAP: usize,
    const DEPOT_CAP: usize,
    I: InterruptControl,
> RegionInit for DepotAllocator<A, S, SLOTS, CAP, DEPOT_CAP, I>
{
    unsafe fn try_init(
        &self,
        phys_base: usize,
        span_len: usize,
        usable: &[crate::allocator::PhysRange],
    ) -> Result<(), crate::InitError> {
        // SAFETY: forwarded unchanged to the backend.
        unsafe { self.backend.try_init(phys_base, span_len, usable) }
    }

    unsafe fn add_usable(&self, base: usize, len: usize) {
        // SAFETY: forwarded unchanged to the backend.
        unsafe { self.backend.add_usable(base, len) };
    }
}

#[cfg(any(feature = "stats", test))]
impl<
    A: crate::AllocatorStats,
    S,
    const SLOTS: usize,
    const CAP: usize,
    const DEPOT_CAP: usize,
    I: InterruptControl,
> crate::AllocatorStats for DepotAllocator<A, S, SLOTS, CAP, DEPOT_CAP, I>
{
    fn total_bytes(&self) -> usize {
        self.backend().total_bytes()
    }

    fn free_bytes(&self) -> usize {
        self.backend().free_bytes() + self.cached_frames() * self.base_frame.bytes()
    }

    fn largest_free_bytes(&self) -> usize {
        self.backend().largest_free_bytes()
    }
}