Add rayon feature flag and implement IntoParIterator

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"

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<G>> {
+        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<G>>;
+
+    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.

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<Q>(&self, key: &Q) -> (usize, u8)
@@ -2718,6 +2749,120 @@ impl<K, V, G> 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<Entry<K, V>>,
+    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<K, V, G> Send for ParIter<'_, K, V, G>
+where
+    K: Sync,
+    V: Sync,
+    G: Sync,
+{
+}
+
+#[cfg(feature = "rayon")]
+unsafe impl<K, V, G> 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<C>(self, consumer: C) -> C::Result
+    where
+        C: rayon::iter::plumbing::UnindexedConsumer<Self::Item>,
+    {
+        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<Self>) {
+        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<F>(self, mut folder: F) -> F
+    where
+        F: rayon::iter::plumbing::Folder<Self::Item>,
+    {
+        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<K, V, S> Drop for HashMap<K, V, S> {
     fn drop(&mut self) {
         let mut raw = *self.table.get_mut();

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>(G);
 
+impl<G> fmt::Debug for MapGuard<G>
+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<G> MapGuard<G> {
     /// Create a new `MapGuard`.
     ///

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<K, S, G> fmt::Debug for HashSetRef<'_, K, S, G>
@@ -860,13 +871,53 @@ 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.
 pub struct Iter<'g, K, G> {
     raw: raw::Iter<'g, K, (), MapGuard<G>>,
 }
 
+#[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<G>>,
+}
+
+#[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<C>(self, consumer: C) -> C::Result
+    where
+        C: rayon::iter::plumbing::UnindexedConsumer<Self::Item>,
+    {
+        self.inner.map(|(k, _)| k).drive_unindexed(consumer)
+    }
+}
+
 impl<'g, K: 'g, G> Iterator for Iter<'g, K, G>
 where
     G: Guard,

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::<usize, usize>(|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::<usize>(|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);
+    });
+}
+