Pack arbitrary bits into a pointer offset in Rust

packed.rs

// TODO: set constraints in no. of bits allowed. currently it's possible to crash a program with invalid bitfield sizes

use std::alloc::{alloc, dealloc, Layout, LayoutError};
use std::ops::Drop;
use std::ptr::NonNull;

fn main() {
    let mut ptr: FlagPointer<2, _> = FlagPointer::new(1234).unwrap();

    println!("value before set = {}", ptr.as_ref());
    println!("bit field before set = {}", ptr.get_bit(0));

    ptr.set_bit_high(0);
    println!();

    println!("value after set = {}", ptr.as_ref());
    println!("bit field after set = {}", ptr.get_bit(0));
}

/// Represent a failure condition upon allocating
/// a [`FlagPointer`] value.
#[derive(Clone, Debug)]
pub enum AllocationError {
    /// Invalid layout resulting from the requested
    /// capacity bits.
    InvalidLayout(LayoutError),
    /// The allocator returned a null pointer.
    NullPtr,
}

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

impl<const CAPACITY: u32, T> FlagPointer<CAPACITY, T> {
    /// The memory layout required to represent `T` along with
    /// up to `CAPACITY` of arbitrary storage bits.
    #[inline]
    fn layout() -> Result<Layout, LayoutError> {
        Layout::new::<T>().align_to(2usize.pow(CAPACITY))
    }

    /// 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);
        }
    }

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

    /// Return an immutable reference to the underlying data.
    pub fn as_ref(&self) -> &T {
        unsafe {
            // SAFETY: Since we masked the bitfield, the pointer
            // dereference is correct.
            &*self.masked_pointer()
        }
    }

    /// Return a mutable reference to the underlying data.
    pub fn as_mut(&mut self) -> &mut T {
        unsafe {
            // SAFETY: Since we masked the bitfield, the pointer
            // dereference is correct.
            &mut *self.masked_pointer()
        }
    }

    /// Allocates a new value of type `T` using the global
    /// allocator.
    pub fn new(value: T) -> Result<Self, AllocationError> {
        let layout = Self::layout().map_err(AllocationError::InvalidLayout)?;
        let offset: *mut T = unsafe {
            // SAFETY: The requested layout for `T` must be
            // correct at this point.
            alloc(layout).cast()
        };
        let pointer = NonNull::new(offset).ok_or(AllocationError::NullPtr)?;
        unsafe {
            // SAFETY: The pointer returned by the global allocator
            // should point to a valid memory location.
            pointer.as_ptr().write(value);
        }
        Ok(FlagPointer { 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);
    }
}

impl<const CAPACITY: u32, T> Drop for FlagPointer<CAPACITY, T> {
    fn drop(&mut self) {
        let layout = Self::layout().unwrap();
        let offset = self.masked_pointer().cast();

        unsafe {
            // SAFETY: The bit field has been properly masked away,
            // and the layout is correct, so the call to `dealloc`
            // should be semantically correct.
            dealloc(offset, layout);
        }
    }
}