diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index a399c1f8..b4505705 100644 --- a/.github/workflows/ci.yml +++ b/.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 diff --git a/Cargo.toml b/Cargo.toml index d099405d..c18dfce2 100644 --- a/Cargo.toml +++ b/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 diff --git a/src/lib.rs b/src/lib.rs index 489e5486..8c44d7bf 100644 --- a/src/lib.rs +++ b/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; diff --git a/src/macros.rs b/src/macros.rs index a7b33508..ca26287b 100644 --- a/src/macros.rs +++ b/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(self, consumer: C) -> C::Result + where + C: UnindexedConsumer, + { + 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 { + 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(self, consumer: C) -> C::Result + where + C: Consumer, + { + self.entries.into_par_iter().map($map_elt).drive(consumer) + } + + fn len(&self) -> usize { + self.entries.len() + } + + fn with_producer(self, callback: CB) -> CB::Output + where + CB: ProducerCallback, + { + self.entries + .into_par_iter() + .map($map_elt) + .with_producer(callback) + } + }; +} diff --git a/src/rayon/mod.rs b/src/rayon/mod.rs index 757a54dd..ebb1ac2d 100644 --- a/src/rayon/mod.rs +++ b/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(self, consumer: C) -> C::Result - where - C: UnindexedConsumer, - { - 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 { - 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(self, consumer: C) -> C::Result - where - C: Consumer, - { - self.entries.into_par_iter().map($map_elt).drive(consumer) - } - - fn len(&self) -> usize { - self.entries.len() - } - - fn with_producer(self, callback: CB) -> CB::Output - where - CB: ProducerCallback, - { - self.entries - .into_par_iter() - .map($map_elt) - .with_producer(callback) - } - }; -} - pub mod map; pub mod set; diff --git a/src/rustc.rs b/src/rustc.rs new file mode 100644 index 00000000..b843858b --- /dev/null +++ b/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 IntoParallelIterator for IndexMap + where + K: Send, + V: Send, + { + type Item = (K, V); + type Iter = IntoParIter; + + fn into_par_iter(self) -> Self::Iter { + IntoParIter { + entries: self.into_entries(), + } + } + } + + pub struct IntoParIter { + entries: Vec>, + } + + impl ParallelIterator for IntoParIter { + type Item = (K, V); + + parallel_iterator_methods!(Bucket::key_value); + } + + impl IndexedParallelIterator for IntoParIter { + indexed_parallel_iterator_methods!(Bucket::key_value); + } + + impl<'a, K, V, S> IntoParallelIterator for &'a IndexMap + 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], + } + + impl<'a, K: Sync, V: Sync> ParallelIterator for ParIter<'a, K, V> { + type Item = (&'a K, &'a V); + + parallel_iterator_methods!(Bucket::refs); + } + + impl IndexedParallelIterator for ParIter<'_, K, V> { + indexed_parallel_iterator_methods!(Bucket::refs); + } + + impl<'a, K, V, S> IntoParallelIterator for &'a mut IndexMap + 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], + } + + 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 IndexedParallelIterator for ParIterMut<'_, K, V> { + indexed_parallel_iterator_methods!(Bucket::ref_mut); + } +} + +mod set { + use super::*; + + impl IntoParallelIterator for IndexSet + where + T: Send, + { + type Item = T; + type Iter = IntoParIter; + + fn into_par_iter(self) -> Self::Iter { + IntoParIter { + entries: self.into_entries(), + } + } + } + + pub struct IntoParIter { + entries: Vec>, + } + + impl ParallelIterator for IntoParIter { + type Item = T; + + parallel_iterator_methods!(Bucket::key); + } + + impl IndexedParallelIterator for IntoParIter { + indexed_parallel_iterator_methods!(Bucket::key); + } + + impl<'a, T, S> IntoParallelIterator for &'a IndexSet + 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], + } + + impl<'a, T: Sync> ParallelIterator for ParIter<'a, T> { + type Item = &'a T; + + parallel_iterator_methods!(Bucket::key_ref); + } + + impl IndexedParallelIterator for ParIter<'_, T> { + indexed_parallel_iterator_methods!(Bucket::key_ref); + } +}