diff --git a/Cargo.toml b/Cargo.toml index aaf7636..37a0d69 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -16,6 +16,7 @@ exclude = ["assets/*"] equivalent = "1" seize = "0.5" serde = { version = "1", optional = true } +rayon = { version = "1", optional = true } [dev-dependencies] rand = "0.8" @@ -30,6 +31,7 @@ serde_json = "1" [features] default = [] serde = ["dep:serde"] +rayon = ["dep:rayon"] [profile.test] inherits = "release" diff --git a/src/map.rs b/src/map.rs index f149a25..286b42c 100644 --- a/src/map.rs +++ b/src/map.rs @@ -1527,6 +1527,15 @@ where } } + /// Returns a parallel iterator over all key-value pairs. + /// + /// Requires the `rayon` feature. + #[cfg(feature = "rayon")] + #[inline] + pub fn par_iter(&self) -> raw::ParIter<'_, K, V, MapGuard> { + self.map.raw.par_iter(&self.guard) + } + /// An iterator visiting all keys in arbitrary order. /// The iterator element type is `&K`. /// @@ -1572,6 +1581,22 @@ where } } +#[cfg(feature = "rayon")] +impl<'a, K, V, S, G> rayon::iter::IntoParallelIterator for &'a HashMapRef<'_, K, V, S, G> +where + K: Hash + Eq + Sync, + V: Sync, + S: BuildHasher, + G: Guard + Sync, +{ + type Item = (&'a K, &'a V); + type Iter = raw::ParIter<'a, K, V, MapGuard>; + + fn into_par_iter(self) -> Self::Iter { + self.map.raw.par_iter(&self.guard) + } +} + /// An iterator over a map's entries. /// /// This struct is created by the [`iter`](HashMap::iter) method on [`HashMap`]. See its documentation for details. diff --git a/src/raw/mod.rs b/src/raw/mod.rs index 79eb825..bd3ead8 100644 --- a/src/raw/mod.rs +++ b/src/raw/mod.rs @@ -3,6 +3,8 @@ mod probe; pub(crate) mod utils; +#[cfg(feature = "rayon")] +use std::fmt; use std::hash::{BuildHasher, Hash}; use std::mem::MaybeUninit; use std::sync::atomic::{AtomicPtr, AtomicU8, AtomicUsize, Ordering}; @@ -1231,6 +1233,35 @@ where Iter { i: 0, guard, table } } + /// Returns a parallel iterator over all key/value pairs in the table. + #[cfg(feature = "rayon")] + #[inline] + pub fn par_iter<'g, G>(&self, guard: &'g G) -> ParIter<'g, K, V, G> + where + G: VerifiedGuard, + { + let root = self.root(guard); + + if root.raw.is_null() { + return ParIter { + start: 0, + end: 0, + guard, + table: root, + }; + } + + let table = self.linearize(root, guard); + let len = table.len(); + + ParIter { + start: 0, + end: len, + guard, + table, + } + } + /// Returns the h1 and h2 hash for the given key. #[inline] fn hash(&self, key: &Q) -> (usize, u8) @@ -2718,6 +2749,120 @@ impl Clone for Iter<'_, K, V, G> { } } +#[cfg(feature = "rayon")] +/// A parallel iterator over the entries of a hash table. +pub struct ParIter<'g, K, V, G> { + start: usize, + end: usize, + table: Table>, + guard: &'g G, +} + +#[cfg(feature = "rayon")] +impl<'g, K, V, G> fmt::Debug for ParIter<'g, K, V, G> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("ParIter") + .field("start", &self.start) + .field("end", &self.end) + .finish() + } +} + +#[cfg(feature = "rayon")] +unsafe impl Send for ParIter<'_, K, V, G> +where + K: Sync, + V: Sync, + G: Sync, +{ +} + +#[cfg(feature = "rayon")] +unsafe impl Sync for ParIter<'_, K, V, G> +where + K: Sync, + V: Sync, + G: Sync, +{ +} + +#[cfg(feature = "rayon")] +impl<'g, K, V, G> rayon::iter::ParallelIterator for ParIter<'g, K, V, G> +where + K: Sync + 'g, + V: Sync + 'g, + G: VerifiedGuard + Sync, +{ + type Item = (&'g K, &'g V); + + fn drive_unindexed(self, consumer: C) -> C::Result + where + C: rayon::iter::plumbing::UnindexedConsumer, + { + rayon::iter::plumbing::bridge_unindexed(self, consumer) + } +} + +#[cfg(feature = "rayon")] +impl<'g, K, V, G> rayon::iter::plumbing::UnindexedProducer for ParIter<'g, K, V, G> +where + K: Sync + 'g, + V: Sync + 'g, + G: VerifiedGuard + Sync, +{ + type Item = (&'g K, &'g V); + + fn split(mut self) -> (Self, Option) { + let len = self.end - self.start; + if len <= 1 { + (self, None) + } else { + let mid = self.start + len / 2; + let right = ParIter { + start: mid, + end: self.end, + table: self.table, + guard: self.guard, + }; + self.end = mid; + (self, Some(right)) + } + } + + fn fold_with(self, mut folder: F) -> F + where + F: rayon::iter::plumbing::Folder, + { + if self.table.raw.is_null() { + return folder; + } + let mut i = self.start; + while i < self.end { + // Safety: i is in-bounds for the table length. + let meta = unsafe { self.table.meta(i) }.load(Ordering::Acquire); + if matches!(meta, meta::EMPTY | meta::TOMBSTONE) { + i += 1; + continue; + } + + let entry = self + .guard + .protect(unsafe { self.table.entry(i) }, Ordering::Acquire) + .unpack(); + + if entry.ptr.is_null() { + i += 1; + continue; + } + + let entry_ref = unsafe { &(*entry.ptr) }; + folder = folder.consume((&entry_ref.key, &entry_ref.value)); + i += 1; + } + folder + } +} + impl Drop for HashMap { fn drop(&mut self) { let mut raw = *self.table.get_mut(); diff --git a/src/raw/utils/mod.rs b/src/raw/utils/mod.rs index ce36dc8..da1dadc 100644 --- a/src/raw/utils/mod.rs +++ b/src/raw/utils/mod.rs @@ -5,7 +5,9 @@ mod tagged; pub use counter::Counter; pub use parker::Parker; +use seize::Guard as _; pub use stack::Stack; +use std::fmt; pub use tagged::{untagged, AtomicPtrFetchOps, StrictProvenance, Tagged, Unpack}; /// A `seize::Guard` that has been verified to belong to a given map. @@ -14,6 +16,17 @@ pub trait VerifiedGuard: seize::Guard {} #[repr(transparent)] pub struct MapGuard(G); +impl fmt::Debug for MapGuard +where + G: seize::Guard, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("MapGuard") + .field("thread_id", &self.thread_id()) + .finish() + } +} + impl MapGuard { /// Create a new `MapGuard`. /// diff --git a/src/set.rs b/src/set.rs index 367776f..92ae6e5 100644 --- a/src/set.rs +++ b/src/set.rs @@ -833,6 +833,17 @@ where raw: self.set.raw.iter(&self.guard), } } + + /// Returns a parallel iterator over all values in arbitrary order. + /// + /// Requires the `rayon` feature. + #[cfg(feature = "rayon")] + #[inline] + pub fn par_iter(&self) -> ParIter<'_, K, G> { + ParIter { + inner: self.set.raw.par_iter(&self.guard), + } + } } impl fmt::Debug for HashSetRef<'_, K, S, G> @@ -860,6 +871,23 @@ where } } +#[cfg(feature = "rayon")] +impl<'a, K, S, G> rayon::iter::IntoParallelIterator for &'a HashSetRef<'_, K, S, G> +where + K: Hash + Eq + Sync, + S: BuildHasher, + G: Guard + Sync, +{ + type Item = &'a K; + type Iter = ParIter<'a, K, G>; + + fn into_par_iter(self) -> Self::Iter { + ParIter { + inner: self.set.raw.par_iter(&self.guard), + } + } +} + /// An iterator over a set's entries. /// /// This struct is created by the [`iter`](HashSet::iter) method on [`HashSet`]. See its documentation for details. @@ -867,6 +895,29 @@ pub struct Iter<'g, K, G> { raw: raw::Iter<'g, K, (), MapGuard>, } +#[cfg(feature = "rayon")] +/// A parallel iterator over a set's entries. +#[derive(Debug)] +pub struct ParIter<'g, K, G> { + inner: raw::ParIter<'g, K, (), MapGuard>, +} + +#[cfg(feature = "rayon")] +impl<'g, K, G> rayon::iter::ParallelIterator for ParIter<'g, K, G> +where + K: Sync + 'g, + G: Guard + Sync, +{ + type Item = &'g K; + + fn drive_unindexed(self, consumer: C) -> C::Result + where + C: rayon::iter::plumbing::UnindexedConsumer, + { + self.inner.map(|(k, _)| k).drive_unindexed(consumer) + } +} + impl<'g, K: 'g, G> Iterator for Iter<'g, K, G> where G: Guard, diff --git a/tests/rayon_iter.rs b/tests/rayon_iter.rs new file mode 100644 index 0000000..ebb7a08 --- /dev/null +++ b/tests/rayon_iter.rs @@ -0,0 +1,68 @@ +#![cfg(feature = "rayon")] + +use rayon::prelude::*; + +mod common; +use common::{with_map, with_set}; + +#[test] +fn hashmap_par_iter() { + if cfg!(papaya_stress) { + return; + } + with_map::(|map| { + let map = map(); + let len = if cfg!(miri) { 100 } else { 10_000 }; + for i in 0..len { + map.pin_owned().insert(i, i + 1); + } + + let mut expected: Vec<_> = (0..len).map(|i| (i, i + 1)).collect(); + + let mut got: Vec<_> = map + .pin_owned() + .par_iter() + .map(|(&k, &v)| (k, v)) + .collect(); + got.sort(); + expected.sort(); + assert_eq!(expected, got); + + let mut via_trait: Vec<_> = (&map.pin_owned()) + .into_par_iter() + .map(|(&k, &v)| (k, v)) + .collect(); + via_trait.sort(); + assert_eq!(expected, via_trait); + }); +} + +#[test] +fn hashset_par_iter() { + if cfg!(papaya_stress) { + return; + } + with_set::(|set| { + let set = set(); + let len = if cfg!(miri) { 100 } else { 10_000 }; + for i in 0..len { + set.pin_owned().insert(i); + } + + let mut expected: Vec<_> = (0..len).collect(); + + let mut got: Vec<_> = set + .pin_owned() + .par_iter() + .copied() + .collect(); + got.sort(); + expected.sort(); + assert_eq!(expected, got); + + let mut via_trait: Vec<_> = (&set.pin_owned()).into_par_iter().copied().collect(); + via_trait.sort(); + assert_eq!(expected, via_trait); + }); +} +