tagged_dispatch/

lib.rs

#![doc = include_str!("../README.md")]
#![cfg_attr(not(feature = "std"), no_std)]

#[cfg(not(feature = "std"))]
extern crate alloc;

use core::marker::PhantomData;

#[cfg(not(feature = "std"))]
use alloc::boxed::Box;
#[cfg(feature = "std")]
use std::boxed::Box;

// Re-export the macro
pub use tagged_dispatch_macros::tagged_dispatch;

// Re-export allocator crates when their features are enabled
#[cfg(feature = "allocator-bumpalo")]
pub use bumpalo;

#[cfg(feature = "allocator-typed-arena")]
pub use typed_arena;

/// The core tagged pointer type used internally.
///
/// Uses the top 7 bits of a 64-bit pointer for type tagging,
/// supporting up to 128 different types while maintaining an 8-byte size.
///
/// # Platform Optimizations
///
/// On Apple Silicon (macOS ARM64), this implementation leverages the hardware's
/// Top Byte Ignore (TBI) feature. TBI allows the processor to automatically
/// ignore the top byte of pointers during memory access, eliminating the need
/// for manual masking operations. This provides a measurable performance
/// improvement by reducing instructions on the critical path of every trait
/// method dispatch.
#[repr(transparent)]
pub struct TaggedPtr<T> {
    ptr: usize,
    _phantom: PhantomData<T>,
}

impl<T> TaggedPtr<T> {
    const TAG_BITS: usize = 7;
    const TAG_SHIFT: usize = 64 - Self::TAG_BITS;
    const TAG_MASK: usize = ((1 << Self::TAG_BITS) - 1) << Self::TAG_SHIFT;
    #[cfg(not(all(target_os = "macos", target_arch = "aarch64")))]
    const PTR_MASK: usize = !Self::TAG_MASK;
    
    /// Maximum number of variants supported (2^7 = 128)
    pub const MAX_VARIANTS: usize = 1 << Self::TAG_BITS;
    
    /// Create a new tagged pointer
    #[inline(always)]
    pub fn new(ptr: *mut T, tag: u8) -> Self {
        debug_assert!(
            tag < Self::MAX_VARIANTS as u8,
            "Tag must be less than 128 (7 bits)"
        );
        
        let addr = ptr as usize;
        debug_assert_eq!(
            addr & Self::TAG_MASK, 
            0, 
            "Pointer already has high bits set!"
        );
        
        Self {
            ptr: addr | ((tag as usize) << Self::TAG_SHIFT),
            _phantom: PhantomData,
        }
    }
    
    /// Get the tag value
    #[inline(always)]
    pub fn tag(&self) -> u8 {
        ((self.ptr & Self::TAG_MASK) >> Self::TAG_SHIFT) as u8
    }
    
    /// Get the untagged pointer.
    ///
    /// # Safety
    /// The returned pointer is only valid if the original pointer passed to `new` is still valid.
    ///
    /// # Platform Optimization
    /// On macOS ARM64 (Apple Silicon), this method leverages the hardware's Top Byte Ignore (TBI)
    /// feature, which automatically masks the top byte during memory access. This eliminates the
    /// need for software masking, providing a performance improvement.
    #[cfg(all(target_os = "macos", target_arch = "aarch64"))]
    #[inline(always)]
    pub fn ptr(&self) -> *mut T {
        self.ptr as *mut T
    }

    /// Get the untagged pointer (standard implementation).
    ///
    /// # Safety
    /// The returned pointer is only valid if the original pointer passed to `new` is still valid.
    #[cfg(not(all(target_os = "macos", target_arch = "aarch64")))]
    #[inline(always)]
    pub fn ptr(&self) -> *mut T {
        // Standard implementation: manually mask off the tag bits
        (self.ptr & Self::PTR_MASK) as *mut T
    }

    /// Get the untagged pointer for deallocation.
    ///
    /// This always masks off the tag bits, even on platforms with TBI support,
    /// because memory allocators require the original untagged pointer.
    ///
    /// # Safety
    /// The returned pointer is only valid if the original pointer passed to `new` is still valid.
    #[doc(hidden)]
    #[inline(always)]
    pub fn untagged_ptr(&self) -> *mut T {
        const PTR_MASK: usize = !(0x7F << 57);
        (self.ptr & PTR_MASK) as *mut T
    }
    
    /// Get a reference to the pointed value.
    ///
    /// # Safety
    /// The caller must ensure that:
    /// - The pointer is valid and points to a properly initialized `T`
    /// - The pointed-to value is not being concurrently mutated
    #[inline(always)]
    pub unsafe fn as_ref(&self) -> &T {
        unsafe { &*self.ptr() }
    }

    /// Get a mutable reference to the pointed value.
    ///
    /// # Safety
    /// The caller must ensure that:
    /// - The pointer is valid and points to a properly initialized `T`
    /// - No other references to the pointed-to value exist
    #[inline(always)]
    pub unsafe fn as_mut(&mut self) -> &mut T {
        unsafe { &mut *self.ptr() }
    }
    
    /// Check if the pointer is null (ignoring the tag)
    #[inline(always)]
    pub fn is_null(&self) -> bool {
        self.ptr() as usize == 0
    }
}

// Safety: TaggedPtr is Send/Sync if T is Send/Sync
unsafe impl<T: Send> Send for TaggedPtr<T> {}
unsafe impl<T: Sync> Sync for TaggedPtr<T> {}

impl<T> Clone for TaggedPtr<T> {
    fn clone(&self) -> Self {
        *self
    }
}

impl<T> Copy for TaggedPtr<T> {}

impl<T> core::fmt::Debug for TaggedPtr<T> {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.debug_struct("TaggedPtr")
            .field("tag", &self.tag())
            .field("ptr", &format_args!("{:p}", self.ptr()))
            .finish()
    }
}

impl<T> core::cmp::PartialEq for TaggedPtr<T> {
    fn eq(&self, other: &Self) -> bool {
        // Compare the raw pointer values (tag + address)
        self.ptr == other.ptr
    }
}

impl<T> core::cmp::Eq for TaggedPtr<T> {}

impl<T> core::cmp::PartialOrd for TaggedPtr<T> {
    fn partial_cmp(&self, other: &Self) -> Option<core::cmp::Ordering> {
        Some(self.cmp(other))
    }
}

impl<T> core::cmp::Ord for TaggedPtr<T> {
    fn cmp(&self, other: &Self) -> core::cmp::Ordering {
        // Compare the raw pointer values (tag is in high bits, so this
        // naturally orders by tag first, then by address)
        self.ptr.cmp(&other.ptr)
    }
}

/// Allocator trait for arena-allocated tagged pointers.
///
/// This trait should be implemented by arena allocators to enable
/// the arena version of tagged dispatch.
///
/// # Example
///
/// ```rust
/// use tagged_dispatch::TaggedAllocator;
/// use bumpalo::Bump;
///
/// // Bumpalo automatically implements TaggedAllocator when the feature is enabled
/// let arena = Bump::new();
/// let ptr = arena.alloc(42);
/// ```
pub trait TaggedAllocator {
    /// Allocate space for a value and return a pointer to it.
    ///
    /// The allocated memory should have the same lifetime as the allocator.
    fn alloc<T>(&self, value: T) -> *mut T;
}

// Implement TaggedAllocator for common arena allocators when their features are enabled

#[cfg(feature = "bumpalo")]
impl TaggedAllocator for bumpalo::Bump {
    #[inline]
    fn alloc<T>(&self, value: T) -> *mut T {
        bumpalo::Bump::alloc(self, value) as *mut T
    }
}

// Note: typed_arena doesn't implement TaggedAllocator directly
// because it can only allocate values of a single type T.
// Instead, the arena builder pattern generates separate arenas
// for each variant type when typed_arena is enabled.

/// Statistics for arena memory usage.
#[derive(Debug, Clone, Copy, Default)]
pub struct ArenaStats {
    /// Total bytes currently allocated
    pub allocated_bytes: usize,
    /// Total capacity of all chunks
    pub chunk_capacity: usize,
}

/// Trait for arena builders generated by the macro.
///
/// Provides memory management capabilities for arena-allocated
/// tagged dispatch types.
pub trait ArenaBuilder<'a>: Sized {
    /// Create a new builder with default settings.
    ///
    /// When both allocators are available, this prefers bumpalo
    /// for its superior flexibility.
    fn new() -> Self;

    /// Reset all allocations, invalidating existing references.
    ///
    /// # Safety
    ///
    /// This invalidates all references previously allocated from this builder.
    /// Using any such references after reset is undefined behavior.
    fn reset(&mut self);

    /// Clear allocations and attempt to reclaim memory.
    ///
    /// More aggressive than reset, this tries to return memory to the OS.
    fn clear(&mut self);

    /// Get current memory usage statistics.
    fn stats(&self) -> ArenaStats;
}

/// A simple box allocator for owned tagged pointers.
///
/// This is used internally by the owned version of tagged dispatch.
pub struct BoxAllocator;

impl TaggedAllocator for BoxAllocator {
    #[inline]
    fn alloc<T>(&self, value: T) -> *mut T {
        Box::into_raw(Box::new(value))
    }
}

// Module with helper utilities
#[doc(hidden)]
pub mod __private {
    pub use core::mem;
    pub use core::ptr;
    pub use core::marker::PhantomData;
}

#[cfg(test)]
mod tests {
    use super::*;
    
    #[test]
    fn test_tag_extraction() {
        let ptr = core::ptr::null_mut::<u32>();
        let tagged = TaggedPtr::new(ptr, 127);
        assert_eq!(tagged.tag(), 127);

        // On macOS ARM64 with TBI, the pointer retains the tag bits
        // because the hardware ignores them automatically
        #[cfg(all(target_os = "macos", target_arch = "aarch64"))]
        {
            // The returned pointer should have the tag in the high byte
            let returned_ptr = tagged.ptr() as usize;
            let expected = ptr as usize | (127usize << TaggedPtr::<u32>::TAG_SHIFT);
            assert_eq!(returned_ptr, expected);
        }

        #[cfg(not(all(target_os = "macos", target_arch = "aarch64")))]
        {
            assert_eq!(tagged.ptr(), ptr);
        }
    }
    
    #[test]
    fn test_tag_preservation() {
        let value = Box::new(42u32);
        let ptr = Box::into_raw(value);

        for tag in 0..128u8 {
            let tagged = TaggedPtr::new(ptr, tag);
            assert_eq!(tagged.tag(), tag);

            // On macOS ARM64 with TBI, the pointer retains the tag bits
            #[cfg(all(target_os = "macos", target_arch = "aarch64"))]
            {
                let returned_ptr = tagged.ptr() as usize;
                let expected = ptr as usize | ((tag as usize) << TaggedPtr::<u32>::TAG_SHIFT);
                assert_eq!(returned_ptr, expected);
            }

            #[cfg(not(all(target_os = "macos", target_arch = "aarch64")))]
            {
                assert_eq!(tagged.ptr(), ptr);
            }
        }

        // Clean up - need to use the original untagged pointer for deallocation
        unsafe { let _ = Box::from_raw(ptr); }
    }
    
    #[test]
    fn test_size() {
        assert_eq!(core::mem::size_of::<TaggedPtr<()>>(), 8);
    }
    
    #[test]
    #[should_panic(expected = "Tag must be less than 128")]
    fn test_tag_overflow() {
        let ptr = core::ptr::null_mut::<u32>();
        let _tagged = TaggedPtr::new(ptr, 128);
    }
    
    #[cfg(feature = "bumpalo")]
    #[test]
    fn test_bumpalo_allocator() {
        use bumpalo::Bump;
        
        let arena = Bump::new();
        let value = 42u32;
        let ptr = arena.alloc(value);
        
        // Should be able to create a tagged pointer with arena allocation
        let tagged = TaggedPtr::new(ptr, 5);
        assert_eq!(tagged.tag(), 5);
        unsafe {
            assert_eq!(*tagged.as_ref(), 42);
        }
    }
}