flat_multimap/rayon/
map.rs

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// TODO: Debug impls
10
11/// Parallel iterator over shared references to entries in a map.
12#[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/// Parallel iterator over shared references to keys in a map.
35#[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/// Parallel iterator over shared references to values in a map.
55#[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
74/// Parallel iterator over mutable references to entries in a map.
75pub 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
96/// Parallel iterator over mutable references to values in a map.
97pub 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
115/// Parallel iterator over entries of a consumed map.
116pub 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
131/// Parallel draining iterator over entries of a map.
132pub 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    /// Visits (potentially in parallel) immutably borrowed keys in an arbitrary order.
149    pub fn par_keys(&self) -> ParKeys<'_, K, V> {
150        ParKeys {
151            inner: unsafe { self.table.par_iter() },
152            marker: PhantomData,
153        }
154    }
155
156    /// Visits (potentially in parallel) immutably borrowed values in an arbitrary order.
157    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    /// Visits (potentially in parallel) mutably borrowed values in an arbitrary order.
167    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    /// Consumes (potentially in parallel) all values in an arbitrary order.
175    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}