Skip to content

Commit 1de6f77

Browse files
cuvipercuviper
committed
Add minimal support for rustc-rayon
1 parent 0af3c46 commit 1de6f77

File tree

6 files changed

+221
-52
lines changed

6 files changed

+221
-52
lines changed

‎.github/workflows/ci.yml

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,8 @@ jobs:
2222
features: serde
2323
- rust: stable
2424
features: rayon
25+
- rust: stable
26+
features: rustc-rayon
2527
- rust: stable
2628
features: std
2729
- rust: beta

‎Cargo.toml

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,10 @@ autocfg = "1"
3636
serde = { version = "1.0", optional = true, default-features = false }
3737
rayon = { version = "1.2", optional = true }
3838

39+
# Internal feature, only used when building as part of rustc,
40+
# not part of the stable interface of this crate.
41+
rustc-rayon = { version = "0.3", optional = true }
42+
3943
[dependencies.hashbrown]
4044
version = "0.11"
4145
default-features = false

‎src/lib.rs

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -103,6 +103,9 @@ pub mod set;
103103
#[cfg(feature = "rayon")]
104104
mod rayon;
105105

106+
#[cfg(feature = "rustc-rayon")]
107+
mod rustc;
108+
106109
pub use crate::equivalent::Equivalent;
107110
pub use crate::map::IndexMap;
108111
pub use crate::set::IndexSet;

‎src/macros.rs

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -122,3 +122,57 @@ macro_rules! double_ended_iterator_methods {
122122
}
123123
};
124124
}
125+
126+
// generate `ParallelIterator` methods by just forwarding to the underlying
127+
// self.entries and mapping its elements.
128+
#[cfg(any(feature = "rayon", feature = "rustc-rayon"))]
129+
macro_rules! parallel_iterator_methods {
130+
// $map_elt is the mapping function from the underlying iterator's element
131+
($map_elt:expr) => {
132+
fn drive_unindexed<C>(self, consumer: C) -> C::Result
133+
where
134+
C: UnindexedConsumer<Self::Item>,
135+
{
136+
self.entries
137+
.into_par_iter()
138+
.map($map_elt)
139+
.drive_unindexed(consumer)
140+
}
141+
142+
// NB: This allows indexed collection, e.g. directly into a `Vec`, but the
143+
// underlying iterator must really be indexed. We should remove this if we
144+
// start having tombstones that must be filtered out.
145+
fn opt_len(&self) -> Option<usize> {
146+
Some(self.entries.len())
147+
}
148+
};
149+
}
150+
151+
// generate `IndexedParallelIterator` methods by just forwarding to the underlying
152+
// self.entries and mapping its elements.
153+
#[cfg(any(feature = "rayon", feature = "rustc-rayon"))]
154+
macro_rules! indexed_parallel_iterator_methods {
155+
// $map_elt is the mapping function from the underlying iterator's element
156+
($map_elt:expr) => {
157+
fn drive<C>(self, consumer: C) -> C::Result
158+
where
159+
C: Consumer<Self::Item>,
160+
{
161+
self.entries.into_par_iter().map($map_elt).drive(consumer)
162+
}
163+
164+
fn len(&self) -> usize {
165+
self.entries.len()
166+
}
167+
168+
fn with_producer<CB>(self, callback: CB) -> CB::Output
169+
where
170+
CB: ProducerCallback<Self::Item>,
171+
{
172+
self.entries
173+
.into_par_iter()
174+
.map($map_elt)
175+
.with_producer(callback)
176+
}
177+
};
178+
}

‎src/rayon/mod.rs

Lines changed: 0 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -4,58 +4,6 @@ use alloc::collections::LinkedList;
44

55
use crate::vec::Vec;
66

7-
// generate `ParallelIterator` methods by just forwarding to the underlying
8-
// self.entries and mapping its elements.
9-
macro_rules! parallel_iterator_methods {
10-
// $map_elt is the mapping function from the underlying iterator's element
11-
($map_elt:expr) => {
12-
fn drive_unindexed<C>(self, consumer: C) -> C::Result
13-
where
14-
C: UnindexedConsumer<Self::Item>,
15-
{
16-
self.entries
17-
.into_par_iter()
18-
.map($map_elt)
19-
.drive_unindexed(consumer)
20-
}
21-
22-
// NB: This allows indexed collection, e.g. directly into a `Vec`, but the
23-
// underlying iterator must really be indexed. We should remove this if we
24-
// start having tombstones that must be filtered out.
25-
fn opt_len(&self) -> Option<usize> {
26-
Some(self.entries.len())
27-
}
28-
};
29-
}
30-
31-
// generate `IndexedParallelIterator` methods by just forwarding to the underlying
32-
// self.entries and mapping its elements.
33-
macro_rules! indexed_parallel_iterator_methods {
34-
// $map_elt is the mapping function from the underlying iterator's element
35-
($map_elt:expr) => {
36-
fn drive<C>(self, consumer: C) -> C::Result
37-
where
38-
C: Consumer<Self::Item>,
39-
{
40-
self.entries.into_par_iter().map($map_elt).drive(consumer)
41-
}
42-
43-
fn len(&self) -> usize {
44-
self.entries.len()
45-
}
46-
47-
fn with_producer<CB>(self, callback: CB) -> CB::Output
48-
where
49-
CB: ProducerCallback<Self::Item>,
50-
{
51-
self.entries
52-
.into_par_iter()
53-
.map($map_elt)
54-
.with_producer(callback)
55-
}
56-
};
57-
}
58-
597
pub mod map;
608
pub mod set;
619

‎src/rustc.rs

Lines changed: 158 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,158 @@
1+
//! Minimal support for `rustc-rayon`, not intended for general use.
2+
3+
use crate::vec::Vec;
4+
use crate::{Bucket, Entries, IndexMap, IndexSet};
5+
6+
use rustc_rayon::iter::plumbing::{Consumer, ProducerCallback, UnindexedConsumer};
7+
use rustc_rayon::iter::{IndexedParallelIterator, IntoParallelIterator, ParallelIterator};
8+
9+
mod map {
10+
use super::*;
11+
12+
impl<K, V, S> IntoParallelIterator for IndexMap<K, V, S>
13+
where
14+
K: Send,
15+
V: Send,
16+
{
17+
type Item = (K, V);
18+
type Iter = IntoParIter<K, V>;
19+
20+
fn into_par_iter(self) -> Self::Iter {
21+
IntoParIter {
22+
entries: self.into_entries(),
23+
}
24+
}
25+
}
26+
27+
pub struct IntoParIter<K, V> {
28+
entries: Vec<Bucket<K, V>>,
29+
}
30+
31+
impl<K: Send, V: Send> ParallelIterator for IntoParIter<K, V> {
32+
type Item = (K, V);
33+
34+
parallel_iterator_methods!(Bucket::key_value);
35+
}
36+
37+
impl<K: Send, V: Send> IndexedParallelIterator for IntoParIter<K, V> {
38+
indexed_parallel_iterator_methods!(Bucket::key_value);
39+
}
40+
41+
impl<'a, K, V, S> IntoParallelIterator for &'a IndexMap<K, V, S>
42+
where
43+
K: Sync,
44+
V: Sync,
45+
{
46+
type Item = (&'a K, &'a V);
47+
type Iter = ParIter<'a, K, V>;
48+
49+
fn into_par_iter(self) -> Self::Iter {
50+
ParIter {
51+
entries: self.as_entries(),
52+
}
53+
}
54+
}
55+
56+
pub struct ParIter<'a, K, V> {
57+
entries: &'a [Bucket<K, V>],
58+
}
59+
60+
impl<'a, K: Sync, V: Sync> ParallelIterator for ParIter<'a, K, V> {
61+
type Item = (&'a K, &'a V);
62+
63+
parallel_iterator_methods!(Bucket::refs);
64+
}
65+
66+
impl<K: Sync, V: Sync> IndexedParallelIterator for ParIter<'_, K, V> {
67+
indexed_parallel_iterator_methods!(Bucket::refs);
68+
}
69+
70+
impl<'a, K, V, S> IntoParallelIterator for &'a mut IndexMap<K, V, S>
71+
where
72+
K: Sync + Send,
73+
V: Send,
74+
{
75+
type Item = (&'a K, &'a mut V);
76+
type Iter = ParIterMut<'a, K, V>;
77+
78+
fn into_par_iter(self) -> Self::Iter {
79+
ParIterMut {
80+
entries: self.as_entries_mut(),
81+
}
82+
}
83+
}
84+
85+
pub struct ParIterMut<'a, K, V> {
86+
entries: &'a mut [Bucket<K, V>],
87+
}
88+
89+
impl<'a, K: Sync + Send, V: Send> ParallelIterator for ParIterMut<'a, K, V> {
90+
type Item = (&'a K, &'a mut V);
91+
92+
parallel_iterator_methods!(Bucket::ref_mut);
93+
}
94+
95+
impl<K: Sync + Send, V: Send> IndexedParallelIterator for ParIterMut<'_, K, V> {
96+
indexed_parallel_iterator_methods!(Bucket::ref_mut);
97+
}
98+
}
99+
100+
mod set {
101+
use super::*;
102+
103+
impl<T, S> IntoParallelIterator for IndexSet<T, S>
104+
where
105+
T: Send,
106+
{
107+
type Item = T;
108+
type Iter = IntoParIter<T>;
109+
110+
fn into_par_iter(self) -> Self::Iter {
111+
IntoParIter {
112+
entries: self.into_entries(),
113+
}
114+
}
115+
}
116+
117+
pub struct IntoParIter<T> {
118+
entries: Vec<Bucket<T, ()>>,
119+
}
120+
121+
impl<T: Send> ParallelIterator for IntoParIter<T> {
122+
type Item = T;
123+
124+
parallel_iterator_methods!(Bucket::key);
125+
}
126+
127+
impl<T: Send> IndexedParallelIterator for IntoParIter<T> {
128+
indexed_parallel_iterator_methods!(Bucket::key);
129+
}
130+
131+
impl<'a, T, S> IntoParallelIterator for &'a IndexSet<T, S>
132+
where
133+
T: Sync,
134+
{
135+
type Item = &'a T;
136+
type Iter = ParIter<'a, T>;
137+
138+
fn into_par_iter(self) -> Self::Iter {
139+
ParIter {
140+
entries: self.as_entries(),
141+
}
142+
}
143+
}
144+
145+
pub struct ParIter<'a, T> {
146+
entries: &'a [Bucket<T, ()>],
147+
}
148+
149+
impl<'a, T: Sync> ParallelIterator for ParIter<'a, T> {
150+
type Item = &'a T;
151+
152+
parallel_iterator_methods!(Bucket::key_ref);
153+
}
154+
155+
impl<T: Sync> IndexedParallelIterator for ParIter<'_, T> {
156+
indexed_parallel_iterator_methods!(Bucket::key_ref);
157+
}
158+
}

0 commit comments

Comments
 (0)