frame-alloc 0.1.0

use core::num::{NonZeroU8, NonZeroUsize};

/// A power-of-two page size, stored as a bit-shift (log2 of bytes).
///
/// # Safety invariant
/// The inner value must be a valid page size for the target architecture.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub struct PageSize(NonZeroU8);

impl PageSize {
    /// Construct a PageSize from a log2 value. Panics if shift is 0 or >= usize::BITS.
    pub const fn from_log2(shift: u8) -> Self {
        assert!(shift > 0 && (shift as u32) < usize::BITS);
        // SAFETY: shift > 0 asserted above
        Self(unsafe { NonZeroU8::new_unchecked(shift) })
    }

    /// The size in bytes.
    pub const fn bytes(&self) -> usize {
        1usize << self.0.get()
    }

    /// The log2 / bit-shift value.
    #[inline]
    pub const fn log2(self) -> u8 {
        self.0.get()
    }

    /// Total bytes for `count` frames of this size.
    /// Returns `None` on overflow.
    #[inline]
    pub const fn total_bytes(self, count: NonZeroUsize) -> Option<usize> {
        self.bytes().checked_mul(count.get())
    }

    /// True if `addr` is naturally aligned to this page size.
    #[inline]
    pub const fn is_aligned(self, addr: usize) -> bool {
        addr & (self.bytes() - 1) == 0
    }

    /// Round `addr` down to the nearest aligned frame base.
    #[inline]
    pub const fn align_down(self, addr: usize) -> usize {
        addr & !(self.bytes() - 1)
    }

    /// Round `addr` up to the next aligned frame base.
    /// Returns `None` on overflow.
    #[inline]
    pub const fn align_up(self, addr: usize) -> Option<usize> {
        let mask = self.bytes() - 1;
        match addr.checked_add(mask) {
            Some(a) => Some(a & !mask),
            None => None,
        }
    }
}