Add minimal support for rustc-rayon

.github/workflows/ci.yml

@@ -22,6 +22,8 @@ jobs:
             features: serde
           - rust: stable
             features: rayon
+          - rust: stable
+            features: rustc-rayon
           - rust: stable
             features: std
           - rust: beta

Cargo.toml

@@ -36,6 +36,10 @@ autocfg = "1"
 serde = { version = "1.0", optional = true, default-features = false }
 rayon = { version = "1.2", optional = true }
 
+# Internal feature, only used when building as part of rustc,
+# not part of the stable interface of this crate.
+rustc-rayon = { version = "0.3", optional = true }
+
 [dependencies.hashbrown]
 version = "0.11"
 default-features = false

src/lib.rs

@@ -103,6 +103,9 @@ pub mod set;
 #[cfg(feature = "rayon")]
 mod rayon;
 
+#[cfg(feature = "rustc-rayon")]
+mod rustc;
+
 pub use crate::equivalent::Equivalent;
 pub use crate::map::IndexMap;
 pub use crate::set::IndexSet;

src/macros.rs

@@ -122,3 +122,57 @@ macro_rules! double_ended_iterator_methods {
         }
     };
 }
+
+// generate `ParallelIterator` methods by just forwarding to the underlying
+// self.entries and mapping its elements.
+#[cfg(any(feature = "rayon", feature = "rustc-rayon"))]
+macro_rules! parallel_iterator_methods {
+    // $map_elt is the mapping function from the underlying iterator's element
+    ($map_elt:expr) => {
+        fn drive_unindexed<C>(self, consumer: C) -> C::Result
+        where
+            C: UnindexedConsumer<Self::Item>,
+        {
+            self.entries
+                .into_par_iter()
+                .map($map_elt)
+                .drive_unindexed(consumer)
+        }
+
+        // NB: This allows indexed collection, e.g. directly into a `Vec`, but the
+        // underlying iterator must really be indexed.  We should remove this if we
+        // start having tombstones that must be filtered out.
+        fn opt_len(&self) -> Option<usize> {
+            Some(self.entries.len())
+        }
+    };
+}
+
+// generate `IndexedParallelIterator` methods by just forwarding to the underlying
+// self.entries and mapping its elements.
+#[cfg(any(feature = "rayon", feature = "rustc-rayon"))]
+macro_rules! indexed_parallel_iterator_methods {
+    // $map_elt is the mapping function from the underlying iterator's element
+    ($map_elt:expr) => {
+        fn drive<C>(self, consumer: C) -> C::Result
+        where
+            C: Consumer<Self::Item>,
+        {
+            self.entries.into_par_iter().map($map_elt).drive(consumer)
+        }
+
+        fn len(&self) -> usize {
+            self.entries.len()
+        }
+
+        fn with_producer<CB>(self, callback: CB) -> CB::Output
+        where
+            CB: ProducerCallback<Self::Item>,
+        {
+            self.entries
+                .into_par_iter()
+                .map($map_elt)
+                .with_producer(callback)
+        }
+    };
+}

src/rayon/mod.rs

@@ -4,58 +4,6 @@ use alloc::collections::LinkedList;
 
 use crate::vec::Vec;
 
-// generate `ParallelIterator` methods by just forwarding to the underlying
-// self.entries and mapping its elements.
-macro_rules! parallel_iterator_methods {
-    // $map_elt is the mapping function from the underlying iterator's element
-    ($map_elt:expr) => {
-        fn drive_unindexed<C>(self, consumer: C) -> C::Result
-        where
-            C: UnindexedConsumer<Self::Item>,
-        {
-            self.entries
-                .into_par_iter()
-                .map($map_elt)
-                .drive_unindexed(consumer)
-        }
-
-        // NB: This allows indexed collection, e.g. directly into a `Vec`, but the
-        // underlying iterator must really be indexed.  We should remove this if we
-        // start having tombstones that must be filtered out.
-        fn opt_len(&self) -> Option<usize> {
-            Some(self.entries.len())
-        }
-    };
-}
-
-// generate `IndexedParallelIterator` methods by just forwarding to the underlying
-// self.entries and mapping its elements.
-macro_rules! indexed_parallel_iterator_methods {
-    // $map_elt is the mapping function from the underlying iterator's element
-    ($map_elt:expr) => {
-        fn drive<C>(self, consumer: C) -> C::Result
-        where
-            C: Consumer<Self::Item>,
-        {
-            self.entries.into_par_iter().map($map_elt).drive(consumer)
-        }
-
-        fn len(&self) -> usize {
-            self.entries.len()
-        }
-
-        fn with_producer<CB>(self, callback: CB) -> CB::Output
-        where
-            CB: ProducerCallback<Self::Item>,
-        {
-            self.entries
-                .into_par_iter()
-                .map($map_elt)
-                .with_producer(callback)
-        }
-    };
-}
-
 pub mod map;
 pub mod set;
 

src/rustc.rs

@@ -0,0 +1,158 @@
+//! Minimal support for `rustc-rayon`, not intended for general use.
+
+use crate::vec::Vec;
+use crate::{Bucket, Entries, IndexMap, IndexSet};
+
+use rustc_rayon::iter::plumbing::{Consumer, ProducerCallback, UnindexedConsumer};
+use rustc_rayon::iter::{IndexedParallelIterator, IntoParallelIterator, ParallelIterator};
+
+mod map {
+    use super::*;
+
+    impl<K, V, S> IntoParallelIterator for IndexMap<K, V, S>
+    where
+        K: Send,
+        V: Send,
+    {
+        type Item = (K, V);
+        type Iter = IntoParIter<K, V>;
+
+        fn into_par_iter(self) -> Self::Iter {
+            IntoParIter {
+                entries: self.into_entries(),
+            }
+        }
+    }
+
+    pub struct IntoParIter<K, V> {
+        entries: Vec<Bucket<K, V>>,
+    }
+
+    impl<K: Send, V: Send> ParallelIterator for IntoParIter<K, V> {
+        type Item = (K, V);
+
+        parallel_iterator_methods!(Bucket::key_value);
+    }
+
+    impl<K: Send, V: Send> IndexedParallelIterator for IntoParIter<K, V> {
+        indexed_parallel_iterator_methods!(Bucket::key_value);
+    }
+
+    impl<'a, K, V, S> IntoParallelIterator for &'a IndexMap<K, V, S>
+    where
+        K: Sync,
+        V: Sync,
+    {
+        type Item = (&'a K, &'a V);
+        type Iter = ParIter<'a, K, V>;
+
+        fn into_par_iter(self) -> Self::Iter {
+            ParIter {
+                entries: self.as_entries(),
+            }
+        }
+    }
+
+    pub struct ParIter<'a, K, V> {
+        entries: &'a [Bucket<K, V>],
+    }
+
+    impl<'a, K: Sync, V: Sync> ParallelIterator for ParIter<'a, K, V> {
+        type Item = (&'a K, &'a V);
+
+        parallel_iterator_methods!(Bucket::refs);
+    }
+
+    impl<K: Sync, V: Sync> IndexedParallelIterator for ParIter<'_, K, V> {
+        indexed_parallel_iterator_methods!(Bucket::refs);
+    }
+
+    impl<'a, K, V, S> IntoParallelIterator for &'a mut IndexMap<K, V, S>
+    where
+        K: Sync + Send,
+        V: Send,
+    {
+        type Item = (&'a K, &'a mut V);
+        type Iter = ParIterMut<'a, K, V>;
+
+        fn into_par_iter(self) -> Self::Iter {
+            ParIterMut {
+                entries: self.as_entries_mut(),
+            }
+        }
+    }
+
+    pub struct ParIterMut<'a, K, V> {
+        entries: &'a mut [Bucket<K, V>],
+    }
+
+    impl<'a, K: Sync + Send, V: Send> ParallelIterator for ParIterMut<'a, K, V> {
+        type Item = (&'a K, &'a mut V);
+
+        parallel_iterator_methods!(Bucket::ref_mut);
+    }
+
+    impl<K: Sync + Send, V: Send> IndexedParallelIterator for ParIterMut<'_, K, V> {
+        indexed_parallel_iterator_methods!(Bucket::ref_mut);
+    }
+}
+
+mod set {
+    use super::*;
+
+    impl<T, S> IntoParallelIterator for IndexSet<T, S>
+    where
+        T: Send,
+    {
+        type Item = T;
+        type Iter = IntoParIter<T>;
+
+        fn into_par_iter(self) -> Self::Iter {
+            IntoParIter {
+                entries: self.into_entries(),
+            }
+        }
+    }
+
+    pub struct IntoParIter<T> {
+        entries: Vec<Bucket<T, ()>>,
+    }
+
+    impl<T: Send> ParallelIterator for IntoParIter<T> {
+        type Item = T;
+
+        parallel_iterator_methods!(Bucket::key);
+    }
+
+    impl<T: Send> IndexedParallelIterator for IntoParIter<T> {
+        indexed_parallel_iterator_methods!(Bucket::key);
+    }
+
+    impl<'a, T, S> IntoParallelIterator for &'a IndexSet<T, S>
+    where
+        T: Sync,
+    {
+        type Item = &'a T;
+        type Iter = ParIter<'a, T>;
+
+        fn into_par_iter(self) -> Self::Iter {
+            ParIter {
+                entries: self.as_entries(),
+            }
+        }
+    }
+
+    pub struct ParIter<'a, T> {
+        entries: &'a [Bucket<T, ()>],
+    }
+
+    impl<'a, T: Sync> ParallelIterator for ParIter<'a, T> {
+        type Item = &'a T;
+
+        parallel_iterator_methods!(Bucket::key_ref);
+    }
+
+    impl<T: Sync> IndexedParallelIterator for ParIter<'_, T> {
+        indexed_parallel_iterator_methods!(Bucket::key_ref);
+    }
+}