1use super::collect;
2use crate::FlatMultimap;
3use hashbrown::raw::rayon::{RawIntoParIter, RawParDrain, RawParIter};
4use rayon::iter::plumbing::UnindexedConsumer;
5use rayon::iter::{FromParallelIterator, IntoParallelIterator, ParallelExtend, ParallelIterator};
6use std::hash::{BuildHasher, Hash};
7use std::marker::PhantomData;
8
9#[derive(Clone)]
13pub struct ParIter<'a, K, V> {
14 inner: RawParIter<(K, V)>,
15 marker: PhantomData<(&'a K, &'a V)>,
16}
17
18impl<'a, K: Sync, V: Sync> ParallelIterator for ParIter<'a, K, V> {
19 type Item = (&'a K, &'a V);
20
21 fn drive_unindexed<C>(self, consumer: C) -> C::Result
22 where
23 C: UnindexedConsumer<(&'a K, &'a V)>,
24 {
25 self.inner
26 .map(|x| unsafe {
27 let r = x.as_ref();
28 (&r.0, &r.1)
29 })
30 .drive_unindexed(consumer)
31 }
32}
33
34#[derive(Clone)]
36pub struct ParKeys<'a, K, V> {
37 inner: RawParIter<(K, V)>,
38 marker: PhantomData<(&'a K, &'a V)>,
39}
40
41impl<'a, K: Sync, V: Sync> ParallelIterator for ParKeys<'a, K, V> {
42 type Item = &'a K;
43
44 fn drive_unindexed<C>(self, consumer: C) -> C::Result
45 where
46 C: UnindexedConsumer<Self::Item>,
47 {
48 self.inner
49 .map(|x| unsafe { &x.as_ref().0 })
50 .drive_unindexed(consumer)
51 }
52}
53
54#[derive(Clone)]
56pub struct ParValues<'a, K, V> {
57 inner: RawParIter<(K, V)>,
58 marker: PhantomData<(&'a K, &'a V)>,
59}
60
61impl<'a, K: Sync, V: Sync> ParallelIterator for ParValues<'a, K, V> {
62 type Item = &'a V;
63
64 fn drive_unindexed<C>(self, consumer: C) -> C::Result
65 where
66 C: UnindexedConsumer<Self::Item>,
67 {
68 self.inner
69 .map(|x| unsafe { &x.as_ref().1 })
70 .drive_unindexed(consumer)
71 }
72}
73
74pub struct ParIterMut<'a, K, V> {
76 inner: RawParIter<(K, V)>,
77 marker: PhantomData<(&'a K, &'a mut V)>,
78}
79
80impl<'a, K: Sync, V: Send> ParallelIterator for ParIterMut<'a, K, V> {
81 type Item = (&'a K, &'a mut V);
82
83 fn drive_unindexed<C>(self, consumer: C) -> C::Result
84 where
85 C: UnindexedConsumer<Self::Item>,
86 {
87 self.inner
88 .map(|x| unsafe {
89 let r = x.as_mut();
90 (&r.0, &mut r.1)
91 })
92 .drive_unindexed(consumer)
93 }
94}
95
96pub struct ParValuesMut<'a, K, V> {
98 inner: RawParIter<(K, V)>,
99 marker: PhantomData<(&'a K, &'a mut V)>,
100}
101
102impl<'a, K: Sync, V: Send> ParallelIterator for ParValuesMut<'a, K, V> {
103 type Item = &'a mut V;
104
105 fn drive_unindexed<C>(self, consumer: C) -> C::Result
106 where
107 C: UnindexedConsumer<Self::Item>,
108 {
109 self.inner
110 .map(|x| unsafe { &mut x.as_mut().1 })
111 .drive_unindexed(consumer)
112 }
113}
114
115pub struct IntoParIter<K, V> {
117 inner: RawIntoParIter<(K, V)>,
118}
119
120impl<K: Send, V: Send> ParallelIterator for IntoParIter<K, V> {
121 type Item = (K, V);
122
123 fn drive_unindexed<C>(self, consumer: C) -> C::Result
124 where
125 C: UnindexedConsumer<Self::Item>,
126 {
127 self.inner.drive_unindexed(consumer)
128 }
129}
130
131pub struct ParDrain<'a, K, V> {
133 inner: RawParDrain<'a, (K, V)>,
134}
135
136impl<K: Send, V: Send> ParallelIterator for ParDrain<'_, K, V> {
137 type Item = (K, V);
138
139 fn drive_unindexed<C>(self, consumer: C) -> C::Result
140 where
141 C: UnindexedConsumer<Self::Item>,
142 {
143 self.inner.drive_unindexed(consumer)
144 }
145}
146
147impl<K: Sync, V: Sync, S> FlatMultimap<K, V, S> {
148 pub fn par_keys(&self) -> ParKeys<'_, K, V> {
150 ParKeys {
151 inner: unsafe { self.table.par_iter() },
152 marker: PhantomData,
153 }
154 }
155
156 pub fn par_values(&self) -> ParValues<'_, K, V> {
158 ParValues {
159 inner: unsafe { self.table.par_iter() },
160 marker: PhantomData,
161 }
162 }
163}
164
165impl<K: Send, V: Send, S> FlatMultimap<K, V, S> {
166 pub fn par_values_mut(&mut self) -> ParValuesMut<'_, K, V> {
168 ParValuesMut {
169 inner: unsafe { self.table.par_iter() },
170 marker: PhantomData,
171 }
172 }
173
174 pub fn par_drain(&mut self) -> ParDrain<'_, K, V> {
176 ParDrain {
177 inner: self.table.par_drain(),
178 }
179 }
180}
181
182impl<K: Send, V: Send, S> IntoParallelIterator for FlatMultimap<K, V, S> {
183 type Item = (K, V);
184 type Iter = IntoParIter<K, V>;
185
186 fn into_par_iter(self) -> Self::Iter {
187 IntoParIter {
188 inner: self.table.into_par_iter(),
189 }
190 }
191}
192
193impl<'a, K: Sync, V: Sync, S> IntoParallelIterator for &'a FlatMultimap<K, V, S> {
194 type Item = (&'a K, &'a V);
195 type Iter = ParIter<'a, K, V>;
196
197 fn into_par_iter(self) -> Self::Iter {
198 ParIter {
199 inner: unsafe { self.table.par_iter() },
200 marker: PhantomData,
201 }
202 }
203}
204
205impl<'a, K: Sync, V: Send, S> IntoParallelIterator for &'a mut FlatMultimap<K, V, S> {
206 type Item = (&'a K, &'a mut V);
207 type Iter = ParIterMut<'a, K, V>;
208
209 fn into_par_iter(self) -> Self::Iter {
210 ParIterMut {
211 inner: unsafe { self.table.par_iter() },
212 marker: PhantomData,
213 }
214 }
215}
216
217impl<K, V, S> FromParallelIterator<(K, V)> for FlatMultimap<K, V, S>
218where
219 K: Eq + Hash + Send,
220 V: Send,
221 S: BuildHasher + Default,
222{
223 fn from_par_iter<P>(par_iter: P) -> Self
224 where
225 P: IntoParallelIterator<Item = (K, V)>,
226 {
227 let mut map = FlatMultimap::default();
228 map.par_extend(par_iter);
229 map
230 }
231}
232
233impl<K, V, S> ParallelExtend<(K, V)> for FlatMultimap<K, V, S>
234where
235 K: Eq + Hash + Send,
236 V: Send,
237 S: BuildHasher,
238{
239 fn par_extend<I>(&mut self, par_iter: I)
240 where
241 I: IntoParallelIterator<Item = (K, V)>,
242 {
243 extend(self, par_iter);
244 }
245}
246
247impl<'a, K, V, S> ParallelExtend<(&'a K, &'a V)> for FlatMultimap<K, V, S>
248where
249 K: Copy + Eq + Hash + Sync,
250 V: Copy + Sync,
251 S: BuildHasher,
252{
253 fn par_extend<I>(&mut self, par_iter: I)
254 where
255 I: IntoParallelIterator<Item = (&'a K, &'a V)>,
256 {
257 extend(self, par_iter);
258 }
259}
260
261fn extend<K, V, S, I>(map: &mut FlatMultimap<K, V, S>, par_iter: I)
262where
263 K: Eq + Hash,
264 S: BuildHasher,
265 I: IntoParallelIterator,
266 FlatMultimap<K, V, S>: Extend<I::Item>,
267{
268 let (list, len) = collect(par_iter);
269
270 map.reserve(len);
271
272 for vec in list {
273 map.extend(vec);
274 }
275}