frame-alloc 0.1.0

use crate::strategies::cpu_id::NoCpuId;
use crate::util::cache::CachePadded;
use crate::{AllocError, CpuId, InitError, PageSize, PhysRange, PhysicalAllocator, RegionInit};
use core::marker::PhantomData;
use core::mem::{ManuallyDrop, MaybeUninit};
use core::num::NonZeroUsize;
use core::ptr;
use core::sync::atomic::AtomicUsize;
use core::sync::atomic::Ordering::Relaxed;

/// Move `[A; N]` into `[CachePadded<A>; N]` in a `const` context.
///
/// `pub(crate)` only so the white-box drop test can reach it; not part of the
/// public API.
pub(crate) const fn pad_cells<A, const N: usize>(items: [A; N]) -> [CachePadded<A>; N] {
    let src = ManuallyDrop::new(items);
    let src_ptr = &src as *const ManuallyDrop<[A; N]> as *const A;
    let mut cells: [MaybeUninit<CachePadded<A>>; N] = [const { MaybeUninit::uninit() }; N];
    let mut i = 0;
    while i < N {
        // SAFETY: `src_ptr.add(i)` is in bounds for `i < N` and is read exactly
        // once across the loop.
        let a = unsafe { ptr::read(src_ptr.add(i)) };
        cells[i] = MaybeUninit::new(CachePadded::new(a));
        i += 1;
    }
    // SAFETY: every cell was initialised above, and `[MaybeUninit<T>; N]` shares
    // layout with `[T; N]`. `cells` is `MaybeUninit`, so its elements are never
    // dropped — no element is dropped twice.
    unsafe {
        ptr::read(&cells as *const [MaybeUninit<CachePadded<A>>; N] as *const [CachePadded<A>; N])
    }
}

/// A wrapper that routes over **`REGIONS` disjoint physical spans**, one
/// unmodified inner allocator `A` per span.
///
/// All regions are expected to share the same backend configuration — base
/// frame, order count, and `max_page`.
///
/// **Allocation** starts on the calling CPU's home region
/// ([`CpuId::current_cpu`] modulo `REGIONS`, or region 0 under the default
/// [`NoCpuId`]) and work-steals round-robin through the other regions on local
/// [`OutOfMemory`](AllocError::OutOfMemory).
///
/// **Deallocation** and [`add_usable`](Self::add_usable) route by physical
/// address to the owning region. An address owned by no region is *dropped*
/// (debug builds panic).
pub struct RegionedAllocator<const REGIONS: usize, A, S = NoCpuId> {
    /// One independent inner allocator per disjoint span, each on its own cache line.
    regions: [CachePadded<A>; REGIONS],
    /// `[base, end)` physical span per region, used to route a physical address
    /// back to its owning region. Written once by `init_at`, read-only after;
    /// `(0, 0)` marks an uninitialised region.
    bounds: [(AtomicUsize, AtomicUsize); REGIONS],
    /// Base frame size; mirrors each inner `A`'s own base frame.
    base_frame: PageSize,
    _selector: PhantomData<fn() -> S>,
}

impl<const REGIONS: usize, A, S> RegionedAllocator<REGIONS, A, S> {
    /// Create a new, uninitialised regioned allocator from `REGIONS` pre-built
    /// inner allocators.
    ///
    /// `base_frame` must match the base frame size of every inner allocator.
    pub const fn new(base_frame: PageSize, regions: [A; REGIONS]) -> Self {
        assert!(REGIONS > 0, "REGIONS must be > 0");
        Self {
            regions: pad_cells(regions),
            bounds: [const { (AtomicUsize::new(0), AtomicUsize::new(0)) }; REGIONS],
            base_frame,
            _selector: PhantomData,
        }
    }

    /// Fallible initialisation of region `idx` over the disjoint physical span
    /// `[phys_base, phys_base + span_len)`, freeing only `usable`. Delegates to
    /// the backend's [`RegionInit::try_init`] and records the span for address
    /// routing.
    ///
    /// Call at most once *successfully* per region index, single-threaded and
    /// before publishing the allocator to concurrent users. On `Err` the region
    /// is left untouched - a corrected retry is permitted.
    ///
    /// # Errors
    ///
    /// * [`InitError::Misaligned`] if `phys_base` is not base-frame aligned;
    /// * [`InitError::InvalidSpan`] if `span_len` is zero, not a base-frame
    ///   multiple, or overflows the address space;
    /// * [`InitError::AlreadyInitialized`] if region `idx` was already initialised;
    /// * [`InitError::OverlapsRegion`] if the span overlaps an initialised region;
    /// * any error returned by the backend's [`RegionInit::try_init`].
    ///
    /// # Panics
    ///
    /// If `idx >= REGIONS` (a structural bug, like any out-of-bounds index).
    ///
    /// # Safety
    ///
    /// The `(phys_base, span_len, usable)` arguments must satisfy the backend's
    /// [`RegionInit::try_init`] safety contract, including the single-threaded
    /// pre-publication requirement. The spans of distinct regions must be
    /// disjoint, and no two `try_init_at` calls may race with each other.
    pub unsafe fn try_init_at(
        &self,
        idx: usize,
        phys_base: usize,
        span_len: usize,
        usable: &[PhysRange],
    ) -> Result<(), InitError>
    where
        A: RegionInit,
    {
        assert!(
            idx < REGIONS,
            "region index {idx} out of range (REGIONS={REGIONS})"
        );
        if !phys_base.is_multiple_of(self.base_frame.bytes()) {
            return Err(InitError::Misaligned {
                required: self.base_frame.bytes(),
            });
        }
        if span_len == 0 || !span_len.is_multiple_of(self.base_frame.bytes()) {
            return Err(InitError::InvalidSpan);
        }
        let span_end = phys_base
            .checked_add(span_len)
            .ok_or(InitError::InvalidSpan)?;

        if self.bounds[idx].0.load(Relaxed) != 0 || self.bounds[idx].1.load(Relaxed) != 0 {
            return Err(InitError::AlreadyInitialized);
        }
        for j in 0..REGIONS {
            if j == idx {
                continue;
            }
            let lo = self.bounds[j].0.load(Relaxed);
            let hi = self.bounds[j].1.load(Relaxed);
            if hi != 0 && phys_base < hi && span_end > lo {
                return Err(InitError::OverlapsRegion { other: j });
            }
        }

        // SAFETY: the caller upholds the backend's `try_init` contract for this
        // span. On `Err` the backend guarantees it is untouched, and we skip the
        // bounds stores below, so the region stays uninitialised and retryable.
        unsafe { self.regions[idx].try_init(phys_base, span_len, usable) }?;

        self.bounds[idx].0.store(phys_base, Relaxed);
        self.bounds[idx].1.store(span_end, Relaxed);
        Ok(())
    }

    /// As [`try_init_at`](Self::try_init_at) but panic on any [`InitError`].
    ///
    /// Call once per region index, single-threaded, before publishing the
    /// allocator to concurrent users. (Same as [`try_init_at`](Self::try_init_at)).
    ///
    /// # Safety
    ///
    /// `idx < REGIONS`. The `(phys_base, span_len, usable)` arguments must satisfy
    /// the backend's [`RegionInit::try_init`] contract, including the
    /// single-threaded pre-publication requirement. The spans of distinct regions
    /// must be disjoint, and no two `init_at`/`try_init_at` calls may race with
    /// each other.
    ///
    /// # Panics
    ///
    /// If `idx >= REGIONS`, or if initialisation returns an [`InitError`].
    pub unsafe fn init_at(
        &self,
        idx: usize,
        phys_base: usize,
        span_len: usize,
        usable: &[PhysRange],
    ) where
        A: RegionInit,
    {
        // SAFETY: forwarded under the caller's trait-level guarantees.
        match unsafe { self.try_init_at(idx, phys_base, span_len, usable) } {
            Ok(()) => {}
            Err(e) => panic!("RegionedAllocator::init_at failed: {e:?}"),
        }
    }

    /// Transition a reserved in-span range to free, routed to the owning region.
    /// The repeatable post-init counterpart of [`RegionInit::add_usable`]
    ///
    /// # Safety
    ///
    /// `[base, base + len)` must lie within a single initialised region's span,
    /// be currently reserved (not already free), exclusively owned, and not
    /// aliased while registered.
    pub unsafe fn add_usable(&self, base: usize, len: usize)
    where
        A: RegionInit,
    {
        let Some(owner) = self.region_of(base) else {
            debug_assert!(false, "add_usable: base {base:#x} not owned by any region");
            return;
        };
        debug_assert!(
            base.checked_add(len)
                .is_some_and(|end| end <= self.bounds[owner].1.load(Relaxed)),
            "add_usable range escapes its owning region"
        );
        // SAFETY: `[base, base + len)` lies in `owner`'s span; the backend's
        // `add_usable` contract is upheld.
        unsafe { self.regions[owner].add_usable(base, len) };
    }

    /// Allocate strictly within region `idx`
    ///
    /// # Panics
    ///
    /// If `idx >= REGIONS`.
    pub fn alloc_in_region(
        &self,
        idx: usize,
        ps: PageSize,
        count: NonZeroUsize,
    ) -> Result<usize, AllocError>
    where
        A: PhysicalAllocator,
    {
        assert!(
            idx < REGIONS,
            "region index {idx} out of range (REGIONS={REGIONS})"
        );
        self.regions[idx].allocate_physical(ps, count)
    }

    /// Allocate from the first region in `chain` that can satisfy the request,
    /// trying them in the given order
    ///
    /// # Panics
    ///
    /// If any index in `chain` is `>= REGIONS`.
    pub fn alloc_in_chain(
        &self,
        chain: &[usize],
        ps: PageSize,
        count: NonZeroUsize,
    ) -> Result<usize, AllocError>
    where
        A: PhysicalAllocator,
    {
        for &idx in chain {
            assert!(
                idx < REGIONS,
                "region index {idx} in chain out of range (REGIONS={REGIONS})"
            );
            match self.regions[idx].allocate_physical(ps, count) {
                Ok(phys) => return Ok(phys),
                Err(AllocError::OutOfMemory) => continue,
                Err(e) => return Err(e),
            }
        }
        Err(AllocError::OutOfMemory)
    }

    /// Index of the region whose span contains `phys`, or `None` if no
    /// initialised region owns it.
    #[inline]
    fn region_of(&self, phys: usize) -> Option<usize> {
        for i in 0..REGIONS {
            let lo = self.bounds[i].0.load(Relaxed);
            let hi = self.bounds[i].1.load(Relaxed);
            if phys >= lo && phys < hi {
                return Some(i);
            }
        }
        None
    }

    /// The inner per-region allocators, for diagnostics.
    #[cfg(any(feature = "stats", test))]
    pub(crate) fn regions(&self) -> impl ExactSizeIterator<Item = &A> {
        self.regions.iter().map(|r| &**r)
    }

    /// The recorded `(base, end)` routing span of region `idx`; `(0, 0)` marks an
    /// uninitialised region. (Test only).
    #[cfg(test)]
    pub(crate) fn region_bounds(&self, idx: usize) -> (usize, usize) {
        (
            self.bounds[idx].0.load(Relaxed),
            self.bounds[idx].1.load(Relaxed),
        )
    }
}

unsafe impl<const REGIONS: usize, A, S> PhysicalAllocator for RegionedAllocator<REGIONS, A, S>
where
    A: PhysicalAllocator,
    S: CpuId,
{
    fn allocate_physical(&self, ps: PageSize, count: NonZeroUsize) -> Result<usize, AllocError> {
        let start = S::current_cpu() % REGIONS;
        for offset in 0..REGIONS {
            let i = (start + offset) % REGIONS;
            match self.regions[i].allocate_physical(ps, count) {
                Ok(phys) => return Ok(phys),
                Err(AllocError::OutOfMemory) => continue,
                Err(e) => return Err(e),
            }
        }
        Err(AllocError::OutOfMemory)
    }

    unsafe fn deallocate_physical(&self, ps: PageSize, count: NonZeroUsize, phys: usize) {
        match self.region_of(phys) {
            // SAFETY: `phys` lies in `owner`'s span; the caller upholds the
            // per-frame contract of the inner allocator.
            Some(owner) => unsafe { self.regions[owner].deallocate_physical(ps, count, phys) },
            None => debug_assert!(
                false,
                "deallocate_physical: phys {phys:#x} not owned by any region"
            ),
        }
    }
}

#[cfg(any(feature = "stats", test))]
impl<const REGIONS: usize, A: crate::AllocatorStats, S> crate::AllocatorStats
    for RegionedAllocator<REGIONS, A, S>
{
    fn total_bytes(&self) -> usize {
        self.regions().map(|r| r.total_bytes()).sum()
    }

    fn free_bytes(&self) -> usize {
        self.regions().map(|r| r.free_bytes()).sum()
    }

    fn largest_free_bytes(&self) -> usize {
        self.regions()
            .map(|r| r.largest_free_bytes())
            .max()
            .unwrap_or(0)
    }
}