hashcons-arena 0.1.0

//! A hash consing arena for efficient interning of values.
//!
//! This arena allows you to intern values such that structurally equal values will yield the
//! same reference, thus saving memory and improving performance in scenarios where many
//! identical values are used.
//!
//! # Features
//!
//! - `sync`: Use a `RwLock`, rather than a `RefCell` to implement interior mutability. This
//! allows the arena to be [`Sync`]. Depends on the `parking_lot` crate.
//!
//! # Example
//!
//! ```rust,ignore
//! use hashcons_arena::HashConsArena;
//!
//! let arena = HashConsArena::new();
//! let a = arena.intern("hello");
//! let b = arena.intern("hello");
//!
//! assert!(a == b); // a and b are the same reference
//! ```
//!
//! # Safety
//!
//! The crate uses `unsafe` in a number of places, and has been tested using miri.
//!

#[cfg(not(feature = "sync"))]
use std::cell::RefCell as RC;
use std::cmp::Ordering;
use std::collections::BTreeMap;
use std::fmt::{self, Debug, Display};
use std::hash::{Hash, Hasher};
use std::mem;
use std::ops::Deref;
use std::pin::Pin;
use std::ptr::NonNull;

use bumpalo::Bump;
use bumpalo::boxed::Box;
use fxhash::hash64;
#[cfg(feature = "sync")]
use parking_lot::RwLock as RC;

#[cfg(feature = "sync")]
mod rc {
    use super::RC;

    #[inline(always)]
    pub(crate) fn read_table<T>(table: &RC<T>) -> parking_lot::RwLockReadGuard<'_, T> {
        table.read()
    }

    #[inline(always)]
    pub(crate) fn write_table<T>(table: &RC<T>) -> parking_lot::RwLockWriteGuard<'_, T> {
        table.write()
    }
}

#[cfg(not(feature = "sync"))]
mod rc {
    use super::RC;

    #[inline(always)]
    pub(crate) fn read_table<T>(table: &RC<T>) -> std::cell::Ref<'_, T> {
        table.borrow()
    }

    #[inline(always)]
    pub(crate) fn write_table<T>(table: &RC<T>) -> std::cell::RefMut<'_, T> {
        table.borrow_mut()
    }
}

/// A hash consing arena that allows for efficient interning of values.
///
/// ## Caveat(s)
///
/// The arena is backed by a [bumpalo](https://crates.io/crates/bumpalo), therefore, objects
/// allocated via the arena, will not have their [`Drop`] implementation called, without manual
/// intervention. Thus, it is relatively easy to leak resources. If you need [`Drop`] to be called,
/// use [`BoxedHashConsArena`] instead.
///
pub struct HashConsArena<T> {
    bump: Bump,
    table: RC<BTreeMap<u64, Vec<NonNull<T>>>>,
}

pub struct HRef<'a, T> {
    ptr: &'a T,
}

impl<'a, T> Clone for HRef<'a, T> {
    fn clone(&self) -> Self {
        Self { ptr: self.ptr }
    }
}

impl<'a, T> Copy for HRef<'a, T> {}

impl<'a, T> HRef<'a, T> {
    pub(crate) fn new(ptr: &'a T) -> Self {
        Self { ptr }
    }

    pub(crate) fn as_ptr(&self) -> *const T {
        self.ptr as *const T
    }

    pub fn as_ref(&self) -> &'a T {
        self.ptr
    }
}

impl<'a, T> Deref for HRef<'a, T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        self.ptr
    }
}

impl<'a, T> PartialEq for HRef<'a, T> {
    fn eq(&self, other: &Self) -> bool {
        std::ptr::eq(self.ptr, other.ptr)
    }
}

impl<'a, T> PartialEq<T> for HRef<'a, T>
where
    T: PartialEq,
{
    fn eq(&self, other: &T) -> bool {
        self.as_ref() == other
    }
}

impl<'a, T> Eq for HRef<'a, T> {}

impl<'a, T> Hash for HRef<'a, T> {
    fn hash<H: Hasher>(&self, state: &mut H) {
        // hash the pointer address, not the value
        self.as_ptr().hash(state);
    }
}

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

impl<'a, T> Ord for HRef<'a, T> {
    fn cmp(&self, other: &Self) -> Ordering {
        self.as_ptr().cmp(&other.as_ptr())
    }
}

impl<'a, T> Debug for HRef<'a, T>
where
    T: Debug,
{
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("HRef")
            .field("ptr", &(self.as_ptr() as usize))
            .field("val", &self.ptr)
            .finish()
    }
}

impl<'a, T> Display for HRef<'a, T>
where
    T: Display,
{
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        self.ptr.fmt(f)
    }
}

impl<T> HashConsArena<T>
where
    T: Eq + Hash,
{
    /// Create a new `HashConsArena`.
    pub fn new() -> Self {
        Self {
            bump: Bump::new(),
            table: RC::new(BTreeMap::new()),
        }
    }

    /// Intern a value, returning a reference that's guaranteed to be the same for structurally
    /// equal values.
    ///
    /// # Arguments
    /// * `value` - The value to intern. It must implement `Eq` and `Hash`.
    ///
    /// # Returns
    /// A `HRef<T>` that points to the interned value. If the value is already interned, it returns
    /// the existing reference.
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// use hashcons_arena::HashConsArena;
    ///
    /// let arena = HashConsArena::new();
    ///
    /// let a = arena.intern("hello");
    /// let b = arena.intern("hello");
    ///
    /// assert!(a == b); // a and b are the same reference
    /// ```
    ///
    pub fn intern<'a>(&'a self, value: T) -> HRef<'a, T> {
        // compute hash of the value
        let hash = hash64(&value);

        // check if we already have this value
        {
            let table = rc::read_table(&self.table);
            if let Some(candidates) = table.get(&hash) {
                for &ptr in candidates {
                    let existing = unsafe { ptr.as_ref() };
                    if *existing == value {
                        // found
                        return HRef::new(existing);
                    }
                }
            }
        } // table borrow ends here

        // not found, allocate new one
        let allocated = &*self.bump.alloc(value);
        let ptr = NonNull::from(allocated);

        // add to hash table
        rc::write_table(&self.table)
            .entry(hash)
            .or_insert_with(Vec::new)
            .push(ptr);

        HRef::new(allocated)
    }

    /// Reset the arena, clearing all interned values.
    pub fn reset(&mut self) {
        rc::write_table(&self.table).clear();
        self.bump.reset();
    }
}

/// A version of `HashConsArena` that ensures interned values are dropped.
pub struct BoxedHashConsArena<T>
where
    T: 'static,
{
    bump: Bump,
    table: RC<BTreeMap<u64, Vec<Pin<Box<'static, T>>>>>,
}

impl<T> BoxedHashConsArena<T>
where
    T: Eq + Hash + 'static,
{
    /// Create a new `BoxedHashConsArena`.
    pub fn new() -> Self {
        Self {
            bump: Bump::new(),
            table: RC::new(BTreeMap::new()),
        }
    }

    /// Intern a value, returning a reference that's guaranteed to be the same for structurally
    /// equal values.
    ///
    /// # Arguments
    ///
    /// * `value` - The value to intern. It must implement `Eq` and `Hash`.
    ///
    /// # Returns
    ///
    /// A `HRef<T>` that points to the interned value. If the value is already interned, it returns
    /// the existing reference.
    ///
    pub fn intern<'a>(&'a self, value: T) -> HRef<'a, T> {
        // compute hash of the value
        let hash = hash64(&value);

        // check if we already have this value
        {
            let table = rc::read_table(&self.table);
            if let Some(candidates) = table.get(&hash) {
                for existing in candidates {
                    if **existing == value {
                        // found - return reference with proper lifetime
                        // safety: the box lives as long as the arena, and we're borrowing
                        // from the arena with lifetime 'a, so this is safe
                        let existing_ref = unsafe { mem::transmute::<&T, &'a T>(existing.deref()) };
                        return HRef::new(existing_ref);
                    }
                }
            }
        } // table borrow ends here

        // not found, allocate new one
        let allocated = Box::pin_in(value, &self.bump);

        // safety: the box will live as long as the bump allocator, and we ensure
        // that the bump allocator is never reset while references exist
        let static_box =
            unsafe { mem::transmute::<Pin<Box<'_, T>>, Pin<Box<'static, T>>>(allocated) };

        // add to hash table first, then get reference from the stored location
        let mut table = rc::write_table(&self.table);

        let entry = table.entry(hash).or_insert_with(Vec::new);

        entry.push(static_box);

        // now get the reference from the box in its final storage location (pleasing miri)
        let stored_box = entry.last().unwrap();
        let allocated_ref = unsafe { mem::transmute::<&T, &'a T>(stored_box.deref()) };

        HRef::new(allocated_ref)
    }

    /// Reset the arena, clearing all interned values.
    pub fn reset(&mut self) {
        rc::write_table(&self.table).clear();
        self.bump.reset();
    }
}

impl<T> Drop for BoxedHashConsArena<T> {
    fn drop(&mut self) {
        // this is necessary to ensure that all boxes are dropped before the bump allocator
        // otherwise we might have dangling references
        rc::write_table(&self.table).clear();
    }
}

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

    #[test]
    fn test_interning() {
        let arena = HashConsArena::new();
        let a = arena.intern("hello");
        let b = arena.intern("hello");

        assert!(a == b); // a and b are the same reference
        assert_eq!(a, "hello");
    }

    #[test]
    fn test_multiple_arenas() {
        let arena1 = HashConsArena::new();
        let arena2 = HashConsArena::new();

        let a1 = arena1.intern("hello");
        let a2 = arena2.intern("hello");

        assert!(a1 != a2); // a1 and a2 are different references
        assert_eq!(a1, "hello");
        assert_eq!(a2, "hello");
    }

    #[test]
    fn test_drop() {
        use std::sync::Arc;
        use std::sync::atomic::AtomicUsize;

        let arena = BoxedHashConsArena::new();
        let drop_ctr = Arc::new(AtomicUsize::new(0));

        struct MyStr {
            value: String,
            drop_ctr: Arc<AtomicUsize>,
        }

        impl MyStr {
            fn new(value: impl Into<String>, drop_ctr: Arc<AtomicUsize>) -> Self {
                Self {
                    value: value.into(),
                    drop_ctr,
                }
            }
        }

        impl PartialEq for MyStr {
            fn eq(&self, other: &Self) -> bool {
                self.value == other.value
            }
        }

        impl Eq for MyStr {}

        impl Hash for MyStr {
            fn hash<H: Hasher>(&self, state: &mut H) {
                self.value.hash(state);
            }
        }

        impl Drop for MyStr {
            fn drop(&mut self) {
                self.drop_ctr
                    .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
            }
        }

        let a = arena.intern(MyStr::new("hello", drop_ctr.clone()));
        let b = arena.intern(MyStr::new("world", drop_ctr.clone()));

        assert!(a != b); // a and b are different references

        // drop the arena, which should clear the table
        drop(arena);

        // check that the drop counter was incremented for both values
        assert_eq!(drop_ctr.load(std::sync::atomic::Ordering::SeqCst), 2);
    }

    #[test]
    fn test_reset() {
        let mut arena = HashConsArena::new();
        let _a = arena.intern("hello");
        assert_eq!(rc::read_table(&arena.table).len(), 1);
        arena.reset();
        assert_eq!(rc::read_table(&arena.table).len(), 0);
    }
}