diff --git a/src/lib.rs b/src/lib.rs index 51998bd..e6d9a8b 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -215,8 +215,8 @@ //! [benchmarks]: https://github.com/ibraheemdev/papaya/blob/master/BENCHMARKS.md #![deny( - missing_debug_implementations, - missing_docs, + // TODO missing_debug_implementations, + // TODO missing_docs, dead_code, unsafe_op_in_unsafe_fn )] @@ -241,3 +241,8 @@ pub use map::{ }; pub use seize::{Guard, LocalGuard, OwnedGuard}; pub use set::{HashSet, HashSetBuilder, HashSetRef}; + +/// A low-level hash table interface. +pub mod table { + pub use crate::raw::table::{HashTable, InsertResult, IntoIter, Iter, IterMut}; +} diff --git a/src/map.rs b/src/map.rs index f149a25..e489fc2 100644 --- a/src/map.rs +++ b/src/map.rs @@ -1,5 +1,5 @@ use crate::raw::utils::MapGuard; -use crate::raw::{self, InsertResult}; +use crate::raw::{map, map::InsertResult}; use crate::Equivalent; use seize::{Collector, Guard, LocalGuard, OwnedGuard}; @@ -14,7 +14,7 @@ use std::marker::PhantomData; /// [`HashMap::guard`] or using the [`HashMap::pin`] API. See the [crate-level documentation](crate#usage) /// for details. pub struct HashMap { - raw: raw::HashMap, + raw: map::HashMap, } // Safety: `HashMap` acts as a single-threaded collection on a single thread. @@ -32,7 +32,7 @@ unsafe impl Send for HashMap {} // Additionally, `HashMap` owns its `seize::Collector` and never exposes it, // so multiple threads cannot be involved in reclamation without sharing the // `HashMap` itself. If this was not true, we would require stricter bounds -// on `HashMap` operations themselves. +// on the `HashMap` operations themselves. unsafe impl Sync for HashMap {} /// A builder for a [`HashMap`]. @@ -131,7 +131,7 @@ impl HashMapBuilder { /// Construct a [`HashMap`] from the builder, using the configured options. pub fn build(self) -> HashMap { HashMap { - raw: raw::HashMap::new(self.capacity, self.hasher, self.collector, self.resize_mode), + raw: map::HashMap::new(self.capacity, self.hasher, self.collector, self.resize_mode), } } } @@ -298,7 +298,7 @@ impl HashMap { /// ``` pub fn with_capacity_and_hasher(capacity: usize, hash_builder: S) -> HashMap { HashMap { - raw: raw::HashMap::new( + raw: map::HashMap::new( capacity, hash_builder, Collector::default(), @@ -1006,6 +1006,35 @@ where } } + /// An iterator visiting all key-value pairs in arbitrary order, with mutable references + /// to the values. The iterator element type is `(&K, &mut V)`. + /// + /// # Examples + /// + /// ``` + /// use papaya::HashMap; + /// + /// let mut map = HashMap::from([ + /// ("a", 1), + /// ("b", 2), + /// ("c", 3), + /// ]); + /// + /// // Update all values. + /// for (_, val) in map.iter_mut() { + /// *val *= 2; + /// } + /// + /// for (key, val) in map.pin().iter() { + /// println!("key: {key} val: {val}"); + /// } + #[inline] + pub fn iter_mut(&mut self) -> IterMut<'_, K, V> { + IterMut { + raw: self.raw.iter_mut(), + } + } + /// An iterator visiting all keys in arbitrary order. /// The iterator element type is `&K`. /// @@ -1071,6 +1100,32 @@ where } } +impl IntoIterator for HashMap { + type Item = (K, V); + type IntoIter = IntoIter; + + /// Creates a consuming iterator, that is, one that moves each key-value + /// pair out of the map in arbitrary order. The map cannot be used after + /// calling this. + /// + /// ``` + /// use papaya::HashMap; + /// + /// let map = HashMap::from([ + /// ("a", 1), + /// ("b", 2), + /// ("c", 3), + /// ]); + /// + /// let v: Vec<(&str, i32)> = map.into_iter().collect(); + /// ``` + fn into_iter(self) -> Self::IntoIter { + IntoIter { + raw: self.raw.into_iter(), + } + } +} + /// An operation to perform on given entry in a [`HashMap`]. /// /// See [`HashMap::compute`] for details. @@ -1165,7 +1220,8 @@ where S: BuildHasher, { fn extend>(&mut self, iter: T) { - // from `hashbrown::HashMap::extend`: + // From `hashbrown::HashMap::extend`: + // // Keys may be already present or show multiple times in the iterator. // Reserve the entire hint lower bound if the map is empty. // Otherwise reserve half the hint (rounded up), so the map @@ -1572,11 +1628,11 @@ where } } -/// An iterator over a map's entries. +/// An iterator over the entries of a `HashMap`. /// /// This struct is created by the [`iter`](HashMap::iter) method on [`HashMap`]. See its documentation for details. pub struct Iter<'g, K, V, G> { - raw: raw::Iter<'g, K, V, MapGuard>, + raw: map::Iter<'g, K, V, MapGuard>, } impl<'g, K: 'g, V: 'g, G> Iterator for Iter<'g, K, V, G> @@ -1591,6 +1647,14 @@ where } } +impl Clone for Iter<'_, K, V, G> { + fn clone(&self) -> Self { + Self { + raw: self.raw.clone(), + } + } +} + impl fmt::Debug for Iter<'_, K, V, G> where K: fmt::Debug, @@ -1598,11 +1662,33 @@ where G: Guard, { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_list() - .entries(Iter { - raw: self.raw.clone(), - }) - .finish() + f.debug_list().entries(self.raw.clone()).finish() + } +} + +/// A mutable iterator over the entries of a `HashMap`. +/// +/// This struct is created by the [`iter_mut`](HashMap::iter_mut) method on [`HashMap`]. See its documentation for details. +pub struct IterMut<'map, K, V> { + raw: map::IterMut<'map, K, V>, +} + +impl<'map, K, V> Iterator for IterMut<'map, K, V> { + type Item = (&'map K, &'map mut V); + + #[inline] + fn next(&mut self) -> Option { + self.raw.next() + } +} + +impl fmt::Debug for IterMut<'_, K, V> +where + K: fmt::Debug, + V: fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_list().entries(self.raw.iter()).finish() } } @@ -1667,3 +1753,32 @@ where f.debug_tuple("Values").field(&self.iter).finish() } } + +/// An owned iterator over the entries of a `HashMap`. +/// +/// This `struct` is created by the [`into_iter`] method on [`HashMap`] +/// (provided by the [`IntoIterator`] trait). See its documentation for more. +/// +/// [`into_iter`]: IntoIterator::into_iter +pub struct IntoIter { + pub(crate) raw: map::IntoIter, +} + +impl Iterator for IntoIter { + type Item = (K, V); + + #[inline] + fn next(&mut self) -> Option { + self.raw.next() + } +} + +impl fmt::Debug for IntoIter +where + K: fmt::Debug, + V: fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_list().entries(self.raw.iter()).finish() + } +} diff --git a/src/raw/alloc.rs b/src/raw/alloc.rs index 3c5792f..eac9870 100644 --- a/src/raw/alloc.rs +++ b/src/raw/alloc.rs @@ -3,7 +3,8 @@ use std::marker::PhantomData; use std::sync::atomic::{AtomicPtr, AtomicU8, Ordering}; use std::{alloc, mem, ptr}; -use super::{probe, State}; +use super::probe; +use super::table::{meta, State}; // A hash-table laid out in a single allocation. // @@ -88,7 +89,7 @@ impl Table { // Initialize the meta table. ptr.add(mem::size_of::>()) .cast::() - .write_bytes(super::meta::EMPTY, len); + .write_bytes(meta::EMPTY, len); } Table { @@ -157,6 +158,25 @@ impl Table { } } + // Returns the entry at the given index. + // + // # Safety + // + // The index must be in-bounds for the length of the table. + #[inline] + pub unsafe fn entry_mut(&mut self, i: usize) -> *mut T { + debug_assert!(i < self.len()); + + // Safety: The caller guarantees the index is in-bounds. + let entry = unsafe { + let meta = self.raw.add(mem::size_of::>()); + let entries = meta.add(self.len()).cast::>(); + &mut *entries.add(i) + }; + + *entry.get_mut() + } + /// Returns the length of the table. #[inline] pub fn len(&self) -> usize { diff --git a/src/raw/map.rs b/src/raw/map.rs new file mode 100644 index 0000000..4b67d2b --- /dev/null +++ b/src/raw/map.rs @@ -0,0 +1,632 @@ +use std::cell::UnsafeCell; +use std::hash::{BuildHasher, Hash}; +use std::marker::PhantomData; +use std::mem::MaybeUninit; +use std::{panic, ptr}; + +use super::utils::{MapGuard, VerifiedGuard}; +use crate::map::{Compute, Operation, ResizeMode}; +use crate::raw::table::{self, Dealloc, HashTable}; +use crate::Equivalent; + +use seize::{Collector, LocalGuard, OwnedGuard}; + +/// A lock-free hash-table. +pub struct HashMap { + table: HashTable, BoxDealloc>, + + /// Hasher for keys. + pub hasher: S, +} + +struct BoxDealloc(PhantomData<(K, V)>); + +impl Dealloc> for BoxDealloc { + unsafe fn dealloc(entry: *mut Entry) { + let _: Box> = unsafe { Box::from_raw(entry) }; + } +} + +// An entry in the hash-table. +#[repr(C, align(8))] // Alignment requirement of `table::Entry`. +pub struct Entry { + /// The key for this entry. + pub key: K, + + /// The value for this entry. + pub value: V, +} + +// The result of an insert operation. +pub enum InsertResult<'g, V> { + /// Inserted the given value. + Inserted(&'g V), + + /// Replaced the given value. + Replaced(&'g V), + + /// Error returned by `try_insert`. + Error { current: &'g V, not_inserted: V }, +} + +impl HashMap { + /// Creates new hash-table with the given options. + #[inline] + pub fn new( + capacity: usize, + hasher: S, + collector: Collector, + resize: ResizeMode, + ) -> HashMap { + HashMap { + hasher, + table: HashTable::new(capacity, collector, resize), + } + } + + /// Returns a guard for this collector + pub fn guard(&self) -> MapGuard> { + self.table.guard() + } + + /// Returns an owned guard for this collector + pub fn owned_guard(&self) -> MapGuard> { + self.table.owned_guard() + } + + /// Verify a guard is valid to use with this map. + #[inline] + pub fn verify<'g, G>(&self, guard: &'g G) -> &'g MapGuard + where + G: seize::Guard, + { + self.table.verify(guard) + } + + /// Returns a reference to the collector. + #[inline] + pub fn collector(&self) -> &Collector { + &self.table.collector + } + + /// Returns the number of entries in the table. + #[inline] + pub fn len(&self) -> usize { + self.table.len() + } +} + +impl HashMap +where + K: Hash + Eq, + S: BuildHasher, +{ + /// Returns a reference to the entry corresponding to the key. + #[inline] + pub fn get<'g, Q>(&self, key: &Q, guard: &'g impl VerifiedGuard) -> Option<(&'g K, &'g V)> + where + Q: Equivalent + Hash + ?Sized, + { + let hash = self.hasher.hash_one(key); + let eq = |entry: *mut Entry| unsafe { key.equivalent(&(*entry).key) }; + + self.table + .find(hash, eq, guard) + .map(|entry| unsafe { (&(*entry).key, &(*entry).value) }) + } + + /// Inserts a key-value pair into the table. + #[inline] + pub fn insert<'g>( + &self, + key: K, + value: V, + replace: bool, + guard: &'g impl VerifiedGuard, + ) -> InsertResult<'g, V> { + let hash = self.hasher.hash_one(&key); + let new_entry = Box::into_raw(Box::new(Entry { key, value })); + let eq = |entry: *mut Entry| unsafe { (*new_entry).key.equivalent(&(*entry).key) }; + + match self + .table + .insert(hash, new_entry, eq, self.hasher(), replace, guard) + { + table::InsertResult::Inserted => InsertResult::Inserted(unsafe { &(*new_entry).value }), + + table::InsertResult::Replaced(entry) => { + InsertResult::Replaced(unsafe { &(*entry).value }) + } + table::InsertResult::Error(entry) => { + let new_entry = unsafe { Box::from_raw(new_entry) }; + + InsertResult::Error { + current: unsafe { &(*entry).value }, + not_inserted: new_entry.value, + } + } + } + } + + /// Removes a key from the map, returning the entry for the key if the key was previously in the map. + #[inline] + pub fn remove<'g, Q>(&self, key: &Q, guard: &'g impl VerifiedGuard) -> Option<(&'g K, &'g V)> + where + Q: Equivalent + Hash + ?Sized, + { + let hash = self.hasher.hash_one(key); + let eq = |entry: *mut Entry| unsafe { key.equivalent(&(*entry).key) }; + + self.table + .remove(hash, eq, self.hasher(), guard) + .map(|entry| unsafe { (&(*entry).key, &(*entry).value) }) + } + + /// Removes a key from the map, returning the entry for the key if the key was previously in the map + /// and the provided closure returns `true` + #[inline] + pub fn remove_if<'g, Q, F>( + &self, + key: &Q, + mut should_remove: F, + guard: &'g impl VerifiedGuard, + ) -> Result, (&'g K, &'g V)> + where + Q: Equivalent + Hash + ?Sized, + F: FnMut(&K, &V) -> bool, + { + let hash = self.hasher.hash_one(key); + let eq = |entry: *mut Entry| unsafe { key.equivalent(&(*entry).key) }; + let should_remove = + |entry: *mut Entry| unsafe { should_remove(&(*entry).key, &(*entry).value) }; + + self.table + .remove_if(hash, eq, should_remove, self.hasher(), guard) + .map(|option| option.map(|entry| unsafe { (&(*entry).key, &(*entry).value) })) + .map_err(|entry| unsafe { (&(*entry).key, &(*entry).value) }) + } + + /// Reserve capacity for `additional` more elements. + #[inline] + pub fn reserve(&self, additional: usize, guard: &impl VerifiedGuard) { + self.table.reserve(additional, self.hasher(), guard); + } + + /// Remove all entries from this table. + #[inline] + pub fn clear(&self, guard: &impl VerifiedGuard) { + self.table.clear(self.hasher(), guard); + } + + /// Retains only the elements specified by the predicate. + #[inline] + pub fn retain(&self, mut f: F, guard: &impl VerifiedGuard) + where + F: FnMut(&K, &V) -> bool, + { + let f = |entry: *mut Entry| unsafe { f(&(*entry).key, &(*entry).value) }; + self.table.retain(f, self.hasher(), guard); + } + + /// Returns an iterator over the keys and values of this table. + #[inline] + pub fn iter<'g, G>(&self, guard: &'g G) -> Iter<'g, K, V, G> + where + G: VerifiedGuard, + { + Iter { + inner: self.table.iter(self.hasher(), guard), + _entries: PhantomData, + } + } + + /// Returns a mutable iterator over the keys and values of this table. + #[inline] + pub fn iter_mut(&mut self) -> IterMut<'_, K, V> { + IterMut { + inner: self.table.iter_mut(), + _entries: PhantomData, + } + } + + #[inline] + fn hasher(&self) -> impl Fn(*mut Entry) -> u64 + '_ { + |entry: *mut Entry| unsafe { self.hasher.hash_one(&(*entry).key) } + } +} + +impl IntoIterator for HashMap { + type Item = (K, V); + type IntoIter = IntoIter; + + fn into_iter(self) -> Self::IntoIter { + IntoIter { + inner: self.table.into_iter(), + _entries: PhantomData, + } + } +} + +/// RMW operations. +impl HashMap +where + K: Hash + Eq, + S: BuildHasher, +{ + /// Tries to insert a key and value computed from a closure into the map, + /// and returns a reference to the value that was inserted. + #[inline] + pub fn try_insert_with<'g, F>( + &self, + key: K, + f: F, + guard: &'g impl VerifiedGuard, + ) -> Result<&'g V, &'g V> + where + F: FnOnce() -> V, + K: 'g, + { + let mut f = Some(f); + let compute = |entry| match entry { + // There is already an existing value. + Some((_, current)) => Operation::Abort(current), + + // Insert the initial value. + // + // Note that this case is guaranteed to be executed at most + // once as insert values are cached, so this can never panic. + // + // TODO: This is no longer cached. + None => Operation::Insert((f.take().unwrap())()), + }; + + match self.compute(key, compute, guard) { + // Failed to insert, return the existing value. + Compute::Aborted(current) => Err(current), + + // Successfully inserted. + Compute::Inserted(_, value) => Ok(value), + + _ => unreachable!(), + } + } + + /// Returns a reference to the value corresponding to the key, or inserts a default value + /// computed from a closure. + #[inline] + pub fn get_or_insert_with<'g, F>(&self, key: K, f: F, guard: &'g impl VerifiedGuard) -> &'g V + where + F: FnOnce() -> V, + K: 'g, + { + match self.try_insert_with(key, f, guard) { + Ok(value) => value, + Err(value) => value, + } + } + + /// Updates an existing entry atomically, returning the value that was inserted. + #[inline] + pub fn update<'g, F>( + &self, + key: K, + mut update: F, + guard: &'g impl VerifiedGuard, + ) -> Option<&'g V> + where + F: FnMut(&V) -> V, + K: 'g, + { + let compute = |entry| match entry { + // There is nothing to update. + None => Operation::Abort(()), + // Perform the update. + Some((_, value)) => Operation::Insert(update(value)), + }; + + match self.compute(key, compute, guard) { + // Return the updated value. + Compute::Updated { + new: (_, value), .. + } => Some(value), + + // There was nothing to update. + Compute::Aborted(_) => None, + + _ => unreachable!(), + } + } + + /// Updates an existing entry or inserts a default value computed from a closure. + #[inline] + pub fn update_or_insert_with<'g, U, F>( + &self, + key: K, + update: U, + f: F, + guard: &'g impl VerifiedGuard, + ) -> &'g V + where + F: FnOnce() -> V, + U: Fn(&V) -> V, + K: 'g, + { + let mut f = Some(f); + + let compute = |entry| match entry { + // Perform the update. + Some((_, value)) => Operation::Insert::<_, ()>(update(value)), + + // Insert the initial value. + // + // Note that this case is guaranteed to be executed at most + // once as insert values are cached, so this can never panic. + // + // TODO: This is no longer cached. + None => Operation::Insert((f.take().unwrap())()), + }; + + match self.compute(key, compute, guard) { + // Return the updated value. + Compute::Updated { + new: (_, value), .. + } => value, + + // Return the value we inserted. + Compute::Inserted(_, value) => value, + + _ => unreachable!(), + } + } + + /// Update an entry with a CAS function. + /// + /// Note that `compute` closure is guaranteed to be called for a `None` input only once, allowing the + /// insertion of values that cannot be cloned or reconstructed. + #[inline] + pub fn compute<'g, F, T>( + &self, + key: K, + mut compute: F, + guard: &'g impl VerifiedGuard, + ) -> Compute<'g, K, V, T> + where + F: FnMut(Option<(&'g K, &'g V)>) -> Operation, + { + let hash = self.hasher.hash_one(&key); + + // Lazy initialize the entry allocation. + let lazy_entry = UnsafeCell::new(LazyEntry::Uninit(key)); + + let eq = |entry: *mut Entry| unsafe { + (*lazy_entry.get()).key().equivalent(&(*entry).key) + }; + + let compute = |entry: Option<*mut Entry>| unsafe { + let entry = entry.map(|entry| (&(*entry).key, &(*entry).value)); + + match compute(entry) { + Operation::Insert(value) => { + let entry = (*lazy_entry.get()).init(); + + (*entry).value.write(value); + + table::Operation::Insert(entry.cast()) + } + Operation::Remove => table::Operation::Remove, + Operation::Abort(value) => table::Operation::Abort(value), + } + }; + + let result = match self.table.compute(hash, eq, compute, self.hasher(), guard) { + table::Compute::Inserted(entry) => unsafe { + Compute::Inserted(&(*entry).key, &(*entry).value) + }, + table::Compute::Updated { old, new } => unsafe { + Compute::Updated { + old: (&(*old).key, &(*old).value), + new: (&(*new).key, &(*new).value), + } + }, + table::Compute::Removed(entry) => unsafe { + Compute::Removed(&(*entry).key, &(*entry).value) + }, + table::Compute::Aborted(value) => Compute::Aborted(value), + }; + + // Deallocate the entry if it was not inserted. + if matches!(result, Compute::Removed(..) | Compute::Aborted(_)) { + if let LazyEntry::Init(entry) = unsafe { &*lazy_entry.get() } { + // Safety: The entry was allocated but not inserted into the map. + let _ = unsafe { Box::from_raw(*entry) }; + } + } + + result + } +} + +/// A lazy initialized `Entry` allocation. +enum LazyEntry { + /// An uninitialized entry, containing just the owned key. + Uninit(K), + + /// An allocated entry. + Init(*mut Entry>), +} + +impl LazyEntry { + /// Returns a reference to the entry's key. + #[inline] + fn key(&self) -> &K { + match self { + LazyEntry::Uninit(key) => key, + LazyEntry::Init(entry) => unsafe { &(**entry).key }, + } + } + + /// Initializes the entry if it has not already been initialized, returning the pointer + /// to the entry allocation. + #[inline] + fn init(&mut self) -> *mut Entry> { + match self { + LazyEntry::Init(entry) => *entry, + LazyEntry::Uninit(key) => { + // Safety: we read the current key with `ptr::read` and overwrite the + // state with `ptr::write`. We also make sure to abort if the allocator + // panics, ensuring the current value is not dropped twice. + unsafe { + let key = ptr::read(key); + let entry = panic::catch_unwind(panic::AssertUnwindSafe(|| { + Box::into_raw(Box::new(Entry { + value: MaybeUninit::uninit(), + key, + })) + })) + .unwrap_or_else(|_| std::process::abort()); + ptr::write(self, LazyEntry::Init(entry)); + entry + } + } + } + } +} + +// An iterator over the keys and values of this table. +pub struct Iter<'g, K, V, G> { + inner: table::Iter<'g, Entry, G>, + _entries: PhantomData<(&'g K, &'g V)>, +} + +impl<'g, K: 'g, V: 'g, G> Iterator for Iter<'g, K, V, G> +where + G: VerifiedGuard, +{ + type Item = (&'g K, &'g V); + + #[inline] + fn next(&mut self) -> Option { + self.inner + .next() + .map(|entry| unsafe { (&(*entry).key, &(*entry).value) }) + } +} + +// Safety: An iterator holds a shared reference to the `HashMap` +// and `Guard`, and outputs shared references to keys and values. +// Thus everything must be `Sync` for the iterator to be `Send` +// or `Sync`. +// +// It is not possible to obtain an owned key, value, or guard +// from an iterator, so `Send` is not a required bound. +unsafe impl Send for Iter<'_, K, V, G> +where + K: Sync, + V: Sync, + G: Sync, +{ +} + +unsafe impl Sync for Iter<'_, K, V, G> +where + K: Sync, + V: Sync, + G: Sync, +{ +} + +impl Clone for Iter<'_, K, V, G> { + #[inline] + fn clone(&self) -> Self { + Iter { + inner: self.inner.clone(), + _entries: PhantomData, + } + } +} + +// A mutable iterator over the keys and values of this table. +pub struct IterMut<'map, K, V> { + inner: table::IterMut>, + // Ensure invariance with respect to `V`. + _entries: PhantomData<(&'map K, &'map mut V)>, +} + +impl<'map, K, V> Iterator for IterMut<'map, K, V> { + type Item = (&'map K, &'map mut V); + + #[inline] + fn next(&mut self) -> Option { + self.inner + .next() + .map(|entry| unsafe { (&(*entry).key, &mut (*entry).value) }) + } +} + +impl<'map, K, V> IterMut<'map, K, V> { + // Returns an immutable iterator over the remaining entries. + pub(crate) fn iter(&self) -> impl Iterator + '_ { + self.inner + .iter() + .map(|entry| unsafe { (&(*entry).key, &(*entry).value) }) + } +} + +// Safety: A mutable iterator does not perform any concurrent access, +// so the normal `Send` and `Sync` rules apply. +unsafe impl Send for IterMut<'_, K, V> +where + K: Send, + V: Send, +{ +} + +unsafe impl Sync for IterMut<'_, K, V> +where + K: Sync, + V: Sync, +{ +} + +// An owned iterator over the keys and values of this table. +pub struct IntoIter { + inner: table::IntoIter, BoxDealloc>, + _entries: PhantomData<(K, V)>, +} + +impl Iterator for IntoIter { + type Item = (K, V); + + #[inline] + fn next(&mut self) -> Option { + self.inner.next().map(|entry| { + let entry = unsafe { ptr::read(entry) }; + (entry.key, entry.value) + }) + } +} + +impl IntoIter { + // Returns an immutable iterator over the remaining entries. + pub(crate) fn iter(&self) -> impl Iterator + '_ { + self.inner + .iter() + .map(|entry| unsafe { (&(*entry).key, &(*entry).value) }) + } +} + +// Safety: An owned iterator does not perform any concurrent access, +// so the normal `Send` and `Sync` rules apply. +unsafe impl Send for IntoIter +where + K: Send, + V: Send, +{ +} + +unsafe impl Sync for IntoIter +where + K: Sync, + V: Sync, +{ +} diff --git a/src/raw/mod.rs b/src/raw/mod.rs index 9aa8f75..e7302c9 100644 --- a/src/raw/mod.rs +++ b/src/raw/mod.rs @@ -1,2840 +1,6 @@ mod alloc; mod probe; +pub(crate) mod map; +pub(crate) mod table; pub(crate) mod utils; - -use std::hash::{BuildHasher, Hash}; -use std::mem::MaybeUninit; -use std::sync::atomic::{AtomicPtr, AtomicU8, AtomicUsize, Ordering}; -use std::sync::Mutex; -use std::{hint, panic, ptr}; - -use self::alloc::{RawTable, Table}; -use self::probe::Probe; -use self::utils::{untagged, AtomicPtrFetchOps, Counter, Parker, StrictProvenance, Tagged}; -use crate::map::{Compute, Operation, ResizeMode}; -use crate::Equivalent; - -use seize::{Collector, LocalGuard, OwnedGuard}; -use utils::{MapGuard, Stack, VerifiedGuard}; - -/// A lock-free hash-table. -pub struct HashMap { - /// A pointer to the root table. - table: AtomicPtr>>, - - /// Collector for memory reclamation. - collector: Collector, - - /// The resize mode, either blocking or incremental. - resize: ResizeMode, - - /// An atomic counter of the number of keys in the table. - count: Counter, - - /// The initial capacity provided to `HashMap::new`. - /// - /// The table is guaranteed to never shrink below this capacity. - initial_capacity: usize, - - /// Hasher for keys. - pub hasher: S, -} - -/// Resize state for the hash-table. -pub struct State { - /// The next table used for resizing. - pub next: AtomicPtr>, - - /// A lock acquired to allocate the next table. - pub allocating: Mutex<()>, - - /// The number of entries that have been copied to the next table. - pub copied: AtomicUsize, - - /// The number of entries that have been claimed by copiers, - /// but not necessarily copied. - pub claim: AtomicUsize, - - /// The status of the resize. - pub status: AtomicU8, - - /// A thread parker for blocking on copy operations. - pub parker: Parker, - - /// Entries whose retirement has been deferred by later tables. - pub deferred: Stack<*mut T>, -} - -impl Default for State { - fn default() -> State { - State { - next: AtomicPtr::new(ptr::null_mut()), - allocating: Mutex::new(()), - copied: AtomicUsize::new(0), - claim: AtomicUsize::new(0), - status: AtomicU8::new(State::PENDING), - parker: Parker::default(), - deferred: Stack::new(), - } - } -} - -impl State<()> { - /// A resize is in-progress. - pub const PENDING: u8 = 0; - - /// The resize has been aborted, continue to the next table. - pub const ABORTED: u8 = 1; - - /// The resize was complete and the table was promoted. - pub const PROMOTED: u8 = 2; -} - -// The result of an insert operation. -pub enum InsertResult<'g, V> { - /// Inserted the given value. - Inserted(&'g V), - - /// Replaced the given value. - Replaced(&'g V), - - /// Error returned by `try_insert`. - Error { current: &'g V, not_inserted: V }, -} - -// The raw result of an insert operation. -pub enum RawInsertResult<'g, K, V> { - /// Inserted the given value. - Inserted(&'g V), - - /// Replaced the given value. - Replaced(&'g V), - - /// Error returned by `try_insert`. - Error { - current: &'g V, - not_inserted: *mut Entry, - }, -} - -// An entry in the hash-table. -#[repr(C, align(8))] // Reserve the lower 3 bits for pointer tagging. -pub struct Entry { - /// The key for this entry. - pub key: K, - - /// The value for this entry. - pub value: V, -} - -impl Entry<(), ()> { - /// The entry is being copied to the new table, no updates are allowed on the old table. - /// - /// This bit is put down to initiate a copy, forcing all writers to complete the resize - /// before making progress. - const COPYING: usize = 0b001; - - /// The entry has been copied to the new table. - /// - /// This bit is put down after a copy completes. Both readers and writers must go to - /// the new table to see the new state of the entry. - /// - /// In blocking mode this is unused. - const COPIED: usize = 0b010; - - /// The entry was copied from a previous table. - /// - /// This bit indicates that an entry may still be accessible from previous tables - /// because the resize is still in progress, and so it is unsafe to reclaim. - /// - /// In blocking mode this is unused. - const BORROWED: usize = 0b100; -} - -impl utils::Unpack for Entry { - /// Mask for an entry pointer, ignoring any tag bits. - const MASK: usize = !(Entry::COPYING | Entry::COPIED | Entry::BORROWED); -} - -impl Entry { - /// A sentinel pointer for a deleted entry. - /// - /// Null pointers are never copied to the new table, so this state is safe to use. - /// Note that tombstone entries may still be marked as `COPYING`, so this state - /// cannot be used for direct equality. - const TOMBSTONE: *mut Entry = Entry::COPIED as _; -} - -/// The status of an entry. -enum EntryStatus { - /// The entry is a tombstone or null (potentially a null copy). - Null, - - /// The entry is being copied. - Copied(Tagged>), - - /// A valid entry. - Value(Tagged>), -} - -impl From>> for EntryStatus { - /// Returns the status for this entry. - #[inline] - fn from(entry: Tagged>) -> Self { - if entry.ptr.is_null() { - EntryStatus::Null - } else if entry.tag() & Entry::COPYING != 0 { - EntryStatus::Copied(entry) - } else { - EntryStatus::Value(entry) - } - } -} - -/// The state of an entry we attempted to update. -enum UpdateStatus { - /// Successfully replaced the given key and value. - Replaced(Tagged>), - - /// A new entry was written before we could update. - Found(EntryStatus), -} - -/// The state of an entry we attempted to insert into. -enum InsertStatus { - /// Successfully inserted the value. - Inserted, - - /// A new entry was written before we could update. - Found(EntryStatus), -} - -impl HashMap { - /// Creates new hash-table with the given options. - #[inline] - pub fn new( - capacity: usize, - hasher: S, - collector: Collector, - resize: ResizeMode, - ) -> HashMap { - // The table is lazily allocated. - if capacity == 0 { - return HashMap { - collector, - resize, - hasher, - initial_capacity: 1, - table: AtomicPtr::new(ptr::null_mut()), - count: Counter::default(), - }; - } - - // Initialize the table and mark it as the root. - let mut table = Table::alloc(probe::entries_for(capacity)); - *table.state_mut().status.get_mut() = State::PROMOTED; - - HashMap { - hasher, - resize, - collector, - initial_capacity: capacity, - table: AtomicPtr::new(table.raw), - count: Counter::default(), - } - } - - /// Returns a guard for this collector - pub fn guard(&self) -> MapGuard> { - // Safety: Created the guard from our collector. - unsafe { MapGuard::new(self.collector().enter()) } - } - - /// Returns an owned guard for this collector - pub fn owned_guard(&self) -> MapGuard> { - // Safety: Created the guard from our collector. - unsafe { MapGuard::new(self.collector().enter_owned()) } - } - - /// Verify a guard is valid to use with this map. - #[inline] - pub fn verify<'g, G>(&self, guard: &'g G) -> &'g MapGuard - where - G: seize::Guard, - { - assert_eq!( - *guard.collector(), - self.collector, - "Attempted to access map with incorrect guard" - ); - - // Safety: Verified the guard above. - unsafe { MapGuard::from_ref(guard) } - } - - /// Returns a reference to the root hash-table. - #[inline] - fn root(&self, guard: &impl VerifiedGuard) -> Table> { - // Load the root table. - let raw = guard.protect(&self.table, Ordering::Acquire); - - // Safety: The root table is either null or a valid table allocation. - unsafe { Table::from_raw(raw) } - } - - /// Returns a reference to the collector. - #[inline] - pub fn collector(&self) -> &Collector { - &self.collector - } - - /// Returns the number of entries in the table. - #[inline] - pub fn len(&self) -> usize { - self.count.sum() - } - - /// Returns true if incremental resizing is enabled. - #[inline] - fn is_incremental(&self) -> bool { - matches!(self.resize, ResizeMode::Incremental(_)) - } -} - -impl HashMap -where - K: Hash + Eq, - S: BuildHasher, -{ - /// Returns a reference to the entry corresponding to the key. - #[inline] - pub fn get<'g, Q>(&self, key: &Q, guard: &'g impl VerifiedGuard) -> Option<(&'g K, &'g V)> - where - Q: Equivalent + Hash + ?Sized, - { - // Load the root table. - let mut table = self.root(guard); - - // The table has not been initialized yet. - if table.raw.is_null() { - return None; - } - - let (h1, h2) = self.hash(key); - - loop { - // Initialize the probe state. - let mut probe = Probe::start(h1, table.mask); - - // Probe until we reach the limit. - 'probe: while probe.len <= table.limit { - // Load the entry metadata first for cheap searches. - // - // Safety: `probe.i` is always in-bounds for the table length. - let meta = unsafe { table.meta(probe.i) }.load(Ordering::Acquire); - - if meta == h2 { - // Load the full entry. - // - // Safety: `probe.i` is always in-bounds for the table length. - let entry = guard - .protect(unsafe { table.entry(probe.i) }, Ordering::Acquire) - .unpack(); - - // The entry was deleted, keep probing. - if entry.ptr.is_null() { - probe.next(table.mask); - continue 'probe; - } - - // Safety: We performed a protected load of the pointer using a verified guard with - // `Acquire` and ensured that it is non-null, meaning it is valid for reads as long - // as we hold the guard. - let entry_ref = unsafe { &(*entry.ptr) }; - - // Check for a full match. - if key.equivalent(&entry_ref.key) { - // The entry was copied to the new table. - // - // In blocking resize mode we do not need to perform self check as all writes block - // until any resizes are complete, making the root table the source of truth for readers. - if entry.tag() & Entry::COPIED != 0 { - break 'probe; - } - - // Found the correct entry, return the key and value. - return Some((&entry_ref.key, &entry_ref.value)); - } - } - - // The key is not in the table. - // - // It also cannot be in the next table because we have not went over the probe limit. - if meta == meta::EMPTY { - return None; - } - - probe.next(table.mask); - } - - // In incremental resize mode, we have to check the next table if we found - // a copied entry or went over the probe limit. - if self.is_incremental() { - if let Some(next) = table.next_table() { - table = next; - continue; - } - } - - // Otherwise, the key is not in the table. - return None; - } - } - - /// Inserts a key-value pair into the table. - #[inline] - pub fn insert<'g>( - &self, - key: K, - value: V, - replace: bool, - guard: &'g impl VerifiedGuard, - ) -> InsertResult<'g, V> { - // Perform the insert. - let raw_result = self.insert_inner(key, value, replace, guard); - - let result = match raw_result { - // Updated an entry. - RawInsertResult::Replaced(value) => InsertResult::Replaced(value), - - // Inserted a new entry. - RawInsertResult::Inserted(value) => { - // Increment the table length. - self.count.get(guard).fetch_add(1, Ordering::Relaxed); - - InsertResult::Inserted(value) - } - - // Failed to insert the entry. - RawInsertResult::Error { - current, - not_inserted, - } => { - // Safety: We allocated this box above and it was not inserted into the table. - let not_inserted = unsafe { Box::from_raw(not_inserted) }; - - InsertResult::Error { - current, - not_inserted: not_inserted.value, - } - } - }; - - result - } - - /// Inserts an entry into the map. - #[inline] - fn insert_inner<'g>( - &self, - key: K, - value: V, - should_replace: bool, - guard: &'g impl VerifiedGuard, - ) -> RawInsertResult<'g, K, V> { - // Allocate the entry to be inserted. - let new_entry = untagged(Box::into_raw(Box::new(Entry { key, value }))); - - // Safety: We just allocated the entry above. - let new_ref = unsafe { &(*new_entry.ptr) }; - - // Load the root table. - let mut table = self.root(guard); - - // Allocate the table if it has not been initialized yet. - if table.raw.is_null() { - table = self.init(None); - } - - let (h1, h2) = self.hash(&new_ref.key); - - let mut help_copy = true; - loop { - // Initialize the probe state. - let mut probe = Probe::start(h1, table.mask); - - // Probe until we reach the limit. - let copying = 'probe: loop { - if probe.len > table.limit { - break None; - } - - // Load the entry metadata first for cheap searches. - // - // Safety: `probe.i` is always in-bounds for the table length. - let meta = unsafe { table.meta(probe.i) }.load(Ordering::Acquire); - - // The entry is empty, try to insert. - let entry = if meta == meta::EMPTY { - // Perform the insertion. - // - // Safety: `probe.i` is always in-bounds for the table length. Additionally, - // `new_entry` was allocated above and never shared. - match unsafe { self.insert_at(probe.i, h2, new_entry.raw, table, guard) } { - // Successfully inserted. - InsertStatus::Inserted => return RawInsertResult::Inserted(&new_ref.value), - - // Lost to a concurrent insert. - // - // If the key matches, we might be able to update the value. - InsertStatus::Found(EntryStatus::Value(found)) - | InsertStatus::Found(EntryStatus::Copied(found)) => found, - - // Otherwise, continue probing. - InsertStatus::Found(EntryStatus::Null) => { - probe.next(table.mask); - continue 'probe; - } - } - } - // Found a potential match. - else if meta == h2 { - // Load the full entry. - // - // Safety: `probe.i` is always in-bounds for the table length. - let entry = guard - .protect(unsafe { table.entry(probe.i) }, Ordering::Acquire) - .unpack(); - - // The entry was deleted, keep probing. - if entry.ptr.is_null() { - probe.next(table.mask); - continue 'probe; - } - - // If the key matches, we might be able to update the value. - entry - } - // Otherwise, continue probing. - else { - probe.next(table.mask); - continue 'probe; - }; - - // Safety: We performed a protected load of the pointer using a verified guard with - // `Acquire` and ensured that it is non-null, meaning it is valid for reads as long - // as we hold the guard. - let entry_ref = unsafe { &(*entry.ptr) }; - - // Check for a full match. - if entry_ref.key != new_ref.key { - probe.next(table.mask); - continue 'probe; - } - - // The entry is being copied to the new table. - if entry.tag() & Entry::COPYING != 0 { - break 'probe Some(probe.i); - } - - // Return an error for calls to `try_insert`. - if !should_replace { - return RawInsertResult::Error { - current: &entry_ref.value, - not_inserted: new_entry.ptr, - }; - } - - // Try to update the value. - // - // Safety: - // - `probe.i` is always in-bounds for the table length - // - `entry` is a valid non-null entry that was inserted into the map. - match unsafe { self.insert_slow(probe.i, entry, new_entry.raw, table, guard) } { - // Successfully performed the update. - UpdateStatus::Replaced(entry) => { - // Safety: `entry` is a valid non-null entry that we found in the map - // before replacing it. - let value = unsafe { &(*entry.ptr).value }; - return RawInsertResult::Replaced(value); - } - - // The entry is being copied. - UpdateStatus::Found(EntryStatus::Copied(_)) => break 'probe Some(probe.i), - - // The entry was deleted before we could update it, continue probing. - UpdateStatus::Found(EntryStatus::Null) => { - probe.next(table.mask); - continue 'probe; - } - - UpdateStatus::Found(EntryStatus::Value(_)) => {} - } - }; - - // Prepare to retry in the next table. - table = self.prepare_retry_insert(copying, &mut help_copy, table, guard); - } - } - - /// The slow-path for `insert`, updating the value. - /// - /// The returned pointer is guaranteed to be non-null and valid for reads. - /// - /// # Safety - /// - /// The safety requirements of `HashMap::update_at` apply. - #[cold] - #[inline(never)] - unsafe fn insert_slow( - &self, - i: usize, - mut entry: Tagged>, - new_entry: *mut Entry, - table: Table>, - guard: &impl VerifiedGuard, - ) -> UpdateStatus { - loop { - // Try to update the value. - // - // Safety: Guaranteed by caller. - match unsafe { self.update_at(i, entry, new_entry, table, guard) } { - // Someone else beat us to the update, retry. - // - // Note that the pointer we find here is a non-null entry that was inserted - // into the map. - UpdateStatus::Found(EntryStatus::Value(found)) => entry = found, - - status => return status, - } - } - } - - /// Prepare to retry an insert operation in the next table. - #[cold] - #[inline(never)] - fn prepare_retry_insert( - &self, - copying: Option, - help_copy: &mut bool, - table: Table>, - guard: &impl VerifiedGuard, - ) -> Table> { - // If went over the probe limit or found a copied entry, trigger a resize. - let mut next_table = self.get_or_alloc_next(None, table); - - let next_table = match self.resize { - // In blocking mode we must complete the resize before proceeding. - ResizeMode::Blocking => self.help_copy(true, &table, guard), - - // In incremental mode we can perform more granular blocking. - ResizeMode::Incremental(_) => { - // Help out with the copy. - if *help_copy { - next_table = self.help_copy(false, &table, guard); - } - - // The entry we want to update is being copied. - if let Some(i) = copying { - // Wait for the entry to be copied. - // - // We could race with the copy to insert into the table. However, - // this entire code path is very rare and likely to complete quickly, - // so blocking allows us to make copies faster. - self.wait_copied(i, &table); - } - - next_table - } - }; - - // Limit incremental copying to once per operation, for more consistent latency. - *help_copy = false; - - // Continue in the new table. - next_table - } - - /// Removes a key from the map, returning the entry for the key if the key was previously in the map. - #[inline] - pub fn remove<'g, Q>(&self, key: &Q, guard: &'g impl VerifiedGuard) -> Option<(&'g K, &'g V)> - where - Q: Equivalent + Hash + ?Sized, - { - #[inline(always)] - fn should_remove(_key: &K, _value: &V) -> bool { - true - } - - // Safety: `should_remove` unconditionally returns `true`. - unsafe { self.remove_if(key, should_remove, guard).unwrap_unchecked() } - } - - /// Removes a key from the map, returning the entry for the key if the key was previously in the map - /// and the provided closure returns `true` - #[inline] - pub fn remove_if<'g, Q, F>( - &self, - key: &Q, - mut should_remove: F, - guard: &'g impl VerifiedGuard, - ) -> Result, (&'g K, &'g V)> - where - Q: Equivalent + Hash + ?Sized, - F: FnMut(&K, &V) -> bool, - { - // Load the root table. - let mut table = self.root(guard); - - // The table has not been initialized yet. - if table.raw.is_null() { - return Ok(None); - } - - let (h1, h2) = self.hash(key); - - let mut help_copy = true; - loop { - // Initialize the probe state. - let mut probe = Probe::start(h1, table.mask); - - // Probe until we reach the limit. - let copying = 'probe: loop { - if probe.len > table.limit { - break None; - } - - // Load the entry metadata first for cheap searches. - // - // Safety: `probe.i` is always in-bounds for the table length. - let meta = unsafe { table.meta(probe.i).load(Ordering::Acquire) }; - - // The key is not in the table. - // It also cannot be in the next table because we have not went over the probe limit. - if meta == meta::EMPTY { - return Ok(None); - } - - // Check for a potential match. - if meta != h2 { - probe.next(table.mask); - continue 'probe; - } - - // Load the full entry. - // - // Safety: `probe.i` is always in-bounds for the table length. - let mut entry = guard - .protect(unsafe { table.entry(probe.i) }, Ordering::Acquire) - .unpack(); - - // The entry was deleted, keep probing. - if entry.ptr.is_null() { - probe.next(table.mask); - continue 'probe; - } - - // Check for a full match. - // - // Safety: We performed a protected load of the pointer using a verified guard with - // `Acquire` and ensured that it is non-null, meaning it is valid for reads as long - // as we hold the guard. - if !key.equivalent(unsafe { &(*entry.ptr).key }) { - probe.next(table.mask); - continue 'probe; - } - - // The entry is being copied to the new table, we have to complete the copy before - // we can remove it. - if entry.tag() & Entry::COPYING != 0 { - break 'probe Some(probe.i); - } - - loop { - // Safety: `entry` is a valid, non-null, protected entry that we found in the map. - let entry_ref = unsafe { &(*entry.ptr) }; - - // Ensure that the entry should be removed. - if !should_remove(&entry_ref.key, &entry_ref.value) { - return Err((&entry_ref.key, &entry_ref.value)); - } - - // Safety: - // - `probe.i` is always in-bounds for the table length - // - `entry` is a valid non-null entry that we found in the map. - let status = - unsafe { self.update_at(probe.i, entry, Entry::TOMBSTONE, table, guard) }; - - match status { - // Successfully removed the entry. - UpdateStatus::Replaced(_entry) => { - // Mark the entry as a tombstone. - // - // Note that this might end up being overwritten by the metadata hash - // if the initial insertion is lagging behind, but we avoid the RMW - // and sacrifice reads in the extremely rare case. - // - // Safety: `probe.i` is always in-bounds for the table length. - unsafe { - table - .meta(probe.i) - .store(meta::TOMBSTONE, Ordering::Release) - }; - - // Decrement the table length. - self.count.get(guard).fetch_sub(1, Ordering::Relaxed); - - // Note that `entry_ref` here is the entry that we just replaced. - return Ok(Some((&entry_ref.key, &entry_ref.value))); - } - - // The entry is being copied to the new table, we have to complete the copy - // before we can remove. - UpdateStatus::Found(EntryStatus::Copied(_)) => break 'probe Some(probe.i), - - // The entry was deleted. - // - // We know that at some point during our execution the key was not in the map. - UpdateStatus::Found(EntryStatus::Null) => return Ok(None), - - // Lost to a concurrent update, retry. - UpdateStatus::Found(EntryStatus::Value(found)) => entry = found, - } - } - }; - - // Prepare to retry in the next table. - table = match self.prepare_retry(copying, &mut help_copy, table, guard) { - Some(table) => table, - - // The search was exhausted. - None => return Ok(None), - } - } - } - - /// Prepare to retry an operation on an existing key in the next table. - /// - /// Returns `None` if the recursive search has been exhausted. - #[cold] - fn prepare_retry( - &self, - copying: Option, - help_copy: &mut bool, - table: Table>, - guard: &impl VerifiedGuard, - ) -> Option>> { - let next_table = match self.resize { - ResizeMode::Blocking => match copying { - // The entry we want to perform the operation on is being copied. - // - // In blocking mode we must complete the resize before proceeding. - Some(_) => self.help_copy(true, &table, guard), - - // If we went over the probe limit, the key is not in the map. - None => return None, - }, - - ResizeMode::Incremental(_) => { - // In incremental resize mode, we always have to check the next table. - let next_table = table.next_table()?; - - // Help out with the copy. - if *help_copy { - self.help_copy(false, &table, guard); - } - - if let Some(i) = copying { - // Wait for the entry to be copied. - // - // We could race with the copy to insert into the table. However, - // this entire code path is very rare and likely to complete quickly, - // so blocking allows us to make copies faster. - self.wait_copied(i, &table); - } - - next_table - } - }; - - // Limit incremental copying to once per operation, for more consistent latency. - *help_copy = false; - - // Continue in the new table. - Some(next_table) - } - - /// Attempts to insert an entry at the given index. - /// - /// In the case of an error, the returned pointer is guaranteed to be - /// protected and valid for reads as long as the guard is held. - /// - /// # Safety - /// - /// The index must be in-bounds for the table. Additionally, `new_entry` must be a - /// valid owned pointer to insert into the map. - #[inline] - unsafe fn insert_at( - &self, - i: usize, - meta: u8, - new_entry: *mut Entry, - table: Table>, - guard: &impl VerifiedGuard, - ) -> InsertStatus { - // Safety: The caller guarantees that `i` is in-bounds. - let entry = unsafe { table.entry(i) }; - let meta_entry = unsafe { table.meta(i) }; - - // Try to claim the empty entry. - let found = match guard.compare_exchange( - entry, - ptr::null_mut(), - new_entry, - Ordering::Release, - Ordering::Acquire, - ) { - // Successfully claimed the entry. - Ok(_) => { - // Update the metadata table. - meta_entry.store(meta, Ordering::Release); - - // Return the value we inserted. - return InsertStatus::Inserted; - } - - // Lost to a concurrent update. - Err(found) => found.unpack(), - }; - - let (meta, status) = match EntryStatus::from(found) { - EntryStatus::Value(_) | EntryStatus::Copied(_) => { - // Safety: We performed a protected load of the pointer using a verified guard - // with `Acquire` and ensured that it is non-null, meaning it is valid for reads - // as long as we hold the guard. - let key = unsafe { &(*found.ptr).key }; - - // An entry was inserted, we have to hash it to get the metadata. - // - // The logic is the same for copied entries here as we have to - // check if the key matches and continue the update in the new table. - let hash = self.hasher.hash_one(key); - (meta::h2(hash), EntryStatus::Value(found)) - } - - // The entry was deleted or null copied. - EntryStatus::Null => (meta::TOMBSTONE, EntryStatus::Null), - }; - - // Ensure the meta table is updated to keep the probe chain alive for readers. - if meta_entry.load(Ordering::Relaxed) == meta::EMPTY { - meta_entry.store(meta, Ordering::Release); - } - - InsertStatus::Found(status) - } - - /// Attempts to replace the value of an existing entry at the given index. - /// - /// In the case of an error, the returned pointer is guaranteed to be - /// protected and valid for reads. - /// - /// # Safety - /// - /// - The index must be in-bounds for the table. - /// - `current` must be a valid non-null entry that was inserted into the map. - /// - `new_entry` must be a valid sentinel or owned pointer to insert into the map. - #[inline] - unsafe fn update_at( - &self, - i: usize, - current: Tagged>, - new_entry: *mut Entry, - table: Table>, - guard: &impl VerifiedGuard, - ) -> UpdateStatus { - // Safety: The caller guarantees that `i` is in-bounds. - let entry = unsafe { table.entry(i) }; - - // Try to perform the update. - let found = match guard.compare_exchange_weak( - entry, - current.raw, - new_entry, - Ordering::Release, - Ordering::Acquire, - ) { - // Successfully updated. - Ok(_) => unsafe { - // Safety: The caller guarantees that `current` is a valid non-null entry that was - // inserted into the map. Additionally, it is now unreachable from this table due - // to the CAS above. - self.defer_retire(current, &table, guard); - - return UpdateStatus::Replaced(current); - }, - - // Lost to a concurrent update. - Err(found) => found.unpack(), - }; - - UpdateStatus::Found(EntryStatus::from(found)) - } - - /// Reserve capacity for `additional` more elements. - #[inline] - pub fn reserve(&self, additional: usize, guard: &impl VerifiedGuard) { - let mut table = self.root(guard); - - // The table has not yet been allocated, initialize it. - if table.raw.is_null() { - table = self.init(Some(probe::entries_for(additional))); - } - - loop { - let capacity = probe::entries_for(self.count.sum().checked_add(additional).unwrap()); - - // We have enough capacity. - if table.len() >= capacity { - return; - } - - // Race to allocate the new table. - self.get_or_alloc_next(Some(capacity), table); - - // Force the copy to complete. - // - // Note that this is not strictly necessary for a `reserve` operation. - table = self.help_copy(true, &table, guard); - } - } - - /// Remove all entries from this table. - #[inline] - pub fn clear(&self, guard: &impl VerifiedGuard) { - // Load the root table. - let mut table = self.root(guard); - - // The table has not been initialized yet. - if table.raw.is_null() { - return; - } - - loop { - // Get a clean copy of the table to delete from. - table = self.linearize(table, guard); - - // Note that this method is not implemented in terms of `retain(|_, _| true)` to avoid - // loading entry metadata, as there is no need to provide consistency with `get`. - let mut copying = false; - - 'probe: for i in 0..table.len() { - // Load the entry to delete. - // - // Safety: `i` is in bounds for the table length. - let mut entry = guard - .protect(unsafe { table.entry(i) }, Ordering::Acquire) - .unpack(); - - loop { - // The entry is empty or already deleted. - if entry.ptr.is_null() { - continue 'probe; - } - - // Found a non-empty entry being copied. - if entry.tag() & Entry::COPYING != 0 { - // Clear every entry in this table that we can, then deal with the copy. - copying = true; - continue 'probe; - } - - // Try to delete the entry. - // - // Safety: `i` is in bounds for the table length. - let result = unsafe { - table.entry(i).compare_exchange( - entry.raw, - Entry::TOMBSTONE, - Ordering::Release, - Ordering::Acquire, - ) - }; - - match result { - // Successfully deleted the entry. - Ok(_) => { - // Update the metadata table. - // - // Safety: `i` is in bounds for the table length. - unsafe { table.meta(i).store(meta::TOMBSTONE, Ordering::Release) }; - - // Decrement the table length. - self.count.get(guard).fetch_sub(1, Ordering::Relaxed); - - // Safety: The caller guarantees that `current` is a valid non-null entry that was - // inserted into the map. Additionally, it is now unreachable from this table due - // to the CAS above. - unsafe { self.defer_retire(entry, &table, guard) }; - continue 'probe; - } - - // Lost to a concurrent update, retry. - Err(found) => entry = found.unpack(), - } - } - } - - // We cleared every entry in this table. - if !copying { - break; - } - - // A resize prevented us from deleting all the entries in this table. - // - // Complete the resize and retry in the new table. - table = self.help_copy(true, &table, guard); - } - } - - /// Retains only the elements specified by the predicate. - #[inline] - pub fn retain(&self, mut f: F, guard: &impl VerifiedGuard) - where - F: FnMut(&K, &V) -> bool, - { - // Load the root table. - let mut table = self.root(guard); - - // The table has not been initialized yet. - if table.raw.is_null() { - return; - } - - loop { - // Get a clean copy of the table to delete from. - table = self.linearize(table, guard); - - let mut copying = false; - 'probe: for i in 0..table.len() { - // Load the entry metadata first to ensure consistency with calls to `get` - // for entries that are retained. - // - // Safety: `i` is in bounds for the table length. - let meta = unsafe { table.meta(i) }.load(Ordering::Acquire); - - // The entry is empty or deleted. - if matches!(meta, meta::EMPTY | meta::TOMBSTONE) { - continue 'probe; - } - - // Load the entry to delete. - // - // Safety: `i` is in bounds for the table length. - let mut entry = guard - .protect(unsafe { table.entry(i) }, Ordering::Acquire) - .unpack(); - - loop { - // The entry is empty or already deleted. - if entry.ptr.is_null() { - continue 'probe; - } - - // Found a non-empty entry being copied. - if entry.tag() & Entry::COPYING != 0 { - // Clear every entry in this table that we can, then deal with the copy. - copying = true; - continue 'probe; - } - - // Safety: We performed a protected load of the pointer using a verified guard with - // `Acquire` and ensured that it is non-null, meaning it is valid for reads as long - // as we hold the guard. - let entry_ref = unsafe { &*entry.ptr }; - - // Should we retain this entry? - if f(&entry_ref.key, &entry_ref.value) { - continue 'probe; - } - - // Try to delete the entry. - // - // Safety: `i` is in bounds for the table length. - let result = unsafe { - table.entry(i).compare_exchange( - entry.raw, - Entry::TOMBSTONE, - Ordering::Release, - Ordering::Acquire, - ) - }; - - match result { - // Successfully deleted the entry. - Ok(_) => { - // Update the metadata table. - // - // Safety: `i` is in bounds for the table length. - unsafe { table.meta(i).store(meta::TOMBSTONE, Ordering::Release) }; - - // Decrement the table length. - self.count.get(guard).fetch_sub(1, Ordering::Relaxed); - - // Safety: The caller guarantees that `current` is a valid non-null entry that was - // inserted into the map. Additionally, it is now unreachable from this table due - // to the CAS above. - unsafe { self.defer_retire(entry, &table, guard) }; - continue 'probe; - } - - // Lost to a concurrent update, retry. - Err(found) => entry = found.unpack(), - } - } - } - - // We cleared every entry in this table. - if !copying { - break; - } - - // A resize prevented us from deleting all the entries in this table. - // - // Complete the resize and retry in the new table. - table = self.help_copy(true, &table, guard); - } - } - - /// Returns an iterator over the keys and values of this table. - #[inline] - pub fn iter<'g, G>(&self, guard: &'g G) -> Iter<'g, K, V, G> - where - G: VerifiedGuard, - { - // Load the root table. - let root = self.root(guard); - - // The table has not been initialized yet, return a dummy iterator. - if root.raw.is_null() { - return Iter { - i: 0, - guard, - table: root, - }; - } - - // Get a clean copy of the table to iterate over. - let table = self.linearize(root, guard); - - Iter { i: 0, guard, table } - } - - /// Returns the h1 and h2 hash for the given key. - #[inline] - fn hash(&self, key: &Q) -> (usize, u8) - where - Q: Hash + ?Sized, - { - let hash = self.hasher.hash_one(key); - (meta::h1(hash), meta::h2(hash)) - } -} - -/// A wrapper around a CAS function that manages the computed state. -struct ComputeState { - /// The CAS function. - compute: F, - - /// A cached insert transition. - insert: Option, - - /// A cached update transition. - update: Option>, -} - -/// A cached update transition. -struct CachedUpdate { - /// The entry that the CAS function was called with. - input: *mut Entry, - - /// The cached result. - output: Operation, -} - -impl<'g, F, K, V, T> ComputeState -where - F: FnMut(Option<(&'g K, &'g V)>) -> Operation, - K: 'g, - V: 'g, -{ - /// Create a new `ComputeState` for the given function. - #[inline] - fn new(compute: F) -> ComputeState { - ComputeState { - compute, - insert: None, - update: None, - } - } - - /// Performs a state transition. - /// - /// # Safety - /// - /// The entry pointer must be valid for reads if provided. - #[inline] - unsafe fn next(&mut self, entry: Option<*mut Entry>) -> Operation { - let Some(entry) = entry else { - // If there is no current entry, perform a transition for the insert. - return match self.insert.take() { - // Use the cached insert. - Some(value) => Operation::Insert(value), - - // Otherwise, compute the value to insert. - None => (self.compute)(None), - }; - }; - - // Otherwise, perform an update transition. - match self.update.take() { - // Used the cached update if the entry has not changed. - Some(CachedUpdate { input, output }) if input == entry => output, - - // Otherwise, compute the value to update. - _ => { - // Safety: The caller guarantees that `entry` is valid for reads. - let entry_ref = unsafe { &*entry }; - (self.compute)(Some((&entry_ref.key, &entry_ref.value))) - } - } - } - - /// Restores the state if an operation fails. - /// - /// This allows the result of the compute closure with a given input to be cached. - /// This is useful at it avoids calling the closure multiple times if an update needs - /// to be retried in a new table. - /// - /// Additionally, update and insert operations are cached separately, although this - /// is not guaranteed in the public API. This means that internal methods can rely on - /// `compute(None)` being called at most once. - #[inline] - fn restore(&mut self, input: Option<*mut Entry>, output: Operation) { - match input { - Some(input) => self.update = Some(CachedUpdate { input, output }), - None => match output { - Operation::Insert(value) => self.insert = Some(value), - _ => unreachable!(), - }, - } - } -} - -/// A lazy initialized `Entry` allocation. -enum LazyEntry { - /// An uninitialized entry, containing just the owned key. - Uninit(K), - - /// An allocated entry. - Init(*mut Entry>), -} - -impl LazyEntry { - /// Returns a reference to the entry's key. - #[inline] - fn key(&self) -> &K { - match self { - LazyEntry::Uninit(key) => key, - LazyEntry::Init(entry) => unsafe { &(**entry).key }, - } - } - - /// Initializes the entry if it has not already been initialized, returning the pointer - /// to the entry allocation. - #[inline] - fn init(&mut self) -> *mut Entry> { - match self { - LazyEntry::Init(entry) => *entry, - LazyEntry::Uninit(key) => { - // Safety: we read the current key with `ptr::read` and overwrite the - // state with `ptr::write`. We also make sure to abort if the allocator - // panics, ensuring the current value is not dropped twice. - unsafe { - let key = ptr::read(key); - let entry = panic::catch_unwind(panic::AssertUnwindSafe(|| { - Box::into_raw(Box::new(Entry { - value: MaybeUninit::uninit(), - key, - })) - })) - .unwrap_or_else(|_| std::process::abort()); - ptr::write(self, LazyEntry::Init(entry)); - entry - } - } - } - } -} - -/// RMW operations. -impl HashMap -where - K: Hash + Eq, - S: BuildHasher, -{ - /// Tries to insert a key and value computed from a closure into the map, - /// and returns a reference to the value that was inserted. - #[inline] - pub fn try_insert_with<'g, F>( - &self, - key: K, - f: F, - guard: &'g impl VerifiedGuard, - ) -> Result<&'g V, &'g V> - where - F: FnOnce() -> V, - K: 'g, - { - let mut f = Some(f); - let compute = |entry| match entry { - // There is already an existing value. - Some((_, current)) => Operation::Abort(current), - - // Insert the initial value. - // - // Note that this case is guaranteed to be executed at most - // once as insert values are cached, so this can never panic. - None => Operation::Insert((f.take().unwrap())()), - }; - - match self.compute(key, compute, guard) { - // Failed to insert, return the existing value. - Compute::Aborted(current) => Err(current), - - // Successfully inserted. - Compute::Inserted(_, value) => Ok(value), - - _ => unreachable!(), - } - } - - /// Returns a reference to the value corresponding to the key, or inserts a default value - /// computed from a closure. - #[inline] - pub fn get_or_insert_with<'g, F>(&self, key: K, f: F, guard: &'g impl VerifiedGuard) -> &'g V - where - F: FnOnce() -> V, - K: 'g, - { - match self.try_insert_with(key, f, guard) { - Ok(value) => value, - Err(value) => value, - } - } - - /// Updates an existing entry atomically, returning the value that was inserted. - #[inline] - pub fn update<'g, F>( - &self, - key: K, - mut update: F, - guard: &'g impl VerifiedGuard, - ) -> Option<&'g V> - where - F: FnMut(&V) -> V, - K: 'g, - { - let compute = |entry| match entry { - // There is nothing to update. - None => Operation::Abort(()), - // Perform the update. - Some((_, value)) => Operation::Insert(update(value)), - }; - - match self.compute(key, compute, guard) { - // Return the updated value. - Compute::Updated { - new: (_, value), .. - } => Some(value), - - // There was nothing to update. - Compute::Aborted(_) => None, - - _ => unreachable!(), - } - } - - /// Updates an existing entry or inserts a default value computed from a closure. - #[inline] - pub fn update_or_insert_with<'g, U, F>( - &self, - key: K, - update: U, - f: F, - guard: &'g impl VerifiedGuard, - ) -> &'g V - where - F: FnOnce() -> V, - U: Fn(&V) -> V, - K: 'g, - { - let mut f = Some(f); - let compute = |entry| match entry { - // Perform the update. - Some((_, value)) => Operation::Insert::<_, ()>(update(value)), - - // Insert the initial value. - // - // Note that this case is guaranteed to be executed at most - // once as insert values are cached, so this can never panic. - None => Operation::Insert((f.take().unwrap())()), - }; - - match self.compute(key, compute, guard) { - // Return the updated value. - Compute::Updated { - new: (_, value), .. - } => value, - - // Return the value we inserted. - Compute::Inserted(_, value) => value, - - _ => unreachable!(), - } - } - - /// Update an entry with a CAS function. - /// - /// Note that `compute` closure is guaranteed to be called for a `None` input only once, allowing the - /// insertion of values that cannot be cloned or reconstructed. - #[inline] - pub fn compute<'g, F, T>( - &self, - key: K, - compute: F, - guard: &'g impl VerifiedGuard, - ) -> Compute<'g, K, V, T> - where - F: FnMut(Option<(&'g K, &'g V)>) -> Operation, - { - // Lazy initialize the entry allocation. - let mut entry = LazyEntry::Uninit(key); - - // Perform the update. - // - // Safety: We just allocated the entry above. - let result = unsafe { self.compute_with(&mut entry, ComputeState::new(compute), guard) }; - - // Deallocate the entry if it was not inserted. - if matches!(result, Compute::Removed(..) | Compute::Aborted(_)) { - if let LazyEntry::Init(entry) = entry { - // Safety: The entry was allocated but not inserted into the map. - let _ = unsafe { Box::from_raw(entry) }; - } - } - - result - } - - /// Update an entry with a CAS function. - /// - /// # Safety - /// - /// The new entry must be a valid owned pointer to insert into the map. - #[inline] - unsafe fn compute_with<'g, F, T>( - &self, - new_entry: &mut LazyEntry, - mut state: ComputeState, - guard: &'g impl VerifiedGuard, - ) -> Compute<'g, K, V, T> - where - F: FnMut(Option<(&'g K, &'g V)>) -> Operation, - { - // Load the root table. - let mut table = self.root(guard); - - // The table has not yet been allocated. - if table.raw.is_null() { - // Compute the value to insert. - // - // Safety: Insert transitions are always sound. - match unsafe { state.next(None) } { - op @ Operation::Insert(_) => state.restore(None, op), - Operation::Remove => panic!("Cannot remove `None` entry."), - Operation::Abort(value) => return Compute::Aborted(value), - } - - // Initialize the table. - table = self.init(None); - } - - let (h1, h2) = self.hash(new_entry.key()); - let mut help_copy = false; - - loop { - // Initialize the probe state. - let mut probe = Probe::start(h1, table.mask); - - // Probe until we reach the limit. - let copying = 'probe: loop { - if probe.len > table.limit { - break 'probe None; - } - - // Load the entry metadata first for cheap searches. - // - // Safety: `probe.i` is always in-bounds for the table length. - let meta = unsafe { table.meta(probe.i) }.load(Ordering::Acquire); - - // The entry is empty. - let mut entry = if meta == meta::EMPTY { - // Compute the value to insert. - // - // Safety: Insert transitions are always sound. - let value = match unsafe { state.next(None) } { - Operation::Insert(value) => value, - Operation::Remove => panic!("Cannot remove `None` entry."), - Operation::Abort(value) => return Compute::Aborted(value), - }; - - let new_entry = new_entry.init(); - // Safety: `new_entry` was just allocated above and is valid for writes. - unsafe { (*new_entry).value = MaybeUninit::new(value) } - - // Attempt to insert. - // - // Safety: `probe.i` is always in-bounds for the table length.Additionally, - // `new_entry` was allocated above and never shared. - match unsafe { self.insert_at(probe.i, h2, new_entry.cast(), table, guard) } { - // Successfully inserted. - InsertStatus::Inserted => { - // Increment the table length. - self.count.get(guard).fetch_add(1, Ordering::Relaxed); - - // Safety: `new_entry` was initialized above. - let new_ref = unsafe { &*new_entry.cast::>() }; - return Compute::Inserted(&new_ref.key, &new_ref.value); - } - - // Lost to a concurrent insert. - // - // If the key matches, we might be able to update the value. - InsertStatus::Found(EntryStatus::Value(found)) - | InsertStatus::Found(EntryStatus::Copied(found)) => { - // Cache the previous value - // - // Safety: `new_entry` was initialized above and was not inserted - // into the map. - let value = unsafe { (*new_entry).value.assume_init_read() }; - state.restore(None, Operation::Insert(value)); - - found - } - - // The entry was removed or invalidated. - InsertStatus::Found(EntryStatus::Null) => { - // Cache the previous value. - // - // Safety: `new_entry` was initialized above and was not inserted - // into the map. - let value = unsafe { (*new_entry).value.assume_init_read() }; - state.restore(None, Operation::Insert(value)); - - // Continue probing. - probe.next(table.mask); - continue 'probe; - } - } - } - // Found a potential match. - else if meta == h2 { - // Load the full entry. - // - // Safety: `probe.i` is always in-bounds for the table length. - let found = guard - .protect(unsafe { table.entry(probe.i) }, Ordering::Acquire) - .unpack(); - - // The entry was deleted, keep probing. - if found.ptr.is_null() { - probe.next(table.mask); - continue 'probe; - } - - // If the key matches, we might be able to update the value. - found - } - // Otherwise, continue probing. - else { - probe.next(table.mask); - continue 'probe; - }; - - // Check for a full match. - // - // Safety: We performed a protected load of the pointer using a verified guard with - // `Acquire` and ensured that it is non-null, meaning it is valid for reads as long - // as we hold the guard. - if unsafe { (*entry.ptr).key != *new_entry.key() } { - probe.next(table.mask); - continue 'probe; - } - - // The entry is being copied to the new table. - if entry.tag() & Entry::COPYING != 0 { - break 'probe Some(probe.i); - } - - loop { - // Compute the value to insert. - // - // Safety: `entry` is valid for reads. - let failure = match unsafe { state.next(Some(entry.ptr)) } { - // The operation was aborted. - Operation::Abort(value) => return Compute::Aborted(value), - - // Update the value. - Operation::Insert(value) => { - let new_entry = new_entry.init(); - - // Safety: `new_entry` was just allocated above and is valid for writes. - unsafe { (*new_entry).value = MaybeUninit::new(value) } - - // Try to perform the update. - // - // Safety: - // - `probe.i` is always in-bounds for the table length - // - `entry` is a valid non-null entry that we found in the map. - // - `new_entry` was initialized above and never shared. - let status = unsafe { - self.update_at(probe.i, entry, new_entry.cast(), table, guard) - }; - - match status { - // Successfully updated. - UpdateStatus::Replaced(entry) => { - // Safety: `entry` is a valid non-null entry that we found in the map - // before replacing it. - let entry_ref = unsafe { &(*entry.ptr) }; - - // Safety: `new_entry` was initialized above. - let new_ref = unsafe { &*new_entry.cast::>() }; - - return Compute::Updated { - old: (&entry_ref.key, &entry_ref.value), - new: (&new_ref.key, &new_ref.value), - }; - } - - // The update failed. - failure => { - // Save the previous value. - // - // Safety: `new_entry` was initialized above and was not inserted - // into the map. - let value = unsafe { (*new_entry).value.assume_init_read() }; - state.restore(Some(entry.ptr), Operation::Insert(value)); - - failure - } - } - } - - // Remove the key from the map. - Operation::Remove => { - // Try to perform the removal. - // - // Safety: - // - `probe.i` is always in-bounds for the table length - // - `entry` is a valid non-null entry that we found in the map. - let status = unsafe { - self.update_at(probe.i, entry, Entry::TOMBSTONE, table, guard) - }; - - match status { - // Successfully removed the entry. - UpdateStatus::Replaced(entry) => { - // Mark the entry as a tombstone. - // - // Note that this might end up being overwritten by the metadata hash - // if the initial insertion is lagging behind, but we avoid the RMW - // and sacrifice reads in the extremely rare case. - unsafe { - table - .meta(probe.i) - .store(meta::TOMBSTONE, Ordering::Release) - }; - - // Decrement the table length. - self.count.get(guard).fetch_sub(1, Ordering::Relaxed); - - // Safety: `entry` is a valid non-null entry that we found in the map - // before replacing it. - let entry_ref = unsafe { &(*entry.ptr) }; - return Compute::Removed(&entry_ref.key, &entry_ref.value); - } - - // The remove failed. - failure => { - // Save the removal operation. - state.restore(Some(entry.ptr), Operation::Remove); - - failure - } - } - } - }; - - match failure { - // The entry is being copied to the new table. - UpdateStatus::Found(EntryStatus::Copied(_)) => break 'probe Some(probe.i), - - // The entry was deleted before we could update it. - // - // We know that at some point during our execution the key was not in the map. - UpdateStatus::Found(EntryStatus::Null) => { - // Compute the next operation. - // - // Safety: Insert transitions are always sound. - match unsafe { state.next(None) } { - Operation::Insert(value) => { - // Save the computed value. - state.restore(None, Operation::Insert(value)); - - // Continue probing to find an empty slot. - probe.next(table.mask); - continue 'probe; - } - Operation::Remove => panic!("Cannot remove `None` entry."), - Operation::Abort(value) => return Compute::Aborted(value), - } - } - - // Someone else beat us to the update, retry. - UpdateStatus::Found(EntryStatus::Value(found)) => entry = found, - - _ => unreachable!(), - } - } - }; - - // Prepare to retry in the next table. - if let Some(next_table) = self.prepare_retry(copying, &mut help_copy, table, guard) { - table = next_table; - continue; - } - - // Otherwise, the key is not in the map. - // - // Safety: Insert transitions are always sound. - match unsafe { state.next(None) } { - // Need to insert into the new table. - op @ Operation::Insert(_) => { - table = self.prepare_retry_insert(None, &mut help_copy, table, guard); - state.restore(None, op); - } - // The operation was aborted. - Operation::Abort(value) => return Compute::Aborted(value), - Operation::Remove => panic!("Cannot remove `None` entry."), - } - } - } -} - -/// Resize operations. -impl HashMap -where - K: Hash + Eq, - S: BuildHasher, -{ - /// Allocate the initial table. - #[cold] - #[inline(never)] - fn init(&self, capacity: Option) -> Table> { - const CAPACITY: usize = 32; - - // Allocate the table and mark it as the root. - let mut new = Table::alloc(capacity.unwrap_or(CAPACITY)); - *new.state_mut().status.get_mut() = State::PROMOTED; - - // Race to write the initial table. - match self.table.compare_exchange( - ptr::null_mut(), - new.raw, - Ordering::Release, - Ordering::Acquire, - ) { - // Successfully initialized the table. - Ok(_) => new, - - // Someone beat us, deallocate our table and use the table that was written. - Err(found) => { - // Safety: We allocated the table above and never shared it. - unsafe { Table::dealloc(new) } - - // Safety: The table was just initialized. - unsafe { Table::from_raw(found) } - } - } - } - - /// Returns the next table, allocating it has not already been created. - #[cold] - #[inline(never)] - fn get_or_alloc_next( - &self, - capacity: Option, - table: Table>, - ) -> Table> { - // Avoid spinning in tests, which can hide race conditions. - const SPIN_ALLOC: usize = if cfg!(any(test, debug_assertions)) { - 1 - } else { - 7 - }; - - // The next table is already allocated. - if let Some(next) = table.next_table() { - return next; - } - - let state = table.state(); - - // Otherwise, try to acquire the allocation lock. - // - // Unlike in `init`, we do not race here to prevent unnecessary allocator pressure. - let _allocating = match state.allocating.try_lock() { - Ok(lock) => lock, - // Someone else is currently allocating. - Err(_) => { - let mut spun = 0; - - // Spin for a bit, waiting for the table to be initialized. - while spun <= SPIN_ALLOC { - for _ in 0..(spun * spun) { - hint::spin_loop(); - } - - // The table was initialized. - if let Some(next) = table.next_table() { - return next; - } - - spun += 1; - } - - // Otherwise, we have to block. - state.allocating.lock().unwrap() - } - }; - - // The table was allocated while we were waiting for the lock. - if let Some(next) = table.next_table() { - return next; - } - - let current_capacity = table.len(); - - // Loading the length here is quite expensive, we may want to consider - // a probabilistic counter to detect high-deletion workloads. - let active_entries = self.len(); - - let next_capacity = match cfg!(papaya_stress) { - // Never grow the table to stress the incremental resizing algorithm. - true => current_capacity, - - // Double the table capacity if we are at least 50% full. - false if active_entries >= (current_capacity >> 1) => current_capacity << 1, - - // Halve the table if we are at most 12.5% full. - // - // This heuristic is intentionally pessimistic as unnecessarily shrinking - // is an expensive operation, but it may change in the future. We also respect - // the initial capacity to give the user a way to retain a strict minimum table - // size. - false if active_entries <= (current_capacity >> 3) => { - self.initial_capacity.max(current_capacity >> 1) - } - - // Otherwise keep the capacity the same. - // - // This can occur due to poor hash distribution or frequent cycling of - // insertions and deletions, in which case we want to avoid continuously - // growing the table. - false => current_capacity, - }; - - let next_capacity = capacity.unwrap_or(next_capacity); - assert!( - next_capacity <= isize::MAX as usize, - "`HashMap` exceeded maximum capacity" - ); - - // Allocate the new table while holding the lock. - let next = Table::alloc(next_capacity); - state.next.store(next.raw, Ordering::Release); - drop(_allocating); - - next - } - - /// Help along with an existing resize operation, returning the new root table. - /// - /// If `copy_all` is `false` in incremental resize mode, this returns the current reference's next - /// table, not necessarily the new root. - #[cold] - #[inline(never)] - fn help_copy( - &self, - copy_all: bool, - table: &Table>, - guard: &impl VerifiedGuard, - ) -> Table> { - match self.resize { - ResizeMode::Blocking => self.help_copy_blocking(table, guard), - ResizeMode::Incremental(chunk) => { - let copied_to = self.help_copy_incremental(chunk, copy_all, guard); - - if !copy_all { - // If we weren't trying to linearize, we have to write to the next table - // even if the copy hasn't completed yet. - return table.next_table().unwrap(); - } - - copied_to - } - } - } - - /// Help along the resize operation until it completes and the next table is promoted. - /// - /// Should only be called on the root table. - fn help_copy_blocking( - &self, - table: &Table>, - guard: &impl VerifiedGuard, - ) -> Table> { - // Load the next table. - let mut next = table.next_table().unwrap(); - - 'copy: loop { - // Make sure we are copying to the correct table. - while next.state().status.load(Ordering::Relaxed) == State::ABORTED { - next = self.get_or_alloc_next(None, next); - } - - // The copy already completed - if self.try_promote(table, &next, 0, guard) { - return next; - } - - let copy_chunk = table.len().min(4096); - - loop { - // Every entry has already been claimed. - if next.state().claim.load(Ordering::Relaxed) >= table.len() { - break; - } - - // Claim a chunk to copy. - let copy_start = next.state().claim.fetch_add(copy_chunk, Ordering::Relaxed); - - // Copy our chunk of entries. - let mut copied = 0; - for i in 0..copy_chunk { - let i = copy_start + i; - - if i >= table.len() { - break; - } - - // Copy the entry. - // - // Safety: We verified that `i` is in-bounds above. - if unsafe { !self.copy_at_blocking(i, table, &next, guard) } { - // This table doesn't have space for the next entry. - // - // Abort the current resize. - // - // Note that the `SeqCst` is necessary to make the store visible - // to threads that are unparked. - next.state().status.store(State::ABORTED, Ordering::SeqCst); - - // Allocate the next table. - let allocated = self.get_or_alloc_next(None, next); - - // Wake anyone waiting for us to finish. - let state = table.state(); - state.parker.unpark(&state.status); - - // Retry in a new table. - next = allocated; - continue 'copy; - } - - copied += 1; - } - - // Are we done? - if self.try_promote(table, &next, copied, guard) { - return next; - } - - // If the resize was aborted while we were copying, continue in the new table. - if next.state().status.load(Ordering::Relaxed) == State::ABORTED { - continue 'copy; - } - } - - let state = next.state(); - // We copied all that we can, wait for the table to be promoted. - for spun in 0.. { - // Avoid spinning in tests, which can hide race conditions. - const SPIN_WAIT: usize = if cfg!(any(test, debug_assertions)) { - 1 - } else { - 7 - }; - - // Note that `Acquire` is necessary here to ensure we see the - // relevant modifications to the root table if see the updated - // state before parking. - // - // Otherwise, `Parker::park` will ensure the necessary synchronization - // when we are unparked. - let status = state.status.load(Ordering::Acquire); - - // If this copy was aborted, we have to retry in the new table. - if status == State::ABORTED { - continue 'copy; - } - - // The copy has completed. - if status == State::PROMOTED { - return next; - } - - // Copy chunks are relatively small and we expect to finish quickly, - // so spin for a bit before resorting to parking. - if spun <= SPIN_WAIT { - for _ in 0..(spun * spun) { - hint::spin_loop(); - } - - continue; - } - - // Park until the table is promoted. - state - .parker - .park(&state.status, |status| status == State::PENDING); - } - } - } - - /// Copy the entry at the given index to the new table. - /// - /// Returns `true` if the entry was copied into the table or `false` if the table was full. - /// - /// # Safety - /// - /// The index must be in-bounds for the table. - unsafe fn copy_at_blocking( - &self, - i: usize, - table: &Table>, - next_table: &Table>, - guard: &impl VerifiedGuard, - ) -> bool { - // Mark the entry as copying. - // - // Safety: The caller guarantees that the index is in-bounds. - // - // Note that we don't need to protect the returned entry here, because - // no one is allowed to retire the entry once we put the `COPYING` bit - // down until it is inserted into the new table. - let entry = unsafe { table.entry(i) } - .fetch_or(Entry::COPYING, Ordering::AcqRel) - .unpack(); - - // The entry is a tombstone. - if entry.raw == Entry::TOMBSTONE { - return true; - } - - // There is nothing to copy, we're done. - if entry.ptr.is_null() { - // Mark as a tombstone so readers avoid having to load the entry. - // - // Safety: The caller guarantees that the index is in-bounds. - unsafe { table.meta(i) }.store(meta::TOMBSTONE, Ordering::Release); - return true; - } - - // Copy the value to the new table. - // - // Safety: We marked the entry as `COPYING`, ensuring that any updates - // or removals wait until we complete the copy, and allowing us to get - // away without a protected load. Additionally, we verified that the - // entry is non-null, meaning that it is valid for reads. - unsafe { - self.insert_copy(entry.ptr.unpack(), false, next_table, guard) - .is_some() - } - } - - /// Help along an in-progress resize incrementally by copying a chunk of entries. - /// - /// Returns the table that was copied to. - fn help_copy_incremental( - &self, - chunk: usize, - block: bool, - guard: &impl VerifiedGuard, - ) -> Table> { - // Always help the highest priority root resize. - let table = self.root(guard); - - // Load the next table. - let Some(next) = table.next_table() else { - // The copy we tried to help was already promoted. - return table; - }; - - loop { - // The copy already completed. - if self.try_promote(&table, &next, 0, guard) { - return next; - } - - loop { - // Every entry has already been claimed. - if next.state().claim.load(Ordering::Relaxed) >= table.len() { - break; - } - - // Claim a chunk to copy. - let copy_start = next.state().claim.fetch_add(chunk, Ordering::Relaxed); - - // Copy our chunk of entries. - let mut copied = 0; - for i in 0..chunk { - let i = copy_start + i; - - if i >= table.len() { - break; - } - - // Copy the entry. - // - // Safety: We verified that `i` is in-bounds above. - unsafe { self.copy_at_incremental(i, &table, &next, guard) }; - copied += 1; - } - - // Update the copy state, and try to promote the table. - // - // Only copy a single chunk if promotion fails, unless we are forced - // to complete the resize. - if self.try_promote(&table, &next, copied, guard) || !block { - return next; - } - } - - // There are no entries that we can copy, block if necessary. - if !block { - return next; - } - - let state = next.state(); - for spun in 0.. { - // Avoid spinning in tests, which can hide race conditions. - const SPIN_WAIT: usize = if cfg!(any(test, debug_assertions)) { - 1 - } else { - 7 - }; - - // The copy has completed. - // - // Note that `Acquire` is necessary here to ensure we see the - // relevant modifications to the root table if see the updated - // state before parking. - // - // Otherwise, `Parker::park` will ensure the necessary synchronization - // when we are unparked. - let status = state.status.load(Ordering::Acquire); - if status == State::PROMOTED { - return next; - } - - // Copy chunks are relatively small and we expect to finish quickly, - // so spin for a bit before resorting to parking. - if spun <= SPIN_WAIT { - for _ in 0..(spun * spun) { - hint::spin_loop(); - } - - continue; - } - - // Park until the table is promoted. - state - .parker - .park(&state.status, |status| status == State::PENDING); - } - } - } - - /// Copy the entry at the given index to the new table. - /// - /// # Safety - /// - /// The index must be in-bounds for the table. - unsafe fn copy_at_incremental( - &self, - i: usize, - table: &Table>, - next_table: &Table>, - guard: &impl VerifiedGuard, - ) { - // Safety: The caller guarantees that the index is in-bounds. - let entry = unsafe { table.entry(i) }; - - // Mark the entry as copying. - let found = entry.fetch_or(Entry::COPYING, Ordering::AcqRel).unpack(); - - // The entry is a tombstone. - if found.raw == Entry::TOMBSTONE { - return; - } - - // There is nothing to copy, we're done. - if found.ptr.is_null() { - // Mark as a tombstone so readers avoid having to load the entry. - // - // Safety: The caller guarantees that the index is in-bounds. - unsafe { table.meta(i) }.store(meta::TOMBSTONE, Ordering::Release); - return; - } - - // Mark the entry as borrowed so writers in the new table know it was copied. - let new_entry = found.map_tag(|addr| addr | Entry::BORROWED); - - // Copy the value to the new table. - // - // Safety: We marked the entry as `COPYING`, ensuring that any updates - // or removals wait until we complete the copy, and allowing us to get - // away without a protected load. Additionally, we verified that the - // entry is non-null, meaning that it is valid for reads. - unsafe { - self.insert_copy(new_entry, true, next_table, guard) - .unwrap(); - } - - // Mark the entry as copied. - let copied = found - .raw - .map_addr(|addr| addr | Entry::COPYING | Entry::COPIED); - - // Note that we already wrote the COPYING bit, so no one is writing to the old - // entry except us. - // - // Note that the `SeqCst` is necessary to make the store visible to threads - // that are unparked. - entry.store(copied, Ordering::SeqCst); - - // Notify any writers that the copy has completed. - table.state().parker.unpark(entry); - } - - // Copy an entry into the table, returning the index it was inserted into. - // - // This is an optimized version of `insert_entry` where the caller is the only writer - // inserting the given key into the new table, as it has already been marked as copying. - // - // # Safety - // - // The new entry must be valid for reads. - unsafe fn insert_copy( - &self, - new_entry: Tagged>, - resize: bool, - table: &Table>, - guard: &impl VerifiedGuard, - ) -> Option<(Table>, usize)> { - // Safety: The new entry is guaranteed to be valid for reads. - let key = unsafe { &(*new_entry.ptr).key }; - - let mut table = *table; - let (h1, h2) = self.hash(key); - - loop { - // Initialize the probe state. - let mut probe = Probe::start(h1, table.mask); - - // Probe until we reach the limit. - while probe.len <= table.limit { - // Safety: `probe.i` is always in-bounds for the table length. - let meta_entry = unsafe { table.meta(probe.i) }; - - // Load the entry metadata first for cheap searches. - let meta = meta_entry.load(Ordering::Acquire); - - // The entry is empty, try to insert. - if meta == meta::EMPTY { - // Safety: `probe.i` is always in-bounds for the table length. - let entry = unsafe { table.entry(probe.i) }; - - // Try to claim the entry. - match guard.compare_exchange( - entry, - ptr::null_mut(), - new_entry.raw, - Ordering::Release, - Ordering::Acquire, - ) { - // Successfully inserted. - Ok(_) => { - // Update the metadata table. - meta_entry.store(h2, Ordering::Release); - return Some((table, probe.i)); - } - Err(found) => { - let found = found.unpack(); - - // The entry was deleted or copied. - let meta = if found.ptr.is_null() { - meta::TOMBSTONE - } else { - // Safety: We performed a protected load of the pointer using a verified guard with - // `Acquire` and ensured that it is non-null, meaning it is valid for reads as long - // as we hold the guard. - let found_ref = unsafe { &(*found.ptr) }; - - // Ensure the meta table is updated to avoid breaking the probe chain. - let hash = self.hasher.hash_one(&found_ref.key); - meta::h2(hash) - }; - - if meta_entry.load(Ordering::Relaxed) == meta::EMPTY { - meta_entry.store(meta, Ordering::Release); - } - } - } - } - - // Continue probing. - probe.next(table.mask); - } - - if !resize { - return None; - } - - // Insert into the next table. - table = self.get_or_alloc_next(None, table); - } - } - - // Update the copy state and attempt to promote a table to the root. - // - // Returns `true` if the table was promoted. - fn try_promote( - &self, - table: &Table>, - next: &Table>, - copied: usize, - guard: &impl VerifiedGuard, - ) -> bool { - let state = next.state(); - - // Update the copy count. - let copied = if copied > 0 { - state.copied.fetch_add(copied, Ordering::AcqRel) + copied - } else { - state.copied.load(Ordering::Acquire) - }; - - // If we copied all the entries in the table, we can try to promote. - if copied == table.len() { - let root = self.table.load(Ordering::Relaxed); - - // Only promote root copies. - // - // We can't promote a nested copy before it's parent has finished, as - // it may not contain all the entries in the table. - if table.raw == root { - // Try to update the root. - if self - .table - .compare_exchange(table.raw, next.raw, Ordering::Release, Ordering::Acquire) - .is_ok() - { - // Successfully promoted the table. - // - // Note that the `SeqCst` is necessary to make the store visible to threads - // that are unparked. - state.status.store(State::PROMOTED, Ordering::SeqCst); - - // Retire the old table. - // - // Safety: `table.raw` is a valid pointer to the table we just copied from. - // Additionally, the CAS above made the previous table unreachable from the - // root pointer, allowing it to be safely retired. - unsafe { - guard.defer_retire(table.raw, |table, collector| { - // Note that we do not drop entries because they have been copied to - // the new root. - drop_table(Table::from_raw(table), collector); - }); - } - } - - // Wake up any writers waiting for the resize to complete. - state.parker.unpark(&state.status); - return true; - } - } - - // Not ready to promote yet. - false - } - - // Completes all pending copies in incremental mode to get a clean copy of the table. - // - // This is necessary for operations like `iter` or `clear`, where entries in multiple tables - // can cause lead to incomplete results. - #[inline] - fn linearize( - &self, - mut table: Table>, - guard: &impl VerifiedGuard, - ) -> Table> { - if self.is_incremental() { - // If we're in incremental resize mode, we need to complete any in-progress resizes to - // ensure we don't miss any entries in the next table. We can't iterate over both because - // we risk returning the same entry twice. - while table.next_table().is_some() { - table = self.help_copy(true, &table, guard); - } - } - - table - } - - // Wait for an incremental copy of a given entry to complete. - #[cold] - #[inline(never)] - fn wait_copied(&self, i: usize, table: &Table>) { - // Avoid spinning in tests, which can hide race conditions. - const SPIN_WAIT: usize = if cfg!(any(test, debug_assertions)) { - 1 - } else { - 5 - }; - - let entry = unsafe { table.entry(i) }; - - // Spin for a short while, waiting for the entry to be copied. - for spun in 0..SPIN_WAIT { - // The entry was copied. - let entry = entry.load(Ordering::Acquire).unpack(); - if entry.tag() & Entry::COPIED != 0 { - return; - } - - for _ in 0..(spun * spun) { - hint::spin_loop(); - } - } - - // Park until the copy completes. - let parker = &table.state().parker; - parker.park(entry, |entry| entry.addr() & Entry::COPIED == 0); - } - - /// Retire an entry that was removed from the current table, but may still be reachable from - /// previous tables. - /// - /// # Safety - /// - /// The entry must be a valid pointer that is unreachable from the current table. Additionally, - /// it is *undefined behavior* to call this method multiple times for the same entry. - #[inline] - unsafe fn defer_retire( - &self, - entry: Tagged>, - table: &Table>, - guard: &impl VerifiedGuard, - ) { - match self.resize { - // Safety: In blocking resize mode, we only ever write to the root table, so the entry - // is inaccessible from all tables. - ResizeMode::Blocking => unsafe { - guard.defer_retire(entry.ptr, seize::reclaim::boxed); - }, - // In incremental resize mode, the entry may be accessible in previous tables. - ResizeMode::Incremental(_) => { - if entry.tag() & Entry::BORROWED == 0 { - // Safety: If the entry is not borrowed, meaning it is not in any previous tables, - // it is inaccessible even if the current table is not root. Thus we can safely retire. - unsafe { guard.defer_retire(entry.ptr, seize::reclaim::boxed) }; - return; - } - - let root = self.root(guard); - - // Check if our table, or any subsequent table, is the root. - let mut next = Some(*table); - while let Some(table) = next { - if table.raw == root.raw { - // Safety: The root table is our table or a table that succeeds ours. - // Thus any previous tables are unreachable from the root, so we can safely retire. - unsafe { guard.defer_retire(entry.ptr, seize::reclaim::boxed) }; - return; - } - - next = table.next_table(); - } - - // Otherwise, we have to wait for the table we are copying from to be reclaimed. - // - // Find the table we are copying from, searching from the root. - let mut prev = root; - - loop { - let next = prev.next_table().unwrap(); - - // Defer the entry to be retired by the table we are copying from. - if next.raw == table.raw { - prev.state().deferred.push(entry.ptr); - return; - } - - prev = next; - } - } - } - } -} - -// An iterator over the keys and values of this table. -pub struct Iter<'g, K, V, G> { - i: usize, - table: Table>, - guard: &'g G, -} - -impl<'g, K: 'g, V: 'g, G> Iterator for Iter<'g, K, V, G> -where - G: VerifiedGuard, -{ - type Item = (&'g K, &'g V); - - #[inline] - fn next(&mut self) -> Option { - // The table has not yet been allocated. - if self.table.raw.is_null() { - return None; - } - - loop { - // Iterated over every entry in the table, we're done. - if self.i >= self.table.len() { - return None; - } - - // Load the entry metadata first to ensure consistency with calls to `get`. - // - // Safety: We verified that `self.i` is in-bounds above. - let meta = unsafe { self.table.meta(self.i) }.load(Ordering::Acquire); - - // The entry is empty or deleted. - if matches!(meta, meta::EMPTY | meta::TOMBSTONE) { - self.i += 1; - continue; - } - - // Load the entry. - // - // Safety: We verified that `self.i` is in-bounds above. - let entry = self - .guard - .protect(unsafe { self.table.entry(self.i) }, Ordering::Acquire) - .unpack(); - - // The entry was deleted. - if entry.ptr.is_null() { - self.i += 1; - continue; - } - - // Safety: We performed a protected load of the pointer using a verified guard with - // `Acquire` and ensured that it is non-null, meaning it is valid for reads as long - // as we hold the guard. - let entry_ref = unsafe { &(*entry.ptr) }; - - self.i += 1; - return Some((&entry_ref.key, &entry_ref.value)); - } - } -} - -// Safety: An iterator holds a shared reference to the HashMap -// and Guard, and outputs shared references to keys and values. -// Thus everything must be `Sync` for the iterator to be `Send` -// or `Sync`. -// -// It is not possible to obtain an owned key, value, or guard -// from an iterator, so `Send` is not a required bound. -unsafe impl Send for Iter<'_, K, V, G> -where - K: Sync, - V: Sync, - G: Sync, -{ -} - -unsafe impl Sync for Iter<'_, K, V, G> -where - K: Sync, - V: Sync, - G: Sync, -{ -} - -impl Clone for Iter<'_, K, V, G> { - #[inline] - fn clone(&self) -> Self { - Iter { - i: self.i, - table: self.table, - guard: self.guard, - } - } -} - -impl Drop for HashMap { - fn drop(&mut self) { - let mut raw = *self.table.get_mut(); - - // Make sure all objects are reclaimed before the collector is dropped. - // - // Dropping a table depends on accessing the collector for deferred retirement, - // using the shared collector pointer that is invalidated by drop. - // - // Safety: We have a unique reference to the collector. - unsafe { self.collector.reclaim_all() }; - - // Drop all nested tables and entries. - while !raw.is_null() { - // Safety: The root and next tables are always valid pointers to a - // table allocation, or null. - let mut table = unsafe { Table::from_raw(raw) }; - - // Read the next table pointer before dropping the current one. - let next = *table.state_mut().next.get_mut(); - - // Safety: We have unique access to the table and do - // not access the entries after this call. - unsafe { drop_entries(table) }; - - // Safety: We have unique access to the table and do - // not access it after this call. - unsafe { drop_table(table, &self.collector) }; - - // Continue for all nested tables. - raw = next; - } - } -} - -// Drop all entries in this table. -// -// # Safety -// -// The table entries must not be accessed after this call. -unsafe fn drop_entries(table: Table>) { - for i in 0..table.len() { - // Safety: `i` is in-bounds and we have unique access to the table. - let entry = unsafe { (*table.entry(i).as_ptr()).unpack() }; - - // The entry was copied, or there is nothing to deallocate. - if entry.ptr.is_null() || entry.tag() & Entry::COPYING != 0 { - continue; - } - - // Drop the entry. - // - // Safety: We verified that the table is non-null and will - // not be accessed after this call. Additionally, we ensured - // that the entry is not copied to avoid double freeing entries - // that may exist in multiple tables. - unsafe { drop(Box::from_raw(entry.ptr)) } - } -} - -// Drop the table allocation. -// -// # Safety -// -// The table must not be accessed after this call. -unsafe fn drop_table(mut table: Table>, collector: &Collector) { - // Drop any entries that were deferred during an incremental resize. - // - // Safety: Entries are deferred after they are made unreachable from the - // next table during a resize from this table. This table must have been accessible - // from the root for any entry to have been deferred. Thus it is being retired now, - // *after* the entry was made inaccessible from the next table. Additionally, for - // this table to have been retired, it also must no longer be accessible from the root, - // meaning that the entry has been totally removed from the map, and can be safely - // retired. - table - .state_mut() - .deferred - .drain(|entry| unsafe { collector.retire(entry, seize::reclaim::boxed) }); - - // Deallocate the table. - // - // Safety: The caller guarantees that the table will not be accessed after this call. - unsafe { Table::dealloc(table) }; -} - -// Entry metadata, inspired by `hashbrown`. -mod meta { - use std::mem; - - // Indicates an empty entry. - pub const EMPTY: u8 = 0x80; - - // Indicates an entry that has been deleted. - pub const TOMBSTONE: u8 = u8::MAX; - - // Returns the primary hash for an entry. - #[inline] - pub fn h1(hash: u64) -> usize { - hash as usize - } - - /// Return a byte of hash metadata, used for cheap searches. - #[inline] - pub fn h2(hash: u64) -> u8 { - const MIN_HASH_LEN: usize = if mem::size_of::() < mem::size_of::() { - mem::size_of::() - } else { - mem::size_of::() - }; - - // Grab the top 7 bits of the hash. - // - // While the hash is normally a full 64-bit value, some hash functions - // (such as fxhash) produce a usize result instead, which means that the - // top 32 bits are 0 on 32-bit platforms. - let top7 = hash >> (MIN_HASH_LEN * 8 - 7); - (top7 & 0x7f) as u8 - } -} diff --git a/src/raw/table.rs b/src/raw/table.rs new file mode 100644 index 0000000..bcfb189 --- /dev/null +++ b/src/raw/table.rs @@ -0,0 +1,2682 @@ +use std::marker::PhantomData; +use std::mem::ManuallyDrop; +use std::sync::atomic::{AtomicPtr, AtomicU8, AtomicUsize, Ordering}; +use std::sync::Mutex; +use std::{hint, panic, ptr}; + +use super::alloc::{RawTable, Table}; +use super::probe::{self, Probe}; +use super::utils::{ + untagged, AtomicPtrFetchOps, Counter, MapGuard, Parker, Stack, StrictProvenance, Tagged, + Unpack, VerifiedGuard, +}; +use crate::map::ResizeMode; + +use seize::{Collector, LocalGuard, OwnedGuard}; + +/// A lock-free hash-table. +pub struct HashTable> { + /// A pointer to the root table. + table: AtomicPtr>, + + /// Collector for memory reclamation. + pub(crate) collector: Collector, + + /// The resize mode, either blocking or incremental. + resize: ResizeMode, + + /// An atomic counter of the number of keys in the table. + count: Counter, + + /// The initial capacity provided to `HashMap::new`. + /// + /// The table is guaranteed to never shrink below this capacity. + initial_capacity: usize, + + _dealloc: PhantomData, +} + +pub trait Dealloc { + unsafe fn dealloc(entry: *mut T); +} + +/// Resize state for the hash-table. +pub struct State { + /// The next table used for resizing. + pub next: AtomicPtr>, + + /// A lock acquired to allocate the next table. + pub allocating: Mutex<()>, + + /// The number of entries that have been copied to the next table. + pub copied: AtomicUsize, + + /// The number of entries that have been claimed by copiers, + /// but not necessarily copied. + pub claim: AtomicUsize, + + /// The status of the resize. + pub status: AtomicU8, + + /// A thread parker for blocking on copy operations. + pub parker: Parker, + + /// Entries whose retirement has been deferred by later tables. + pub deferred: Stack<*mut T>, +} + +impl Default for State { + fn default() -> State { + State { + next: AtomicPtr::new(ptr::null_mut()), + allocating: Mutex::new(()), + copied: AtomicUsize::new(0), + claim: AtomicUsize::new(0), + status: AtomicU8::new(State::PENDING), + parker: Parker::default(), + deferred: Stack::new(), + } + } +} + +impl State<()> { + /// A resize is in-progress. + pub const PENDING: u8 = 0; + + /// The resize has been aborted, continue to the next table. + pub const ABORTED: u8 = 1; + + /// The resize was complete and the table was promoted. + pub const PROMOTED: u8 = 2; +} + +// The raw result of an insert operation. +pub enum InsertResult { + /// Inserted the entry. + Inserted, + + /// Replaced the given entry. + Replaced(*mut T), + + /// Did not insert due to the given existing entry. + Error(*mut T), +} + +// An entry in the hash-table. +pub struct Entry; + +impl Entry { + /// The entry is being copied to the new table, no updates are allowed on the old table. + /// + /// This bit is put down to initiate a copy, forcing all writers to complete the resize + /// before making progress. + const COPYING: usize = 0b001; + + /// The entry has been copied to the new table. + /// + /// This bit is put down after a copy completes. Both readers and writers must go to + /// the new table to see the new state of the entry. + /// + /// In blocking mode this is unused. + const COPIED: usize = 0b010; + + /// The entry was copied from a previous table. + /// + /// This bit indicates that an entry may still be accessible from previous tables + /// because the resize is still in progress, and so it is unsafe to reclaim. + /// + /// In blocking mode this is unused. + const BORROWED: usize = 0b100; +} + +impl Unpack for Entry { + /// Mask for an entry pointer, ignoring any tag bits. + const MASK: usize = !(Entry::COPYING | Entry::COPIED | Entry::BORROWED); +} + +impl Entry { + /// A sentinel pointer for a deleted entry. + /// + /// Null pointers are never copied to the new table, so this state is safe to use. + /// Note that tombstone entries may still be marked as `COPYING`, so this state + /// cannot be used for direct equality. + const TOMBSTONE: *mut () = Entry::COPIED as _; + + unsafe fn reclaim>(entry: *mut T, _collector: &Collector) { + unsafe { D::dealloc(entry) }; + } +} + +/// The status of an entry. +enum EntryStatus { + /// The entry is a tombstone or null (potentially a null copy). + Null, + + /// The entry is being copied. + Copied(Tagged), + + /// A valid entry. + Value(Tagged), +} + +impl From> for EntryStatus { + /// Returns the status for this entry. + #[inline] + fn from(entry: Tagged) -> Self { + if entry.ptr.is_null() { + EntryStatus::Null + } else if entry.tag() & Entry::COPYING != 0 { + EntryStatus::Copied(entry) + } else { + EntryStatus::Value(entry) + } + } +} + +/// The state of an entry we attempted to update. +enum UpdateStatus { + /// Successfully replaced the given key and value. + Replaced(Tagged), + + /// A new entry was written before we could update. + Found(EntryStatus), +} + +/// The state of an entry we attempted to insert into. +enum InsertStatus { + /// Successfully inserted the value. + Inserted, + + /// A new entry was written before we could update. + Found(EntryStatus), +} + +impl> HashTable { + /// Creates new hash-table with the given options. + #[inline] + pub fn new(capacity: usize, collector: Collector, resize: ResizeMode) -> HashTable { + // The table is lazily allocated. + if capacity == 0 { + return HashTable { + collector, + resize, + initial_capacity: 1, + table: AtomicPtr::new(ptr::null_mut()), + count: Counter::default(), + _dealloc: PhantomData, + }; + } + + // Initialize the table and mark it as the root. + let mut table = Table::alloc(probe::entries_for(capacity)); + *table.state_mut().status.get_mut() = State::PROMOTED; + + HashTable { + resize, + collector, + initial_capacity: capacity, + table: AtomicPtr::new(table.raw), + count: Counter::default(), + _dealloc: PhantomData, + } + } + + /// Returns a guard for this collector + pub fn guard(&self) -> MapGuard> { + // Safety: Created the guard from our collector. + unsafe { MapGuard::new(self.collector().enter()) } + } + + /// Returns an owned guard for this collector + pub fn owned_guard(&self) -> MapGuard> { + // Safety: Created the guard from our collector. + unsafe { MapGuard::new(self.collector().enter_owned()) } + } + + /// Verify a guard is valid to use with this map. + #[inline] + pub fn verify<'g, G>(&self, guard: &'g G) -> &'g MapGuard + where + G: seize::Guard, + { + assert_eq!( + *guard.collector(), + self.collector, + "Attempted to access map with incorrect guard" + ); + + // Safety: Verified the guard above. + unsafe { MapGuard::from_ref(guard) } + } + + /// Returns a reference to the root hash-table. + #[inline] + fn root(&self, guard: &impl VerifiedGuard) -> Table { + // Load the root table. + let raw = guard.protect(&self.table, Ordering::Acquire); + + // Safety: The root table is either null or a valid table allocation. + unsafe { Table::from_raw(raw) } + } + + /// Returns a reference to the collector. + #[inline] + pub fn collector(&self) -> &Collector { + &self.collector + } + + /// Returns the number of entries in the table. + #[inline] + pub fn len(&self) -> usize { + self.count.sum() + } + + /// Returns `true` if the table is empty. Otherwise returns `false`. + #[inline] + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + /// Returns true if incremental resizing is enabled. + #[inline] + fn is_incremental(&self) -> bool { + matches!(self.resize, ResizeMode::Incremental(_)) + } +} + +impl> HashTable { + /// Returns a reference to the entry corresponding to the key. + #[inline] + pub fn find( + &self, + hash: u64, + mut eq: impl FnMut(*mut T) -> bool, + guard: &impl VerifiedGuard, + ) -> Option<*mut T> { + // Load the root table. + let mut table = self.root(guard); + + // The table has not been initialized yet. + if table.raw.is_null() { + return None; + } + + let (h1, h2) = meta::split(hash); + + loop { + // Initialize the probe state. + let mut probe = Probe::start(h1, table.mask); + + // Probe until we reach the limit. + 'probe: while probe.len <= table.limit { + // Load the entry metadata first for cheap searches. + // + // Safety: `probe.i` is always in-bounds for the table length. + let meta = unsafe { table.meta(probe.i) }.load(Ordering::Acquire); + + if meta == h2 { + // Load the full entry. + // + // Safety: `probe.i` is always in-bounds for the table length. + let entry = guard + .protect(unsafe { table.entry(probe.i) }, Ordering::Acquire) + .unpack::(); + + // The entry was deleted, keep probing. + if entry.ptr.is_null() { + probe.next(table.mask); + continue 'probe; + } + + // Check for a full match. + // + // Safety: We performed a protected load of the pointer using a verified guard with + // `Acquire` and ensured that it is non-null, meaning it is valid for reads as long + // as we hold the guard. + if eq(entry.ptr) { + // The entry was copied to the new table. + // + // In blocking resize mode we do not need to perform self check as all writes block + // until any resizes are complete, making the root table the source of truth for readers. + if entry.tag() & Entry::COPIED != 0 { + break 'probe; + } + + return Some(entry.ptr); + } + } + + // The key is not in the table. + // + // It also cannot be in the next table because we have not went over the probe limit. + if meta == meta::EMPTY { + return None; + } + + probe.next(table.mask); + } + + // In incremental resize mode, we have to check the next table if we found + // a copied entry or went over the probe limit. + if self.is_incremental() { + if let Some(next) = table.next_table() { + table = next; + continue; + } + } + + // Otherwise, the key is not in the table. + return None; + } + } + + /// Inserts an entry into the map. + #[inline] + pub fn insert( + &self, + hash: u64, + entry: *mut T, + mut eq: impl FnMut(*mut T) -> bool, + hasher: impl Fn(*mut T) -> u64, + should_replace: bool, + guard: &impl VerifiedGuard, + ) -> InsertResult { + // Allocate the entry to be inserted. + let new_entry = untagged::<_, Entry>(entry); + + // Load the root table. + let mut table = self.root(guard); + + // Allocate the table if it has not been initialized yet. + if table.raw.is_null() { + table = self.init(None); + } + + let (h1, h2) = meta::split(hash); + + let mut help_copy = true; + loop { + // Initialize the probe state. + let mut probe = Probe::start(h1, table.mask); + + // Probe until we reach the limit. + let copying = 'probe: loop { + if probe.len > table.limit { + break None; + } + + // Load the entry metadata first for cheap searches. + // + // Safety: `probe.i` is always in-bounds for the table length. + let meta = unsafe { table.meta(probe.i) }.load(Ordering::Acquire); + + // The entry is empty, try to insert. + let entry = if meta == meta::EMPTY { + // Perform the insertion. + // + // Safety: `probe.i` is always in-bounds for the table length. Additionally, + // `new_entry` was allocated above and never shared. + match unsafe { + self.insert_at(probe.i, h2, new_entry.raw, table, &hasher, guard) + } { + // Successfully inserted. + InsertStatus::Inserted => { + // Increment the table length. + self.count.get(guard).fetch_add(1, Ordering::Relaxed); + return InsertResult::Inserted; + } + + // Lost to a concurrent insert. + // + // If the key matches, we might be able to update the value. + InsertStatus::Found(EntryStatus::Value(found)) + | InsertStatus::Found(EntryStatus::Copied(found)) => found, + + // Otherwise, continue probing. + InsertStatus::Found(EntryStatus::Null) => { + probe.next(table.mask); + continue 'probe; + } + } + } + // Found a potential match. + else if meta == h2 { + // Load the full entry. + // + // Safety: `probe.i` is always in-bounds for the table length. + let entry = guard + .protect(unsafe { table.entry(probe.i) }, Ordering::Acquire) + .unpack(); + + // The entry was deleted, keep probing. + if entry.ptr.is_null() { + probe.next(table.mask); + continue 'probe; + } + + // If the key matches, we might be able to update the value. + entry + } + // Otherwise, continue probing. + else { + probe.next(table.mask); + continue 'probe; + }; + + // Check for a full match. + // + // Safety: We performed a protected load of the pointer using a verified guard with + // `Acquire` and ensured that it is non-null, meaning it is valid for reads as long + // as we hold the guard. + if !eq(entry.ptr) { + probe.next(table.mask); + continue 'probe; + } + + // The entry is being copied to the new table. + if entry.tag() & Entry::COPYING != 0 { + break 'probe Some(probe.i); + } + + // Return an error for calls to `try_insert`. + if !should_replace { + return InsertResult::Error(entry.ptr); + } + + // Try to update the value. + // + // Safety: + // - `probe.i` is always in-bounds for the table length + // - `entry` is a valid non-null entry that was inserted into the map. + match unsafe { self.insert_slow(probe.i, entry, new_entry.raw, table, guard) } { + // Successfully performed the update. + UpdateStatus::Replaced(entry) => { + // Safety: `entry` is a valid non-null entry that we found in the map + // before replacing it. + return InsertResult::Replaced(entry.ptr); + } + + // The entry is being copied. + UpdateStatus::Found(EntryStatus::Copied(_)) => break 'probe Some(probe.i), + + // The entry was deleted before we could update it, continue probing. + UpdateStatus::Found(EntryStatus::Null) => { + probe.next(table.mask); + continue 'probe; + } + + UpdateStatus::Found(EntryStatus::Value(_)) => {} + } + }; + + // Prepare to retry in the next table. + table = self.prepare_retry_insert(copying, &mut help_copy, table, &hasher, guard); + } + } + + /// The slow-path for `insert`, updating the value. + /// + /// The returned pointer is guaranteed to be non-null and valid for reads. + /// + /// # Safety + /// + /// The safety requirements of `HashMap::update_at` apply. + #[cold] + #[inline(never)] + unsafe fn insert_slow( + &self, + i: usize, + mut entry: Tagged, + new_entry: *mut T, + table: Table, + guard: &impl VerifiedGuard, + ) -> UpdateStatus { + loop { + // Try to update the value. + // + // Safety: Guaranteed by caller. + match unsafe { self.update_at(i, entry, new_entry, table, guard) } { + // Someone else beat us to the update, retry. + // + // Note that the pointer we find here is a non-null entry that was inserted + // into the map. + UpdateStatus::Found(EntryStatus::Value(found)) => entry = found, + + status => return status, + } + } + } + + /// Prepare to retry an insert operation in the next table. + #[cold] + #[inline(never)] + fn prepare_retry_insert( + &self, + copying: Option, + help_copy: &mut bool, + table: Table, + hasher: impl Fn(*mut T) -> u64, + guard: &impl VerifiedGuard, + ) -> Table { + // If went over the probe limit or found a copied entry, trigger a resize. + let mut next_table = self.get_or_alloc_next(None, table); + + let next_table = match self.resize { + // In blocking mode we must complete the resize before proceeding. + ResizeMode::Blocking => self.help_copy(true, &table, hasher, guard), + + // In incremental mode we can perform more granular blocking. + ResizeMode::Incremental(_) => { + // Help out with the copy. + if *help_copy { + next_table = self.help_copy(false, &table, hasher, guard); + } + + // The entry we want to update is being copied. + if let Some(i) = copying { + // Wait for the entry to be copied. + // + // We could race with the copy to insert into the table. However, + // this entire code path is very rare and likely to complete quickly, + // so blocking allows us to make copies faster. + self.wait_copied(i, &table); + } + + next_table + } + }; + + // Limit incremental copying to once per operation, for more consistent latency. + *help_copy = false; + + // Continue in the new table. + next_table + } + + /// Removes a key from the map, returning the entry for the key if the key was previously in the map. + #[inline] + pub fn remove( + &self, + hash: u64, + eq: impl FnMut(*mut T) -> bool, + hasher: impl Fn(*mut T) -> u64, + guard: &impl VerifiedGuard, + ) -> Option<*mut T> { + #[inline(always)] + fn should_remove(_entry: *mut T) -> bool { + true + } + + // Safety: `should_remove` unconditionally returns `true`. + unsafe { + self.remove_if(hash, eq, should_remove::, hasher, guard) + .unwrap_unchecked() + } + } + + /// Removes a key from the map, returning the entry for the key if the key was previously in the map + /// and the provided closure returns `true` + #[inline] + pub fn remove_if( + &self, + hash: u64, + mut eq: impl FnMut(*mut T) -> bool, + mut should_remove: impl FnMut(*mut T) -> bool, + hasher: impl Fn(*mut T) -> u64, + guard: &impl VerifiedGuard, + ) -> Result, *mut T> { + // Load the root table. + let mut table = self.root(guard); + + // The table has not been initialized yet. + if table.raw.is_null() { + return Ok(None); + } + + let (h1, h2) = meta::split(hash); + + let mut help_copy = true; + loop { + // Initialize the probe state. + let mut probe = Probe::start(h1, table.mask); + + // Probe until we reach the limit. + let copying = 'probe: loop { + if probe.len > table.limit { + break None; + } + + // Load the entry metadata first for cheap searches. + // + // Safety: `probe.i` is always in-bounds for the table length. + let meta = unsafe { table.meta(probe.i).load(Ordering::Acquire) }; + + // The key is not in the table. + // It also cannot be in the next table because we have not went over the probe limit. + if meta == meta::EMPTY { + return Ok(None); + } + + // Check for a potential match. + if meta != h2 { + probe.next(table.mask); + continue 'probe; + } + + // Load the full entry. + // + // Safety: `probe.i` is always in-bounds for the table length. + let mut entry = guard + .protect(unsafe { table.entry(probe.i) }, Ordering::Acquire) + .unpack(); + + // The entry was deleted, keep probing. + if entry.ptr.is_null() { + probe.next(table.mask); + continue 'probe; + } + + // Check for a full match. + // + // Safety: We performed a protected load of the pointer using a verified guard with + // `Acquire` and ensured that it is non-null, meaning it is valid for reads as long + // as we hold the guard. + if !eq(entry.ptr) { + probe.next(table.mask); + continue 'probe; + } + + // The entry is being copied to the new table, we have to complete the copy before + // we can remove it. + if entry.tag() & Entry::COPYING != 0 { + break 'probe Some(probe.i); + } + + loop { + // Ensure that the entry should be removed. + // + // Safety: `entry` is a valid, non-null, protected entry that we found in the map. + if !should_remove(entry.ptr) { + return Err(entry.ptr); + } + + // Safety: + // - `probe.i` is always in-bounds for the table length + // - `entry` is a valid non-null entry that we found in the map. + let status = unsafe { + self.update_at(probe.i, entry, Entry::TOMBSTONE.cast(), table, guard) + }; + + match status { + // Successfully removed the entry. + UpdateStatus::Replaced(entry) => { + // Mark the entry as a tombstone. + // + // Note that this might end up being overwritten by the metadata hash + // if the initial insertion is lagging behind, but we avoid the RMW + // and sacrifice reads in the extremely rare case. + // + // Safety: `probe.i` is always in-bounds for the table length. + unsafe { + table + .meta(probe.i) + .store(meta::TOMBSTONE, Ordering::Release) + }; + + // Decrement the table length. + self.count.get(guard).fetch_sub(1, Ordering::Relaxed); + + // Safety: `entry` is the non-null entry that we just replaced. + return Ok(Some(entry.ptr)); + } + + // The entry is being copied to the new table, we have to complete the copy + // before we can remove. + UpdateStatus::Found(EntryStatus::Copied(_)) => break 'probe Some(probe.i), + + // The entry was deleted. + // + // We know that at some point during our execution the key was not in the map. + UpdateStatus::Found(EntryStatus::Null) => return Ok(None), + + // Lost to a concurrent update, retry. + UpdateStatus::Found(EntryStatus::Value(found)) => entry = found, + } + } + }; + + // Prepare to retry in the next table. + table = match self.prepare_retry(copying, &mut help_copy, table, &hasher, guard) { + Some(table) => table, + + // The search was exhausted. + None => return Ok(None), + } + } + } + + /// Prepare to retry an operation on an existing key in the next table. + /// + /// Returns `None` if the recursive search has been exhausted. + #[cold] + fn prepare_retry( + &self, + copying: Option, + help_copy: &mut bool, + table: Table, + hasher: impl Fn(*mut T) -> u64, + guard: &impl VerifiedGuard, + ) -> Option> { + let next_table = match self.resize { + ResizeMode::Blocking => match copying { + // The entry we want to perform the operation on is being copied. + // + // In blocking mode we must complete the resize before proceeding. + Some(_) => self.help_copy(true, &table, hasher, guard), + + // If we went over the probe limit, the key is not in the map. + None => return None, + }, + + ResizeMode::Incremental(_) => { + // In incremental resize mode, we always have to check the next table. + let next_table = table.next_table()?; + + // Help out with the copy. + if *help_copy { + self.help_copy(false, &table, hasher, guard); + } + + if let Some(i) = copying { + // Wait for the entry to be copied. + // + // We could race with the copy to insert into the table. However, + // this entire code path is very rare and likely to complete quickly, + // so blocking allows us to make copies faster. + self.wait_copied(i, &table); + } + + next_table + } + }; + + // Limit incremental copying to once per operation, for more consistent latency. + *help_copy = false; + + // Continue in the new table. + Some(next_table) + } + + /// Attempts to insert an entry at the given index. + /// + /// In the case of an error, the returned pointer is guaranteed to be + /// protected and valid for reads as long as the guard is held. + /// + /// # Safety + /// + /// The index must be in-bounds for the table. Additionally, `new_entry` must be a + /// valid owned pointer to insert into the map. + #[inline] + unsafe fn insert_at( + &self, + i: usize, + meta: u8, + new_entry: *mut T, + table: Table, + hasher: impl Fn(*mut T) -> u64, + guard: &impl VerifiedGuard, + ) -> InsertStatus { + // Safety: The caller guarantees that `i` is in-bounds. + let entry = unsafe { table.entry(i) }; + let meta_entry = unsafe { table.meta(i) }; + + // Try to claim the empty entry. + let found = match guard.compare_exchange( + entry, + ptr::null_mut(), + new_entry, + Ordering::Release, + Ordering::Acquire, + ) { + // Successfully claimed the entry. + Ok(_) => { + // Update the metadata table. + meta_entry.store(meta, Ordering::Release); + + // Return the value we inserted. + return InsertStatus::Inserted; + } + + // Lost to a concurrent update. + Err(found) => found.unpack(), + }; + + let (meta, status) = match EntryStatus::from(found) { + EntryStatus::Value(_) | EntryStatus::Copied(_) => { + // An entry was inserted, we have to hash it to get the metadata. + // + // The logic is the same for copied entries here as we have to + // check if the key matches and continue the update in the new table. + // + // Safety: We performed a protected load of the pointer using a verified guard + // with `Acquire` and ensured that it is non-null, meaning it is valid for reads + // as long as we hold the guard. + let hash = hasher(found.ptr); + (meta::h2(hash), EntryStatus::Value(found)) + } + + // The entry was deleted or null copied. + EntryStatus::Null => (meta::TOMBSTONE, EntryStatus::Null), + }; + + // Ensure the meta table is updated to keep the probe chain alive for readers. + if meta_entry.load(Ordering::Relaxed) == meta::EMPTY { + meta_entry.store(meta, Ordering::Release); + } + + InsertStatus::Found(status) + } + + /// Attempts to replace the value of an existing entry at the given index. + /// + /// In the case of an error, the returned pointer is guaranteed to be + /// protected and valid for reads. + /// + /// # Safety + /// + /// - The index must be in-bounds for the table. + /// - `current` must be a valid non-null entry that was inserted into the map. + /// - `new_entry` must be a valid sentinel or owned pointer to insert into the map. + #[inline] + unsafe fn update_at( + &self, + i: usize, + current: Tagged, + new_entry: *mut T, + table: Table, + guard: &impl VerifiedGuard, + ) -> UpdateStatus { + // Safety: The caller guarantees that `i` is in-bounds. + let entry = unsafe { table.entry(i) }; + + // Try to perform the update. + let found = match guard.compare_exchange_weak( + entry, + current.raw, + new_entry, + Ordering::Release, + Ordering::Acquire, + ) { + // Successfully updated. + Ok(_) => unsafe { + // Safety: The caller guarantees that `current` is a valid non-null entry that was + // inserted into the map. Additionally, it is now unreachable from this table due + // to the CAS above. + self.defer_retire(current, &table, guard); + + return UpdateStatus::Replaced(current); + }, + + // Lost to a concurrent update. + Err(found) => found.unpack(), + }; + + UpdateStatus::Found(EntryStatus::from(found)) + } + + /// Reserve capacity for `additional` more elements. + #[inline] + pub fn reserve( + &self, + additional: usize, + hasher: impl Fn(*mut T) -> u64, + guard: &impl VerifiedGuard, + ) { + let mut table = self.root(guard); + + // The table has not yet been allocated, initialize it. + if table.raw.is_null() { + table = self.init(Some(probe::entries_for(additional))); + } + + loop { + let capacity = probe::entries_for(self.count.sum().checked_add(additional).unwrap()); + + // We have enough capacity. + if table.len() >= capacity { + return; + } + + // Race to allocate the new table. + self.get_or_alloc_next(Some(capacity), table); + + // Force the copy to complete. + // + // Note that this is not strictly necessary for a `reserve` operation. + table = self.help_copy(true, &table, &hasher, guard); + } + } + + /// Remove all entries from this table. + #[inline] + pub fn clear(&self, hasher: impl Fn(*mut T) -> u64, guard: &impl VerifiedGuard) { + // Load the root table. + let mut table = self.root(guard); + + // The table has not been initialized yet. + if table.raw.is_null() { + return; + } + + loop { + // Get a clean copy of the table to delete from. + table = self.linearize(table, &hasher, guard); + + // Note that this method is not implemented in terms of `retain(|_, _| true)` to avoid + // loading entry metadata, as there is no need to provide consistency with `get`. + let mut copying = false; + + 'probe: for i in 0..table.len() { + // Load the entry to delete. + // + // Safety: `i` is in bounds for the table length. + let mut entry = guard + .protect(unsafe { table.entry(i) }, Ordering::Acquire) + .unpack(); + + loop { + // The entry is empty or already deleted. + if entry.ptr.is_null() { + continue 'probe; + } + + // Found a non-empty entry being copied. + if entry.tag() & Entry::COPYING != 0 { + // Clear every entry in this table that we can, then deal with the copy. + copying = true; + continue 'probe; + } + + // Try to delete the entry. + // + // Safety: `i` is in bounds for the table length. + let result = unsafe { + table.entry(i).compare_exchange( + entry.raw, + Entry::TOMBSTONE.cast(), + Ordering::Release, + Ordering::Acquire, + ) + }; + + match result { + // Successfully deleted the entry. + Ok(_) => { + // Update the metadata table. + // + // Safety: `i` is in bounds for the table length. + unsafe { table.meta(i).store(meta::TOMBSTONE, Ordering::Release) }; + + // Decrement the table length. + self.count.get(guard).fetch_sub(1, Ordering::Relaxed); + + // Safety: The caller guarantees that `current` is a valid non-null entry that was + // inserted into the map. Additionally, it is now unreachable from this table due + // to the CAS above. + unsafe { self.defer_retire(entry, &table, guard) }; + continue 'probe; + } + + // Lost to a concurrent update, retry. + Err(found) => entry = found.unpack(), + } + } + } + + // We cleared every entry in this table. + if !copying { + break; + } + + // A resize prevented us from deleting all the entries in this table. + // + // Complete the resize and retry in the new table. + table = self.help_copy(true, &table, &hasher, guard); + } + } + + /// Retains only the elements specified by the predicate. + #[inline] + pub fn retain(&self, mut f: F, hasher: impl Fn(*mut T) -> u64, guard: &impl VerifiedGuard) + where + F: FnMut(*mut T) -> bool, + { + // Load the root table. + let mut table = self.root(guard); + + // The table has not been initialized yet. + if table.raw.is_null() { + return; + } + + loop { + // Get a clean copy of the table to delete from. + table = self.linearize(table, &hasher, guard); + + let mut copying = false; + 'probe: for i in 0..table.len() { + // Load the entry metadata first to ensure consistency with calls to `get` + // for entries that are retained. + // + // Safety: `i` is in bounds for the table length. + let meta = unsafe { table.meta(i) }.load(Ordering::Acquire); + + // The entry is empty or deleted. + if matches!(meta, meta::EMPTY | meta::TOMBSTONE) { + continue 'probe; + } + + // Load the entry to delete. + // + // Safety: `i` is in bounds for the table length. + let mut entry = guard + .protect(unsafe { table.entry(i) }, Ordering::Acquire) + .unpack(); + + loop { + // The entry is empty or already deleted. + if entry.ptr.is_null() { + continue 'probe; + } + + // Found a non-empty entry being copied. + if entry.tag() & Entry::COPYING != 0 { + // Clear every entry in this table that we can, then deal with the copy. + copying = true; + continue 'probe; + } + + // Should we retain this entry? + // + // Safety: We performed a protected load of the pointer using a verified guard with + // `Acquire` and ensured that it is non-null, meaning it is valid for reads as long + // as we hold the guard. + if f(entry.ptr) { + continue 'probe; + } + + // Try to delete the entry. + // + // Safety: `i` is in bounds for the table length. + let result = unsafe { + table.entry(i).compare_exchange( + entry.raw, + Entry::TOMBSTONE.cast(), + Ordering::Release, + Ordering::Acquire, + ) + }; + + match result { + // Successfully deleted the entry. + Ok(_) => { + // Update the metadata table. + // + // Safety: `i` is in bounds for the table length. + unsafe { table.meta(i).store(meta::TOMBSTONE, Ordering::Release) }; + + // Decrement the table length. + self.count.get(guard).fetch_sub(1, Ordering::Relaxed); + + // Safety: The caller guarantees that `current` is a valid non-null entry that was + // inserted into the map. Additionally, it is now unreachable from this table due + // to the CAS above. + unsafe { self.defer_retire(entry, &table, guard) }; + continue 'probe; + } + + // Lost to a concurrent update, retry. + Err(found) => entry = found.unpack(), + } + } + } + + // We cleared every entry in this table. + if !copying { + break; + } + + // A resize prevented us from deleting all the entries in this table. + // + // Complete the resize and retry in the new table. + table = self.help_copy(true, &table, &hasher, guard); + } + } + + /// Returns an iterator over the entries of this table. + #[inline] + pub fn iter<'g, G>(&self, hasher: impl Fn(*mut T) -> u64, guard: &'g G) -> Iter<'g, T, G> + where + G: VerifiedGuard, + { + // Load the root table. + let root = self.root(guard); + + // The table has not been initialized yet, return a dummy iterator. + if root.raw.is_null() { + return Iter { + i: 0, + guard, + table: root, + }; + } + + // Get a clean copy of the table to iterate over. + let table = self.linearize(root, &hasher, guard); + + Iter { i: 0, guard, table } + } + + /// Returns a mutable iterator over the entries of this table. + #[inline] + pub fn iter_mut(&mut self) -> IterMut { + // Safety: The root table is either null or a valid table allocation. + let table = unsafe { Table::from_raw(*self.table.get_mut()) }; + + // The table has not been initialized yet, return a dummy iterator. + if table.raw.is_null() { + return IterMut { i: 0, table }; + } + + IterMut { i: 0, table } + } +} + +impl> IntoIterator for HashTable { + type Item = *mut T; + type IntoIter = IntoIter; + + /// Returns an owned iterator over the entries of this table. + fn into_iter(self) -> Self::IntoIter { + let mut map = ManuallyDrop::new(self); + + // Safety: The root table is either null or a valid table allocation. + let table = unsafe { Table::from_raw(*map.table.get_mut()) }; + + // The table has not been initialized yet, return a dummy iterator. + if table.raw.is_null() { + return IntoIter { i: 0, map, table }; + } + + IntoIter { i: 0, map, table } + } +} + +// An operation to perform on given entry in a [`HashTable`]. +#[derive(Debug, PartialEq, Eq)] +pub enum Operation { + /// Insert the given value. + Insert(*mut T), + + /// Remove the entry from the map. + Remove, + + /// Abort the operation with the given value. + Abort(A), +} + +#[derive(Debug, PartialEq, Eq)] +pub enum Compute { + /// The given entry was inserted. + Inserted(*mut T), + + /// The entry was updated. + Updated { + /// The entry that was replaced. + old: *mut T, + + /// The entry that was inserted. + new: *mut T, + }, + + /// The given entry was removed. + Removed(*mut T), + + /// The operation was aborted with the given value. + Aborted(A), +} + +/// RMW operations. +impl> HashTable { + /// Update an entry with a CAS function. + /// + /// Note that `compute` closure is guaranteed to be called for a `None` input only once, allowing the + /// insertion of values that cannot be cloned or reconstructed. + #[inline] + pub fn compute( + &self, + hash: u64, + mut eq: impl FnMut(*mut T) -> bool, + mut compute: F, + hasher: impl Fn(*mut T) -> u64, + guard: &impl VerifiedGuard, + ) -> Compute + where + F: FnMut(Option<*mut T>) -> Operation, + { + // Load the root table. + let mut table = self.root(guard); + + // The table has not yet been allocated. + if table.raw.is_null() { + // Compute the value to insert. + match compute(None) { + Operation::Insert(_) => {} + Operation::Remove => panic!("Cannot remove `None` entry."), + Operation::Abort(value) => return Compute::Aborted(value), + } + + // Initialize the table. + table = self.init(None); + } + + let (h1, h2) = meta::split(hash); + let mut help_copy = false; + + loop { + // Initialize the probe state. + let mut probe = Probe::start(h1, table.mask); + + // Probe until we reach the limit. + let copying = 'probe: loop { + if probe.len > table.limit { + break 'probe None; + } + + // Load the entry metadata first for cheap searches. + // + // Safety: `probe.i` is always in-bounds for the table length. + let meta = unsafe { table.meta(probe.i) }.load(Ordering::Acquire); + + // The entry is empty. + let mut entry = if meta == meta::EMPTY { + // Compute the entry to insert. + let new_entry = match compute(None) { + Operation::Insert(new_entry) => new_entry, + Operation::Remove => panic!("Cannot remove `None` entry."), + Operation::Abort(value) => return Compute::Aborted(value), + }; + + // Attempt to insert. + // + // Safety: `probe.i` is always in-bounds for the table length.Additionally, + // `new_entry` was allocated above and never shared. + match unsafe { self.insert_at(probe.i, h2, new_entry, table, &hasher, guard) } { + // Successfully inserted. + InsertStatus::Inserted => { + // Increment the table length. + self.count.get(guard).fetch_add(1, Ordering::Relaxed); + + // Safety: `new_entry` was initialized above. + return Compute::Inserted(new_entry); + } + + // Lost to a concurrent insert. + // + // If the key matches, we might be able to update the value. + InsertStatus::Found(EntryStatus::Value(found)) + | InsertStatus::Found(EntryStatus::Copied(found)) => found, + + // The entry was removed or invalidated. + InsertStatus::Found(EntryStatus::Null) => { + // Continue probing. + probe.next(table.mask); + continue 'probe; + } + } + } + // Found a potential match. + else if meta == h2 { + // Load the full entry. + // + // Safety: `probe.i` is always in-bounds for the table length. + let found = guard + .protect(unsafe { table.entry(probe.i) }, Ordering::Acquire) + .unpack(); + + // The entry was deleted, keep probing. + if found.ptr.is_null() { + probe.next(table.mask); + continue 'probe; + } + + // If the key matches, we might be able to update the value. + found + } + // Otherwise, continue probing. + else { + probe.next(table.mask); + continue 'probe; + }; + + // Check for a full match. + // + // Safety: We performed a protected load of the pointer using a verified guard with + // `Acquire` and ensured that it is non-null, meaning it is valid for reads as long + // as we hold the guard. + if !eq(entry.ptr) { + probe.next(table.mask); + continue 'probe; + } + + // The entry is being copied to the new table. + if entry.tag() & Entry::COPYING != 0 { + break 'probe Some(probe.i); + } + + loop { + // Compute the value to insert. + // + // Safety: `entry` is valid for reads. + let failure = match compute(Some(entry.ptr)) { + // The operation was aborted. + Operation::Abort(value) => return Compute::Aborted(value), + + // Update the value. + Operation::Insert(new_entry) => { + // Try to perform the update. + // + // Safety: + // - `probe.i` is always in-bounds for the table length + // - `entry` is a valid non-null entry that we found in the map. + // - `new_entry` was initialized above and never shared. + let status = + unsafe { self.update_at(probe.i, entry, new_entry, table, guard) }; + + match status { + // Successfully updated. + UpdateStatus::Replaced(entry) => { + return Compute::Updated { + // Safety: `entry` is a valid non-null entry that we found in the map + // before replacing it. + old: entry.ptr, + // Safety: `new_entry` was initialized above. + new: new_entry, + }; + } + + // The update failed. + failure => failure, + } + } + + // Remove the key from the map. + Operation::Remove => { + // Try to perform the removal. + // + // Safety: + // - `probe.i` is always in-bounds for the table length + // - `entry` is a valid non-null entry that we found in the map. + let status = unsafe { + self.update_at( + probe.i, + entry, + Entry::TOMBSTONE.cast(), + table, + guard, + ) + }; + + match status { + // Successfully removed the entry. + UpdateStatus::Replaced(entry) => { + // Mark the entry as a tombstone. + // + // Note that this might end up being overwritten by the metadata hash + // if the initial insertion is lagging behind, but we avoid the RMW + // and sacrifice reads in the extremely rare case. + unsafe { + table + .meta(probe.i) + .store(meta::TOMBSTONE, Ordering::Release) + }; + + // Decrement the table length. + self.count.get(guard).fetch_sub(1, Ordering::Relaxed); + + // Safety: `entry` is a valid non-null entry that we found in the map + // before replacing it. + return Compute::Removed(entry.ptr); + } + + // The remove failed. + failure => failure, + } + } + }; + + match failure { + // The entry is being copied to the new table. + UpdateStatus::Found(EntryStatus::Copied(_)) => break 'probe Some(probe.i), + + // The entry was deleted before we could update it. + // + // We know that at some point during our execution the key was not in the map. + UpdateStatus::Found(EntryStatus::Null) => { + // Compute the next operation. + match compute(None) { + Operation::Insert(_) => { + // Continue probing to find an empty slot. + probe.next(table.mask); + continue 'probe; + } + Operation::Remove => panic!("Cannot remove `None` entry."), + Operation::Abort(value) => return Compute::Aborted(value), + } + } + + // Someone else beat us to the update, retry. + UpdateStatus::Found(EntryStatus::Value(found)) => entry = found, + + _ => unreachable!(), + } + } + }; + + // Prepare to retry in the next table. + if let Some(next_table) = + self.prepare_retry(copying, &mut help_copy, table, &hasher, guard) + { + table = next_table; + continue; + } + + // Otherwise, we have exhausted our search, and the key is not in the map. + // + // Check if we should attempt to resize and insert into a new table. + match compute(None) { + // Need to insert into the new table. + Operation::Insert(_) => { + table = self.prepare_retry_insert(None, &mut help_copy, table, &hasher, guard); + } + // The operation was aborted. + Operation::Abort(value) => return Compute::Aborted(value), + Operation::Remove => panic!("Cannot remove `None` entry."), + } + } + } +} + +/// Resize operations. +impl> HashTable { + /// Allocate the initial table. + #[cold] + #[inline(never)] + fn init(&self, capacity: Option) -> Table { + const CAPACITY: usize = 32; + + // Allocate the table and mark it as the root. + let mut new = Table::alloc(capacity.unwrap_or(CAPACITY)); + *new.state_mut().status.get_mut() = State::PROMOTED; + + // Race to write the initial table. + match self.table.compare_exchange( + ptr::null_mut(), + new.raw, + Ordering::Release, + Ordering::Acquire, + ) { + // Successfully initialized the table. + Ok(_) => new, + + // Someone beat us, deallocate our table and use the table that was written. + Err(found) => { + // Safety: We allocated the table above and never shared it. + unsafe { Table::dealloc(new) } + + // Safety: The table was just initialized. + unsafe { Table::from_raw(found) } + } + } + } + + /// Returns the next table, allocating it has not already been created. + #[cold] + #[inline(never)] + fn get_or_alloc_next(&self, capacity: Option, table: Table) -> Table { + // Avoid spinning in tests, which can hide race conditions. + const SPIN_ALLOC: usize = if cfg!(any(test, debug_assertions)) { + 1 + } else { + 7 + }; + + // The next table is already allocated. + if let Some(next) = table.next_table() { + return next; + } + + let state = table.state(); + + // Otherwise, try to acquire the allocation lock. + // + // Unlike in `init`, we do not race here to prevent unnecessary allocator pressure. + let _allocating = match state.allocating.try_lock() { + Ok(lock) => lock, + // Someone else is currently allocating. + Err(_) => { + let mut spun = 0; + + // Spin for a bit, waiting for the table to be initialized. + while spun <= SPIN_ALLOC { + for _ in 0..(spun * spun) { + hint::spin_loop(); + } + + // The table was initialized. + if let Some(next) = table.next_table() { + return next; + } + + spun += 1; + } + + // Otherwise, we have to block. + state.allocating.lock().unwrap() + } + }; + + // The table was allocated while we were waiting for the lock. + if let Some(next) = table.next_table() { + return next; + } + + let current_capacity = table.len(); + + // Loading the length here is quite expensive, we may want to consider + // a probabilistic counter to detect high-deletion workloads. + let active_entries = self.len(); + + let next_capacity = match cfg!(papaya_stress) { + // Never grow the table to stress the incremental resizing algorithm. + true => current_capacity, + + // Double the table capacity if we are at least 50% full. + false if active_entries >= (current_capacity >> 1) => current_capacity << 1, + + // Halve the table if we are at most 12.5% full. + // + // This heuristic is intentionally pessimistic as unnecessarily shrinking + // is an expensive operation, but it may change in the future. We also respect + // the initial capacity to give the user a way to retain a strict minimum table + // size. + false if active_entries <= (current_capacity >> 3) => { + self.initial_capacity.max(current_capacity >> 1) + } + + // Otherwise keep the capacity the same. + // + // This can occur due to poor hash distribution or frequent cycling of + // insertions and deletions, in which case we want to avoid continuously + // growing the table. + false => current_capacity, + }; + + let next_capacity = capacity.unwrap_or(next_capacity); + assert!( + next_capacity <= isize::MAX as usize, + "`HashMap` exceeded maximum capacity" + ); + + // Allocate the new table while holding the lock. + let next = Table::alloc(next_capacity); + state.next.store(next.raw, Ordering::Release); + drop(_allocating); + + next + } + + /// Help along with an existing resize operation, returning the new root table. + /// + /// If `copy_all` is `false` in incremental resize mode, this returns the current reference's next + /// table, not necessarily the new root. + #[cold] + #[inline(never)] + fn help_copy( + &self, + copy_all: bool, + table: &Table, + hasher: impl Fn(*mut T) -> u64, + guard: &impl VerifiedGuard, + ) -> Table { + match self.resize { + ResizeMode::Blocking => self.help_copy_blocking(table, hasher, guard), + ResizeMode::Incremental(chunk) => { + let copied_to = self.help_copy_incremental(chunk, copy_all, hasher, guard); + + if !copy_all { + // If we weren't trying to linearize, we have to write to the next table + // even if the copy hasn't completed yet. + return table.next_table().unwrap(); + } + + copied_to + } + } + } + + /// Help along the resize operation until it completes and the next table is promoted. + /// + /// Should only be called on the root table. + fn help_copy_blocking( + &self, + table: &Table, + hasher: impl Fn(*mut T) -> u64, + guard: &impl VerifiedGuard, + ) -> Table { + // Load the next table. + let mut next = table.next_table().unwrap(); + + 'copy: loop { + // Make sure we are copying to the correct table. + while next.state().status.load(Ordering::Relaxed) == State::ABORTED { + next = self.get_or_alloc_next(None, next); + } + + // The copy already completed + if self.try_promote(table, &next, 0, guard) { + return next; + } + + let copy_chunk = table.len().min(4096); + + loop { + // Every entry has already been claimed. + if next.state().claim.load(Ordering::Relaxed) >= table.len() { + break; + } + + // Claim a chunk to copy. + let copy_start = next.state().claim.fetch_add(copy_chunk, Ordering::Relaxed); + + // Copy our chunk of entries. + let mut copied = 0; + for i in 0..copy_chunk { + let i = copy_start + i; + + if i >= table.len() { + break; + } + + // Copy the entry. + // + // Safety: We verified that `i` is in-bounds above. + if unsafe { !self.copy_at_blocking(i, table, &next, &hasher, guard) } { + // This table doesn't have space for the next entry. + // + // Abort the current resize. + // + // Note that the `SeqCst` is necessary to make the store visible + // to threads that are unparked. + next.state().status.store(State::ABORTED, Ordering::SeqCst); + + // Allocate the next table. + let allocated = self.get_or_alloc_next(None, next); + + // Wake anyone waiting for us to finish. + let state = table.state(); + state.parker.unpark(&state.status); + + // Retry in a new table. + next = allocated; + continue 'copy; + } + + copied += 1; + } + + // Are we done? + if self.try_promote(table, &next, copied, guard) { + return next; + } + + // If the resize was aborted while we were copying, continue in the new table. + if next.state().status.load(Ordering::Relaxed) == State::ABORTED { + continue 'copy; + } + } + + let state = next.state(); + // We copied all that we can, wait for the table to be promoted. + for spun in 0.. { + // Avoid spinning in tests, which can hide race conditions. + const SPIN_WAIT: usize = if cfg!(any(test, debug_assertions)) { + 1 + } else { + 7 + }; + + // Note that `Acquire` is necessary here to ensure we see the + // relevant modifications to the root table if see the updated + // state before parking. + // + // Otherwise, `Parker::park` will ensure the necessary synchronization + // when we are unparked. + let status = state.status.load(Ordering::Acquire); + + // If this copy was aborted, we have to retry in the new table. + if status == State::ABORTED { + continue 'copy; + } + + // The copy has completed. + if status == State::PROMOTED { + return next; + } + + // Copy chunks are relatively small and we expect to finish quickly, + // so spin for a bit before resorting to parking. + if spun <= SPIN_WAIT { + for _ in 0..(spun * spun) { + hint::spin_loop(); + } + + continue; + } + + // Park until the table is promoted. + state + .parker + .park(&state.status, |status| status == State::PENDING); + } + } + } + + /// Copy the entry at the given index to the new table. + /// + /// Returns `true` if the entry was copied into the table or `false` if the table was full. + /// + /// # Safety + /// + /// The index must be in-bounds for the table. + unsafe fn copy_at_blocking( + &self, + i: usize, + table: &Table, + next_table: &Table, + hasher: impl Fn(*mut T) -> u64, + guard: &impl VerifiedGuard, + ) -> bool { + // Mark the entry as copying. + // + // Safety: The caller guarantees that the index is in-bounds. + // + // Note that we don't need to protect the returned entry here, because + // no one is allowed to retire the entry once we put the `COPYING` bit + // down until it is inserted into the new table. + let entry = unsafe { table.entry(i) } + .fetch_or(Entry::COPYING, Ordering::AcqRel) + .unpack::(); + + // The entry is a tombstone. + if entry.raw.cast() == Entry::TOMBSTONE { + return true; + } + + // There is nothing to copy, we're done. + if entry.ptr.is_null() { + // Mark as a tombstone so readers avoid having to load the entry. + // + // Safety: The caller guarantees that the index is in-bounds. + unsafe { table.meta(i) }.store(meta::TOMBSTONE, Ordering::Release); + return true; + } + + // Copy the value to the new table. + // + // Safety: We marked the entry as `COPYING`, ensuring that any updates + // or removals wait until we complete the copy, and allowing us to get + // away without a protected load. Additionally, we verified that the + // entry is non-null, meaning that it is valid for reads. + unsafe { + self.insert_copy(entry.ptr.unpack(), false, next_table, hasher, guard) + .is_some() + } + } + + /// Help along an in-progress resize incrementally by copying a chunk of entries. + /// + /// Returns the table that was copied to. + fn help_copy_incremental( + &self, + chunk: usize, + block: bool, + hasher: impl Fn(*mut T) -> u64, + guard: &impl VerifiedGuard, + ) -> Table { + // Always help the highest priority root resize. + let table = self.root(guard); + + // Load the next table. + let Some(next) = table.next_table() else { + // The copy we tried to help was already promoted. + return table; + }; + + loop { + // The copy already completed. + if self.try_promote(&table, &next, 0, guard) { + return next; + } + + loop { + // Every entry has already been claimed. + if next.state().claim.load(Ordering::Relaxed) >= table.len() { + break; + } + + // Claim a chunk to copy. + let copy_start = next.state().claim.fetch_add(chunk, Ordering::Relaxed); + + // Copy our chunk of entries. + let mut copied = 0; + for i in 0..chunk { + let i = copy_start + i; + + if i >= table.len() { + break; + } + + // Copy the entry. + // + // Safety: We verified that `i` is in-bounds above. + unsafe { self.copy_at_incremental(i, &table, &next, &hasher, guard) }; + copied += 1; + } + + // Update the copy state, and try to promote the table. + // + // Only copy a single chunk if promotion fails, unless we are forced + // to complete the resize. + if self.try_promote(&table, &next, copied, guard) || !block { + return next; + } + } + + // There are no entries that we can copy, block if necessary. + if !block { + return next; + } + + let state = next.state(); + for spun in 0.. { + // Avoid spinning in tests, which can hide race conditions. + const SPIN_WAIT: usize = if cfg!(any(test, debug_assertions)) { + 1 + } else { + 7 + }; + + // The copy has completed. + // + // Note that `Acquire` is necessary here to ensure we see the + // relevant modifications to the root table if see the updated + // state before parking. + // + // Otherwise, `Parker::park` will ensure the necessary synchronization + // when we are unparked. + let status = state.status.load(Ordering::Acquire); + if status == State::PROMOTED { + return next; + } + + // Copy chunks are relatively small and we expect to finish quickly, + // so spin for a bit before resorting to parking. + if spun <= SPIN_WAIT { + for _ in 0..(spun * spun) { + hint::spin_loop(); + } + + continue; + } + + // Park until the table is promoted. + state + .parker + .park(&state.status, |status| status == State::PENDING); + } + } + } + + /// Copy the entry at the given index to the new table. + /// + /// # Safety + /// + /// The index must be in-bounds for the table. + unsafe fn copy_at_incremental( + &self, + i: usize, + table: &Table, + next_table: &Table, + hasher: impl Fn(*mut T) -> u64, + guard: &impl VerifiedGuard, + ) { + // Safety: The caller guarantees that the index is in-bounds. + let entry = unsafe { table.entry(i) }; + + // Mark the entry as copying. + let found = entry.fetch_or(Entry::COPYING, Ordering::AcqRel).unpack(); + + // The entry is a tombstone. + if found.raw.cast() == Entry::TOMBSTONE { + return; + } + + // There is nothing to copy, we're done. + if found.ptr.is_null() { + // Mark as a tombstone so readers avoid having to load the entry. + // + // Safety: The caller guarantees that the index is in-bounds. + unsafe { table.meta(i) }.store(meta::TOMBSTONE, Ordering::Release); + return; + } + + // Mark the entry as borrowed so writers in the new table know it was copied. + let new_entry = found.map_tag(|addr| addr | Entry::BORROWED); + + // Copy the value to the new table. + // + // Safety: We marked the entry as `COPYING`, ensuring that any updates + // or removals wait until we complete the copy, and allowing us to get + // away without a protected load. Additionally, we verified that the + // entry is non-null, meaning that it is valid for reads. + unsafe { + self.insert_copy(new_entry, true, next_table, hasher, guard) + .unwrap(); + } + + // Mark the entry as copied. + let copied = found + .raw + .map_addr(|addr| addr | Entry::COPYING | Entry::COPIED); + + // Note that we already wrote the COPYING bit, so no one is writing to the old + // entry except us. + // + // Note that the `SeqCst` is necessary to make the store visible to threads + // that are unparked. + entry.store(copied, Ordering::SeqCst); + + // Notify any writers that the copy has completed. + table.state().parker.unpark(entry); + } + + // Copy an entry into the table, returning the index it was inserted into. + // + // This is an optimized version of `insert_entry` where the caller is the only writer + // inserting the given key into the new table, as it has already been marked as copying. + // + // # Safety + // + // The new entry must be valid for reads. + unsafe fn insert_copy( + &self, + new_entry: Tagged, + resize: bool, + table: &Table, + hasher: impl Fn(*mut T) -> u64, + guard: &impl VerifiedGuard, + ) -> Option<(Table, usize)> { + let mut table = *table; + + // Safety: The new entry is guaranteed to be valid for reads. + let hash = hasher(new_entry.ptr); + + let (h1, h2) = meta::split(hash); + + loop { + // Initialize the probe state. + let mut probe = Probe::start(h1, table.mask); + + // Probe until we reach the limit. + while probe.len <= table.limit { + // Safety: `probe.i` is always in-bounds for the table length. + let meta_entry = unsafe { table.meta(probe.i) }; + + // Load the entry metadata first for cheap searches. + let meta = meta_entry.load(Ordering::Acquire); + + // The entry is empty, try to insert. + if meta == meta::EMPTY { + // Safety: `probe.i` is always in-bounds for the table length. + let entry = unsafe { table.entry(probe.i) }; + + // Try to claim the entry. + match guard.compare_exchange( + entry, + ptr::null_mut(), + new_entry.raw, + Ordering::Release, + Ordering::Acquire, + ) { + // Successfully inserted. + Ok(_) => { + // Update the metadata table. + meta_entry.store(h2, Ordering::Release); + return Some((table, probe.i)); + } + Err(found) => { + let found = found.unpack::(); + + // The entry was deleted or copied. + let meta = if found.ptr.is_null() { + meta::TOMBSTONE + } else { + // Ensure the meta table is updated to avoid breaking the probe chain. + // + // Safety: We performed a protected load of the pointer using a verified guard with + // `Acquire` and ensured that it is non-null, meaning it is valid for reads as long + // as we hold the guard. + let hash = hasher(found.ptr); + meta::h2(hash) + }; + + if meta_entry.load(Ordering::Relaxed) == meta::EMPTY { + meta_entry.store(meta, Ordering::Release); + } + } + } + } + + // Continue probing. + probe.next(table.mask); + } + + if !resize { + return None; + } + + // Insert into the next table. + table = self.get_or_alloc_next(None, table); + } + } + + // Update the copy state and attempt to promote a table to the root. + // + // Returns `true` if the table was promoted. + fn try_promote( + &self, + table: &Table, + next: &Table, + copied: usize, + guard: &impl VerifiedGuard, + ) -> bool { + let state = next.state(); + + // Update the copy count. + let copied = if copied > 0 { + state.copied.fetch_add(copied, Ordering::AcqRel) + copied + } else { + state.copied.load(Ordering::Acquire) + }; + + // If we copied all the entries in the table, we can try to promote. + if copied == table.len() { + let root = self.table.load(Ordering::Relaxed); + + // Only promote root copies. + // + // We can't promote a nested copy before it's parent has finished, as + // it may not contain all the entries in the table. + if table.raw == root { + // Try to update the root. + if self + .table + .compare_exchange(table.raw, next.raw, Ordering::Release, Ordering::Acquire) + .is_ok() + { + // Successfully promoted the table. + // + // Note that the `SeqCst` is necessary to make the store visible to threads + // that are unparked. + state.status.store(State::PROMOTED, Ordering::SeqCst); + + // Retire the old table. + // + // Safety: `table.raw` is a valid pointer to the table we just copied from. + // Additionally, the CAS above made the previous table unreachable from the + // root pointer, allowing it to be safely retired. + unsafe { + guard.defer_retire(table.raw, |table, collector| { + // Note that we do not drop entries because they have been copied to + // the new root. + drop_table::<_, D>(Table::from_raw(table), collector); + }); + } + } + + // Wake up any writers waiting for the resize to complete. + state.parker.unpark(&state.status); + return true; + } + } + + // Not ready to promote yet. + false + } + + // Completes all pending copies in incremental mode to get a clean copy of the table. + // + // This is necessary for operations like `iter` or `clear`, where entries in multiple tables + // can cause lead to incomplete results. + #[inline] + fn linearize( + &self, + mut table: Table, + hasher: impl Fn(*mut T) -> u64, + guard: &impl VerifiedGuard, + ) -> Table { + if self.is_incremental() { + // If we're in incremental resize mode, we need to complete any in-progress resizes to + // ensure we don't miss any entries in the next table. We can't iterate over both because + // we risk returning the same entry twice. + while table.next_table().is_some() { + table = self.help_copy(true, &table, &hasher, guard); + } + } + + table + } + + // Wait for an incremental copy of a given entry to complete. + #[cold] + #[inline(never)] + fn wait_copied(&self, i: usize, table: &Table) { + // Avoid spinning in tests, which can hide race conditions. + const SPIN_WAIT: usize = if cfg!(any(test, debug_assertions)) { + 1 + } else { + 5 + }; + + let entry = unsafe { table.entry(i) }; + + // Spin for a short while, waiting for the entry to be copied. + for spun in 0..SPIN_WAIT { + // The entry was copied. + let entry = entry.load(Ordering::Acquire).unpack::(); + if entry.tag() & Entry::COPIED != 0 { + return; + } + + for _ in 0..(spun * spun) { + hint::spin_loop(); + } + } + + // Park until the copy completes. + let parker = &table.state().parker; + parker.park(entry, |entry| entry.addr() & Entry::COPIED == 0); + } + + /// Retire an entry that was removed from the current table, but may still be reachable from + /// previous tables. + /// + /// # Safety + /// + /// The entry must be a valid pointer that is unreachable from the current table. Additionally, + /// it is *undefined behavior* to call this method multiple times for the same entry. + #[inline] + unsafe fn defer_retire( + &self, + entry: Tagged, + table: &Table, + guard: &impl VerifiedGuard, + ) { + match self.resize { + // Safety: In blocking resize mode, we only ever write to the root table, so the entry + // is inaccessible from all tables. + ResizeMode::Blocking => unsafe { + guard.defer_retire(entry.ptr, Entry::reclaim::); + }, + // In incremental resize mode, the entry may be accessible in previous tables. + ResizeMode::Incremental(_) => { + if entry.tag() & Entry::BORROWED == 0 { + // Safety: If the entry is not borrowed, meaning it is not in any previous tables, + // it is inaccessible even if the current table is not root. Thus we can safely retire. + unsafe { guard.defer_retire(entry.ptr, Entry::reclaim::) }; + return; + } + + let root = self.root(guard); + + // Check if our table, or any subsequent table, is the root. + let mut next = Some(*table); + while let Some(table) = next { + if table.raw == root.raw { + // Safety: The root table is our table or a table that succeeds ours. + // Thus any previous tables are unreachable from the root, so we can safely retire. + unsafe { guard.defer_retire(entry.ptr, Entry::reclaim::) }; + return; + } + + next = table.next_table(); + } + + // Otherwise, we have to wait for the table we are copying from to be reclaimed. + // + // Find the table we are copying from, searching from the root. + let mut prev = root; + + loop { + let next = prev.next_table().unwrap(); + + // Defer the entry to be retired by the table we are copying from. + if next.raw == table.raw { + prev.state().deferred.push(entry.ptr); + return; + } + + prev = next; + } + } + } + } +} + +// An iterator over the keys and values of this table. +pub struct Iter<'g, T, G> { + i: usize, + table: Table, + guard: &'g G, +} + +impl<'g, T, G> Iterator for Iter<'g, T, G> +where + G: VerifiedGuard, +{ + type Item = *mut T; + + #[inline] + fn next(&mut self) -> Option { + // The table has not yet been allocated. + if self.table.raw.is_null() { + return None; + } + + loop { + // Iterated over every entry in the table, we're done. + if self.i >= self.table.len() { + return None; + } + + // Load the entry metadata first to ensure consistency with calls to `get`. + // + // Safety: We verified that `self.i` is in-bounds above. + let meta = unsafe { self.table.meta(self.i) }.load(Ordering::Acquire); + + // The entry is empty or deleted. + if matches!(meta, meta::EMPTY | meta::TOMBSTONE) { + self.i += 1; + continue; + } + + // Load the entry. + // + // Safety: We verified that `self.i` is in-bounds above. + let entry = self + .guard + .protect(unsafe { self.table.entry(self.i) }, Ordering::Acquire) + .unpack::(); + + // The entry was deleted. + if entry.ptr.is_null() { + self.i += 1; + continue; + } + + self.i += 1; + + // Safety: We performed a protected load of the pointer using a verified guard with + // `Acquire` and ensured that it is non-null, meaning it is valid for reads as long + // as we hold the guard. + return Some(entry.ptr); + } + } +} + +// Safety: An iterator simply yields pointers to the keys and values. +// Dereferencing those pointers is up to the caller. Additionally, +// an iterator holds a shared reference to the guard, so the +// guard must be `Sync` for the iterator to be `Send` or `Sync`. +unsafe impl Send for Iter<'_, T, G> where G: Sync {} +unsafe impl Sync for Iter<'_, T, G> where G: Sync {} + +impl Clone for Iter<'_, T, G> { + #[inline] + fn clone(&self) -> Self { + Iter { + i: self.i, + table: self.table, + guard: self.guard, + } + } +} + +// A mutable iterator over the keys and values of this table. +pub struct IterMut { + i: usize, + table: Table, +} + +impl Iterator for IterMut { + type Item = *mut T; + + #[inline] + fn next(&mut self) -> Option { + self.next_raw() + } +} + +impl IterMut { + // Note that this method is guaranteed to only yield entry pointers that are valid for reads + // and writes. + #[inline] + fn next_raw(&mut self) -> Option<*mut T> { + // The table has not yet been allocated. + if self.table.raw.is_null() { + return None; + } + + loop { + // Iterated over every entry in the table, proceed to a nested resize if there is one. + if self.i >= self.table.len() { + if let Some(next_table) = self.table.next_table() { + self.i = 0; + self.table = next_table; + continue; + } + + // Otherwise, we're done. + return None; + } + + // Load the entry. + // + // Safety: We verified that `self.i` is in-bounds above. + let entry = unsafe { self.table.entry_mut(self.i) }.unpack::(); + + // The entry was deleted. + if entry.ptr.is_null() { + self.i += 1; + continue; + } + + // The entry was copied, we'll yield it when iterating over the table it was copied to. + // + // We have a mutable reference to the table, so there are no concurrent removals that + // can lead to us yielding duplicate entries. + if entry.tag() & Entry::COPIED != 0 { + self.i += 1; + continue; + } + + self.i += 1; + + // Safety: We ensured the entry is non-null. + return Some(entry.ptr); + } + } + + // Returns an immutable iterator over the remaining entries. + pub(crate) fn iter(&self) -> impl Iterator + '_ { + // Note that we still hold on to the mutable reference to the table, + // so we can iterate without synchronization. + IterMut { + i: self.i, + table: self.table, + } + } +} + +impl Clone for IterMut { + #[inline] + fn clone(&self) -> Self { + IterMut { + i: self.i, + table: self.table, + } + } +} + +// Safety: A mutable iterator does not perform any concurrent access and +// simply yields pointers to the keys and values. Dereferencing those +// pointers is up to the caller. +unsafe impl Send for IterMut {} +unsafe impl Sync for IterMut {} + +// An owned iterator over the keys and values of this table. +pub struct IntoIter> { + i: usize, + table: Table, + map: ManuallyDrop>, +} + +impl> Iterator for IntoIter { + type Item = *mut T; + + #[inline] + fn next(&mut self) -> Option { + // The table has not yet been allocated. + if self.table.raw.is_null() { + return None; + } + + loop { + // Iterated over every entry in the table, proceed to a nested resize if there is one. + if self.i >= self.table.len() { + if let Some(next_table) = self.table.next_table() { + // Drop the previous table. + // + // Safety: We have unique access to the table and do + // not access it after this call. + unsafe { drop_table::<_, D>(self.table, &self.map.collector) }; + + // Reset the iterator for the next table. + self.i = 0; + self.table = next_table; + continue; + } + + // Otherwise, we're done. + return None; + } + + // Load the entry. + // + // Safety: We verified that `self.i` is in-bounds above. + let entry = unsafe { self.table.entry_mut(self.i) }.unpack::(); + + // The entry was deleted. + if entry.ptr.is_null() { + self.i += 1; + continue; + } + + // The entry was copied, we'll yield it when iterating over the table it was copied to. + // + // We own the table, so there are no concurrent removals that can lead to us yielding + // duplicate entries. + if entry.tag() & Entry::COPIED != 0 { + self.i += 1; + continue; + } + + self.i += 1; + + // Safety: We ensured the entry is non-null. Additionally, we own the map + // and ensure not to drop already yielded entries. + return Some(entry.ptr); + } + } +} + +impl> IntoIter { + // Returns an immutable iterator over the remaining entries. + pub(crate) fn iter(&self) -> impl Iterator + '_ { + // `IntoIter` owns the `HashTable` and do not expose any access + // except through `&mut self` so we can use a mutable iterator. + IterMut { + i: self.i, + table: self.table, + } + } +} + +impl> Drop for IntoIter { + fn drop(&mut self) { + // Drop the remaining elements that have not been yielded. + drop_parts::<_, D>(self.i, self.table.raw, &mut self.map.collector); + } +} + +// Safety: An owned iterator does not perform any concurrent access and +// simply yields pointers to the keys and values. Dereferencing those +// pointers is up to the caller. +unsafe impl> Send for IntoIter {} +unsafe impl> Sync for IntoIter {} + +impl> Drop for HashTable { + fn drop(&mut self) { + drop_parts::<_, D>(0, *self.table.get_mut(), &mut self.collector); + } +} + +// Drop the elements of a `HashMap`. +fn drop_parts>( + mut start: usize, + mut raw: *mut RawTable, + collector: &mut Collector, +) { + // Make sure all objects are reclaimed before the collector is dropped. + // + // Dropping a table depends on accessing the collector for deferred retirement, + // using the shared collector pointer that is invalidated by drop. + // + // Safety: We have a unique reference to the collector. + unsafe { collector.reclaim_all() }; + + // Drop all nested tables and entries. + while !raw.is_null() { + // Safety: The root and next tables are always valid pointers to a + // table allocation, or null. + let mut table = unsafe { Table::from_raw(raw) }; + + // Read the next table pointer before dropping the current one. + let next = *table.state_mut().next.get_mut(); + + // Safety: We have unique access to the table and do + // not access the entries after this call. + unsafe { drop_entries::<_, D>(start, table) }; + + // Safety: We have unique access to the table and do + // not access it after this call. + unsafe { drop_table::<_, D>(table, collector) }; + + // Continue for all nested tables. + raw = next; + start = 0; + } +} + +// Drop all entries in this table. +// +// # Safety +// +// The table entries must not be accessed after this call. +unsafe fn drop_entries>(start: usize, table: Table) { + for i in start..table.len() { + // Safety: `i` is in-bounds and we have unique access to the table. + let entry = unsafe { (*table.entry(i).as_ptr()).unpack::() }; + + // The entry was copied, or there is nothing to deallocate. + if entry.ptr.is_null() || entry.tag() & Entry::COPYING != 0 { + continue; + } + + // Drop the entry. + // + // Safety: We verified that the table is non-null and will + // not be accessed after this call. Additionally, we ensured + // that the entry is not copied to avoid double freeing entries + // that may exist in multiple tables. + unsafe { D::dealloc(entry.ptr) }; + } +} + +// Drop the table allocation. +// +// # Safety +// +// The table must not be accessed after this call. +unsafe fn drop_table>(mut table: Table, collector: &Collector) { + // Drop any entries that were deferred during an incremental resize. + // + // Safety: Entries are deferred after they are made unreachable from the + // next table during a resize from this table. This table must have been accessible + // from the root for any entry to have been deferred. Thus it is being retired now, + // *after* the entry was made inaccessible from the next table. Additionally, for + // this table to have been retired, it also must no longer be accessible from the root, + // meaning that the entry has been totally removed from the map, and can be safely + // retired. + table + .state_mut() + .deferred + .drain(|entry| unsafe { collector.retire(entry, Entry::reclaim::) }); + + // Deallocate the table. + // + // Safety: The caller guarantees that the table will not be accessed after this call. + unsafe { Table::dealloc(table) }; +} + +// Entry metadata, inspired by `hashbrown`. +pub mod meta { + use std::mem; + + // Indicates an empty entry. + pub const EMPTY: u8 = 0x80; + + // Indicates an entry that has been deleted. + pub const TOMBSTONE: u8 = u8::MAX; + + // Returns the primary hash and secondary for an entry. + #[inline] + pub fn split(hash: u64) -> (usize, u8) { + (h1(hash), h2(hash)) + } + + // Returns the primary hash for an entry. + #[inline] + pub fn h1(hash: u64) -> usize { + hash as usize + } + + /// Return a byte of hash metadata, used for cheap searches. + #[inline] + pub fn h2(hash: u64) -> u8 { + const MIN_HASH_LEN: usize = if mem::size_of::() < mem::size_of::() { + mem::size_of::() + } else { + mem::size_of::() + }; + + // Grab the top 7 bits of the hash. + // + // While the hash is normally a full 64-bit value, some hash functions + // (such as fxhash) produce a usize result instead, which means that the + // top 32 bits are 0 on 32-bit platforms. + let top7 = hash >> (MIN_HASH_LEN * 8 - 7); + (top7 & 0x7f) as u8 + } +} diff --git a/src/raw/utils/counter.rs b/src/raw/utils/counter.rs index e9ab8de..5fa6d56 100644 --- a/src/raw/utils/counter.rs +++ b/src/raw/utils/counter.rs @@ -1,7 +1,6 @@ -use std::sync::{ - atomic::{AtomicIsize, Ordering}, - OnceLock, -}; +use std::fmt; +use std::sync::atomic::{AtomicIsize, Ordering}; +use std::sync::OnceLock; use super::CachePadded; @@ -59,3 +58,9 @@ impl Counter { .unwrap_or(0) } } + +impl fmt::Debug for Counter { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_tuple("Counter").field(&self.sum()).finish() + } +} diff --git a/src/raw/utils/tagged.rs b/src/raw/utils/tagged.rs index c92de4b..0b14638 100644 --- a/src/raw/utils/tagged.rs +++ b/src/raw/utils/tagged.rs @@ -1,5 +1,9 @@ -use std::mem::align_of; -use std::sync::atomic::{AtomicPtr, Ordering}; +use std::{ + marker::PhantomData, + sync::atomic::{AtomicPtr, Ordering}, +}; + +use crate::raw::table::Entry; // Polyfill for the unstable strict-provenance APIs. #[allow(clippy::missing_safety_doc)] @@ -7,19 +11,21 @@ use std::sync::atomic::{AtomicPtr, Ordering}; pub unsafe trait StrictProvenance: Sized { fn addr(self) -> usize; fn map_addr(self, f: impl FnOnce(usize) -> usize) -> Self; - fn unpack(self) -> Tagged + fn unpack(self) -> Tagged where - T: Unpack; + U: Unpack; + + // This constant, if used, will fail to compile if `T` doesn't have an alignment + // that guarantees all valid pointers have zero in the bits excluded by `T::MASK`. + // + // TODO: Make this generic over `T`. + const ASSERT_ALIGNMENT: () = assert!(align_of::() > !Entry::MASK); } // Unpack a tagged pointer. pub trait Unpack: Sized { // A mask for the pointer tag bits. const MASK: usize; - - // This constant, if used, will fail to compile if T doesn't have an alignment - // that guarantees all valid pointers have zero in the bits excluded by T::MASK. - const ASSERT_ALIGNMENT: () = assert!(align_of::() > !Self::MASK); } unsafe impl StrictProvenance for *mut T { @@ -34,44 +40,49 @@ unsafe impl StrictProvenance for *mut T { } #[inline(always)] - fn unpack(self) -> Tagged + fn unpack(self) -> Tagged where - T: Unpack, + U: Unpack, { - let () = T::ASSERT_ALIGNMENT; + let () = Self::ASSERT_ALIGNMENT; + Tagged { raw: self, - ptr: self.map_addr(|addr| addr & T::MASK), + ptr: self.map_addr(|addr| addr & Entry::MASK), + _unpack: PhantomData, } } } // An unpacked tagged pointer. -pub struct Tagged { +pub struct Tagged { // The raw tagged pointer. pub raw: *mut T, // The untagged pointer. pub ptr: *mut T, + + _unpack: PhantomData, } // Creates a `Tagged` from an untagged pointer. #[inline] -pub fn untagged(value: *mut T) -> Tagged { +pub fn untagged(value: *mut T) -> Tagged { Tagged { raw: value, ptr: value, + _unpack: PhantomData, } } -impl Tagged +impl Tagged where - T: Unpack, + U: Unpack, { // Returns the tag portion of this pointer. #[inline] pub fn tag(self) -> usize { - self.raw.addr() & !T::MASK + self.raw.addr() & !U::MASK } // Maps the tag of this pointer. @@ -80,13 +91,14 @@ where Tagged { raw: self.raw.map_addr(f), ptr: self.ptr, + _unpack: PhantomData, } } } -impl Copy for Tagged {} +impl Copy for Tagged {} -impl Clone for Tagged { +impl Clone for Tagged { fn clone(&self) -> Self { *self } diff --git a/src/set.rs b/src/set.rs index 367776f..62e0580 100644 --- a/src/set.rs +++ b/src/set.rs @@ -1,9 +1,7 @@ -use crate::raw::utils::MapGuard; -use crate::raw::{self, InsertResult}; -use crate::Equivalent; +use crate::{Equivalent, HashMap, HashMapRef}; use seize::{Collector, Guard, LocalGuard, OwnedGuard}; -use crate::map::ResizeMode; +use crate::map::{self, ResizeMode}; use std::collections::hash_map::RandomState; use std::fmt; use std::hash::{BuildHasher, Hash}; @@ -15,7 +13,7 @@ use std::marker::PhantomData; /// [`HashSet::guard`] or using the [`HashSet::pin`] API. See the [crate-level documentation](crate#usage) /// for details. pub struct HashSet { - raw: raw::HashMap, + map: HashMap, } // Safety: We only ever hand out &K through shared references to the map, @@ -120,7 +118,12 @@ impl HashSetBuilder { /// Construct a [`HashSet`] from the builder, using the configured options. pub fn build(self) -> HashSet { HashSet { - raw: raw::HashMap::new(self.capacity, self.hasher, self.collector, self.resize_mode), + map: HashMap::builder() + .capacity(self.capacity) + .hasher(self.hasher) + .collector(self.collector) + .resize_mode(self.resize_mode) + .build(), } } } @@ -244,12 +247,7 @@ impl HashSet { /// ``` pub fn with_capacity_and_hasher(capacity: usize, hash_builder: S) -> HashSet { HashSet { - raw: raw::HashMap::new( - capacity, - hash_builder, - Collector::default(), - ResizeMode::default(), - ), + map: HashMap::with_capacity_and_hasher(capacity, hash_builder), } } @@ -260,8 +258,8 @@ impl HashSet { #[inline] pub fn pin(&self) -> HashSetRef<'_, K, S, LocalGuard<'_>> { HashSetRef { - guard: self.raw.guard(), set: self, + map_ref: self.map.pin(), } } @@ -276,8 +274,8 @@ impl HashSet { #[inline] pub fn pin_owned(&self) -> HashSetRef<'_, K, S, OwnedGuard<'_>> { HashSetRef { - guard: self.raw.owned_guard(), set: self, + map_ref: self.map.pin_owned(), } } @@ -287,7 +285,7 @@ impl HashSet { /// See the [crate-level documentation](crate#usage) for details. #[inline] pub fn guard(&self) -> LocalGuard<'_> { - self.raw.collector().enter() + self.map.guard() } /// Returns an owned guard for use with this set. @@ -300,7 +298,7 @@ impl HashSet { /// See the [crate-level documentation](crate#usage) for details. #[inline] pub fn owned_guard(&self) -> OwnedGuard<'_> { - self.raw.collector().enter_owned() + self.map.owned_guard() } } @@ -324,7 +322,7 @@ where /// ``` #[inline] pub fn len(&self) -> usize { - self.raw.len() + self.map.len() } /// Returns `true` if the set is empty. Otherwise returns `false`. @@ -369,7 +367,7 @@ where where Q: Equivalent + Hash + ?Sized, { - self.get(key, self.raw.verify(guard)).is_some() + self.map.get(key, guard).is_some() } /// Returns a reference to the value corresponding to the key. @@ -396,7 +394,7 @@ where where Q: Equivalent + Hash + ?Sized, { - match self.raw.get(key, self.raw.verify(guard)) { + match self.map.get_key_value(key, guard) { Some((key, _)) => Some(key), None => None, } @@ -427,15 +425,10 @@ where /// ``` #[inline] pub fn insert(&self, key: K, guard: &impl Guard) -> bool { - match self.raw.insert(key, (), true, self.raw.verify(guard)) { - InsertResult::Inserted(_) => true, - InsertResult::Replaced(_) => false, - InsertResult::Error { .. } => unreachable!(), - } + self.map.insert(key, (), guard).is_none() } - /// Removes a key from the set, returning the value at the key if the key - /// was previously in the set. + /// Removes a value from the set. Returns whether the value was present in the set. /// /// The key may be any borrowed form of the set's key type, but /// [`Hash`] and [`Eq`] on the borrowed form *must* match those for @@ -456,10 +449,7 @@ where where Q: Equivalent + Hash + ?Sized, { - match self.raw.remove(key, self.raw.verify(guard)) { - Some((_, _)) => true, - None => false, - } + self.map.remove(key, guard).is_some() } /// Tries to reserve capacity for `additional` more elements to be inserted @@ -484,7 +474,7 @@ where /// ``` #[inline] pub fn reserve(&self, additional: usize, guard: &impl Guard) { - self.raw.reserve(additional, self.raw.verify(guard)) + self.map.reserve(additional, guard) } /// Clears the set, removing all values. @@ -506,7 +496,7 @@ where /// ``` #[inline] pub fn clear(&self, guard: &impl Guard) { - self.raw.clear(self.raw.verify(guard)) + self.map.clear(guard) } /// Retains only the elements specified by the predicate. @@ -537,7 +527,7 @@ where where F: FnMut(&K) -> bool, { - self.raw.retain(|k, _| f(k), self.raw.verify(guard)) + self.map.retain(|k, _| f(k), guard) } /// An iterator visiting all values in arbitrary order. @@ -566,7 +556,35 @@ where G: Guard, { Iter { - raw: self.raw.iter(self.raw.verify(guard)), + inner: self.map.iter(guard), + } + } +} + +impl IntoIterator for HashSet { + type Item = K; + type IntoIter = IntoIter; + + /// Creates a consuming iterator, that is, one that moves each value out + /// of the set in arbitrary order. The set cannot be used after calling + /// this. + /// + /// # Examples + /// + /// ``` + /// use papaya::HashSet; + /// + /// let set = HashSet::from([ + /// "a", + /// "b", + /// "c" + /// ]); + /// + /// let v: Vec<&str> = set.into_iter().collect(); + /// ``` + fn into_iter(self) -> Self::IntoIter { + IntoIter { + inner: self.map.into_iter(), } } } @@ -577,14 +595,7 @@ where S: BuildHasher, { fn eq(&self, other: &Self) -> bool { - if self.len() != other.len() { - return false; - } - - let (guard1, guard2) = (&self.guard(), &other.guard()); - - let mut iter = self.iter(guard1); - iter.all(|key| other.get(key, guard2).is_some()) + self.map.eq(&other.map) } } @@ -612,24 +623,8 @@ where S: BuildHasher, { fn extend>(&mut self, iter: T) { - // from `hashbrown::HashSet::extend`: - // Keys may be already present or show multiple times in the iterator. - // Reserve the entire hint lower bound if the set is empty. - // Otherwise reserve half the hint (rounded up), so the set - // will only resize twice in the worst case. - let iter = iter.into_iter(); - let reserve = if self.is_empty() { - iter.size_hint().0 - } else { - (iter.size_hint().0 + 1) / 2 - }; - - let guard = self.guard(); - self.reserve(reserve, &guard); - - for key in iter { - self.insert(key, &guard); - } + let mut map = &self.map; + map.extend(iter.into_iter().map(|key| (key, ()))); } } @@ -658,27 +653,8 @@ where S: BuildHasher + Default, { fn from_iter>(iter: T) -> Self { - let mut iter = iter.into_iter(); - - if let Some(key) = iter.next() { - let (lower, _) = iter.size_hint(); - let set = HashSet::with_capacity_and_hasher(lower.saturating_add(1), S::default()); - - // Ideally we could use an unprotected guard here. However, `insert` - // returns references to values that were replaced and retired, so - // we need a "real" guard. A `raw_insert` method that strictly returns - // pointers would fix this. - { - let set = set.pin(); - set.insert(key); - for key in iter { - set.insert(key); - } - } - - set - } else { - Self::default() + HashSet { + map: HashMap::from_iter(iter.into_iter().map(|key| (key, ()))), } } } @@ -689,20 +665,9 @@ where S: BuildHasher + Clone, { fn clone(&self) -> HashSet { - let other = HashSet::builder() - .capacity(self.len()) - .hasher(self.raw.hasher.clone()) - .collector(seize::Collector::new()) - .build(); - - { - let (guard1, guard2) = (&self.guard(), &other.guard()); - for key in self.iter(guard1) { - other.insert(key.clone(), guard2); - } + HashSet { + map: self.map.clone(), } - - other } } @@ -711,8 +676,8 @@ where /// This type is created with [`HashSet::pin`] and can be used to easily access a [`HashSet`] /// without explicitly managing a guard. See the [crate-level documentation](crate#usage) for details. pub struct HashSetRef<'set, K, S, G> { - guard: MapGuard, set: &'set HashSet, + map_ref: HashMapRef<'set, K, (), S, G>, } impl<'set, K, S, G> HashSetRef<'set, K, S, G> @@ -732,7 +697,7 @@ where /// See [`HashSet::len`] for details. #[inline] pub fn len(&self) -> usize { - self.set.raw.len() + self.map_ref.len() } /// Returns `true` if the set is empty. Otherwise returns `false`. @@ -740,7 +705,7 @@ where /// See [`HashSet::is_empty`] for details. #[inline] pub fn is_empty(&self) -> bool { - self.len() == 0 + self.map_ref.is_empty() } /// Returns `true` if the set contains a value for the specified key. @@ -751,7 +716,7 @@ where where Q: Equivalent + Hash + ?Sized, { - self.get(key).is_some() + self.map_ref.get(key).is_some() } /// Returns a reference to the value corresponding to the key. @@ -762,7 +727,7 @@ where where Q: Equivalent + Hash + ?Sized, { - match self.set.raw.get(key, &self.guard) { + match self.map_ref.get_key_value(key) { Some((k, _)) => Some(k), None => None, } @@ -773,11 +738,7 @@ where /// See [`HashSet::insert`] for details. #[inline] pub fn insert(&self, key: K) -> bool { - match self.set.raw.insert(key, (), true, &self.guard) { - InsertResult::Inserted(_) => true, - InsertResult::Replaced(_) => false, - InsertResult::Error { .. } => unreachable!(), - } + self.map_ref.insert(key, ()).is_none() } /// Removes a key from the set, returning the value at the key if the key @@ -789,10 +750,7 @@ where where Q: Equivalent + Hash + ?Sized, { - match self.set.raw.remove(key, &self.guard) { - Some((_, _)) => true, - None => false, - } + self.map_ref.remove(key).is_some() } /// Clears the set, removing all values. @@ -800,7 +758,7 @@ where /// See [`HashSet::clear`] for details. #[inline] pub fn clear(&self) { - self.set.raw.clear(&self.guard) + self.map_ref.clear() } /// Retains only the elements specified by the predicate. @@ -811,7 +769,7 @@ where where F: FnMut(&K) -> bool, { - self.set.raw.retain(|k, _| f(k), &self.guard) + self.map_ref.retain(|k, _| f(k)) } /// Tries to reserve capacity for `additional` more elements to be inserted @@ -820,7 +778,7 @@ where /// See [`HashSet::reserve`] for details. #[inline] pub fn reserve(&self, additional: usize) { - self.set.raw.reserve(additional, &self.guard) + self.map_ref.reserve(additional) } /// An iterator visiting all values in arbitrary order. @@ -830,7 +788,7 @@ where #[inline] pub fn iter(&self) -> Iter<'_, K, G> { Iter { - raw: self.set.raw.iter(&self.guard), + inner: self.map_ref.iter(), } } } @@ -864,7 +822,7 @@ where /// /// This struct is created by the [`iter`](HashSet::iter) method on [`HashSet`]. See its documentation for details. pub struct Iter<'g, K, G> { - raw: raw::Iter<'g, K, (), MapGuard>, + inner: map::Iter<'g, K, (), G>, } impl<'g, K: 'g, G> Iterator for Iter<'g, K, G> @@ -875,7 +833,15 @@ where #[inline] fn next(&mut self) -> Option { - self.raw.next().map(|(k, _)| k) + self.inner.next().map(|(k, _)| k) + } +} + +impl Clone for Iter<'_, K, G> { + fn clone(&self) -> Self { + Self { + inner: self.inner.clone(), + } } } @@ -885,10 +851,34 @@ where G: Guard, { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_list() - .entries(Iter { - raw: self.raw.clone(), - }) - .finish() + f.debug_list().entries(self.clone()).finish() + } +} + +/// An owned iterator over the entries of a `HashSet`. +/// +/// This `struct` is created by the [`into_iter`] method on [`HashSet`] +/// (provided by the [`IntoIterator`] trait). See its documentation for more. +/// +/// [`into_iter`]: IntoIterator::into_iter +pub struct IntoIter { + inner: map::IntoIter, +} + +impl Iterator for IntoIter { + type Item = K; + + #[inline] + fn next(&mut self) -> Option { + self.inner.next().map(|(k, _)| k) + } +} + +impl fmt::Debug for IntoIter +where + K: fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_list().entries(self.inner.raw.iter()).finish() } } diff --git a/tests/basic.rs b/tests/basic.rs index 5e48414..d704136 100644 --- a/tests/basic.rs +++ b/tests/basic.rs @@ -776,7 +776,7 @@ fn len() { let map = map(); let len = if cfg!(miri) { 100 } else { 10_000 }; for i in 0..len { - map.pin().insert(i, i + 1); + assert_eq!(map.pin().insert(i, i + 1), None); } assert_eq!(map.len(), len); }); @@ -802,6 +802,57 @@ fn iter() { }); } +#[test] +fn iter_mut() { + if cfg!(papaya_stress) { + return; + } + + with_map::(|map| { + let mut map = map(); + let len = if cfg!(miri) { 100 } else { 10_000 }; + for i in 0..len { + assert_eq!(map.pin().insert(i, i + 1), None); + } + + let v: Vec<_> = (0..len).map(|i| (i, i + 1)).collect(); + let mut got: Vec<_> = map.iter_mut().map(|(&k, &mut v)| (k, v)).collect(); + got.sort(); + assert_eq!(v, got); + }); +} + +#[test] +fn into_iter() { + if cfg!(papaya_stress) { + return; + } + + with_map::(|map| { + let map = map(); + let len = if cfg!(miri) { 100 } else { 10_000 }; + for i in 0..len { + assert_eq!(map.pin().insert(i, (i + 1).to_string()), None); + } + + let v: Vec<_> = (0..len).map(|i| (i, (i + 1).to_string())).collect(); + let mut got: Vec<_> = map.into_iter().collect(); + got.sort(); + assert_eq!(v, got); + }); + + with_map::(|map| { + let map = map(); + let len = if cfg!(miri) { 100 } else { 10_000 }; + for i in 0..len { + assert_eq!(map.pin().insert(i, (i + 1).to_string()), None); + } + + let got: Vec<_> = map.into_iter().take(len / 2).collect(); + assert_eq!(got.len(), len / 2); + }); +} + #[test] fn retain_empty() { with_map::(|map| {