1//! Rayon extensions for `HashTable`.
2
3use super::raw::{RawIntoParIter, RawParDrain, RawParIter};
4use crate::hash_table::HashTable;
5use crate::raw::{Allocator, Global};
6use core::fmt;
7use core::marker::PhantomData;
8use rayon::iter::plumbing::UnindexedConsumer;
9use rayon::iter::{IntoParallelIterator, ParallelIterator};
10
11/// Parallel iterator over shared references to entries in a map.
12///
13/// This iterator is created by the [`par_iter`] method on [`HashTable`]
14/// (provided by the [`IntoParallelRefIterator`] trait).
15/// See its documentation for more.
16///
17/// [`par_iter`]: /hashbrown/struct.HashTable.html#method.par_iter
18/// [`HashTable`]: /hashbrown/struct.HashTable.html
19/// [`IntoParallelRefIterator`]: https://docs.rs/rayon/1.0/rayon/iter/trait.IntoParallelRefIterator.html
20pub struct ParIter<'a, T> {
21 inner: RawParIter<T>,
22 marker: PhantomData<&'a T>,
23}
24
25impl<'a, T: Sync> ParallelIterator for ParIter<'a, T> {
26 type Item = &'a T;
27
28 #[cfg_attr(feature = "inline-more", inline)]
29 fn drive_unindexed<C>(self, consumer: C) -> C::Result
30 where
31 C: UnindexedConsumer<Self::Item>,
32 {
33 self.inner
34 .map(|x: Bucket| unsafe { x.as_ref() })
35 .drive_unindexed(consumer)
36 }
37}
38
39impl<T> Clone for ParIter<'_, T> {
40 #[cfg_attr(feature = "inline-more", inline)]
41 fn clone(&self) -> Self {
42 Self {
43 inner: self.inner.clone(),
44 marker: PhantomData,
45 }
46 }
47}
48
49impl<T: fmt::Debug> fmt::Debug for ParIter<'_, T> {
50 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
51 let iter: impl Iterator = unsafe { self.inner.iter() }.map(|x: Bucket| unsafe { x.as_ref() });
52 f.debug_list().entries(iter).finish()
53 }
54}
55
56/// Parallel iterator over mutable references to entries in a map.
57///
58/// This iterator is created by the [`par_iter_mut`] method on [`HashTable`]
59/// (provided by the [`IntoParallelRefMutIterator`] trait).
60/// See its documentation for more.
61///
62/// [`par_iter_mut`]: /hashbrown/struct.HashTable.html#method.par_iter_mut
63/// [`HashTable`]: /hashbrown/struct.HashTable.html
64/// [`IntoParallelRefMutIterator`]: https://docs.rs/rayon/1.0/rayon/iter/trait.IntoParallelRefMutIterator.html
65pub struct ParIterMut<'a, T> {
66 inner: RawParIter<T>,
67 marker: PhantomData<&'a mut T>,
68}
69
70impl<'a, T: Send> ParallelIterator for ParIterMut<'a, T> {
71 type Item = &'a mut T;
72
73 #[cfg_attr(feature = "inline-more", inline)]
74 fn drive_unindexed<C>(self, consumer: C) -> C::Result
75 where
76 C: UnindexedConsumer<Self::Item>,
77 {
78 self.inner
79 .map(|x: Bucket| unsafe { x.as_mut() })
80 .drive_unindexed(consumer)
81 }
82}
83
84impl<T: fmt::Debug> fmt::Debug for ParIterMut<'_, T> {
85 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
86 ParIter {
87 inner: self.inner.clone(),
88 marker: PhantomData,
89 }
90 .fmt(f)
91 }
92}
93
94/// Parallel iterator over entries of a consumed map.
95///
96/// This iterator is created by the [`into_par_iter`] method on [`HashTable`]
97/// (provided by the [`IntoParallelIterator`] trait).
98/// See its documentation for more.
99///
100/// [`into_par_iter`]: /hashbrown/struct.HashTable.html#method.into_par_iter
101/// [`HashTable`]: /hashbrown/struct.HashTable.html
102/// [`IntoParallelIterator`]: https://docs.rs/rayon/1.0/rayon/iter/trait.IntoParallelIterator.html
103pub struct IntoParIter<T, A: Allocator = Global> {
104 inner: RawIntoParIter<T, A>,
105}
106
107impl<T: Send, A: Allocator + Send> ParallelIterator for IntoParIter<T, A> {
108 type Item = T;
109
110 #[cfg_attr(feature = "inline-more", inline)]
111 fn drive_unindexed<C>(self, consumer: C) -> C::Result
112 where
113 C: UnindexedConsumer<Self::Item>,
114 {
115 self.inner.drive_unindexed(consumer)
116 }
117}
118
119impl<T: fmt::Debug, A: Allocator> fmt::Debug for IntoParIter<T, A> {
120 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
121 ParIter {
122 inner: unsafe { self.inner.par_iter() },
123 marker: PhantomData,
124 }
125 .fmt(f)
126 }
127}
128
129/// Parallel draining iterator over entries of a map.
130///
131/// This iterator is created by the [`par_drain`] method on [`HashTable`].
132/// See its documentation for more.
133///
134/// [`par_drain`]: /hashbrown/struct.HashTable.html#method.par_drain
135/// [`HashTable`]: /hashbrown/struct.HashTable.html
136pub struct ParDrain<'a, T, A: Allocator = Global> {
137 inner: RawParDrain<'a, T, A>,
138}
139
140impl<T: Send, A: Allocator + Sync> ParallelIterator for ParDrain<'_, T, A> {
141 type Item = T;
142
143 #[cfg_attr(feature = "inline-more", inline)]
144 fn drive_unindexed<C>(self, consumer: C) -> C::Result
145 where
146 C: UnindexedConsumer<Self::Item>,
147 {
148 self.inner.drive_unindexed(consumer)
149 }
150}
151
152impl<T: fmt::Debug, A: Allocator> fmt::Debug for ParDrain<'_, T, A> {
153 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
154 ParIter {
155 inner: unsafe { self.inner.par_iter() },
156 marker: PhantomData,
157 }
158 .fmt(f)
159 }
160}
161
162impl<T: Send, A: Allocator> HashTable<T, A> {
163 /// Consumes (potentially in parallel) all values in an arbitrary order,
164 /// while preserving the map's allocated memory for reuse.
165 #[cfg_attr(feature = "inline-more", inline)]
166 pub fn par_drain(&mut self) -> ParDrain<'_, T, A> {
167 ParDrain {
168 inner: self.raw.par_drain(),
169 }
170 }
171}
172
173impl<T: Send, A: Allocator + Send> IntoParallelIterator for HashTable<T, A> {
174 type Item = T;
175 type Iter = IntoParIter<T, A>;
176
177 #[cfg_attr(feature = "inline-more", inline)]
178 fn into_par_iter(self) -> Self::Iter {
179 IntoParIter {
180 inner: self.raw.into_par_iter(),
181 }
182 }
183}
184
185impl<'a, T: Sync, A: Allocator> IntoParallelIterator for &'a HashTable<T, A> {
186 type Item = &'a T;
187 type Iter = ParIter<'a, T>;
188
189 #[cfg_attr(feature = "inline-more", inline)]
190 fn into_par_iter(self) -> Self::Iter {
191 ParIter {
192 inner: unsafe { self.raw.par_iter() },
193 marker: PhantomData,
194 }
195 }
196}
197
198impl<'a, T: Send, A: Allocator> IntoParallelIterator for &'a mut HashTable<T, A> {
199 type Item = &'a mut T;
200 type Iter = ParIterMut<'a, T>;
201
202 #[cfg_attr(feature = "inline-more", inline)]
203 fn into_par_iter(self) -> Self::Iter {
204 ParIterMut {
205 inner: unsafe { self.raw.par_iter() },
206 marker: PhantomData,
207 }
208 }
209}
210
211#[cfg(test)]
212mod test_par_table {
213 use alloc::vec::Vec;
214 use core::sync::atomic::{AtomicUsize, Ordering};
215
216 use rayon::prelude::*;
217
218 use crate::{hash_map::make_hash, hash_table::HashTable, DefaultHashBuilder};
219
220 #[test]
221 fn test_iterate() {
222 let hasher = DefaultHashBuilder::default();
223 let mut a = HashTable::new();
224 for i in 0..32 {
225 a.insert_unique(make_hash(&hasher, &i), i, |x| make_hash(&hasher, x));
226 }
227 let observed = AtomicUsize::new(0);
228 a.par_iter().for_each(|k| {
229 observed.fetch_or(1 << *k, Ordering::Relaxed);
230 });
231 assert_eq!(observed.into_inner(), 0xFFFF_FFFF);
232 }
233
234 #[test]
235 fn test_move_iter() {
236 let hasher = DefaultHashBuilder::default();
237 let hs = {
238 let mut hs = HashTable::new();
239
240 hs.insert_unique(make_hash(&hasher, &'a'), 'a', |x| make_hash(&hasher, x));
241 hs.insert_unique(make_hash(&hasher, &'b'), 'b', |x| make_hash(&hasher, x));
242
243 hs
244 };
245
246 let v = hs.into_par_iter().collect::<Vec<char>>();
247 assert!(v == ['a', 'b'] || v == ['b', 'a']);
248 }
249}
250