Implement rough symbol interning infra

Author: Veykril

.typos.toml

@@ -14,6 +14,8 @@ extend-ignore-re = [
     "\\w*\\.{3,4}\\w*",
     '"flate2"',
     "raison d'être",
+    "inout",
+    "optin"
 ]
 
 [default.extend-words]

Cargo.lock

@@ -18,6 +18,7 @@ dashmap.workspace = true
 hashbrown.workspace = true
 rustc-hash.workspace = true
 triomphe.workspace = true
+sptr = "0.3.2"
 
 [lints]
-workspace = true
+workspace = true

crates/intern/Cargo.toml

@@ -20,6 +20,9 @@ type Guard<T> = dashmap::RwLockWriteGuard<
     HashMap<Arc<T>, SharedValue<()>, BuildHasherDefault<FxHasher>>,
 >;
 
+mod symbol;
+pub use self::symbol::{symbols, Symbol};
+
 pub struct Interned<T: Internable + ?Sized> {
     arc: Arc<T>,
 }

crates/intern/src/lib.rs

@@ -0,0 +1,291 @@
+//! Attempt at flexible symbol interning, allowing to intern and free strings at runtime while also
+//! supporting
+
+use std::{
+    borrow::Borrow,
+    fmt,
+    hash::{BuildHasherDefault, Hash, Hasher},
+    mem,
+    ptr::NonNull,
+    sync::OnceLock,
+};
+
+use dashmap::{DashMap, SharedValue};
+use hashbrown::{hash_map::RawEntryMut, HashMap};
+use rustc_hash::FxHasher;
+use sptr::Strict;
+use triomphe::Arc;
+
+pub mod symbols;
+
+const _: () = assert!(std::mem::size_of::<Box<str>>() == std::mem::size_of::<&str>());
+const _: () = assert!(std::mem::align_of::<Box<str>>() == std::mem::align_of::<&str>());
+
+const _: () = assert!(std::mem::size_of::<Arc<Box<str>>>() == std::mem::size_of::<&&str>());
+const _: () = assert!(std::mem::align_of::<Arc<Box<str>>>() == std::mem::align_of::<&&str>());
+
+#[derive(PartialEq, Eq, Hash, Copy, Clone, Debug)]
+struct TaggedArcPtr {
+    packed: NonNull<*const str>,
+}
+
+unsafe impl Send for TaggedArcPtr {}
+unsafe impl Sync for TaggedArcPtr {}
+
+impl TaggedArcPtr {
+    const BOOL_BITS: usize = true as usize;
+
+    const fn non_arc(r: &&str) -> Self {
+        Self {
+            // SAFETY: The pointer is non-null as it is derived from a reference
+            // Additionally, the tag is empty as we aren't encoding an arc, so we do not need to
+            // handle tagging here. (we couldn't do it if we wanted to as we are in a const context)
+            packed: unsafe {
+                NonNull::new_unchecked((r as *const &str).cast::<*const str>().cast_mut())
+            },
+        }
+    }
+
+    fn arc(arc: Arc<Box<str>>) -> Self {
+        Self {
+            packed: Self::pack_arc(
+                // Safety: `Arc::into_raw`` always returns a non null pointer
+                unsafe { NonNull::new_unchecked(Arc::into_raw(arc).cast_mut().cast()) },
+            ),
+        }
+    }
+
+    /// Retrieves the tag.
+    #[inline]
+    pub(crate) fn try_as_arc_owned(self) -> Option<Arc<Box<str>>> {
+        // Unpack the tag, according to the `self.packed` encoding scheme
+        let tag = Strict::addr(self.packed.as_ptr()) & Self::BOOL_BITS;
+        if tag != 0 {
+            // Safety: We checked that the tag is true, so we are pointing to the data offset of an `Arc`
+            Some(unsafe { Arc::from_raw(self.pointer().as_ptr().cast::<Box<str>>()) })
+        } else {
+            None
+        }
+    }
+
+    #[inline]
+    const fn pack_arc(ptr: NonNull<*const str>) -> NonNull<*const str> {
+        let packed_tag = true as usize;
+
+        // can't use this due to trait methods not being const
+        // unsafe {
+        //     // Safety: The pointer is derived from a non-null
+        //     NonNull::new_unchecked(Strict::map_addr(ptr.as_ptr(), |addr| {
+        //         // Safety:
+        //         // - The pointer is `NonNull` => it's address is `NonZero<usize>`
+        //         // - `P::BITS` least significant bits are always zero (`Pointer` contract)
+        //         // - `T::BITS <= P::BITS` (from `Self::ASSERTION`)
+        //         //
+        //         // Thus `addr >> T::BITS` is guaranteed to be non-zero.
+        //         //
+        //         // `{non_zero} | packed_tag` can't make the value zero.
+
+        //         (addr >> Self::BOOL_BITS) | packed_tag
+        //     }))
+        // }
+
+        let self_addr = unsafe { core::mem::transmute::<*const _, usize>(ptr.as_ptr()) };
+        let addr = self_addr | packed_tag;
+
+        // FIXME(strict_provenance_magic): I am magic and should be a compiler intrinsic.
+        //
+        // In the mean-time, this operation is defined to be "as if" it was
+        // a wrapping_offset, so we can emulate it as such. This should properly
+        // restore pointer provenance even under today's compiler.
+        let dest_addr = addr as isize;
+        let offset = dest_addr.wrapping_sub(self_addr as isize);
+
+        unsafe { NonNull::new_unchecked(ptr.as_ptr().cast::<u8>().wrapping_offset(offset).cast()) }
+    }
+
+    #[inline]
+    pub(crate) fn pointer(self) -> NonNull<*const str> {
+        // SAFETY: The resulting pointer is guaranteed to be NonNull as we only modify the niche bytes
+        unsafe {
+            NonNull::new_unchecked(Strict::map_addr(self.packed.as_ptr(), |addr| {
+                addr & !Self::BOOL_BITS
+            }))
+        }
+    }
+
+    #[inline]
+    pub(crate) fn as_str(&self) -> &str {
+        // SAFETY: We always point to a pointer to a str
+        unsafe { *self.pointer().as_ptr().cast::<&str>() }
+    }
+}
+
+#[derive(PartialEq, Eq, Hash, Clone, Debug)]
+pub struct Symbol {
+    repr: TaggedArcPtr,
+}
+const _: () = assert!(std::mem::size_of::<Symbol>() == std::mem::size_of::<NonNull<()>>());
+const _: () = assert!(std::mem::align_of::<Symbol>() == std::mem::align_of::<NonNull<()>>());
+
+static MAP: OnceLock<DashMap<SymbolProxy, (), BuildHasherDefault<FxHasher>>> = OnceLock::new();
+
+impl Symbol {
+    pub fn intern(s: &str) -> Self {
+        let (mut shard, hash) = Self::select_shard(s);
+        // Atomically,
+        // - check if `obj` is already in the map
+        //   - if so, clone its `Arc` and return it
+        //   - if not, box it up, insert it, and return a clone
+        // This needs to be atomic (locking the shard) to avoid races with other thread, which could
+        // insert the same object between us looking it up and inserting it.
+        match shard.raw_entry_mut().from_key_hashed_nocheck(hash, s) {
+            RawEntryMut::Occupied(occ) => Self { repr: increase_arc_refcount(occ.key().0) },
+            RawEntryMut::Vacant(vac) => Self {
+                repr: increase_arc_refcount(
+                    vac.insert_hashed_nocheck(
+                        hash,
+                        SymbolProxy(TaggedArcPtr::arc(Arc::new(Box::<str>::from(s)))),
+                        SharedValue::new(()),
+                    )
+                    .0
+                     .0,
+                ),
+            },
+        }
+    }
+
+    pub fn as_str(&self) -> &str {
+        self.repr.as_str()
+    }
+
+    #[inline]
+    fn select_shard(
+        s: &str,
+    ) -> (
+        dashmap::RwLockWriteGuard<
+            'static,
+            HashMap<SymbolProxy, SharedValue<()>, BuildHasherDefault<FxHasher>>,
+        >,
+        u64,
+    ) {
+        let storage = MAP.get_or_init(symbols::prefill);
+        let hash = {
+            let mut hasher = std::hash::BuildHasher::build_hasher(storage.hasher());
+            s.hash(&mut hasher);
+            hasher.finish()
+        };
+        let shard_idx = storage.determine_shard(hash as usize);
+        let shard = &storage.shards()[shard_idx];
+        (shard.write(), hash)
+    }
+
+    #[cold]
+    fn drop_slow(arc: &Arc<Box<str>>) {
+        let (mut shard, hash) = Self::select_shard(arc);
+
+        if Arc::count(arc) != 2 {
+            // Another thread has interned another copy
+            return;
+        }
+
+        match shard.raw_entry_mut().from_key_hashed_nocheck::<str>(hash, arc.as_ref()) {
+            RawEntryMut::Occupied(occ) => occ.remove_entry(),
+            RawEntryMut::Vacant(_) => unreachable!(),
+        }
+        .0
+         .0
+        .try_as_arc_owned()
+        .unwrap();
+
+        // Shrink the backing storage if the shard is less than 50% occupied.
+        if shard.len() * 2 < shard.capacity() {
+            shard.shrink_to_fit();
+        }
+    }
+}
+
+impl Drop for Symbol {
+    #[inline]
+    fn drop(&mut self) {
+        let Some(arc) = self.repr.try_as_arc_owned() else {
+            return;
+        };
+        // When the last `Ref` is dropped, remove the object from the global map.
+        if Arc::count(&arc) == 2 {
+            // Only `self` and the global map point to the object.
+
+            Self::drop_slow(&arc);
+        }
+        // decrement the ref count
+        drop(arc);
+    }
+}
+
+fn increase_arc_refcount(repr: TaggedArcPtr) -> TaggedArcPtr {
+    let Some(arc) = repr.try_as_arc_owned() else {
+        return repr;
+    };
+    // increase the ref count
+    mem::forget(arc.clone());
+    mem::forget(arc);
+    repr
+}
+
+impl fmt::Display for Symbol {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        self.as_str().fmt(f)
+    }
+}
+
+// only exists so we can use `from_key_hashed_nocheck` with a &str
+#[derive(Debug, PartialEq, Eq)]
+struct SymbolProxy(TaggedArcPtr);
+
+impl Hash for SymbolProxy {
+    fn hash<H: Hasher>(&self, state: &mut H) {
+        self.0.as_str().hash(state);
+    }
+}
+
+impl Borrow<str> for SymbolProxy {
+    fn borrow(&self) -> &str {
+        self.0.as_str()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn smoke_test() {
+        Symbol::intern("isize");
+        let base_len = MAP.get().unwrap().len();
+        let hello = Symbol::intern("hello");
+        let world = Symbol::intern("world");
+        let bang = Symbol::intern("!");
+        let q = Symbol::intern("?");
+        assert_eq!(MAP.get().unwrap().len(), base_len + 4);
+        let bang2 = Symbol::intern("!");
+        assert_eq!(MAP.get().unwrap().len(), base_len + 4);
+        drop(bang2);
+        assert_eq!(MAP.get().unwrap().len(), base_len + 4);
+        drop(q);
+        assert_eq!(MAP.get().unwrap().len(), base_len + 3);
+        let default = Symbol::intern("default");
+        assert_eq!(MAP.get().unwrap().len(), base_len + 3);
+        assert_eq!(
+            "hello default world!",
+            format!("{} {} {}{}", hello.as_str(), default.as_str(), world.as_str(), bang.as_str())
+        );
+        drop(default);
+        assert_eq!(
+            "hello world!",
+            format!("{} {}{}", hello.as_str(), world.as_str(), bang.as_str())
+        );
+        drop(hello);
+        drop(world);
+        drop(bang);
+        assert_eq!(MAP.get().unwrap().len(), base_len);
+    }
+}