peckr 0.6.1

//! Raw packed pointers.

use core::ops::{Deref, DerefMut};
use core::ptr::NonNull;

/// Pointer type with the ability to store up to `CAPACITY`
/// arbitrary bits inline with the underlying memory offset.
///
/// The variance guarantees are the same as that of
/// a [`NonNull`] pointer.
#[derive(Clone, Debug, Hash, Ord, PartialOrd, Eq, PartialEq)]
#[repr(transparent)]
pub struct PackedPtr<const CAPACITY: u32, T> {
    pointer: NonNull<T>,
}

impl<const CAPACITY: u32, T> AsRef<T> for PackedPtr<CAPACITY, T> {
    fn as_ref(&self) -> &T {
        unsafe {
            // SAFETY: Since we masked the bitfield, the pointer
            // dereference is correct.
            &*self.masked_pointer()
        }
    }
}

impl<const CAPACITY: u32, T> Deref for PackedPtr<CAPACITY, T> {
    type Target = T;

    #[inline]
    fn deref(&self) -> &Self::Target {
        self.as_ref()
    }
}

impl<const CAPACITY: u32, T> AsMut<T> for PackedPtr<CAPACITY, T> {
    fn as_mut(&mut self) -> &mut T {
        unsafe {
            // SAFETY: Since we masked the bitfield, the pointer
            // dereference is correct.
            &mut *self.masked_pointer()
        }
    }
}

impl<const CAPACITY: u32, T> DerefMut for PackedPtr<CAPACITY, T> {
    #[inline]
    fn deref_mut(&mut self) -> &mut Self::Target {
        self.as_mut()
    }
}

impl<const CAPACITY: u32, T> PackedPtr<CAPACITY, T> {
    /// Check an access to a bitfield index. Panic if
    /// the access is incorrect.
    #[inline(always)]
    fn check_bitfield_bounds(index: u32) {
        if index >= CAPACITY {
            panic!("Cannot index bits higher than {}", CAPACITY);
        }
    }

    /// Return the mask that may be used to erase the bitfield
    /// set by the user.
    #[inline(always)]
    pub(crate) const fn mask() -> usize {
        (1 << CAPACITY) - 1
    }

    /// Retrieve the offset in memory that was originally
    /// allocated, by masking away the bitfield that has
    /// been set by the user.
    #[inline]
    pub(crate) fn masked_pointer(&self) -> *mut T {
        (self.pointer.as_ptr() as usize & !Self::mask()) as *mut _
    }

    /// Constructs a new [`PackedPtr`], without performing any
    /// checks on the source pointer.
    ///
    /// ## Safety
    ///
    /// The caller must guarantee that the pointer contains at
    /// least `CAPACITY` free storage bits, starting from the
    /// base offset.
    #[inline]
    pub unsafe fn new_unchecked(pointer: NonNull<T>) -> Self {
        Self { pointer }
    }

    /// Constructs a new [`PackedPtr`], without performing any
    /// checks on the source pointer.
    ///
    /// ## Safety
    ///
    /// The caller must guarantee that the pointer contains at
    /// least `CAPACITY` free storage bits, starting from the
    /// base offset, and that the source pointer is non-null.
    #[inline]
    pub unsafe fn new_unchecked_raw(pointer: *mut T) -> Self {
        Self::new_unchecked(NonNull::new_unchecked(pointer))
    }

    /// Get the value of the bit at `index`.
    ///
    /// If the index higher than or equal to `CAPACITY`,
    /// panic.
    #[inline]
    pub fn get_bit(&self, index: u32) -> bool {
        Self::check_bitfield_bounds(index);
        self.pointer.as_ptr() as usize & (1 << index) != 0
    }

    /// Set the value of the bit at `index`.
    ///
    /// If the index higher than or equal to `CAPACITY`,
    /// panic.
    pub fn set_bit(&mut self, index: u32, value: bool) {
        Self::check_bitfield_bounds(index);

        let old_off = self.pointer.as_ptr() as usize;
        let new_off = if value {
            old_off | (1usize << index)
        } else {
            old_off & !(1usize << index)
        };

        self.pointer = unsafe {
            // SAFETY: The updated pointer is still non-null.
            NonNull::new_unchecked(new_off as *mut _)
        };
    }

    /// Set the bit at `index` to 1.
    ///
    /// If the index higher than or equal to `CAPACITY`,
    /// panic.
    #[inline]
    pub fn set_bit_high(&mut self, index: u32) {
        self.set_bit(index, true);
    }

    /// Set the bit at `index` to 0.
    ///
    /// If the index higher than or equal to `CAPACITY`,
    /// panic.
    #[inline]
    pub fn set_bit_low(&mut self, index: u32) {
        self.set_bit(index, false);
    }

    /// Get the stored bit field.
    #[inline]
    pub fn bitfield(&self) -> usize {
        self.pointer.as_ptr() as usize & Self::mask()
    }

    /// Set the stored bit field.
    #[inline]
    pub fn set_bitfield(&mut self, bitfield: usize) {
        let mask = Self::mask();

        let old_off = self.pointer.as_ptr() as usize & !mask;
        let new_off = old_off | (bitfield & mask);

        self.pointer = unsafe {
            // SAFETY: The updated pointer is still non-null.
            NonNull::new_unchecked(new_off as *mut _)
        };
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    /// 32-bit integer with at least 4 free bits in its memory offsets.
    #[repr(align(16))]
    #[derive(Debug, Copy, Clone, Eq, PartialEq)]
    struct IntCap4(i32);

    /// Test that even after modifying the pointer's offset, we are
    /// still able to dereference the underlying data. Additionally,
    /// test that our changes are still queriable even after updating
    /// the pointed memory location.
    #[test]
    fn pointer_deref() {
        let mut value = IntCap4(1234);

        let mut ptr: PackedPtr<1, IntCap4> =
            unsafe { PackedPtr::new_unchecked_raw(&mut value as *mut _) };
        assert_eq!(*ptr.as_ref(), IntCap4(1234));
        assert_eq!(ptr.get_bit(0), false);

        ptr.set_bit_high(0);
        assert_eq!(*ptr.as_ref(), IntCap4(1234));
        assert_eq!(ptr.get_bit(0), true);

        *ptr.as_mut() = IntCap4(5678);
        assert_eq!(*ptr.as_ref(), IntCap4(5678));
        assert_eq!(ptr.get_bit(0), true);

        ptr.set_bit_low(0);
        assert_eq!(*ptr.as_ref(), IntCap4(5678));
        assert_eq!(ptr.get_bit(0), false);
    }

    /// Test manipulating the entire bit field of a pointer.
    #[test]
    fn manipulate_entire_bitfield() {
        let mut value = IntCap4(1234);

        let mut ptr: PackedPtr<4, IntCap4> =
            unsafe { PackedPtr::new_unchecked_raw(&mut value as *mut _) };
        assert_eq!(ptr.bitfield(), 0b0000);
        assert_eq!(*ptr.as_ref(), IntCap4(1234));

        ptr.set_bitfield(0b1011);
        assert_eq!(ptr.bitfield(), 0b1011);
        assert_eq!(*ptr.as_ref(), IntCap4(1234));

        ptr.set_bit_low(3);
        assert_eq!(ptr.bitfield(), 0b0011);
        assert_eq!(*ptr.as_ref(), IntCap4(1234));

        ptr.set_bitfield(0b1001);
        assert_eq!(ptr.bitfield(), 0b1001);
        assert_eq!(*ptr.as_ref(), IntCap4(1234));
    }
}