]>
Commit | Line | Data |
---|---|---|
48663c56 XL |
1 | use crate::raw::Bucket; |
2 | use crate::raw::{RawIterRange, RawTable}; | |
3 | use crate::scopeguard::guard; | |
4 | use alloc::alloc::dealloc; | |
5 | use core::marker::PhantomData; | |
6 | use core::mem; | |
7 | use core::ptr::NonNull; | |
8 | use rayon::iter::{ | |
9 | plumbing::{self, Folder, UnindexedConsumer, UnindexedProducer}, | |
10 | ParallelIterator, | |
11 | }; | |
12 | ||
13 | /// Parallel iterator which returns a raw pointer to every full bucket in the table. | |
14 | pub struct RawParIter<T> { | |
15 | iter: RawIterRange<T>, | |
16 | } | |
17 | ||
18 | impl<T> ParallelIterator for RawParIter<T> { | |
19 | type Item = Bucket<T>; | |
20 | ||
e74abb32 | 21 | #[cfg_attr(feature = "inline-more", inline)] |
48663c56 XL |
22 | fn drive_unindexed<C>(self, consumer: C) -> C::Result |
23 | where | |
24 | C: UnindexedConsumer<Self::Item>, | |
25 | { | |
26 | let producer = ParIterProducer { iter: self.iter }; | |
27 | plumbing::bridge_unindexed(producer, consumer) | |
28 | } | |
29 | } | |
30 | ||
31 | /// Producer which returns a `Bucket<T>` for every element. | |
32 | struct ParIterProducer<T> { | |
33 | iter: RawIterRange<T>, | |
34 | } | |
35 | ||
36 | impl<T> UnindexedProducer for ParIterProducer<T> { | |
37 | type Item = Bucket<T>; | |
38 | ||
e74abb32 | 39 | #[cfg_attr(feature = "inline-more", inline)] |
48663c56 XL |
40 | fn split(self) -> (Self, Option<Self>) { |
41 | let (left, right) = self.iter.split(); | |
42 | let left = ParIterProducer { iter: left }; | |
43 | let right = right.map(|right| ParIterProducer { iter: right }); | |
44 | (left, right) | |
45 | } | |
46 | ||
e74abb32 | 47 | #[cfg_attr(feature = "inline-more", inline)] |
48663c56 XL |
48 | fn fold_with<F>(self, folder: F) -> F |
49 | where | |
50 | F: Folder<Self::Item>, | |
51 | { | |
52 | folder.consume_iter(self.iter) | |
53 | } | |
54 | } | |
55 | ||
56 | /// Parallel iterator which consumes a table and returns elements. | |
57 | pub struct RawIntoParIter<T> { | |
58 | table: RawTable<T>, | |
59 | } | |
60 | ||
61 | impl<T: Send> ParallelIterator for RawIntoParIter<T> { | |
62 | type Item = T; | |
63 | ||
e74abb32 | 64 | #[cfg_attr(feature = "inline-more", inline)] |
48663c56 XL |
65 | fn drive_unindexed<C>(self, consumer: C) -> C::Result |
66 | where | |
67 | C: UnindexedConsumer<Self::Item>, | |
68 | { | |
69 | let iter = unsafe { self.table.iter().iter }; | |
70 | let _guard = guard(self.table.into_alloc(), |alloc| { | |
71 | if let Some((ptr, layout)) = *alloc { | |
72 | unsafe { | |
73 | dealloc(ptr.as_ptr(), layout); | |
74 | } | |
75 | } | |
76 | }); | |
77 | let producer = ParDrainProducer { iter }; | |
78 | plumbing::bridge_unindexed(producer, consumer) | |
79 | } | |
80 | } | |
81 | ||
82 | /// Parallel iterator which consumes elements without freeing the table storage. | |
83 | pub struct RawParDrain<'a, T> { | |
84 | // We don't use a &'a mut RawTable<T> because we want RawParDrain to be | |
85 | // covariant over T. | |
86 | table: NonNull<RawTable<T>>, | |
87 | marker: PhantomData<&'a RawTable<T>>, | |
88 | } | |
89 | ||
90 | unsafe impl<T> Send for RawParDrain<'_, T> {} | |
91 | ||
92 | impl<T: Send> ParallelIterator for RawParDrain<'_, T> { | |
93 | type Item = T; | |
94 | ||
e74abb32 | 95 | #[cfg_attr(feature = "inline-more", inline)] |
48663c56 XL |
96 | fn drive_unindexed<C>(self, consumer: C) -> C::Result |
97 | where | |
98 | C: UnindexedConsumer<Self::Item>, | |
99 | { | |
100 | let _guard = guard(self.table, |table| unsafe { | |
101 | table.as_mut().clear_no_drop() | |
102 | }); | |
103 | let iter = unsafe { self.table.as_ref().iter().iter }; | |
104 | mem::forget(self); | |
105 | let producer = ParDrainProducer { iter }; | |
106 | plumbing::bridge_unindexed(producer, consumer) | |
107 | } | |
108 | } | |
109 | ||
110 | impl<T> Drop for RawParDrain<'_, T> { | |
111 | fn drop(&mut self) { | |
112 | // If drive_unindexed is not called then simply clear the table. | |
113 | unsafe { self.table.as_mut().clear() } | |
114 | } | |
115 | } | |
116 | ||
117 | /// Producer which will consume all elements in the range, even if it is dropped | |
118 | /// halfway through. | |
119 | struct ParDrainProducer<T> { | |
120 | iter: RawIterRange<T>, | |
121 | } | |
122 | ||
123 | impl<T: Send> UnindexedProducer for ParDrainProducer<T> { | |
124 | type Item = T; | |
125 | ||
e74abb32 | 126 | #[cfg_attr(feature = "inline-more", inline)] |
48663c56 XL |
127 | fn split(self) -> (Self, Option<Self>) { |
128 | let (left, right) = self.iter.clone().split(); | |
129 | mem::forget(self); | |
130 | let left = ParDrainProducer { iter: left }; | |
131 | let right = right.map(|right| ParDrainProducer { iter: right }); | |
132 | (left, right) | |
133 | } | |
134 | ||
e74abb32 | 135 | #[cfg_attr(feature = "inline-more", inline)] |
48663c56 XL |
136 | fn fold_with<F>(mut self, mut folder: F) -> F |
137 | where | |
138 | F: Folder<Self::Item>, | |
139 | { | |
140 | // Make sure to modify the iterator in-place so that any remaining | |
141 | // elements are processed in our Drop impl. | |
142 | while let Some(item) = self.iter.next() { | |
143 | folder = folder.consume(unsafe { item.read() }); | |
144 | if folder.full() { | |
145 | return folder; | |
146 | } | |
147 | } | |
148 | ||
149 | // If we processed all elements then we don't need to run the drop. | |
150 | mem::forget(self); | |
151 | folder | |
152 | } | |
153 | } | |
154 | ||
155 | impl<T> Drop for ParDrainProducer<T> { | |
e74abb32 | 156 | #[cfg_attr(feature = "inline-more", inline)] |
48663c56 XL |
157 | fn drop(&mut self) { |
158 | // Drop all remaining elements | |
159 | if mem::needs_drop::<T>() { | |
160 | while let Some(item) = self.iter.next() { | |
161 | unsafe { | |
162 | item.drop(); | |
163 | } | |
164 | } | |
165 | } | |
166 | } | |
167 | } | |
168 | ||
169 | impl<T> RawTable<T> { | |
170 | /// Returns a parallel iterator over the elements in a `RawTable`. | |
e74abb32 | 171 | #[cfg_attr(feature = "inline-more", inline)] |
48663c56 XL |
172 | pub fn par_iter(&self) -> RawParIter<T> { |
173 | RawParIter { | |
174 | iter: unsafe { self.iter().iter }, | |
175 | } | |
176 | } | |
177 | ||
178 | /// Returns a parallel iterator over the elements in a `RawTable`. | |
e74abb32 | 179 | #[cfg_attr(feature = "inline-more", inline)] |
48663c56 XL |
180 | pub fn into_par_iter(self) -> RawIntoParIter<T> { |
181 | RawIntoParIter { table: self } | |
182 | } | |
183 | ||
184 | /// Returns a parallel iterator which consumes all elements of a `RawTable` | |
185 | /// without freeing its memory allocation. | |
e74abb32 | 186 | #[cfg_attr(feature = "inline-more", inline)] |
48663c56 XL |
187 | pub fn par_drain(&mut self) -> RawParDrain<'_, T> { |
188 | RawParDrain { | |
189 | table: NonNull::from(self), | |
190 | marker: PhantomData, | |
191 | } | |
192 | } | |
193 | } |