crates/intern/src/symbol.rs - third_party/github.com/rust-lang/rust-analyzer - Git at Google

 //! Attempt at flexible symbol interning, allowing to intern and free strings at runtime while also
 //! supporting compile time declaration of symbols that will never be freed.

 use std::{
     fmt,
     hash::{BuildHasher, BuildHasherDefault, Hash},
     mem::{self, ManuallyDrop},
     ptr::NonNull,
     sync::OnceLock,
 };

 use dashmap::{DashMap, SharedValue};
 use hashbrown::raw::RawTable;
 use rustc_hash::FxHasher;
 use triomphe::Arc;

 pub mod symbols;

 // some asserts for layout compatibility
 const _: () = assert!(size_of::<Box<str>>() == size_of::<&str>());
 const _: () = assert!(align_of::<Box<str>>() == align_of::<&str>());

 const _: () = assert!(size_of::<Arc<Box<str>>>() == size_of::<&&str>());
 const _: () = assert!(align_of::<Arc<Box<str>>>() == align_of::<&&str>());

 const _: () = assert!(size_of::<*const *const str>() == size_of::<TaggedArcPtr>());
 const _: () = assert!(align_of::<*const *const str>() == align_of::<TaggedArcPtr>());

 const _: () = assert!(size_of::<Arc<Box<str>>>() == size_of::<TaggedArcPtr>());
 const _: () = assert!(align_of::<Arc<Box<str>>>() == align_of::<TaggedArcPtr>());

 /// A pointer that points to a pointer to a `str`, it may be backed as a `&'static &'static str` or
 /// `Arc<Box<str>>` but its size is that of a thin pointer. The active variant is encoded as a tag
 /// in the LSB of the alignment niche.
 // Note, Ideally this would encode a `ThinArc<str>` and `ThinRef<str>`/`ThinConstPtr<str>` instead of the double indirection.
 #[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: &'static &'static str) -> Self {
         assert!(align_of::<&'static &'static str>().trailing_zeros() as usize > Self::BOOL_BITS);
         // SAFETY: The pointer is non-null as it is derived from a reference
         // Ideally we would call out to `pack_arc` but for a `false` tag, unfortunately the
         // packing stuff requires reading out the pointer to an integer which is not supported
         // in const contexts, so here we make use of the fact that for the non-arc version the
         // tag is false (0) and thus does not need touching the actual pointer value.ext)

         let packed =
             unsafe { NonNull::new_unchecked((r as *const &str).cast::<*const str>().cast_mut()) };
         Self { packed }
     }

     fn arc(arc: Arc<Box<str>>) -> Self {
         assert!(align_of::<&'static &'static str>().trailing_zeros() as usize > Self::BOOL_BITS);
         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.
     ///
     /// # Safety
     ///
     /// You can only drop the `Arc` if the instance is dropped.
     #[inline]
     pub(crate) unsafe fn try_as_arc_owned(self) -> Option<ManuallyDrop<Arc<Box<str>>>> {
         // Unpack the tag from the alignment niche
         let tag = self.packed.as_ptr().addr() & Self::BOOL_BITS;
         if tag != 0 {
             // Safety: We checked that the tag is non-zero -> true, so we are pointing to the data offset of an `Arc`
             Some(ManuallyDrop::new(unsafe {
                 Arc::from_raw(self.pointer().as_ptr().cast::<Box<str>>())
             }))
         } else {
             None
         }
     }

     #[inline]
     fn pack_arc(ptr: NonNull<*const str>) -> NonNull<*const str> {
         let packed_tag = true as usize;

         unsafe {
             // Safety: The pointer is derived from a non-null and bit-oring it with true (1) will
             // not make it null.
             NonNull::new_unchecked(ptr.as_ptr().map_addr(|addr| addr | packed_tag))
         }
     }

     #[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(self.packed.as_ptr().map_addr(|addr| addr & !Self::BOOL_BITS))
         }
     }

     #[inline]
     pub(crate) fn as_str(&self) -> &str {
         // SAFETY: We always point to a pointer to a str no matter what variant is active
         unsafe { *self.pointer().as_ptr().cast::<&str>() }
     }
 }

 #[derive(PartialEq, Eq, Hash)]
 pub struct Symbol {
     repr: TaggedArcPtr,
 }

 impl fmt::Debug for Symbol {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         self.as_str().fmt(f)
     }
 }

 const _: () = assert!(size_of::<Symbol>() == size_of::<NonNull<()>>());
 const _: () = assert!(align_of::<Symbol>() == align_of::<NonNull<()>>());

 type Map = DashMap<Symbol, (), BuildHasherDefault<FxHasher>>;
 static MAP: OnceLock<Map> = OnceLock::new();

 impl Symbol {
     pub fn intern(s: &str) -> Self {
         let storage = MAP.get_or_init(symbols::prefill);
         let (mut shard, hash) = Self::select_shard(storage, s);
         // Atomically,
         // - check if `obj` is already in the map
         //   - if so, copy out its entry, conditionally bumping the backing Arc and return it
         //   - if not, put it into a box and then into an Arc, insert it, bump the ref-count and return the copy
         // 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.
         let bucket = match shard.find_or_find_insert_slot(
             hash,
             |(other, _)| other.as_str() == s,
             |(x, _)| Self::hash(storage, x.as_str()),
         ) {
             Ok(bucket) => bucket,
             // SAFETY: The slot came from `find_or_find_insert_slot()`, and the table wasn't modified since then.
             Err(insert_slot) => unsafe {
                 shard.insert_in_slot(
                     hash,
                     insert_slot,
                     (
                         Symbol { repr: TaggedArcPtr::arc(Arc::new(Box::<str>::from(s))) },
                         SharedValue::new(()),
                     ),
                 )
             },
         };
         // SAFETY: We just retrieved/inserted this bucket.
         unsafe { bucket.as_ref().0.clone() }
     }

     pub fn integer(i: usize) -> Self {
         match i {
             0 => symbols::INTEGER_0,
             1 => symbols::INTEGER_1,
             2 => symbols::INTEGER_2,
             3 => symbols::INTEGER_3,
             4 => symbols::INTEGER_4,
             5 => symbols::INTEGER_5,
             6 => symbols::INTEGER_6,
             7 => symbols::INTEGER_7,
             8 => symbols::INTEGER_8,
             9 => symbols::INTEGER_9,
             10 => symbols::INTEGER_10,
             11 => symbols::INTEGER_11,
             12 => symbols::INTEGER_12,
             13 => symbols::INTEGER_13,
             14 => symbols::INTEGER_14,
             15 => symbols::INTEGER_15,
             i => Symbol::intern(&format!("{i}")),
         }
     }

     pub fn empty() -> Self {
         symbols::__empty
     }

     #[inline]
     pub fn as_str(&self) -> &str {
         self.repr.as_str()
     }

     #[inline]
     fn select_shard(
         storage: &'static Map,
         s: &str,
     ) -> (dashmap::RwLockWriteGuard<'static, RawTable<(Symbol, SharedValue<()>)>>, u64) {
         let hash = Self::hash(storage, s);
         let shard_idx = storage.determine_shard(hash as usize);
         let shard = &storage.shards()[shard_idx];
         (shard.write(), hash)
     }

     #[inline]
     fn hash(storage: &'static Map, s: &str) -> u64 {
         storage.hasher().hash_one(s)
     }

     #[cold]
     fn drop_slow(arc: &Arc<Box<str>>) {
         let storage = MAP.get_or_init(symbols::prefill);
         let (mut shard, hash) = Self::select_shard(storage, arc);

         match Arc::count(arc) {
             0 | 1 => unreachable!(),
             2 => (),
             _ => {
                 // Another thread has interned another copy
                 return;
             }
         }

         let s = &***arc;
         let (ptr, _) = shard.remove_entry(hash, |(x, _)| x.as_str() == s).unwrap();
         let ptr = ManuallyDrop::new(ptr);
         // SAFETY: We're dropping, we have ownership.
         ManuallyDrop::into_inner(unsafe { ptr.repr.try_as_arc_owned().unwrap() });
         debug_assert_eq!(Arc::count(arc), 1);

         // Shrink the backing storage if the shard is less than 50% occupied.
         if shard.len() * 2 < shard.capacity() {
             let len = shard.len();
             shard.shrink_to(len, |(x, _)| Self::hash(storage, x.as_str()));
         }
     }
 }

 impl Drop for Symbol {
     #[inline]
     fn drop(&mut self) {
         // SAFETY: We're dropping, we have ownership.
         let Some(arc) = (unsafe { 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
         ManuallyDrop::into_inner(arc);
     }
 }

 impl Clone for Symbol {
     fn clone(&self) -> Self {
         Self { repr: increase_arc_refcount(self.repr) }
     }
 }

 fn increase_arc_refcount(repr: TaggedArcPtr) -> TaggedArcPtr {
     // SAFETY: We're not dropping the `Arc`.
     let Some(arc) = (unsafe { repr.try_as_arc_owned() }) else {
         return repr;
     };
     // increase the ref count
     mem::forget(Arc::clone(&arc));
     repr
 }

 impl fmt::Display for Symbol {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         self.as_str().fmt(f)
     }
 }

 #[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 more_worlds = world.clone();
         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");
         let many_worlds = world.clone();
         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(many_worlds);
         drop(more_worlds);
         drop(hello);
         drop(world);
         drop(bang);
         assert_eq!(MAP.get().unwrap().len(), base_len);
     }
 }
	//! Attempt at flexible symbol interning, allowing to intern and free strings at runtime while also
	//! supporting compile time declaration of symbols that will never be freed.

	use std::{
	fmt,
	hash::{BuildHasher, BuildHasherDefault, Hash},
	mem::{self, ManuallyDrop},
	ptr::NonNull,
	sync::OnceLock,
	};

	use dashmap::{DashMap, SharedValue};
	use hashbrown::raw::RawTable;
	use rustc_hash::FxHasher;
	use triomphe::Arc;

	pub mod symbols;

	// some asserts for layout compatibility
	const _: () = assert!(size_of::<Box<str>>() == size_of::<&str>());
	const _: () = assert!(align_of::<Box<str>>() == align_of::<&str>());

	const _: () = assert!(size_of::<Arc<Box<str>>>() == size_of::<&&str>());
	const _: () = assert!(align_of::<Arc<Box<str>>>() == align_of::<&&str>());

	const _: () = assert!(size_of::<const const str>() == size_of::<TaggedArcPtr>());
	const _: () = assert!(align_of::<const const str>() == align_of::<TaggedArcPtr>());

	const _: () = assert!(size_of::<Arc<Box<str>>>() == size_of::<TaggedArcPtr>());
	const _: () = assert!(align_of::<Arc<Box<str>>>() == align_of::<TaggedArcPtr>());

	/// A pointer that points to a pointer to a `str`, it may be backed as a `&'static &'static str` or
	/// `Arc<Box<str>>` but its size is that of a thin pointer. The active variant is encoded as a tag
	/// in the LSB of the alignment niche.
	// Note, Ideally this would encode a `ThinArc<str>` and `ThinRef<str>`/`ThinConstPtr<str>` instead of the double indirection.
	#[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: &'static &'static str) -> Self {
	assert!(align_of::<&'static &'static str>().trailing_zeros() as usize > Self::BOOL_BITS);
	// SAFETY: The pointer is non-null as it is derived from a reference
	// Ideally we would call out to `pack_arc` but for a `false` tag, unfortunately the
	// packing stuff requires reading out the pointer to an integer which is not supported
	// in const contexts, so here we make use of the fact that for the non-arc version the
	// tag is false (0) and thus does not need touching the actual pointer value.ext)

	let packed =
	unsafe { NonNull::new_unchecked((r as const &str).cast::<const str>().cast_mut()) };
	Self { packed }
	}

	fn arc(arc: Arc<Box<str>>) -> Self {
	assert!(align_of::<&'static &'static str>().trailing_zeros() as usize > Self::BOOL_BITS);
	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.
	///
	/// # Safety
	///
	/// You can only drop the `Arc` if the instance is dropped.
	#[inline]
	pub(crate) unsafe fn try_as_arc_owned(self) -> Option<ManuallyDrop<Arc<Box<str>>>> {
	// Unpack the tag from the alignment niche
	let tag = self.packed.as_ptr().addr() & Self::BOOL_BITS;
	if tag != 0 {
	// Safety: We checked that the tag is non-zero -> true, so we are pointing to the data offset of an `Arc`
	Some(ManuallyDrop::new(unsafe {
	Arc::from_raw(self.pointer().as_ptr().cast::<Box<str>>())
	}))
	} else {
	None
	}
	}

	#[inline]
	fn pack_arc(ptr: NonNull<const str>) -> NonNull<const str> {
	let packed_tag = true as usize;

	unsafe {
	// Safety: The pointer is derived from a non-null and bit-oring it with true (1) will
	// not make it null.
	NonNull::new_unchecked(ptr.as_ptr().map_addr(\|addr\| addr \| packed_tag))
	}
	}

	#[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(self.packed.as_ptr().map_addr(\|addr\| addr & !Self::BOOL_BITS))
	}
	}

	#[inline]
	pub(crate) fn as_str(&self) -> &str {
	// SAFETY: We always point to a pointer to a str no matter what variant is active
	unsafe { *self.pointer().as_ptr().cast::<&str>() }
	}
	}

	#[derive(PartialEq, Eq, Hash)]
	pub struct Symbol {
	repr: TaggedArcPtr,
	}

	impl fmt::Debug for Symbol {
	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
	self.as_str().fmt(f)
	}
	}

	const _: () = assert!(size_of::<Symbol>() == size_of::<NonNull<()>>());
	const _: () = assert!(align_of::<Symbol>() == align_of::<NonNull<()>>());

	type Map = DashMap<Symbol, (), BuildHasherDefault<FxHasher>>;
	static MAP: OnceLock<Map> = OnceLock::new();

	impl Symbol {
	pub fn intern(s: &str) -> Self {
	let storage = MAP.get_or_init(symbols::prefill);
	let (mut shard, hash) = Self::select_shard(storage, s);
	// Atomically,
	// - check if `obj` is already in the map
	// - if so, copy out its entry, conditionally bumping the backing Arc and return it
	// - if not, put it into a box and then into an Arc, insert it, bump the ref-count and return the copy
	// 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.
	let bucket = match shard.find_or_find_insert_slot(
	hash,
	\|(other, _)\| other.as_str() == s,
	\|(x, _)\| Self::hash(storage, x.as_str()),
	) {
	Ok(bucket) => bucket,
	// SAFETY: The slot came from `find_or_find_insert_slot()`, and the table wasn't modified since then.
	Err(insert_slot) => unsafe {
	shard.insert_in_slot(
	hash,
	insert_slot,
	(
	Symbol { repr: TaggedArcPtr::arc(Arc::new(Box::<str>::from(s))) },
	SharedValue::new(()),
	),
	)
	},
	};
	// SAFETY: We just retrieved/inserted this bucket.
	unsafe { bucket.as_ref().0.clone() }
	}

	pub fn integer(i: usize) -> Self {
	match i {
	0 => symbols::INTEGER_0,
	1 => symbols::INTEGER_1,
	2 => symbols::INTEGER_2,
	3 => symbols::INTEGER_3,
	4 => symbols::INTEGER_4,
	5 => symbols::INTEGER_5,
	6 => symbols::INTEGER_6,
	7 => symbols::INTEGER_7,
	8 => symbols::INTEGER_8,
	9 => symbols::INTEGER_9,
	10 => symbols::INTEGER_10,
	11 => symbols::INTEGER_11,
	12 => symbols::INTEGER_12,
	13 => symbols::INTEGER_13,
	14 => symbols::INTEGER_14,
	15 => symbols::INTEGER_15,
	i => Symbol::intern(&format!("{i}")),
	}
	}

	pub fn empty() -> Self {
	symbols::__empty
	}

	#[inline]
	pub fn as_str(&self) -> &str {
	self.repr.as_str()
	}

	#[inline]
	fn select_shard(
	storage: &'static Map,
	s: &str,
	) -> (dashmap::RwLockWriteGuard<'static, RawTable<(Symbol, SharedValue<()>)>>, u64) {
	let hash = Self::hash(storage, s);
	let shard_idx = storage.determine_shard(hash as usize);
	let shard = &storage.shards()[shard_idx];
	(shard.write(), hash)
	}

	#[inline]
	fn hash(storage: &'static Map, s: &str) -> u64 {
	storage.hasher().hash_one(s)
	}

	#[cold]
	fn drop_slow(arc: &Arc<Box<str>>) {
	let storage = MAP.get_or_init(symbols::prefill);
	let (mut shard, hash) = Self::select_shard(storage, arc);

	match Arc::count(arc) {
	0 \| 1 => unreachable!(),
	2 => (),
	_ => {
	// Another thread has interned another copy
	return;
	}
	}

	let s = &***arc;
	let (ptr, _) = shard.remove_entry(hash, \|(x, _)\| x.as_str() == s).unwrap();
	let ptr = ManuallyDrop::new(ptr);
	// SAFETY: We're dropping, we have ownership.
	ManuallyDrop::into_inner(unsafe { ptr.repr.try_as_arc_owned().unwrap() });
	debug_assert_eq!(Arc::count(arc), 1);

	// Shrink the backing storage if the shard is less than 50% occupied.
	if shard.len() * 2 < shard.capacity() {
	let len = shard.len();
	shard.shrink_to(len, \|(x, _)\| Self::hash(storage, x.as_str()));
	}
	}
	}

	impl Drop for Symbol {
	#[inline]
	fn drop(&mut self) {
	// SAFETY: We're dropping, we have ownership.
	let Some(arc) = (unsafe { 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
	ManuallyDrop::into_inner(arc);
	}
	}

	impl Clone for Symbol {
	fn clone(&self) -> Self {
	Self { repr: increase_arc_refcount(self.repr) }
	}
	}

	fn increase_arc_refcount(repr: TaggedArcPtr) -> TaggedArcPtr {
	// SAFETY: We're not dropping the `Arc`.
	let Some(arc) = (unsafe { repr.try_as_arc_owned() }) else {
	return repr;
	};
	// increase the ref count
	mem::forget(Arc::clone(&arc));
	repr
	}

	impl fmt::Display for Symbol {
	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
	self.as_str().fmt(f)
	}
	}

	#[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 more_worlds = world.clone();
	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");
	let many_worlds = world.clone();
	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(many_worlds);
	drop(more_worlds);
	drop(hello);
	drop(world);
	drop(bang);
	assert_eq!(MAP.get().unwrap().len(), base_len);
	}
	}