diff --git a/crates/intern/src/lib.rs b/crates/intern/src/lib.rs index d17eddf1560..dabbf3a38b5 100644 --- a/crates/intern/src/lib.rs +++ b/crates/intern/src/lib.rs @@ -9,7 +9,7 @@ }; use dashmap::{DashMap, SharedValue}; -use hashbrown::HashMap; +use hashbrown::{hash_map::RawEntryMut, HashMap}; use once_cell::sync::OnceCell; use rustc_hash::FxHasher; use triomphe::Arc; @@ -26,59 +26,61 @@ pub struct Interned { impl Interned { pub fn new(obj: T) -> Self { - match Interned::lookup(&obj) { - Ok(this) => this, - Err(shard) => { - let arc = Arc::new(obj); - Self::alloc(arc, shard) - } - } - } -} - -impl Interned { - fn lookup(obj: &T) -> Result> { - let storage = T::storage().get(); - let shard_idx = storage.determine_map(obj); - let shard = &storage.shards()[shard_idx]; - let shard = shard.write(); - + let (mut shard, hash) = Self::select(&obj); // 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. - - // FIXME: avoid double lookup/hashing by using raw entry API (once stable, or when - // hashbrown can be plugged into dashmap) - match shard.get_key_value(obj) { - Some((arc, _)) => Ok(Self { arc: arc.clone() }), - None => Err(shard), + match shard.raw_entry_mut().from_key_hashed_nocheck(hash as u64, &obj) { + RawEntryMut::Occupied(occ) => Self { arc: occ.key().clone() }, + RawEntryMut::Vacant(vac) => Self { + arc: vac + .insert_hashed_nocheck(hash as u64, Arc::new(obj), SharedValue::new(())) + .0 + .clone(), + }, } } - - fn alloc(arc: Arc, mut shard: Guard) -> Self { - let arc2 = arc.clone(); - - shard.insert(arc2, SharedValue::new(())); - - Self { arc } - } } impl Interned { pub fn new_str(s: &str) -> Self { - match Interned::lookup(s) { - Ok(this) => this, - Err(shard) => { - let arc = Arc::::from(s); - Self::alloc(arc, shard) - } + let (mut shard, hash) = Self::select(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 as u64, s) { + RawEntryMut::Occupied(occ) => Self { arc: occ.key().clone() }, + RawEntryMut::Vacant(vac) => Self { + arc: vac + .insert_hashed_nocheck(hash as u64, Arc::from(s), SharedValue::new(())) + .0 + .clone(), + }, } } } +impl Interned { + #[inline] + fn select(obj: &T) -> (Guard, u64) { + let storage = T::storage().get(); + let hash = { + let mut hasher = std::hash::BuildHasher::build_hasher(storage.hasher()); + obj.hash(&mut hasher); + hasher.finish() + }; + let shard_idx = storage.determine_shard(hash as usize); + let shard = &storage.shards()[shard_idx]; + (shard.write(), hash) + } +} + impl Drop for Interned { #[inline] fn drop(&mut self) { @@ -94,20 +96,17 @@ fn drop(&mut self) { impl Interned { #[cold] fn drop_slow(&mut self) { - let storage = T::storage().get(); - let shard_idx = storage.determine_map(&self.arc); - let shard = &storage.shards()[shard_idx]; - let mut shard = shard.write(); + let (mut shard, hash) = Self::select(&self.arc); - // FIXME: avoid double lookup - let (arc, _) = shard.get_key_value(&self.arc).expect("interned value removed prematurely"); - - if Arc::count(arc) != 2 { + if Arc::count(&self.arc) != 2 { // Another thread has interned another copy return; } - shard.remove(&self.arc); + match shard.raw_entry_mut().from_key_hashed_nocheck(hash, &self.arc) { + RawEntryMut::Occupied(occ) => occ.remove(), + RawEntryMut::Vacant(_) => unreachable!(), + }; // Shrink the backing storage if the shard is less than 50% occupied. if shard.len() * 2 < shard.capacity() {