diff --git a/Cargo.lock b/Cargo.lock index 874687ad..791ac476 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2,6 +2,16 @@ # It is not intended for manual editing. version = 3 +[[package]] +name = "atomic-wait" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a55b94919229f2c42292fd71ffa4b75e83193bffdd77b1e858cd55fd2d0b0ea8" +dependencies = [ + "libc", + "windows-sys 0.42.0", +] + [[package]] name = "cfg-if" version = "1.0.0" @@ -12,6 +22,7 @@ checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd" name = "dashmap" version = "6.1.0" dependencies = [ + "atomic-wait", "cfg-if", "seize", ] @@ -25,11 +36,26 @@ checksum = "b4668fb0ea861c1df094127ac5f1da3409a82116a4ba74fca2e58ef927159bb3" [[package]] name = "seize" version = "0.4.5" -source = "git+https://github.com/xacrimon/seize?rev=0f41a677bccb60bdd8f3fb66a9d3e0744a5a99e8#0f41a677bccb60bdd8f3fb66a9d3e0744a5a99e8" +source = "git+https://github.com/xacrimon/seize?rev=3525d40db2ea0deeee273f13700f1211d5c7d102#3525d40db2ea0deeee273f13700f1211d5c7d102" dependencies = [ "cfg-if", "libc", - "windows-sys", + "windows-sys 0.48.0", +] + +[[package]] +name = "windows-sys" +version = "0.42.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a3e1820f08b8513f676f7ab6c1f99ff312fb97b553d30ff4dd86f9f15728aa7" +dependencies = [ + "windows_aarch64_gnullvm 0.42.2", + "windows_aarch64_msvc 0.42.2", + "windows_i686_gnu 0.42.2", + "windows_i686_msvc 0.42.2", + "windows_x86_64_gnu 0.42.2", + "windows_x86_64_gnullvm 0.42.2", + "windows_x86_64_msvc 0.42.2", ] [[package]] @@ -47,51 +73,93 @@ version = "0.48.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9a2fa6e2155d7247be68c096456083145c183cbbbc2764150dda45a87197940c" dependencies = [ - "windows_aarch64_gnullvm", - "windows_aarch64_msvc", - "windows_i686_gnu", - "windows_i686_msvc", - "windows_x86_64_gnu", - "windows_x86_64_gnullvm", - "windows_x86_64_msvc", + "windows_aarch64_gnullvm 0.48.5", + "windows_aarch64_msvc 0.48.5", + "windows_i686_gnu 0.48.5", + "windows_i686_msvc 0.48.5", + "windows_x86_64_gnu 0.48.5", + "windows_x86_64_gnullvm 0.48.5", + "windows_x86_64_msvc 0.48.5", ] +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "597a5118570b68bc08d8d59125332c54f1ba9d9adeedeef5b99b02ba2b0698f8" + [[package]] name = "windows_aarch64_gnullvm" version = "0.48.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2b38e32f0abccf9987a4e3079dfb67dcd799fb61361e53e2882c3cbaf0d905d8" +[[package]] +name = "windows_aarch64_msvc" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e08e8864a60f06ef0d0ff4ba04124db8b0fb3be5776a5cd47641e942e58c4d43" + [[package]] name = "windows_aarch64_msvc" version = "0.48.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dc35310971f3b2dbbf3f0690a219f40e2d9afcf64f9ab7cc1be722937c26b4bc" +[[package]] +name = "windows_i686_gnu" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c61d927d8da41da96a81f029489353e68739737d3beca43145c8afec9a31a84f" + [[package]] name = "windows_i686_gnu" version = "0.48.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a75915e7def60c94dcef72200b9a8e58e5091744960da64ec734a6c6e9b3743e" +[[package]] +name = "windows_i686_msvc" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "44d840b6ec649f480a41c8d80f9c65108b92d89345dd94027bfe06ac444d1060" + [[package]] name = "windows_i686_msvc" version = "0.48.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8f55c233f70c4b27f66c523580f78f1004e8b5a8b659e05a4eb49d4166cca406" +[[package]] +name = "windows_x86_64_gnu" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8de912b8b8feb55c064867cf047dda097f92d51efad5b491dfb98f6bbb70cb36" + [[package]] name = "windows_x86_64_gnu" version = "0.48.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "53d40abd2583d23e4718fddf1ebec84dbff8381c07cae67ff7768bbf19c6718e" +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "26d41b46a36d453748aedef1486d5c7a85db22e56aff34643984ea85514e94a3" + [[package]] name = "windows_x86_64_gnullvm" version = "0.48.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0b7b52767868a23d5bab768e390dc5f5c55825b6d30b86c844ff2dc7414044cc" +[[package]] +name = "windows_x86_64_msvc" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9aec5da331524158c6d1a4ac0ab1541149c0b9505fde06423b02f5ef0106b9f0" + [[package]] name = "windows_x86_64_msvc" version = "0.48.5" diff --git a/Cargo.toml b/Cargo.toml index 1aa1e955..82355745 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -14,8 +14,9 @@ keywords = ["atomic", "concurrent", "hashmap"] categories = ["concurrency", "algorithms", "data-structures"] [dependencies] -seize = { git = "https://github.com/xacrimon/seize", rev = "0f41a677bccb60bdd8f3fb66a9d3e0744a5a99e8" } cfg-if = "1.0.0" +seize = { git = "https://github.com/xacrimon/seize", rev = "3525d40db2ea0deeee273f13700f1211d5c7d102" } +atomic-wait = "1.1.0" [package.metadata.docs.rs] features = [] diff --git a/src/lib.rs b/src/lib.rs index 3f716165..6d659282 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1 +1,2 @@ mod raw; +mod map; diff --git a/src/map.rs b/src/map.rs new file mode 100644 index 00000000..ae844454 --- /dev/null +++ b/src/map.rs @@ -0,0 +1,1576 @@ +use crate::raw::{self, InsertResult}; +use seize::{Collector, Guard, LocalGuard, OwnedGuard}; + +use std::borrow::Borrow; +use std::collections::hash_map::RandomState; +use std::fmt; +use std::hash::{BuildHasher, Hash}; +use std::marker::PhantomData; +use std::iter::FromIterator; + +/// A concurrent hash table. +/// +/// Most hash table operations require a [`Guard`](crate::Guard), which can be acquired through +/// [`HashMap::guard`] or using the [`HashMap::pin`] API. See the [crate-level documentation](crate#usage) +/// for details. +pub struct HashMap { + raw: raw::HashMap, +} + +// Safety: We only ever hand out &K/V through shared references to the map, +// so normal Send/Sync rules apply. We never expose owned or mutable references +// to keys or values. +unsafe impl Send for HashMap {} +unsafe impl Sync for HashMap {} + +/// A builder for a [`HashMap`]. +/// +/// # Examples +/// +/// ```rust +/// use papaya::{HashMap, ResizeMode}; +/// use seize::Collector; +/// use std::collections::hash_map::RandomState; +/// +/// let map: HashMap = HashMap::builder() +/// // Set the initial capacity. +/// .capacity(2048) +/// // Set the hasher. +/// .hasher(RandomState::new()) +/// // Set the resize mode. +/// .resize_mode(ResizeMode::Blocking) +/// // Set a custom garbage collector. +/// .collector(Collector::new().batch_size(128)) +/// // Construct the hash map. +/// .build(); +/// ``` +pub struct HashMapBuilder { + hasher: S, + capacity: usize, + collector: Collector, + resize_mode: ResizeMode, + _kv: PhantomData<(K, V)>, +} + +impl HashMapBuilder { + /// Set the hash builder used to hash keys. + /// + /// Warning: `hash_builder` is normally randomly generated, and is designed + /// to allow HashMaps to be resistant to attacks that cause many collisions + /// and very poor performance. Setting it manually using this function can + /// expose a DoS attack vector. + /// + /// The `hash_builder` passed should implement the [`BuildHasher`] trait for + /// the HashMap to be useful, see its documentation for details. + pub fn hasher(self, hasher: S) -> HashMapBuilder { + HashMapBuilder { + hasher, + capacity: self.capacity, + collector: self.collector, + resize_mode: self.resize_mode, + _kv: PhantomData, + } + } +} + +impl HashMapBuilder { + /// Set the initial capacity of the map. + /// + /// The table should be able to hold at least `capacity` elements before resizing. + /// However, the capacity is an estimate, and the table may prematurely resize due + /// to poor hash distribution. If `capacity` is 0, the hash map will not allocate. + pub fn capacity(self, capacity: usize) -> HashMapBuilder { + HashMapBuilder { + capacity, + hasher: self.hasher, + collector: self.collector, + resize_mode: self.resize_mode, + _kv: PhantomData, + } + } + + /// Set the resizing mode of the map. See [`ResizeMode`] for details. + pub fn resize_mode(self, resize_mode: ResizeMode) -> Self { + HashMapBuilder { + resize_mode, + hasher: self.hasher, + capacity: self.capacity, + collector: self.collector, + _kv: PhantomData, + } + } + + /// Set the [`seize::Collector`] used for garbage collection. + /// + /// This method may be useful when you want more control over garbage collection. + /// + /// Note that all `Guard` references used to access the map must be produced by + /// the provided `collector`. + pub fn collector(self, collector: Collector) -> Self { + HashMapBuilder { + collector, + hasher: self.hasher, + capacity: self.capacity, + resize_mode: self.resize_mode, + _kv: PhantomData, + } + } + + /// 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), + } + } +} + +impl fmt::Debug for HashMapBuilder { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("HashMapBuilder") + .field("capacity", &self.capacity) + .field("collector", &self.collector) + .field("resize_mode", &self.resize_mode) + .finish() + } +} + +/// Resize behavior for a [`HashMap`]. +/// +/// Hash maps must resize when the underlying table becomes full, migrating all key and value pairs +/// to a new table. This type allows you to configure the resizing behavior when passed to +/// [`HashMapBuilder::resize_mode`]. +#[derive(Debug)] +pub enum ResizeMode { + /// Writers copy a constant number of key/value pairs to the new table before making + /// progress. + /// + /// Incremental resizes avoids latency spikes that can occur when insert operations have + /// to resize a large table. However, they reduce parallelism during the resize and so can reduce + /// overall throughput. Incremental resizing also means reads or write operations during an + /// in-progress resize may have to search both the current and new table before succeeding, trading + /// off median latency during a resize for tail latency. + /// + /// This is the default resize mode, with a chunk size of `64`. + Incremental(usize), + /// All writes to the map must wait till the resize completes before making progress. + /// + /// Blocking resizes tend to be better in terms of throughput, especially in setups with + /// multiple writers that can perform the resize in parallel. However, they can lead to latency + /// spikes for insert operations that have to resize large tables. + /// + /// If insert latency is not a concern, such as if the keys in your map are stable, enabling blocking + /// resizes may yield better performance. + /// + /// Blocking resizing may also be a better option if you rely heavily on iteration or similar + /// operations, as they require completing any in-progress resizes for consistency. + Blocking, +} + +impl Default for ResizeMode { + fn default() -> Self { + // Incremental resizing is a good default for most workloads as it avoids + // unexpected latency spikes. + ResizeMode::Incremental(64) + } +} + +impl HashMap { + /// Creates an empty `HashMap`. + /// + /// The hash map is initially created with a capacity of 0, so it will not allocate + /// until it is first inserted into. + /// + /// # Examples + /// + /// ``` + /// use papaya::HashMap; + /// let map: HashMap<&str, i32> = HashMap::new(); + /// ``` + pub fn new() -> HashMap { + HashMap::with_capacity_and_hasher(0, RandomState::new()) + } + + /// Creates an empty `HashMap` with the specified capacity. + /// + /// The table should be able to hold at least `capacity` elements before resizing. + /// However, the capacity is an estimate, and the table may prematurely resize due + /// to poor hash distribution. If `capacity` is 0, the hash map will not allocate. + /// + /// # Examples + /// + /// ``` + /// use papaya::HashMap; + /// let map: HashMap<&str, i32> = HashMap::with_capacity(10); + /// ``` + pub fn with_capacity(capacity: usize) -> HashMap { + HashMap::with_capacity_and_hasher(capacity, RandomState::new()) + } + + /// Returns a builder for a `HashMap`. + /// + /// The builder can be used for more complex configuration, such as using + /// a custom [`Collector`], or [`ResizeMode`]. + pub fn builder() -> HashMapBuilder { + HashMapBuilder { + capacity: 0, + hasher: RandomState::default(), + collector: Collector::new(), + resize_mode: ResizeMode::default(), + _kv: PhantomData, + } + } +} + +impl Default for HashMap +where + S: Default, +{ + fn default() -> Self { + HashMap::with_hasher(S::default()) + } +} + +impl HashMap { + /// Creates an empty `HashMap` which will use the given hash builder to hash + /// keys. + /// + /// Warning: `hash_builder` is normally randomly generated, and is designed + /// to allow HashMaps to be resistant to attacks that cause many collisions + /// and very poor performance. Setting it manually using this function can + /// expose a DoS attack vector. + /// + /// The `hash_builder` passed should implement the [`BuildHasher`] trait for + /// the HashMap to be useful, see its documentation for details. + /// + /// # Examples + /// + /// ``` + /// use papaya::HashMap; + /// use std::hash::RandomState; + /// + /// let s = RandomState::new(); + /// let map = HashMap::with_hasher(s); + /// map.pin().insert(1, 2); + /// ``` + pub fn with_hasher(hash_builder: S) -> HashMap { + HashMap::with_capacity_and_hasher(0, hash_builder) + } + + /// Creates an empty `HashMap` with at least the specified capacity, using + /// `hash_builder` to hash the keys. + /// + /// The table should be able to hold at least `capacity` elements before resizing. + /// However, the capacity is an estimate, and the table may prematurely resize due + /// to poor hash distribution. If `capacity` is 0, the hash map will not allocate. + /// + /// Warning: `hash_builder` is normally randomly generated, and is designed + /// to allow HashMaps to be resistant to attacks that cause many collisions + /// and very poor performance. Setting it manually using this function can + /// expose a DoS attack vector. + /// + /// The `hasher` passed should implement the [`BuildHasher`] trait for + /// the HashMap to be useful, see its documentation for details. + /// + /// # Examples + /// + /// ``` + /// use papaya::HashMap; + /// use std::hash::RandomState; + /// + /// let s = RandomState::new(); + /// let map = HashMap::with_capacity_and_hasher(10, s); + /// map.pin().insert(1, 2); + /// ``` + pub fn with_capacity_and_hasher(capacity: usize, hash_builder: S) -> HashMap { + HashMap { + raw: raw::HashMap::new( + capacity, + hash_builder, + Collector::default(), + ResizeMode::default(), + ), + } + } + + /// Returns a pinned reference to the map. + /// + /// The returned reference manages a guard internally, preventing garbage collection + /// for as long as it is held. See the [crate-level documentation](crate#usage) for details. + #[inline] + pub fn pin(&self) -> HashMapRef<'_, K, V, S, LocalGuard<'_>> { + HashMapRef { + guard: self.guard(), + map: self, + } + } + + /// Returns a pinned reference to the map. + /// + /// Unlike [`HashMap::pin`], the returned reference implements `Send` and `Sync`, + /// allowing it to be held across `.await` points in work-stealing schedulers. + /// This is especially useful for iterators. + /// + /// The returned reference manages a guard internally, preventing garbage collection + /// for as long as it is held. See the [crate-level documentation](crate#usage) for details. + #[inline] + pub fn pin_owned(&self) -> HashMapRef<'_, K, V, S, OwnedGuard<'_>> { + HashMapRef { + guard: self.owned_guard(), + map: self, + } + } + + /// Returns a guard for use with this map. + /// + /// Note that holding on to a guard prevents garbage collection. + /// See the [crate-level documentation](crate#usage) for details. + #[inline] + pub fn guard(&self) -> LocalGuard<'_> { + self.raw.collector().enter() + } + + /// Returns an owned guard for use with this map. + /// + /// Owned guards implement `Send` and `Sync`, allowing them to be held across + /// `.await` points in work-stealing schedulers. This is especially useful + /// for iterators. + /// + /// Note that holding on to a guard prevents garbage collection. + /// See the [crate-level documentation](crate#usage) for details. + #[inline] + pub fn owned_guard(&self) -> OwnedGuard<'_> { + self.raw.collector().enter_owned() + } +} + +impl HashMap +where + K: Hash + Eq, + S: BuildHasher, +{ + /// Returns the number of entries in the map. + /// + /// # Examples + /// + /// ``` + /// use papaya::HashMap; + /// + /// let map = HashMap::new(); + /// + /// map.pin().insert(1, "a"); + /// map.pin().insert(2, "b"); + /// assert!(map.len() == 2); + /// ``` + #[inline] + pub fn len(&self) -> usize { + self.raw.len() + } + + /// Returns `true` if the map is empty. Otherwise returns `false`. + /// + /// # Examples + /// + /// ``` + /// use papaya::HashMap; + /// + /// let map = HashMap::new(); + /// assert!(map.is_empty()); + /// map.pin().insert("a", 1); + /// assert!(!map.is_empty()); + /// ``` + #[inline] + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + /// Returns `true` if the map contains a value for the specified key. + /// + /// The key may be any borrowed form of the map's key type, but + /// [`Hash`] and [`Eq`] on the borrowed form *must* match those for + /// the key type. + /// + /// [`Eq`]: std::cmp::Eq + /// [`Hash`]: std::hash::Hash + /// + /// + /// # Examples + /// + /// ``` + /// use papaya::HashMap; + /// + /// let map = HashMap::new(); + /// let m = map.pin(); + /// m.insert(1, "a"); + /// assert_eq!(m.contains_key(&1), true); + /// assert_eq!(m.contains_key(&2), false); + /// ``` + #[inline] + pub fn contains_key(&self, key: &Q, guard: &impl Guard) -> bool + where + K: Borrow, + Q: Hash + Eq + ?Sized, + { + self.get(key, guard).is_some() + } + + /// Returns a reference to the value corresponding to the key. + /// + /// The key may be any borrowed form of the map's key type, but + /// [`Hash`] and [`Eq`] on the borrowed form *must* match those for + /// the key type. + /// + /// [`Eq`]: std::cmp::Eq + /// [`Hash`]: std::hash::Hash + /// + /// # Examples + /// + /// ``` + /// use papaya::HashMap; + /// + /// let map = HashMap::new(); + /// let m = map.pin(); + /// m.insert(1, "a"); + /// assert_eq!(m.get(&1), Some(&"a")); + /// assert_eq!(m.get(&2), None); + /// ``` + #[inline] + pub fn get<'g, Q>(&self, key: &Q, guard: &'g impl Guard) -> Option<&'g V> + where + K: Borrow + 'g, + Q: Hash + Eq + ?Sized, + { + match self.raw.root(guard).get(key, guard) { + Some((_, v)) => Some(v), + None => None, + } + } + + /// Returns the key-value pair corresponding to the supplied key. + /// + /// The supplied key may be any borrowed form of the map's key type, but + /// [`Hash`] and [`Eq`] on the borrowed form *must* match those for + /// the key type. + /// + /// [`Eq`]: std::cmp::Eq + /// [`Hash`]: std::hash::Hash + /// + /// # Examples + /// + /// ``` + /// use papaya::HashMap; + /// + /// let map = HashMap::new(); + /// let m = map.pin(); + /// m.insert(1, "a"); + /// assert_eq!(m.get_key_value(&1), Some((&1, &"a"))); + /// assert_eq!(m.get_key_value(&2), None); + /// ``` + #[inline] + pub fn get_key_value<'g, Q>(&self, key: &Q, guard: &'g impl Guard) -> Option<(&'g K, &'g V)> + where + K: Borrow, + Q: Hash + Eq + ?Sized, + { + self.raw.root(guard).get(key, guard) + } + + /// Inserts a key-value pair into the map. + /// + /// If the map did not have this key present, [`None`] is returned. + /// + /// If the map did have this key present, the value is updated, and the old + /// value is returned. The key is not updated, though; this matters for + /// types that can be `==` without being identical. See the [standard library + /// documentation] for details. + /// + /// [standard library documentation]: https://doc.rust-lang.org/std/collections/index.html#insert-and-complex-keys + /// + /// # Examples + /// + /// ``` + /// use papaya::HashMap; + /// + /// let map = HashMap::new(); + /// assert_eq!(map.pin().insert(37, "a"), None); + /// assert_eq!(map.pin().is_empty(), false); + /// + /// // note: you can also re-use a map pin like so: + /// let m = map.pin(); + /// + /// m.insert(37, "b"); + /// assert_eq!(m.insert(37, "c"), Some(&"b")); + /// assert_eq!(m.get(&37), Some(&"c")); + /// ``` + #[inline] + pub fn insert<'g>(&self, key: K, value: V, guard: &'g impl Guard) -> Option<&'g V> { + match self.raw.root(guard).insert(key, value, true, guard) { + InsertResult::Inserted(_) => None, + InsertResult::Replaced(value) => Some(value), + InsertResult::Error { .. } => unreachable!(), + } + } + + /// Tries to insert a key-value pair into the map, and returns + /// a reference to the value that was inserted. + /// + /// If the map already had this key present, nothing is updated, and + /// an error containing the existing value is returned. + /// + /// # Examples + /// + /// ``` + /// use papaya::HashMap; + /// + /// let map = HashMap::new(); + /// let m = map.pin(); + /// assert_eq!(m.try_insert(37, "a").unwrap(), &"a"); + /// + /// let err = m.try_insert(37, "b").unwrap_err(); + /// assert_eq!(err.current, &"a"); + /// assert_eq!(err.not_inserted, "b"); + /// ``` + #[inline] + pub fn try_insert<'g>( + &self, + key: K, + value: V, + guard: &'g impl Guard, + ) -> Result<&'g V, OccupiedError<'g, V>> { + match self.raw.root(guard).insert(key, value, false, guard) { + InsertResult::Inserted(value) => Ok(value), + InsertResult::Error { + current, + not_inserted, + } => Err(OccupiedError { + current, + not_inserted, + }), + InsertResult::Replaced(_) => unreachable!(), + } + } + + /// Returns a reference to the value corresponding to the key, or inserts a default value. + /// + /// If the given key is present, the corresponding value is returned. If it is not present, + /// the provided `value` is inserted, and a reference to the newly inserted value is returned. + /// + /// # Examples + /// + /// ``` + /// use papaya::HashMap; + /// + /// let map = HashMap::new(); + /// assert_eq!(map.pin().get_or_insert("a", 3), &3); + /// assert_eq!(map.pin().get_or_insert("a", 6), &3); + /// ``` + #[inline] + pub fn get_or_insert<'g>(&self, key: K, value: V, guard: &'g impl Guard) -> &'g V { + // Note that we use `try_insert` instead of `compute` or `get_or_insert_with` here, as it + // allows us to avoid the closure indirection. + match self.try_insert(key, value, guard) { + Ok(inserted) => inserted, + Err(OccupiedError { current, .. }) => current, + } + } + + /// Returns a reference to the value corresponding to the key, or inserts a default value + /// computed from a closure. + /// + /// If the given key is present, the corresponding value is returned. If it is not present, + /// the value computed from `f` is inserted, and a reference to the newly inserted value is + /// returned. + /// + /// + /// # Examples + /// + /// ``` + /// use papaya::HashMap; + /// + /// let map = HashMap::new(); + /// assert_eq!(map.pin().get_or_insert_with("a", || 3), &3); + /// assert_eq!(map.pin().get_or_insert_with("a", || 6), &3); + /// ``` + #[inline] + pub fn get_or_insert_with<'g, F>(&self, key: K, f: F, guard: &'g impl Guard) -> &'g V + where + F: FnOnce() -> V, + K: 'g, + { + self.raw.root(guard).get_or_insert_with(key, f, guard) + } + + /// Updates an existing entry atomically. + /// + /// If the value for the specified `key` is present, the new value is computed and stored the + /// using the provided update function, and the new value is returned. Otherwise, `None` + /// is returned. + /// + /// + /// The update function is given the current value associated with the given key and returns the + /// new value to be stored. The operation is applied atomically only if the state of the entry remains + /// the same, meaning that it is not concurrently modified in any way. If the entry is + /// modified, the operation is retried with the new entry, similar to a traditional [compare-and-swap](https://en.wikipedia.org/wiki/Compare-and-swap) + /// operation. + /// + /// Note that the `update` function should be pure as it may be called multiple times, and the output + /// for a given entry may be memoized across retries. + /// + /// + /// # Examples + /// + /// ``` + /// use papaya::HashMap; + /// + /// let map = HashMap::new(); + /// map.pin().insert("a", 1); + /// assert_eq!(map.pin().get(&"a"), Some(&1)); + /// + /// map.pin().update("a", |v| v + 1); + /// assert_eq!(map.pin().get(&"a"), Some(&2)); + /// ``` + #[inline] + pub fn update<'g, F>(&self, key: K, update: F, guard: &'g impl Guard) -> Option<&'g V> + where + F: Fn(&V) -> V, + K: 'g, + { + self.raw.root(guard).update(key, update, guard) + } + + /// Updates an existing entry or inserts a default value. + /// + /// If the value for the specified `key` is present, the new value is computed and stored the + /// using the provided update function, and the new value is returned. Otherwise, the provided + /// `value` is inserted into the map, and a reference to the newly inserted value is returned. + /// + /// See [`HashMap::update`] for details about how atomic updates are performed. + /// + /// # Examples + /// + /// ``` + /// use papaya::HashMap; + /// + /// let map = HashMap::new(); + /// assert_eq!(*map.pin().update_or_insert("a", |i| i + 1, 0), 0); + /// assert_eq!(*map.pin().update_or_insert("a", |i| i + 1, 0), 1); + /// ``` + #[inline] + pub fn update_or_insert<'g, F>( + &self, + key: K, + update: F, + value: V, + guard: &'g impl Guard, + ) -> &'g V + where + F: Fn(&V) -> V, + K: 'g, + { + self.update_or_insert_with(key, update, || value, guard) + } + + /// Updates an existing entry or inserts a default value computed from a closure. + /// + /// If the value for the specified `key` is present, the new value is computed and stored the + /// using the provided update function, and the new value is returned. Otherwise, the value + /// computed by `f` is inserted into the map, and a reference to the newly inserted value is + /// returned. + /// + /// See [`HashMap::update`] for details about how atomic updates are performed. + /// + /// # Examples + /// + /// ``` + /// use papaya::HashMap; + /// + /// let map = HashMap::new(); + /// assert_eq!(*map.pin().update_or_insert_with("a", |i| i + 1, || 0), 0); + /// assert_eq!(*map.pin().update_or_insert_with("a", |i| i + 1, || 0), 1); + /// ``` + #[inline] + pub fn update_or_insert_with<'g, U, F>( + &self, + key: K, + update: U, + f: F, + guard: &'g impl Guard, + ) -> &'g V + where + F: FnOnce() -> V, + U: Fn(&V) -> V, + K: 'g, + { + self.raw + .root(guard) + .update_or_insert_with(key, update, f, guard) + } + + /// Updates an entry with a compare-and-swap (CAS) function. + /// + /// This method allows you to perform complex operations on the map atomically. The `compute` + /// closure is given the current state of the entry and returns the operation that should be + /// performed. The operation is applied atomically only if the state of the entry remains the same, + /// meaning it is not concurrently modified in any way. + /// + /// Note that the `compute` function should be pure as it may be called multiple times, and + /// the output for a given entry may be memoized across retries. + /// + /// In most cases you can avoid this method and instead use a higher-level atomic operation. + /// See the [crate-level documentation](crate#atomic-operations) for details. + /// + /// # Examples + /// + /// ```rust + /// use papaya::{HashMap, Operation, Compute}; + /// + /// let map = HashMap::new(); + /// let map = map.pin(); + /// + /// let compute = |entry| match entry { + /// // Remove the value if it is even. + /// Some((_key, value)) if value % 2 == 0 => { + /// Operation::Remove + /// } + /// + /// // Increment the value if it is odd. + /// Some((_key, value)) => { + /// Operation::Insert(value + 1) + /// } + /// + /// // Do nothing if the key does not exist + /// None => Operation::Abort(()), + /// }; + /// + /// assert_eq!(map.compute('A', compute), Compute::Aborted(())); + /// + /// map.insert('A', 1); + /// assert_eq!(map.compute('A', compute), Compute::Updated { + /// old: (&'A', &1), + /// new: (&'A', &2), + /// }); + /// assert_eq!(map.compute('A', compute), Compute::Removed(&'A', &2)); + /// ``` + #[inline] + pub fn compute<'g, F, T>( + &self, + key: K, + compute: F, + guard: &'g impl Guard, + ) -> Compute<'g, K, V, T> + where + F: FnMut(Option<(&'g K, &'g V)>) -> Operation, + { + self.raw.root(guard).compute(key, compute, guard) + } + + /// Removes a key from the map, returning the value at the key if the key + /// was previously in the map. + /// + /// The key may be any borrowed form of the map's key type, but + /// [`Hash`] and [`Eq`] on the borrowed form *must* match those for + /// the key type. + /// + /// # Examples + /// + /// ``` + /// use papaya::HashMap; + /// + /// let map = HashMap::new(); + /// map.pin().insert(1, "a"); + /// assert_eq!(map.pin().remove(&1), Some(&"a")); + /// assert_eq!(map.pin().remove(&1), None); + /// ``` + #[inline] + pub fn remove<'g, Q>(&self, key: &Q, guard: &'g impl Guard) -> Option<&'g V> + where + K: Borrow + 'g, + Q: Hash + Eq + ?Sized, + { + match self.raw.root(guard).remove(key, guard) { + Some((_, value)) => Some(value), + None => None, + } + } + + /// Removes a key from the map, returning the stored key and value if the + /// key was previously in the map. + /// + /// The key may be any borrowed form of the map's key type, but + /// [`Hash`] and [`Eq`] on the borrowed form *must* match those for + /// the key type. + /// + /// # Examples + /// + /// ``` + /// use papaya::HashMap; + /// + /// let map = HashMap::new(); + /// map.pin().insert(1, "a"); + /// assert_eq!(map.pin().get(&1), Some(&"a")); + /// assert_eq!(map.pin().remove_entry(&1), Some((&1, &"a"))); + /// assert_eq!(map.pin().remove(&1), None); + /// ``` + #[inline] + pub fn remove_entry<'g, Q>(&self, key: &Q, guard: &'g impl Guard) -> Option<(&'g K, &'g V)> + where + K: Borrow, + Q: Hash + Eq + ?Sized, + { + self.raw.root(guard).remove(key, guard) + } + + /// Tries to reserve capacity for `additional` more elements to be inserted + /// in the `HashMap`. + /// + /// After calling this method, the table should be able to hold at least `capacity` elements + /// before resizing. However, the capacity is an estimate, and the table may prematurely resize + /// due to poor hash distribution. The collection may also reserve more space to avoid frequent + /// reallocations. + /// + /// # Panics + /// + /// Panics if the new allocation size overflows `usize`. + /// + /// # Examples + /// + /// ``` + /// use papaya::HashMap; + /// + /// let map: HashMap<&str, i32> = HashMap::new(); + /// map.pin().reserve(10); + /// ``` + #[inline] + pub fn reserve(&self, additional: usize, guard: &impl Guard) { + self.raw.root(guard).reserve(additional, guard); + } + + /// Clears the map, removing all key-value pairs. + /// + /// Note that this method will block until any in-progress resizes are + /// completed before proceeding. See the [consistency](crate#consistency) + /// section for details. + /// + /// # Examples + /// + /// ``` + /// use papaya::HashMap; + /// + /// let map = HashMap::new(); + /// + /// map.pin().insert(1, "a"); + /// map.pin().clear(); + /// assert!(map.pin().is_empty()); + /// ``` + #[inline] + pub fn clear(&self, guard: &impl Guard) { + self.raw.root(guard).clear(guard) + } + + /// Retains only the elements specified by the predicate. + /// + /// In other words, remove all pairs `(k, v)` for which `f(&k, &v)` returns `false`. + /// The elements are visited in unsorted (and unspecified) order. + /// + /// Note the function may be called more than once for a given key if its value is + /// concurrently modified during removal. + /// + /// Additionally, this method will block until any in-progress resizes are + /// completed before proceeding. See the [consistency](crate#consistency) + /// section for details. + /// + /// # Examples + /// + /// ``` + /// use papaya::HashMap; + /// + /// let mut map: HashMap = (0..8).map(|x| (x, x * 10)).collect(); + /// map.pin().retain(|&k, _| k % 2 == 0); + /// assert_eq!(map.len(), 4); + /// ``` + #[inline] + pub fn retain(&mut self, f: F, guard: &impl Guard) + where + F: FnMut(&K, &V) -> bool, + { + self.raw.root(guard).retain(f, guard) + } + + /// An iterator visiting all key-value pairs in arbitrary order. + /// The iterator element type is `(&K, &V)`. + /// + /// Note that this method will block until any in-progress resizes are + /// completed before proceeding. See the [consistency](crate#consistency) + /// section for details. + /// + /// # Examples + /// + /// ``` + /// use papaya::HashMap; + /// + /// let map = HashMap::from([ + /// ("a", 1), + /// ("b", 2), + /// ("c", 3), + /// ]); + /// + /// for (key, val) in map.pin().iter() { + /// println!("key: {key} val: {val}"); + /// } + #[inline] + pub fn iter<'g, G>(&self, guard: &'g G) -> Iter<'g, K, V, G> + where + G: Guard, + { + Iter { + raw: self.raw.root(guard).iter(guard), + } + } + + /// An iterator visiting all keys in arbitrary order. + /// The iterator element type is `&K`. + /// + /// Note that this method will block until any in-progress resizes are + /// completed before proceeding. See the [consistency](crate#consistency) + /// section for details. + /// + /// # Examples + /// + /// ``` + /// use papaya::HashMap; + /// + /// let map = HashMap::from([ + /// ("a", 1), + /// ("b", 2), + /// ("c", 3), + /// ]); + /// + /// for key in map.pin().keys() { + /// println!("{key}"); + /// } + /// ``` + #[inline] + pub fn keys<'g, G>(&self, guard: &'g G) -> Keys<'g, K, V, G> + where + G: Guard, + { + Keys { + iter: self.iter(guard), + } + } + + /// An iterator visiting all values in arbitrary order. + /// The iterator element type is `&V`. + /// + /// Note that this method will block until any in-progress resizes are + /// completed before proceeding. See the [consistency](crate#consistency) + /// section for details. + /// + /// # Examples + /// + /// ``` + /// use papaya::HashMap; + /// + /// let map = HashMap::from([ + /// ("a", 1), + /// ("b", 2), + /// ("c", 3), + /// ]); + /// + /// for value in map.pin().values() { + /// println!("{value}"); + /// } + /// ``` + #[inline] + pub fn values<'g, G>(&self, guard: &'g G) -> Values<'g, K, V, G> + where + G: Guard, + { + Values { + iter: self.iter(guard), + } + } +} + +/// An operation to perform on given entry in a [`HashMap`]. +/// +/// See [`HashMap::compute`] for details. +#[derive(Debug, PartialEq, Eq)] +pub enum Operation { + /// Insert the given value. + Insert(V), + + /// Remove the entry from the map. + Remove, + + /// Abort the operation with the given value. + Abort(T), +} + +/// The result of a [`compute`](HashMap::compute) operation. +/// +/// Contains information about the [`Operation`] that was performed. +#[derive(Debug, PartialEq, Eq)] +pub enum Compute<'g, K, V, T> { + /// The given entry was inserted. + Inserted(&'g K, &'g V), + + /// The entry was updated. + Updated { + /// The entry that was replaced. + old: (&'g K, &'g V), + + /// The entry that was inserted. + new: (&'g K, &'g V), + }, + + /// The given entry was removed. + Removed(&'g K, &'g V), + + /// The operation was aborted with the given value. + Aborted(T), +} + +/// An error returned by [`try_insert`](HashMap::try_insert) when the key already exists. +/// +/// Contains the existing value, and the value that was not inserted. +#[derive(Debug, PartialEq, Eq)] +pub struct OccupiedError<'a, V: 'a> { + /// The value in the map that was already present. + pub current: &'a V, + /// The value which was not inserted, because the entry was already occupied. + pub not_inserted: V, +} + +impl PartialEq for HashMap +where + K: Hash + Eq, + V: PartialEq, + 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, value)| other.get(key, guard2).map_or(false, |v| *value == *v)) + } +} + +impl Eq for HashMap +where + K: Hash + Eq, + V: Eq, + S: BuildHasher, +{ +} + +impl fmt::Debug for HashMap +where + K: Hash + Eq + fmt::Debug, + V: fmt::Debug, + S: BuildHasher, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let guard = self.guard(); + f.debug_map().entries(self.iter(&guard)).finish() + } +} + +impl Extend<(K, V)> for &HashMap +where + K: Hash + Eq, + S: BuildHasher, +{ + fn extend>(&mut self, iter: T) { + // 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 + // 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, value) in iter { + self.insert(key, value, &guard); + } + } +} + +impl<'a, K, V, S> Extend<(&'a K, &'a V)> for &HashMap +where + K: Copy + Hash + Eq, + V: Copy, + S: BuildHasher, +{ + fn extend>(&mut self, iter: T) { + self.extend(iter.into_iter().map(|(&key, &value)| (key, value))); + } +} + +impl From<[(K, V); N]> for HashMap +where + K: Hash + Eq, +{ + fn from(arr: [(K, V); N]) -> Self { + HashMap::from_iter(arr) + } +} + +impl FromIterator<(K, V)> for HashMap +where + K: Hash + Eq, + S: BuildHasher + Default, +{ + fn from_iter>(iter: T) -> Self { + let mut iter = iter.into_iter(); + + if let Some((key, value)) = iter.next() { + let (lower, _) = iter.size_hint(); + let map = HashMap::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 map = map.pin(); + map.insert(key, value); + for (key, value) in iter { + map.insert(key, value); + } + } + + map + } else { + Self::default() + } + } +} + +impl Clone for HashMap +where + K: Clone + Hash + Eq, + V: Clone, + S: BuildHasher + Clone, +{ + fn clone(&self) -> HashMap { + let other = HashMap::builder() + .capacity(self.len()) + .hasher(self.raw.hasher.clone()) + .collector(self.raw.collector().clone()) + .build(); + + { + let (guard1, guard2) = (&self.guard(), &other.guard()); + for (key, value) in self.iter(guard1) { + other.insert(key.clone(), value.clone(), guard2); + } + } + + other + } +} + +/// A pinned reference to a [`HashMap`]. +/// +/// This type is created with [`HashMap::pin`] and can be used to easily access a [`HashMap`] +/// without explicitly managing a guard. See the [crate-level documentation](crate#usage) for details. +pub struct HashMapRef<'map, K, V, S, G> { + guard: G, + map: &'map HashMap, +} + +impl<'map, K, V, S, G> HashMapRef<'map, K, V, S, G> +where + K: Hash + Eq, + S: BuildHasher, + G: Guard, +{ + /// Returns a reference to the inner [`HashMap`]. + #[inline] + pub fn map(&self) -> &'map HashMap { + self.map + } + + /// Returns the number of entries in the map. + /// + /// See [`HashMap::len`] for details. + #[inline] + pub fn len(&self) -> usize { + self.map.raw.len() + } + + /// Returns `true` if the map is empty. Otherwise returns `false`. + /// + /// See [`HashMap::is_empty`] for details. + #[inline] + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + /// Returns `true` if the map contains a value for the specified key. + /// + /// See [`HashMap::contains_key`] for details. + #[inline] + pub fn contains_key(&self, key: &Q) -> bool + where + K: Borrow, + Q: Hash + Eq + ?Sized, + { + self.get(key).is_some() + } + + /// Returns a reference to the value corresponding to the key. + /// + /// See [`HashMap::get`] for details. + #[inline] + pub fn get(&self, key: &Q) -> Option<&V> + where + K: Borrow, + Q: Hash + Eq + ?Sized, + { + match self.root().get(key, &self.guard) { + Some((_, v)) => Some(v), + None => None, + } + } + + /// Returns the key-value pair corresponding to the supplied key. + /// + /// See [`HashMap::get_key_value`] for details. + #[inline] + pub fn get_key_value(&self, key: &Q) -> Option<(&K, &V)> + where + K: Borrow, + Q: Hash + Eq + ?Sized, + { + self.root().get(key, &self.guard) + } + + /// Inserts a key-value pair into the map. + /// + /// See [`HashMap::insert`] for details. + #[inline] + pub fn insert(&self, key: K, value: V) -> Option<&V> { + match self.root().insert(key, value, true, &self.guard) { + InsertResult::Inserted(_) => None, + InsertResult::Replaced(value) => Some(value), + InsertResult::Error { .. } => unreachable!(), + } + } + + /// Tries to insert a key-value pair into the map, and returns + /// a reference to the value that was inserted. + /// + /// See [`HashMap::try_insert`] for details. + #[inline] + pub fn try_insert(&self, key: K, value: V) -> Result<&V, OccupiedError<'_, V>> { + match self.root().insert(key, value, false, &self.guard) { + InsertResult::Inserted(value) => Ok(value), + InsertResult::Error { + current, + not_inserted, + } => Err(OccupiedError { + current, + not_inserted, + }), + InsertResult::Replaced(_) => unreachable!(), + } + } + + /// Returns a reference to the value corresponding to the key, or inserts a default value. + /// + /// See [`HashMap::get_or_insert`] for details. + #[inline] + pub fn get_or_insert(&self, key: K, value: V) -> &V { + // Note that we use `try_insert` instead of `compute` or `get_or_insert_with` here, as it + // allows us to avoid the closure indirection. + match self.try_insert(key, value) { + Ok(inserted) => inserted, + Err(OccupiedError { current, .. }) => current, + } + } + + /// Returns a reference to the value corresponding to the key, or inserts a default value + /// computed from a closure. + /// + /// See [`HashMap::get_or_insert_with`] for details. + #[inline] + pub fn get_or_insert_with(&self, key: K, f: F) -> &V + where + F: FnOnce() -> V, + { + self.root().get_or_insert_with(key, f, &self.guard) + } + + /// Updates an existing entry atomically. + /// + /// See [`HashMap::update`] for details. + #[inline] + pub fn update(&self, key: K, update: F) -> Option<&V> + where + F: Fn(&V) -> V, + { + self.root().update(key, update, &self.guard) + } + + /// Updates an existing entry or inserts a default value. + /// + /// See [`HashMap::update_or_insert`] for details. + #[inline] + pub fn update_or_insert(&self, key: K, update: F, value: V) -> &V + where + F: Fn(&V) -> V, + { + self.update_or_insert_with(key, update, || value) + } + + /// Updates an existing entry or inserts a default value computed from a closure. + /// + /// See [`HashMap::update_or_insert_with`] for details. + #[inline] + pub fn update_or_insert_with(&self, key: K, update: U, f: F) -> &V + where + F: FnOnce() -> V, + U: Fn(&V) -> V, + { + self.root() + .update_or_insert_with(key, update, f, &self.guard) + } + + // Updates an entry with a compare-and-swap (CAS) function. + // + /// See [`HashMap::compute`] for details. + #[inline] + pub fn compute<'g, F, T>(&'g self, key: K, compute: F) -> Compute<'g, K, V, T> + where + F: FnMut(Option<(&'g K, &'g V)>) -> Operation, + { + self.root().compute(key, compute, &self.guard) + } + + /// Removes a key from the map, returning the value at the key if the key + /// was previously in the map. + /// + /// See [`HashMap::remove`] for details. + #[inline] + pub fn remove(&self, key: &Q) -> Option<&V> + where + K: Borrow, + Q: Hash + Eq + ?Sized, + { + match self.root().remove(key, &self.guard) { + Some((_, value)) => Some(value), + None => None, + } + } + + /// Removes a key from the map, returning the stored key and value if the + /// key was previously in the map. + /// + /// See [`HashMap::remove_entry`] for details. + #[inline] + pub fn remove_entry(&self, key: &Q) -> Option<(&K, &V)> + where + K: Borrow, + Q: Hash + Eq + ?Sized, + { + self.root().remove(key, &self.guard) + } + + /// Clears the map, removing all key-value pairs. + /// + /// See [`HashMap::clear`] for details. + #[inline] + pub fn clear(&self) { + self.root().clear(&self.guard) + } + + /// Retains only the elements specified by the predicate. + /// + /// See [`HashMap::retain`] for details. + #[inline] + pub fn retain(&mut self, f: F) + where + F: FnMut(&K, &V) -> bool, + { + self.root().retain(f, &self.guard) + } + + /// Tries to reserve capacity for `additional` more elements to be inserted + /// in the map. + /// + /// See [`HashMap::reserve`] for details. + #[inline] + pub fn reserve(&self, additional: usize) { + self.root().reserve(additional, &self.guard) + } + + /// An iterator visiting all key-value pairs in arbitrary order. + /// The iterator element type is `(&K, &V)`. + /// + /// See [`HashMap::iter`] for details. + #[inline] + pub fn iter(&self) -> Iter<'_, K, V, G> { + Iter { + raw: self.root().iter(&self.guard), + } + } + + /// An iterator visiting all keys in arbitrary order. + /// The iterator element type is `&K`. + /// + /// See [`HashMap::keys`] for details. + #[inline] + pub fn keys(&self) -> Keys<'_, K, V, G> { + Keys { iter: self.iter() } + } + + /// An iterator visiting all values in arbitrary order. + /// The iterator element type is `&V`. + /// + /// See [`HashMap::values`] for details. + #[inline] + pub fn values(&self) -> Values<'_, K, V, G> { + Values { iter: self.iter() } + } + + #[inline] + fn root(&self) -> raw::HashMapRef<'_, K, V, S> { + // Safety: A `HashMapRef` can only be created through `HashMap::pin` or + // `HashMap::pin_owned`, so we know the guard belongs to our collector. + unsafe { self.map.raw.root_unchecked(&self.guard) } + } +} + +impl fmt::Debug for HashMapRef<'_, K, V, S, G> +where + K: Hash + Eq + fmt::Debug, + V: fmt::Debug, + S: BuildHasher, + G: Guard, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_map().entries(self.iter()).finish() + } +} + +impl<'a, K, V, S, G> IntoIterator for &'a HashMapRef<'_, K, V, S, G> +where + K: Hash + Eq, + S: BuildHasher, + G: Guard, +{ + type Item = (&'a K, &'a V); + type IntoIter = Iter<'a, K, V, G>; + + fn into_iter(self) -> Self::IntoIter { + self.iter() + } +} + +/// An iterator over a map's entries. +/// +/// 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, G>, +} + +impl<'g, K: 'g, V: 'g, G> Iterator for Iter<'g, K, V, G> +where + G: Guard, +{ + type Item = (&'g K, &'g V); + + #[inline] + fn next(&mut self) -> Option { + self.raw.next() + } +} + +impl fmt::Debug for Iter<'_, K, V, G> +where + K: fmt::Debug, + V: fmt::Debug, + G: Guard, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_list() + .entries(Iter { + raw: self.raw.clone(), + }) + .finish() + } +} + +/// An iterator over a map's keys. +/// +/// This struct is created by the [`keys`](HashMap::keys) method on [`HashMap`]. See its documentation for details. +pub struct Keys<'g, K, V, G> { + iter: Iter<'g, K, V, G>, +} + +impl<'g, K: 'g, V: 'g, G> Iterator for Keys<'g, K, V, G> +where + G: Guard, +{ + type Item = &'g K; + + #[inline] + fn next(&mut self) -> Option { + let (key, _) = self.iter.next()?; + Some(key) + } +} + +impl fmt::Debug for Keys<'_, K, V, G> +where + K: fmt::Debug, + V: fmt::Debug, + G: Guard, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_tuple("Keys").field(&self.iter).finish() + } +} + +/// An iterator over a map's values. +/// +/// This struct is created by the [`values`](HashMap::values) method on [`HashMap`]. See its documentation for details. +pub struct Values<'g, K, V, G> { + iter: Iter<'g, K, V, G>, +} + +impl<'g, K: 'g, V: 'g, G> Iterator for Values<'g, K, V, G> +where + G: Guard, +{ + type Item = &'g V; + + #[inline] + fn next(&mut self) -> Option { + let (_, value) = self.iter.next()?; + Some(value) + } +} + +impl fmt::Debug for Values<'_, K, V, G> +where + K: fmt::Debug, + V: fmt::Debug, + G: Guard, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_tuple("Values").field(&self.iter).finish() + } +} diff --git a/src/raw/alloc.rs b/src/raw/alloc.rs new file mode 100644 index 00000000..647bb318 --- /dev/null +++ b/src/raw/alloc.rs @@ -0,0 +1,195 @@ +use std::alloc::Layout; +use std::marker::PhantomData; +use std::sync::atomic::{AtomicPtr, AtomicU8}; +use std::{alloc, mem, ptr}; + +use seize::Collector; + +use super::{probe, State}; + +// A hash-table laid out in a single allocation. +#[repr(transparent)] +pub struct RawTable(u8); + +// Safety: `seize::Link` is the first field (see `TableLayout`). +unsafe impl seize::AsLink for RawTable {} + +// The layout of the table allocation. +#[repr(C)] +struct TableLayout { + link: seize::Link, + mask: usize, + limit: usize, + capacity: usize, + state: State, + meta: [AtomicU8; 0], + entries: [AtomicPtr<()>; 0], +} + +// Manages a table allocation. +#[repr(C)] +pub struct Table { + // Mask for the table length. + pub mask: usize, + // The probe limit. + pub limit: usize, + // The raw table pointer. + pub raw: *mut RawTable, + // The true (padded) table capacity. + capacity: usize, + _t: PhantomData, +} + +impl Copy for Table {} + +impl Clone for Table { + fn clone(&self) -> Self { + *self + } +} + +impl Table { + // Allocate a table with the provided length. + pub fn alloc(len: usize, collector: &Collector) -> Table { + assert!(len.is_power_of_two()); + assert!(mem::align_of::() % mem::align_of::<*mut T>() == 0); + + // Pad the meta table to fulfill the alignment requirement of an entry. + let capacity = (len + mem::align_of::<*mut T>() - 1) & !(mem::align_of::<*mut T>() - 1); + let mask = len - 1; + let limit = probe::limit(len); + + // Temporary guard needed to allocate the table. + let guard = collector.enter(); + + unsafe { + let layout = Self::layout(capacity); + + // Allocate the table, zeroing the entries. + let ptr = alloc::alloc_zeroed(layout); + if ptr.is_null() { + alloc::handle_alloc_error(layout); + } + + // Write the table state. + ptr.cast::().write(TableLayout { + link: collector.link(&guard), + mask, + limit, + capacity, + state: State { + collector, + ..State::default() + }, + meta: [], + entries: [], + }); + + // Initialize the meta table. + ptr.add(mem::size_of::()) + .cast::() + .write_bytes(super::meta::EMPTY, capacity); + + Table { + mask, + limit, + capacity, + raw: ptr.cast::(), + _t: PhantomData, + } + } + } + + // Creates a `Table` from a raw pointer. + #[inline] + pub unsafe fn from_raw(raw: *mut RawTable) -> Table { + if raw.is_null() { + return Table { + raw, + mask: 0, + limit: 0, + capacity: 0, + _t: PhantomData, + }; + } + + let layout = unsafe { &*raw.cast::() }; + + Table { + raw, + mask: layout.mask, + limit: layout.limit, + capacity: layout.capacity, + _t: PhantomData, + } + } + + // Returns the metadata entry at the given index. + #[inline] + pub unsafe fn meta(&self, i: usize) -> &AtomicU8 { + debug_assert!(i < self.capacity); + &*self + .raw + .add(mem::size_of::()) + .add(i * mem::size_of::()) + .cast::() + } + + // Returns the entry at the given index. + #[inline] + pub unsafe fn entry(&self, i: usize) -> &AtomicPtr { + let offset = mem::size_of::() + + mem::size_of::() * self.capacity + + i * mem::size_of::>(); + + debug_assert!(i < self.capacity); + &*self.raw.add(offset).cast::>() + } + + /// Returns the length of the table. + #[inline] + pub fn len(&self) -> usize { + self.mask + 1 + } + + // Returns a reference to the table state. + #[inline] + pub fn state(&self) -> &State { + unsafe { &(*self.raw.cast::()).state } + } + + // Returns a mutable reference to the table state. + #[inline] + pub fn state_mut(&mut self) -> &mut State { + unsafe { &mut (*self.raw.cast::()).state } + } + + // Deallocate the table. + pub unsafe fn dealloc(table: Table) { + let layout = Self::layout(table.capacity); + ptr::drop_in_place(table.raw.cast::()); + unsafe { alloc::dealloc(table.raw.cast::(), layout) } + } + + // The table layout used for allocation. + fn layout(capacity: usize) -> Layout { + let size = mem::size_of::() + + (mem::size_of::() * capacity) // meta + + (mem::size_of::() * capacity); // entries + Layout::from_size_align(size, mem::align_of::()).unwrap() + } +} + +#[test] +fn layout() { + unsafe { + let collector = seize::Collector::new(); + let table: Table = Table::alloc(4, &collector); + let table: Table = Table::from_raw(table.raw); + assert_eq!(table.mask, 3); + assert_eq!(table.len(), 4); + // The capacity is padded for pointer alignment. + assert_eq!(table.capacity, 8); + Table::dealloc(table); + } +} diff --git a/src/raw/mod.rs b/src/raw/mod.rs index a918f958..d77eeeca 100644 --- a/src/raw/mod.rs +++ b/src/raw/mod.rs @@ -1,2 +1,2482 @@ mod utils; mod parker; +mod alloc; +mod probe; + +use std::borrow::Borrow; +use std::hash::{BuildHasher, Hash}; +use std::marker::PhantomData; +use std::mem::{self, MaybeUninit}; +use std::sync::atomic::{fence, AtomicPtr, AtomicU32, AtomicUsize, Ordering}; +use std::sync::Mutex; +use std::{hint, panic, ptr}; + +use self::alloc::RawTable; +use probe::Probe; +use utils::{untagged, AtomicPtrFetchOps, Counter, Shared, StrictProvenance, Tagged}; +use crate::map::{Compute, Operation, ResizeMode}; +use parker::Parker; + +use seize::{AsLink, Collector, Guard, Link}; + +// A lock-free hash-table. +pub struct HashMap { + // A pointer to the root table. + table: AtomicPtr, + // Collector for memory reclamation. + // + // The collector is allocated as it's aliased by each table, + // in case it needs to be accessed during reclamation. + collector: Shared, + // The resize mode, either blocking or incremental. + resize: ResizeMode, + // The number of keys in the table. + count: Counter, + // Hasher for keys. + pub hasher: S, + _kv: PhantomData<(K, V)>, +} + +// The hash-table allocation. +type Table = self::alloc::Table>; + +// Hash-table state. +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: AtomicU32, + // A thread parker for blocking on copy operations. + pub parker: Parker, + // Entries whose retirement has been deferred by later tables. + pub deferred: seize::Deferred, + // A pointer to the root collector, valid as long as the map is alive. + pub collector: *const Collector, +} + +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: AtomicU32::new(State::PENDING), + parker: Parker::default(), + deferred: seize::Deferred::new(), + collector: ptr::null(), + } + } +} + +impl State { + // A resize is in-progress. + pub const PENDING: u32 = 0; + // The resize has been aborted, continue to the next table. + pub const ABORTED: u32 = 1; + // The resize was complete and the table was promoted. + pub const PROMOTED: u32 = 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: Tagged>, + not_inserted: *mut Entry, + }, +} + +// An entry in the hash-table. +#[repr(C)] +pub struct Entry { + pub link: Link, + pub key: K, + pub value: V, +} + +// Safety: repr(C) and seize::Link is the first field +unsafe impl AsLink for 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; + + // Reclaims an entry. + #[inline] + unsafe fn reclaim(link: *mut Link) { + let entry: *mut Entry = link.cast(); + let _entry = unsafe { Box::from_raw(entry) }; + } +} + +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 { + #[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 { + let collector = Shared::from(collector); + + // The table is lazily allocated. + if capacity == 0 { + return HashMap { + collector, + resize, + hasher, + table: AtomicPtr::new(ptr::null_mut()), + count: Counter::default(), + _kv: PhantomData, + }; + } + + // Initialize the table and mark it as the root. + let mut table = Table::::alloc(probe::entries_for(capacity), &collector); + *table.state_mut().status.get_mut() = State::PROMOTED; + + HashMap { + hasher, + resize, + collector, + table: AtomicPtr::new(table.raw), + count: Counter::default(), + _kv: PhantomData, + } + } + + // Returns a reference to the root hash-table. + #[inline] + pub fn root<'g>(&self, guard: &'g impl Guard) -> HashMapRef<'g, K, V, S> { + assert!( + guard.belongs_to(&self.collector), + "accessed map with incorrect guard" + ); + + // Safety: verified that the guard belongs to our collector. + unsafe { self.root_unchecked(guard) } + } + + // Returns a reference to the root hash-table. + // + // # Safety + // + // The guard must belong to this table's collector. + #[inline] + pub unsafe fn root_unchecked<'g>(&self, guard: &'g impl Guard) -> HashMapRef<'g, K, V, S> { + // Load the root table. + let raw = guard.protect(&self.table, Ordering::Acquire); + let table = unsafe { Table::::from_raw(raw) }; + + // Safety: The caller guarantees that the guard comes from our collector, so + // &'g Guard implies &'g self. This makes bounds a little nicer for users. + unsafe { + mem::transmute::, HashMapRef<'g, K, V, S>>(self.as_ref(table)) + } + } + + // 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(_)) + } + + // Returns a reference to the given table. + #[inline] + fn as_ref(&self, table: Table) -> HashMapRef<'_, K, V, S> { + HashMapRef { table, root: self } + } +} + +// A reference to the root table, or an arbitrarily nested table migration. +pub struct HashMapRef<'a, K, V, S> { + table: Table, + root: &'a HashMap, +} + +// Hash-table operations. +impl<'root, K, V, S> HashMapRef<'root, K, V, S> +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 Guard) -> Option<(&'g K, &'g V)> + where + K: Borrow, + Q: Hash + Eq + ?Sized, + { + // The table has not yet been allocated. + if self.table.raw.is_null() { + return None; + } + + // Initialize the probe state. + let (h1, h2) = self.hash(key); + let mut probe = Probe::start(h1, self.table.mask); + + // Probe until we reach the limit. + while probe.len <= self.table.limit { + // Load the entry metadata first for cheap searches. + let meta = unsafe { self.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 None; + } + + // Check for a potential match. + if meta != h2 { + probe.next(self.table.mask); + continue; + } + + // Load the full entry. + let entry = + unsafe { guard.protect(self.table.entry(probe.i), Ordering::Acquire) }.unpack(); + + // The entry was deleted, keep probing. + if entry.ptr.is_null() { + probe.next(self.table.mask); + continue; + } + + // Check for a full match. + if unsafe { (*entry.ptr).key.borrow() } == key { + // The entry was copied to the new table. + // + // In blocking resize mode we do not need to perform this 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; + } + + // Found the correct entry, return the key and value. + return unsafe { Some((&(*entry.ptr).key, &(*entry.ptr).value)) }; + } + + probe.next(self.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.root.is_incremental() { + if let Some(next) = self.next_table_ref() { + // Retry in the new table. + return next.get(key, guard); + } + } + + // Otherwise, the key is not in the table. + None + } + + // Inserts a key-value pair into the table. + #[inline] + pub fn insert<'g>( + &mut self, + key: K, + value: V, + replace: bool, + guard: &'g impl Guard, + ) -> InsertResult<'g, V> { + // Allocate the entry to be inserted. + let entry = Box::into_raw(Box::new(Entry { + key, + value, + link: self.root.collector.link(guard), + })); + + // Perform the insert. + // + // Safety: We just allocated the entry above. + let result = unsafe { self.insert_with(untagged(entry), replace, true, guard) }; + let result = match result { + RawInsertResult::Inserted(value) => InsertResult::Inserted(value), + RawInsertResult::Replaced(value) => InsertResult::Replaced(value), + RawInsertResult::Error { + current, + not_inserted, + } => { + let current = unsafe { &(*current.ptr).value }; + + // 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, + } + } + }; + + // Increment the length if we inserted a new entry. + if matches!(result, InsertResult::Inserted(_)) { + self.root + .count + .get(guard.thread().id()) + .fetch_add(1, Ordering::Relaxed); + } + + result + } + + // Inserts an entry into the map. + // + // # Safety + // + // The new entry must be a valid pointer. + #[inline] + unsafe fn insert_with<'g>( + &mut self, + new_entry: Tagged>, + should_replace: bool, + help_copy: bool, + guard: &'g impl Guard, + ) -> RawInsertResult<'g, K, V> { + // Allocate the table if it has not been initialized yet. + if self.table.raw.is_null() { + self.init(None); + } + + // Safety: The new entry is guaranteed to be valid by the caller. + let new_ref = unsafe { &*(new_entry).ptr }; + + // Initialize the probe state. + let (h1, h2) = self.hash(&new_ref.key); + let mut probe = Probe::start(h1, self.table.mask); + + // Probe until we reach the limit. + let copying = 'probe: loop { + if probe.len > self.table.limit { + break None; + } + + // Load the entry metadata first for cheap searches. + let meta = unsafe { self.table.meta(probe.i) }.load(Ordering::Acquire); + + // The entry is empty, try to insert. + let mut entry = if meta == meta::EMPTY { + match self.insert_at(probe.i, h2, new_entry.raw, 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(self.table.mask); + continue 'probe; + } + } + } + // Found a potential match. + else if meta == h2 { + // Load the full entry. + let found = guard + .protect(self.table.entry(probe.i), Ordering::Acquire) + .unpack(); + + // The entry was deleted, keep probing. + if found.ptr.is_null() { + probe.next(self.table.mask); + continue 'probe; + } + + // If the key matches, we might be able to update the value. + found + } + // Otherwise, continue probing. + else { + probe.next(self.table.mask); + continue 'probe; + }; + + // Check for a full match. + if unsafe { (*entry.ptr).key != new_ref.key } { + probe.next(self.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, + not_inserted: new_entry.ptr, + }; + } + + loop { + // Try to update the value. + match self.update_at(probe.i, entry, new_entry.raw, guard) { + // Successfully updated. + UpdateStatus::Replaced(entry) => { + return RawInsertResult::Replaced(&(*entry.ptr).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(self.table.mask); + continue 'probe; + } + + // Someone else beat us to the update, retry. + UpdateStatus::Found(EntryStatus::Value(found)) => entry = found, + } + } + }; + + // If went over the probe limit or found a copied entry, trigger a resize. + let mut next_table = self.get_or_alloc_next(None); + + let next_table = match self.root.resize { + ResizeMode::Blocking => { + // In blocking mode we must complete the resize before proceeding. + self.help_copy(guard, false) + } + + ResizeMode::Incremental(_) => { + // Help out with the copy. + if help_copy { + next_table = self.help_copy(guard, false); + } + + // 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); + } + + next_table + } + }; + + // Insert into the next table. + self.as_ref(next_table) + .insert_with(new_entry, should_replace, false, guard) + } + + // 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: ?Sized>(&self, key: &Q, guard: &'g impl Guard) -> Option<(&'g K, &'g V)> + where + K: Borrow, + Q: Hash + Eq, + { + self.remove_inner(key, true, guard) + } + + // Removes a key from the map, returning the entry for the key if the key was previously in the map. + // + // This is a recursive helper for `remove_entry`. + #[inline] + pub fn remove_inner<'g, Q: ?Sized>( + &self, + key: &Q, + help_copy: bool, + guard: &'g impl Guard, + ) -> Option<(&'g K, &'g V)> + where + K: Borrow, + Q: Hash + Eq, + { + if self.table.raw.is_null() { + return None; + } + + // Initialize the probe state. + let (h1, h2) = self.hash(key); + let mut probe = Probe::start(h1, self.table.mask); + + // Probe until we reach the limit. + let copying = 'probe: loop { + if probe.len > self.table.limit { + break None; + } + + // Load the entry metadata first for cheap searches. + let meta = unsafe { self.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 None; + } + + // Check for a potential match. + if meta != h2 { + probe.next(self.table.mask); + continue 'probe; + } + + // Load the full entry. + let mut entry = + unsafe { guard.protect(self.table.entry(probe.i), Ordering::Acquire) }.unpack(); + + // The entry was deleted, keep probing. + if entry.ptr.is_null() { + probe.next(self.table.mask); + continue 'probe; + } + + // Check for a full match. + if unsafe { (*entry.ptr).key.borrow() != key } { + probe.next(self.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 { + match unsafe { self.update_at(probe.i, entry, Entry::TOMBSTONE, guard) } { + // 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 { + self.table + .meta(probe.i) + .store(meta::TOMBSTONE, Ordering::Release) + }; + + // Decrement the table length. + let count = self.root.count.get(guard.thread().id()); + count.fetch_sub(1, Ordering::Relaxed); + + let entry = unsafe { &(*entry.ptr) }; + return Some((&entry.key, &entry.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 None, + + // Lost to a concurrent update, retry. + UpdateStatus::Found(EntryStatus::Value(found)) => entry = found, + } + } + }; + + match self.root.resize { + ResizeMode::Blocking => match copying { + // The entry we want to remove is being copied. + Some(_) => { + // In blocking mode we must complete the resize before proceeding. + let next_table = self.help_copy(guard, false); + + // Continue in the new table. + return self.as_ref(next_table).remove_inner(key, help_copy, guard); + } + // If we went over the probe limit, the key is not in this table. + None => None, + }, + + ResizeMode::Incremental(_) => { + // The entry we want to remove is being copied. + if let Some(i) = copying { + let next_table = self.next_table_ref().unwrap(); + + // Help out with the copy. + if help_copy { + self.help_copy(guard, false); + } + + // Wait for the entry to be copied. + self.wait_copied(i); + + // Continue in the new table. + return next_table.remove_inner(key, false, guard); + } + + // In incremental resize mode, we have to check the next table if we found + // a copied entry or went over the probe limit. + if let Some(next_table) = self.next_table_ref() { + // Help out with the copy. + if help_copy { + self.help_copy(guard, false); + } + + // Continue in the new table. + return next_table.remove_inner(key, false, guard); + } + + // Otherwise, the key is not in the table. + None + } + } + } + + // Attempts to insert an entry at the given index. + #[inline] + unsafe fn insert_at( + &self, + i: usize, + meta: u8, + new_entry: *mut Entry, + guard: &impl Guard, + ) -> InsertStatus { + let entry = unsafe { self.table.entry(i) }; + + // Try to claim the empty entry. + let found = match entry.compare_exchange( + ptr::null_mut(), + new_entry, + Ordering::Release, + Ordering::Acquire, + ) { + // Successfully claimed the entry. + Ok(_) => { + // Update the metadata table. + unsafe { self.table.meta(i).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(_) => { + // Protect the entry before accessing it. + let found = guard.protect(entry, Ordering::Acquire).unpack(); + + // Re-check the entry status. + match EntryStatus::from(found) { + EntryStatus::Value(found) | EntryStatus::Copied(found) => { + // 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.root.hasher.hash_one(&(*found.ptr).key); + (meta::h2(hash), EntryStatus::Value(found)) + } + + // The entry was deleted or null copied. + EntryStatus::Null => (meta::TOMBSTONE, EntryStatus::Null), + } + } + + // 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 self.table.meta(i).load(Ordering::Relaxed) == meta::EMPTY { + self.table.meta(i).store(meta, Ordering::Release); + } + + InsertStatus::Found(status) + } + + // Attempts to replace the value of an existing entry at the given index. + #[inline] + unsafe fn update_at( + &self, + i: usize, + current: Tagged>, + new_entry: *mut Entry, + guard: &impl Guard, + ) -> UpdateStatus { + let entry = unsafe { self.table.entry(i) }; + + // Try to perform the update. + let found = match entry.compare_exchange_weak( + current.raw, + new_entry, + Ordering::Release, + Ordering::Acquire, + ) { + // Successfully updated. + Ok(_) => unsafe { + // Safety: The old value is now unreachable from this table. + self.defer_retire(current, guard); + return UpdateStatus::Replaced(current); + }, + + // Lost to a concurrent update. + Err(found) => found.unpack(), + }; + + let status = match EntryStatus::from(found) { + EntryStatus::Value(_) => { + // Protect the entry before accessing it. + let found = guard.protect(entry, Ordering::Acquire).unpack(); + + // Re-check the entry status. + EntryStatus::from(found) + } + + // The entry was copied. + // + // We don't need to protect the entry as we never access it, + // we wait for it to be copied and continue in the new table. + EntryStatus::Copied(entry) => EntryStatus::Copied(entry), + + // The entry was deleted. + removed => removed, + }; + + UpdateStatus::Found(status) + } + + // Reserve capacity for `additional` more elements. + #[inline] + pub fn reserve(&mut self, additional: usize, guard: &impl Guard) { + // The table has not yet been allocated, try to initialize it. + if self.table.raw.is_null() && self.init(Some(probe::entries_for(additional))) { + return; + } + + loop { + let capacity = + probe::entries_for(self.root.count.sum().checked_add(additional).unwrap()); + + // We have enough capacity. + if self.table.len() >= capacity { + return; + } + + // Race to allocate the new table. + self.get_or_alloc_next(Some(capacity)); + + // Force the copy to complete. + // + // Note that this is not strictly necessary for a `reserve` operation. + self.table = self.help_copy(guard, true); + } + } + + // Remove all entries from this table. + #[inline] + pub fn clear(&mut self, guard: &impl Guard) { + if self.table.raw.is_null() { + return; + } + + // Get a clean copy of the table to delete from. + self.linearize(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..self.table.len() { + // Load the entry to delete. + let mut entry = + unsafe { guard.protect(self.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. + let result = unsafe { + self.table.entry(i).compare_exchange( + entry.raw, + Entry::TOMBSTONE, + Ordering::Release, + Ordering::Acquire, + ) + }; + + match result { + // Successfully deleted the entry. + Ok(_) => unsafe { + // Update the metadata table. + self.table.meta(i).store(meta::TOMBSTONE, Ordering::Release); + + // Decrement the table length. + let count = self.root.count.get(guard.thread().id()); + count.fetch_sub(1, Ordering::Relaxed); + + // Safety: We just removed the old value from this table. + self.defer_retire(entry, guard); + + continue 'probe; + }, + + // Lost to a concurrent update, retry. + Err(found) => entry = found.unpack(), + } + } + } + + // A resize prevented us from deleting all the entries in this table. + // + // Complete the resize and retry in the new table. + if copying { + let next_table = self.help_copy(guard, true); + return self.as_ref(next_table).clear(guard); + } + } + + // Retains only the elements specified by the predicate.. + #[inline] + pub fn retain(&mut self, mut f: F, guard: &impl Guard) + where + F: FnMut(&K, &V) -> bool, + { + if self.table.raw.is_null() { + return; + } + + // Get a clean copy of the table to delete from. + self.linearize(guard); + + let mut copying = false; + 'probe: for i in 0..self.table.len() { + // Load the entry metadata first to ensure consistency with calls to `get` + // for entries that are retained. + let meta = unsafe { self.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. + let mut entry = + unsafe { guard.protect(self.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; + } + + let entry_ref = unsafe { &*entry.raw }; + + // Should we retain this entry? + if f(&entry_ref.key, &entry_ref.value) { + continue 'probe; + } + + // Try to delete the entry. + let result = unsafe { + self.table.entry(i).compare_exchange( + entry.raw, + Entry::TOMBSTONE, + Ordering::Release, + Ordering::Acquire, + ) + }; + + match result { + // Successfully deleted the entry. + Ok(_) => unsafe { + // Update the metadata table. + self.table.meta(i).store(meta::TOMBSTONE, Ordering::Release); + + // Decrement the table length. + let count = self.root.count.get(guard.thread().id()); + count.fetch_sub(1, Ordering::Relaxed); + + // Safety: We just removed the old value from this table. + self.defer_retire(entry, guard); + + continue 'probe; + }, + + // Lost to a concurrent update, retry. + Err(found) => entry = found.unpack(), + } + } + } + + // A resize prevented us from deleting all the entries in this table. + // + // Complete the resize and retry in the new table. + if copying { + let next_table = self.help_copy(guard, true); + return self.as_ref(next_table).retain(f, guard); + } + } + + // Returns an iterator over the keys and values of this table. + #[inline] + pub fn iter<'g, G>(&mut self, guard: &'g G) -> Iter<'g, K, V, G> + where + G: Guard, + { + if self.table.raw.is_null() { + return Iter { + i: 0, + guard, + table: self.table, + }; + } + + // Get a clean copy of the table to iterate over. + self.linearize(guard); + + Iter { + i: 0, + guard, + table: self.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.root.hasher.hash_one(key); + (meta::h1(hash), meta::h2(hash)) + } +} + +// A wrapper around a CAS function that manages the computed state. +struct ComputeState { + compute: F, + insert: Option, + update: Option>, +} + +struct CachedUpdate { + input: *mut Entry, + 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. + #[inline] + fn next(&mut self, entry: Option<*mut Entry>) -> Operation { + match entry { + Some(entry) => match self.update.take() { + Some(CachedUpdate { input, output }) if input == entry => output, + _ => { + let entry = unsafe { &*entry }; + (self.compute)(Some((&entry.key, &entry.value))) + } + }, + None => match self.insert.take() { + Some(value) => Operation::Insert(value), + None => (self.compute)(None), + }, + } + } + + // Restores the state if an operation fails. + #[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. + #[inline] + fn init(&mut self, collector: &Collector, guard: &impl Guard) -> *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 { + link: collector.link(guard), + value: MaybeUninit::uninit(), + key, + })) + })) + .unwrap_or_else(|_| std::process::abort()); + ptr::write(self, LazyEntry::Init(entry)); + entry + } + } + } + } +} + +// Update operations. +impl<'root, K, V, S> HashMapRef<'root, K, V, S> +where + K: Hash + Eq, + S: BuildHasher, +{ + /// 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>(&mut self, key: K, f: F, guard: &'g impl Guard) -> &'g V + where + F: FnOnce() -> V, + K: 'g, + { + let mut f = Some(f); + let compute = |entry| match entry { + // Return the existing value. + Some((_, current)) => Operation::Abort(current), + // Insert the initial value. + None => Operation::Insert((f.take().unwrap())()), + }; + + match self.compute(key, compute, guard) { + Compute::Aborted(value) => value, + Compute::Inserted(_, value) => value, + _ => unreachable!(), + } + } + + // Updates an existing entry atomically, returning the value that was inserted. + #[inline] + pub fn update<'g, F>(&mut self, key: K, mut update: F, guard: &'g impl Guard) -> Option<&'g V> + where + F: FnMut(&V) -> V, + K: 'g, + { + let compute = |entry| match entry { + Some((_, value)) => Operation::Insert(update(value)), + None => Operation::Abort(()), + }; + + match self.compute(key, compute, guard) { + Compute::Updated { + new: (_, value), .. + } => Some(value), + 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>( + &mut self, + key: K, + update: U, + f: F, + guard: &'g impl Guard, + ) -> &'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. + None => Operation::Insert((f.take().unwrap())()), + }; + + match self.compute(key, compute, guard) { + Compute::Updated { + new: (_, value), .. + } => value, + 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>( + &mut self, + key: K, + compute: F, + guard: &'g impl Guard, + ) -> 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), true, guard) }; + + // Deallocate the entry if it was not inserted. + if matches!(result, Compute::Removed(..) | Compute::Aborted(_)) { + if let LazyEntry::Init(entry) = entry { + // Safety: We allocated this box above and it was 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 pointer. + #[inline] + unsafe fn compute_with<'g, F, T>( + &mut self, + new_entry: &mut LazyEntry, + mut state: ComputeState, + help_copy: bool, + guard: &'g impl Guard, + ) -> Compute<'g, K, V, T> + where + F: FnMut(Option<(&'g K, &'g V)>) -> Operation, + { + // The table has not yet been allocated. + if self.table.raw.is_null() { + // Compute the value to insert. + match 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. + self.init(None); + } + + // Initialize the probe state. + let (h1, h2) = self.hash(new_entry.key()); + let mut probe = Probe::start(h1, self.table.mask); + + // Probe until we reach the limit. + let copying = 'probe: loop { + if probe.len > self.table.limit { + break None; + } + + // Load the entry metadata first for cheap searches. + let meta = unsafe { self.table.meta(probe.i) }.load(Ordering::Acquire); + + // The entry is empty. + let mut entry = if meta == meta::EMPTY { + // Compute the value to insert. + let value = match 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(&self.root.collector, guard); + unsafe { (*new_entry).value = MaybeUninit::new(value) } + + // Attempt to insert. + match self.insert_at(probe.i, h2, new_entry.cast(), guard) { + // Successfully inserted. + InsertStatus::Inserted => { + // Increment the table length. + let count = self.root.count.get(guard.thread().id()); + count.fetch_add(1, Ordering::Relaxed); + + let new = unsafe { &*new_entry.cast::>() }; + return Compute::Inserted(&new.key, &new.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)) => { + // Save the previous value in case the update fails. + 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) => { + // Save the previous value. + let value = unsafe { (*new_entry).value.assume_init_read() }; + state.restore(None, Operation::Insert(value)); + + // Continue probing. + probe.next(self.table.mask); + continue 'probe; + } + } + } + // Found a potential match. + else if meta == h2 { + // Load the full entry. + let found = guard + .protect(self.table.entry(probe.i), Ordering::Acquire) + .unpack(); + + // The entry was deleted, keep probing. + if found.ptr.is_null() { + probe.next(self.table.mask); + continue 'probe; + } + + // If the key matches, we might be able to update the value. + found + } + // Otherwise, continue probing. + else { + probe.next(self.table.mask); + continue 'probe; + }; + + // Check for a full match. + if unsafe { (*entry.ptr).key != *new_entry.key() } { + probe.next(self.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. + let failure = match 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(&self.root.collector, guard); + unsafe { (*new_entry).value = MaybeUninit::new(value) } + + // Try to perform the update. + match self.update_at(probe.i, entry, new_entry.cast(), guard) { + // Successfully updated. + UpdateStatus::Replaced(entry) => { + let old = unsafe { &(*entry.ptr) }; + let new = unsafe { &*new_entry.cast::>() }; + + return Compute::Updated { + old: (&old.key, &old.value), + new: (&new.key, &new.value), + }; + } + + // The update failed. + failure => { + // Save the previous value. + let value = unsafe { (*new_entry).value.assume_init_read() }; + state.restore(Some(entry.ptr), Operation::Insert(value)); + + failure + } + } + } + Operation::Remove => { + match unsafe { self.update_at(probe.i, entry, Entry::TOMBSTONE, guard) } { + // 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 { + self.table + .meta(probe.i) + .store(meta::TOMBSTONE, Ordering::Release) + }; + + // Decrement the table length. + let count = self.root.count.get(guard.thread().id()); + count.fetch_sub(1, Ordering::Relaxed); + + let entry = unsafe { &(*entry.ptr) }; + return Compute::Removed(&entry.key, &entry.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. + match 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(self.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!(), + } + } + }; + + match self.root.resize { + ResizeMode::Blocking => match copying { + // The entry we want to update is being copied. + Some(_) => { + // In blocking mode we must complete the resize before proceeding. + let next_table = self.help_copy(guard, false); + + // Continue in the new table. + return self + .as_ref(next_table) + .compute_with(new_entry, state, help_copy, guard); + } + // If we went over the probe limit, the key is not in this table. + None => {} + }, + + ResizeMode::Incremental(_) => { + // The entry we want to update is being copied. + if let Some(i) = copying { + let mut next_table = self.next_table_ref().unwrap(); + + // Help out with the copy. + if help_copy { + self.help_copy(guard, false); + } + + // Wait for the entry to be copied. + // + // We could perform an update based on the copied entry and + // race to insert it into the table here instead. 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); + + // Continue in the new table. + return next_table.compute_with(new_entry, state, false, guard); + } + + // In incremental resize mode, we have to check the next table if we found + // a copied entry or went over the probe limit. + if let Some(mut next_table) = self.next_table_ref() { + // Help out with the copy. + if help_copy { + self.help_copy(guard, false); + } + + // Continue in the new table. + return next_table.compute_with(new_entry, state, false, guard); + } + } + } + + // Otherwise, the key is not in the table. + match state.next(None) { + // Need to insert into the new table. + op @ Operation::Insert(_) => { + // Trigger a resize. + self.get_or_alloc_next(None); + + // Help out with the resize. + let next_table = self.help_copy(guard, false); + + state.restore(None, op); + return self + .as_ref(next_table) + .compute_with(new_entry, state, false, guard); + } + Operation::Remove => panic!("Cannot remove `None` entry."), + Operation::Abort(value) => return Compute::Aborted(value), + } + } +} + +// Resize operations. +impl<'root, K, V, S> HashMapRef<'root, K, V, S> +where + K: Hash + Eq, + S: BuildHasher, +{ + // Returns a reference to the given table. + #[inline] + fn as_ref(&self, table: Table) -> HashMapRef<'root, K, V, S> { + HashMapRef { + table, + root: self.root, + } + } + + // Returns a reference to the next table, if it has already been created. + #[inline] + fn next_table_ref(&self) -> Option> { + let state = self.table.state(); + let next = state.next.load(Ordering::Acquire); + + if !next.is_null() { + return unsafe { Some(self.as_ref(Table::from_raw(next))) }; + } + + None + } + + // Allocate the initial table. + #[cold] + fn init(&mut self, capacity: Option) -> bool { + const CAPACITY: usize = 32; + + // Allocate the table and mark it as the root. + let mut table = Table::::alloc(capacity.unwrap_or(CAPACITY), &self.root.collector); + *table.state_mut().status.get_mut() = State::PROMOTED; + + // Race to write the initial table. + match self.root.table.compare_exchange( + ptr::null_mut(), + table.raw, + Ordering::Release, + Ordering::Acquire, + ) { + // Successfully initialized the table. + Ok(_) => { + self.table = table; + true + } + + // Someone beat us, deallocate our table and use the table that was written. + Err(found) => { + unsafe { Table::dealloc(table) } + self.table = unsafe { Table::from_raw(found) }; + false + } + } + } + + // Returns the next table, allocating it has not already been created. + #[cold] + fn get_or_alloc_next(&self, capacity: Option) -> Table { + // Avoid spinning in tests, which can hide race conditions. + const SPIN_ALLOC: usize = if cfg!(any(test, debug_assertions)) { + 1 + } else { + 7 + }; + + let state = self.table.state(); + let next = state.next.load(Ordering::Acquire); + + // The next table is already allocated. + if !next.is_null() { + return unsafe { Table::from_raw(next) }; + } + + // 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(); + } + + let next = state.next.load(Ordering::Acquire); + if !next.is_null() { + // The table was initialized. + return unsafe { Table::from_raw(next) }; + } + + spun += 1; + } + + // Otherwise, we have to block. + state.allocating.lock().unwrap() + } + }; + + let next = state.next.load(Ordering::Acquire); + if !next.is_null() { + // The table was allocated while we were waiting for the lock. + return unsafe { Table::from_raw(next) }; + } + + let next_capacity = match cfg!(papaya_stress) { + // Never grow the table to stress the incremental resizing algorithm. + true => self.table.len(), + // Double the table capacity if we are at least 50% full. + // + // Loading the length here is quite expensive, we may want to consider + // a probabilistic counter to detect high-deletion workloads. + false if self.root.len() >= (self.table.len() >> 1) => self.table.len() << 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 => self.table.len(), + }; + + 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, &self.root.collector); + 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] + fn help_copy(&self, guard: &impl Guard, copy_all: bool) -> Table { + match self.root.resize { + ResizeMode::Blocking => self.help_copy_blocking(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 self.next_table_ref().unwrap().table; + } + + 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, guard: &impl Guard) -> Table { + // Load the next table. + let next = self.table.state().next.load(Ordering::Acquire); + debug_assert!(!next.is_null()); + let mut next = unsafe { Table::::from_raw(next) }; + + 'copy: loop { + // Make sure we are copying to the correct table. + while next.state().status.load(Ordering::Relaxed) == State::ABORTED { + next = self.as_ref(next).get_or_alloc_next(None); + } + + // The copy already completed + if self.try_promote(next, 0, guard) { + return next; + } + + let copy_chunk = self.table.len().min(4096); + + loop { + // Every entry has already been claimed. + if next.state().claim.load(Ordering::Relaxed) >= self.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 >= self.table.len() { + break; + } + + // Copy the entry. + if !self.copy_at_blocking(i, next, guard) { + // This table doesn't have space for the next entry. + // + // Abort the current resize. + next.state().status.store(State::ABORTED, Ordering::Relaxed); + + // Allocate the next table. + let allocated = self.as_ref(next).get_or_alloc_next(None); + + // Wake anyone waiting for us to finish. + atomic_wait::wake_all(&next.state().status); + + // Retry in a new table. + next = allocated; + continue 'copy; + } + + copied += 1; + } + + // Are we done? + if self.try_promote(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; + } + } + + // 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 + }; + + let status = next.state().status.load(Ordering::Relaxed); + + // 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 { + fence(Ordering::Acquire); + 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. + atomic_wait::wait(&next.state().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. + #[inline] + fn copy_at_blocking(&self, i: usize, next_table: Table, guard: &impl Guard) -> bool { + // Mark the entry as copying. + let entry = unsafe { + self.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. + unsafe { self.table.meta(i).store(meta::TOMBSTONE, Ordering::Release) }; + return true; + } + + // Copy the value to the new table. + unsafe { + self.as_ref(next_table) + .insert_copy(entry.ptr.unpack(), false, 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 Guard) -> Table { + // Always help the highest priority root resize. + let root = self.root.root(guard); + if self.table.raw != root.table.raw { + return root.help_copy_incremental(chunk, block, guard); + } + + // Load the next table. + let next = self.table.state().next.load(Ordering::Acquire); + + // The copy we tried to help was already promoted. + if next.is_null() { + return self.table; + } + + let next = unsafe { Table::::from_raw(next) }; + + loop { + // The copy already completed. + if self.try_promote(next, 0, guard) { + return next; + } + + loop { + // Every entry has already been claimed. + if next.state().claim.load(Ordering::Relaxed) >= self.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 >= self.table.len() { + break; + } + + // Copy the entry. + self.copy_at_incremental(i, 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(next, copied, guard) || !block { + return next; + } + } + + // There are no entries that we can copy, block if necessary. + if !block { + return next; + } + + 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. + let status = next.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. + atomic_wait::wait(&next.state().status, State::PENDING); + } + } + } + + // Copy the entry at the given index to the new table. + #[inline] + fn copy_at_incremental(&self, i: usize, next_table: Table, guard: &impl Guard) { + // Mark the entry as copying. + let entry = unsafe { + self.table + .entry(i) + .fetch_or(Entry::COPYING, Ordering::AcqRel) + .unpack() + }; + + // The entry is a tombstone. + if entry.raw == Entry::TOMBSTONE { + return; + } + + // 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. + unsafe { self.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 = entry.map_tag(|addr| addr | Entry::BORROWED); + + // Copy the value to the new table. + unsafe { + self.as_ref(next_table) + .insert_copy(new_entry, true, guard) + .unwrap(); + } + + // Mark the entry as copied. + let copied = entry + .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. + unsafe { self.table.entry(i).store(copied, Ordering::SeqCst) }; + + // Notify any writers that the copy has completed. + unsafe { self.table.state().parker.unpark(self.table.entry(i)) }; + } + + // 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. + #[inline] + unsafe fn insert_copy( + &self, + new_entry: Tagged>, + resize: bool, + guard: &impl Guard, + ) -> Option<(Table, usize)> { + // Safety: The new entry is guaranteed to be valid by the caller. + let key = unsafe { &(*new_entry.ptr).key }; + + // Initialize the probe state. + let (h1, h2) = self.hash(key); + let mut probe = Probe::start(h1, self.table.mask); + + // Probe until we reach the limit. + while probe.len <= self.table.limit { + // Load the entry metadata first for cheap searches. + let meta = unsafe { self.table.meta(probe.i) }.load(Ordering::Acquire); + + // The entry is empty, try to insert. + if meta == meta::EMPTY { + let entry = unsafe { self.table.entry(probe.i) }; + + // Try to claim the entry. + match entry.compare_exchange( + ptr::null_mut(), + new_entry.raw, + Ordering::Release, + Ordering::Acquire, + ) { + // Successfully inserted. + Ok(_) => { + // Update the metadata table. + unsafe { self.table.meta(probe.i).store(h2, Ordering::Release) }; + return Some((self.table, probe.i)); + } + Err(found) => { + // The entry was deleted or copied. + let meta = if found.unpack().ptr.is_null() { + meta::TOMBSTONE + } else { + // Protect the entry before accessing it. + let found = guard.protect(entry, Ordering::Acquire).unpack(); + + // Recheck the pointer. + if found.ptr.is_null() { + meta::TOMBSTONE + } else { + // Ensure the meta table is updated to avoid breaking the probe chain. + let hash = self.root.hasher.hash_one(&(*found.ptr).key); + meta::h2(hash) + } + }; + + if self.table.meta(probe.i).load(Ordering::Relaxed) == meta::EMPTY { + self.table.meta(probe.i).store(meta, Ordering::Release); + } + } + } + } + + // Continue probing. + probe.next(self.table.mask); + } + + if !resize { + return None; + } + + // Insert into the next table. + let next_table = self.get_or_alloc_next(None); + self.as_ref(next_table) + .insert_copy(new_entry, resize, guard) + } + + // Update the copy state and attempt to promote a table to the root. + // + // Returns `true` if the table was promoted. + #[inline] + fn try_promote(&self, next: Table, copied: usize, guard: &impl Guard) -> bool { + // Update the copy count. + let copied = if copied > 0 { + next.state().copied.fetch_add(copied, Ordering::AcqRel) + copied + } else { + next.state().copied.load(Ordering::Acquire) + }; + + // If we copied all the entries in the table, we can try to promote. + if copied == self.table.len() { + let root = self.root.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 self.table.raw == root { + // Try to update the root. + if self + .root + .table + .compare_exchange( + self.table.raw, + next.raw, + Ordering::Release, + Ordering::Acquire, + ) + .is_ok() + { + // Successfully promoted the table. + next.state() + .status + .store(State::PROMOTED, Ordering::Release); + + unsafe { + // Retire the old table. + // + // Note that we do not drop entries because they have been copied to the + // new root. + guard.defer_retire(self.table.raw, |link| { + let raw: *mut RawTable = link.cast(); + let table = Table::::from_raw(raw); + drop_table::(table); + }); + } + } + + // Wake up any writers waiting for the resize to complete. + atomic_wait::wake_all(&next.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(&mut self, guard: &impl Guard) { + if self.root.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 self.next_table_ref().is_some() { + self.table = self.help_copy(guard, true); + } + } + } + + // Wait for an incremental copy of a given entry to complete. + #[inline] + fn wait_copied(&self, i: usize) { + // Avoid spinning in tests, which can hide race conditions. + const SPIN_WAIT: usize = if cfg!(any(test, debug_assertions)) { + 1 + } else { + 5 + }; + + let parker = &self.table.state().parker; + let entry = unsafe { self.table.entry(i) }; + + // Spin for a short while, waiting for the entry to be copied. + for spun in 0..SPIN_WAIT { + for _ in 0..(spun * spun) { + hint::spin_loop(); + } + + // The entry was copied. + let entry = entry.load(Ordering::Acquire).unpack(); + if entry.tag() & Entry::COPIED != 0 { + return; + } + } + + // Park until the copy completes. + 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 unreachable from the current table. + #[inline] + unsafe fn defer_retire(&self, entry: Tagged>, guard: &impl Guard) { + match self.root.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 we are not the root. Thus we can safely retire. + unsafe { guard.defer_retire(entry.ptr, Entry::reclaim::) }; + return; + } + + let root = self.root.table.load(Ordering::Relaxed); + + // Check if our table, or any subsequent table, is the root. + let mut next = Some(self.clone()); + while let Some(map) = next { + if map.table.raw == root { + // Safety: The root table is our table or a table that succeeds ours. + // Thus any previous tables are unreachable and we can safely retire. + unsafe { guard.defer_retire(entry.ptr, Entry::reclaim::) }; + return; + } + + next = map.next_table_ref(); + } + + // 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. + fence(Ordering::Acquire); + let mut prev = self.as_ref(unsafe { Table::::from_raw(root) }); + + loop { + let next = prev.next_table_ref().unwrap(); + + // Defer the entry to be retired by the table we are copying from. + if next.table.raw == self.table.raw { + prev.table.state().deferred.defer(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: Guard, +{ + 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`. + 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. + let entry = unsafe { + self.guard + .protect(self.table.entry(self.i), Ordering::Acquire) + .unpack() + }; + + // The entry was deleted. + if entry.ptr.is_null() { + self.i += 1; + continue; + } + + // Read the key and value. + let entry = unsafe { (&(*entry.ptr).key, &(*entry.ptr).value) }; + + self.i += 1; + return Some(entry); + } + } +} + +// 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/Sync. +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 Clone for HashMapRef<'_, K, V, S> { + #[inline] + fn clone(&self) -> Self { + HashMapRef { + table: self.table, + root: self.root, + } + } +} + +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. + unsafe { self.collector.reclaim_all() }; + + // Drop all nested tables and entries. + while !raw.is_null() { + let mut table = unsafe { Table::::from_raw(raw) }; + let next = *table.state_mut().next.get_mut(); + unsafe { drop_entries::(table) }; + unsafe { drop_table::(table) }; + raw = next; + } + } +} + +// Drop all entries in this table. +unsafe fn drop_entries(table: Table) { + for i in 0..table.len() { + 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. + unsafe { Entry::reclaim::(entry.ptr.cast()) } + } +} + +// Drop the table allocation. +unsafe fn drop_table(mut table: Table) { + // Safety: `drop_table` is being called from `reclaim_all` in `Drop` or + // a table is being reclaimed by our thread. In both cases, the collector + // is still alive and safe to access through the state pointer. + let collector = unsafe { &*table.state().collector }; + + // Drop any entries that were deferred during an incremental resize. + // + // Safety: A deferred entry was retired after it was made unreachable + // from the next table during a resize. Because our table was still accessible + // for this entry to be deferred, our table must have been retired *after* the + // entry was made accessible in the next table. Now that our table is being reclaimed, + // the entry has thus been totally removed from the map, and can be safely retired. + unsafe { + table + .state_mut() + .deferred + .retire_all(collector, Entry::reclaim::) + } + + // Deallocate the table. + 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/probe.rs b/src/raw/probe.rs new file mode 100644 index 00000000..2805e339 --- /dev/null +++ b/src/raw/probe.rs @@ -0,0 +1,47 @@ +// A quadratic probe sequence. +#[derive(Default)] +pub struct Probe { + // The current index in the probe sequence. + pub i: usize, + // The current length of the probe sequence. + pub len: usize, + // The current quadratic stride. + stride: usize, +} + +impl Probe { + // Initialize the probe sequence. + #[inline] + pub fn start(hash: usize, mask: usize) -> Probe { + Probe { + i: hash & mask, + len: 0, + stride: 0, + } + } + + // Increment the probe sequence. + #[inline] + pub fn next(&mut self, mask: usize) { + self.len += 1; + self.stride += 1; + self.i = (self.i + self.stride) & mask; + } +} + +// The maximum probe length for table operations. +// +// Estimating a load factor for the hash-table based on probe lengths allows +// the hash-table to avoid loading the length every insert, which is a source +// of contention. +pub fn limit(capacity: usize) -> usize { + // 5 * log2(capacity): Testing shows this gives us a ~85% load factor. + 5 * ((usize::BITS as usize) - (capacity.leading_zeros() as usize) - 1) +} + +// Returns an estimate of the number of entries needed to hold `capacity` elements. +pub fn entries_for(capacity: usize) -> usize { + // We should rarely resize before 75%. + let capacity = capacity.checked_mul(8).expect("capacity overflow") / 6; + capacity.next_power_of_two() +}