Merge pull request #116 from cuviper/rayon

Add optional rayon support

Cargo.toml

@@ -23,6 +23,7 @@ ahash = "0.3.8"
 serde = { version = "1.0.114", optional = true, features = ["derive"] }
 cfg-if = "1.0.0"
 hashbrown = { version = "0.8.0", optional = true }
+rayon = { version = "1.5.0", optional = true }
 
 [package.metadata.docs.rs]
-features = ["raw-api", "serde"]
+features = ["rayon", "raw-api", "serde"]

src/lib.rs

@@ -13,6 +13,12 @@ pub mod setref;
 mod t;
 mod util;
 
+#[cfg(feature = "rayon")]
+pub mod rayon {
+    pub mod map;
+    pub mod set;
+}
+
 use ahash::RandomState;
 use cfg_if::cfg_if;
 use core::borrow::Borrow;

src/rayon/map.rs

@@ -0,0 +1,228 @@
+use crate::lock::RwLock;
+use crate::mapref::multiple::{RefMulti, RefMutMulti};
+use crate::util;
+use crate::{DashMap, HashMap};
+use ahash::RandomState;
+use core::hash::{BuildHasher, Hash};
+use rayon::iter::plumbing::UnindexedConsumer;
+use rayon::iter::{FromParallelIterator, IntoParallelIterator, ParallelExtend, ParallelIterator};
+
+cfg_if::cfg_if! {
+    if #[cfg(feature = "no_std")] {
+        use alloc::{boxed::Box, sync::Arc, vec::Vec};
+    } else {
+        use std::sync::Arc;
+    }
+}
+
+impl<K, V, S> ParallelExtend<(K, V)> for DashMap<K, V, S>
+where
+    K: Send + Sync + Eq + Hash,
+    V: Send + Sync,
+    S: Send + Sync + Clone + BuildHasher,
+{
+    fn par_extend<I>(&mut self, par_iter: I)
+    where
+        I: IntoParallelIterator<Item = (K, V)>,
+    {
+        (&*self).par_extend(par_iter);
+    }
+}
+
+// Since we don't actually need mutability, we can implement this on a
+// reference, similar to `io::Write for &File`.
+impl<K, V, S> ParallelExtend<(K, V)> for &'_ DashMap<K, V, S>
+where
+    K: Send + Sync + Eq + Hash,
+    V: Send + Sync,
+    S: Send + Sync + Clone + BuildHasher,
+{
+    fn par_extend<I>(&mut self, par_iter: I)
+    where
+        I: IntoParallelIterator<Item = (K, V)>,
+    {
+        let &mut map = self;
+        par_iter.into_par_iter().for_each(move |(key, value)| {
+            map.insert(key, value);
+        });
+    }
+}
+
+impl<K, V, S> FromParallelIterator<(K, V)> for DashMap<K, V, S>
+where
+    K: Send + Sync + Eq + Hash,
+    V: Send + Sync,
+    S: Send + Sync + Clone + Default + BuildHasher,
+{
+    fn from_par_iter<I>(par_iter: I) -> Self
+    where
+        I: IntoParallelIterator<Item = (K, V)>,
+    {
+        let map = Self::default();
+        (&map).par_extend(par_iter);
+        map
+    }
+}
+
+// Implementation note: while the shards will iterate in parallel, we flatten
+// sequentially within each shard (`flat_map_iter`), because the standard
+// `HashMap` only implements `ParallelIterator` by collecting to a `Vec` first.
+// There is real parallel support in the `hashbrown/rayon` feature, but we don't
+// always use that map.
+
+impl<K, V, S> IntoParallelIterator for DashMap<K, V, S>
+where
+    K: Send + Eq + Hash,
+    V: Send,
+    S: Send + Clone + BuildHasher,
+{
+    type Iter = OwningIter<K, V, S>;
+    type Item = (K, V);
+
+    fn into_par_iter(self) -> Self::Iter {
+        OwningIter {
+            shards: self.shards,
+        }
+    }
+}
+
+pub struct OwningIter<K, V, S = RandomState> {
+    shards: Box<[RwLock<HashMap<K, V, S>>]>,
+}
+
+impl<K, V, S> ParallelIterator for OwningIter<K, V, S>
+where
+    K: Send + Eq + Hash,
+    V: Send,
+    S: Send + Clone + BuildHasher,
+{
+    type Item = (K, V);
+
+    fn drive_unindexed<C>(self, consumer: C) -> C::Result
+    where
+        C: UnindexedConsumer<Self::Item>,
+    {
+        Vec::from(self.shards)
+            .into_par_iter()
+            .flat_map_iter(|shard| {
+                shard
+                    .into_inner()
+                    .into_iter()
+                    .map(|(k, v)| (k, v.into_inner()))
+            })
+            .drive_unindexed(consumer)
+    }
+}
+
+// This impl also enables `IntoParallelRefIterator::par_iter`
+impl<'a, K, V, S> IntoParallelIterator for &'a DashMap<K, V, S>
+where
+    K: Send + Sync + Eq + Hash,
+    V: Send + Sync,
+    S: Send + Sync + Clone + BuildHasher,
+{
+    type Iter = Iter<'a, K, V, S>;
+    type Item = RefMulti<'a, K, V, S>;
+
+    fn into_par_iter(self) -> Self::Iter {
+        Iter {
+            shards: &self.shards,
+        }
+    }
+}
+
+pub struct Iter<'a, K, V, S = RandomState> {
+    shards: &'a [RwLock<HashMap<K, V, S>>],
+}
+
+impl<'a, K, V, S> ParallelIterator for Iter<'a, K, V, S>
+where
+    K: Send + Sync + Eq + Hash,
+    V: Send + Sync,
+    S: Send + Sync + Clone + BuildHasher,
+{
+    type Item = RefMulti<'a, K, V, S>;
+
+    fn drive_unindexed<C>(self, consumer: C) -> C::Result
+    where
+        C: UnindexedConsumer<Self::Item>,
+    {
+        self.shards
+            .into_par_iter()
+            .flat_map_iter(|shard| {
+                let guard = shard.read();
+                let sref: &'a HashMap<K, V, S> = unsafe { util::change_lifetime_const(&*guard) };
+
+                let guard = Arc::new(guard);
+                sref.iter().map(move |(k, v)| {
+                    let guard = Arc::clone(&guard);
+                    RefMulti::new(guard, k, v.get())
+                })
+            })
+            .drive_unindexed(consumer)
+    }
+}
+
+// This impl also enables `IntoParallelRefMutIterator::par_iter_mut`
+impl<'a, K, V, S> IntoParallelIterator for &'a mut DashMap<K, V, S>
+where
+    K: Send + Sync + Eq + Hash,
+    V: Send + Sync,
+    S: Send + Sync + Clone + BuildHasher,
+{
+    type Iter = IterMut<'a, K, V, S>;
+    type Item = RefMutMulti<'a, K, V, S>;
+
+    fn into_par_iter(self) -> Self::Iter {
+        IterMut {
+            shards: &self.shards,
+        }
+    }
+}
+
+impl<'a, K, V, S> DashMap<K, V, S>
+where
+    K: Send + Sync + Eq + Hash,
+    V: Send + Sync,
+    S: Send + Sync + Clone + BuildHasher,
+{
+    // Unlike `IntoParallelRefMutIterator::par_iter_mut`, we only _need_ `&self`.
+    pub fn par_iter_mut(&self) -> IterMut<'_, K, V, S> {
+        IterMut {
+            shards: &self.shards,
+        }
+    }
+}
+
+pub struct IterMut<'a, K, V, S = RandomState> {
+    shards: &'a [RwLock<HashMap<K, V, S>>],
+}
+
+impl<'a, K, V, S> ParallelIterator for IterMut<'a, K, V, S>
+where
+    K: Send + Sync + Eq + Hash,
+    V: Send + Sync,
+    S: Send + Sync + Clone + BuildHasher,
+{
+    type Item = RefMutMulti<'a, K, V, S>;
+
+    fn drive_unindexed<C>(self, consumer: C) -> C::Result
+    where
+        C: UnindexedConsumer<Self::Item>,
+    {
+        self.shards
+            .into_par_iter()
+            .flat_map_iter(|shard| {
+                let mut guard = shard.write();
+                let sref: &'a mut HashMap<K, V, S> =
+                    unsafe { util::change_lifetime_mut(&mut *guard) };
+
+                let guard = Arc::new(guard);
+                sref.iter_mut().map(move |(k, v)| {
+                    let guard = Arc::clone(&guard);
+                    RefMutMulti::new(guard, k, v.get_mut())
+                })
+            })
+            .drive_unindexed(consumer)
+    }
+}

src/rayon/set.rs

@@ -0,0 +1,121 @@
+use crate::setref::multiple::RefMulti;
+use crate::DashSet;
+use ahash::RandomState;
+use core::hash::{BuildHasher, Hash};
+use rayon::iter::plumbing::UnindexedConsumer;
+use rayon::iter::{FromParallelIterator, IntoParallelIterator, ParallelExtend, ParallelIterator};
+
+impl<K, S> ParallelExtend<K> for DashSet<K, S>
+where
+    K: Send + Sync + Eq + Hash,
+    S: Send + Sync + Clone + BuildHasher,
+{
+    fn par_extend<I>(&mut self, par_iter: I)
+    where
+        I: IntoParallelIterator<Item = K>,
+    {
+        (&*self).par_extend(par_iter);
+    }
+}
+
+// Since we don't actually need mutability, we can implement this on a
+// reference, similar to `io::Write for &File`.
+impl<K, S> ParallelExtend<K> for &'_ DashSet<K, S>
+where
+    K: Send + Sync + Eq + Hash,
+    S: Send + Sync + Clone + BuildHasher,
+{
+    fn par_extend<I>(&mut self, par_iter: I)
+    where
+        I: IntoParallelIterator<Item = K>,
+    {
+        let &mut set = self;
+        par_iter.into_par_iter().for_each(move |key| {
+            set.insert(key);
+        });
+    }
+}
+
+impl<K, S> FromParallelIterator<K> for DashSet<K, S>
+where
+    K: Send + Sync + Eq + Hash,
+    S: Send + Sync + Clone + Default + BuildHasher,
+{
+    fn from_par_iter<I>(par_iter: I) -> Self
+    where
+        I: IntoParallelIterator<Item = K>,
+    {
+        let set = Self::default();
+        (&set).par_extend(par_iter);
+        set
+    }
+}
+
+impl<K, S> IntoParallelIterator for DashSet<K, S>
+where
+    K: Send + Eq + Hash,
+    S: Send + Clone + BuildHasher,
+{
+    type Iter = OwningIter<K, S>;
+    type Item = K;
+
+    fn into_par_iter(self) -> Self::Iter {
+        OwningIter {
+            inner: self.inner.into_par_iter(),
+        }
+    }
+}
+
+pub struct OwningIter<K, S = RandomState> {
+    inner: super::map::OwningIter<K, (), S>,
+}
+
+impl<K, S> ParallelIterator for OwningIter<K, S>
+where
+    K: Send + Eq + Hash,
+    S: Send + Clone + BuildHasher,
+{
+    type Item = K;
+
+    fn drive_unindexed<C>(self, consumer: C) -> C::Result
+    where
+        C: UnindexedConsumer<Self::Item>,
+    {
+        self.inner.map(|(k, _)| k).drive_unindexed(consumer)
+    }
+}
+
+// This impl also enables `IntoParallelRefIterator::par_iter`
+impl<'a, K, S> IntoParallelIterator for &'a DashSet<K, S>
+where
+    K: Send + Sync + Eq + Hash,
+    S: Send + Sync + Clone + BuildHasher,
+{
+    type Iter = Iter<'a, K, S>;
+    type Item = RefMulti<'a, K, S>;
+
+    fn into_par_iter(self) -> Self::Iter {
+        Iter {
+            inner: (&self.inner).into_par_iter(),
+        }
+    }
+}
+
+pub struct Iter<'a, K, S = RandomState> {
+    inner: super::map::Iter<'a, K, (), S>,
+}
+
+impl<'a, K, S> ParallelIterator for Iter<'a, K, S>
+where
+    K: Send + Sync + Eq + Hash,
+    S: Send + Sync + Clone + BuildHasher,
+{
+    type Item = RefMulti<'a, K, S>;
+
+    fn drive_unindexed<C>(self, consumer: C) -> C::Result
+    where
+        C: UnindexedConsumer<Self::Item>,
+    {
+        self.inner.map(RefMulti::new).drive_unindexed(consumer)
+    }
+}

src/set.rs

@@ -18,7 +18,7 @@ use core::iter::FromIterator;
 /// [`DashMap`]: struct.DashMap.html
 
 pub struct DashSet<K, S = RandomState> {
-    inner: DashMap<K, (), S>,
+    pub(crate) inner: DashMap<K, (), S>,
 }
 
 impl<K: Eq + Hash + fmt::Debug, S: BuildHasher + Clone> fmt::Debug for DashSet<K, S> {