/rust/registry/src/index.crates.io-1949cf8c6b5b557f/roaring-0.11.4/src/treemap/multiops.rs
Line | Count | Source |
1 | | use alloc::collections::{binary_heap::PeekMut, BTreeMap, BinaryHeap}; |
2 | | use core::{borrow::Borrow, cmp::Ordering, mem}; |
3 | | |
4 | | use crate::{MultiOps, RoaringBitmap, RoaringTreemap}; |
5 | | |
6 | | #[cfg(not(feature = "std"))] |
7 | | use alloc::vec::Vec; |
8 | | |
9 | | impl<I> MultiOps<RoaringTreemap> for I |
10 | | where |
11 | | I: IntoIterator<Item = RoaringTreemap>, |
12 | | { |
13 | | type Output = RoaringTreemap; |
14 | | |
15 | 0 | fn union(self) -> Self::Output { |
16 | 0 | try_simple_multi_op_owned::<_, _, UnionOp>( |
17 | 0 | self.into_iter().map(Ok::<_, core::convert::Infallible>), |
18 | | ) |
19 | 0 | .unwrap() |
20 | 0 | } |
21 | | |
22 | 0 | fn intersection(self) -> Self::Output { |
23 | 0 | try_ordered_multi_op_owned::<_, _, IntersectionOp>( |
24 | 0 | self.into_iter().map(Ok::<_, core::convert::Infallible>), |
25 | | ) |
26 | 0 | .unwrap() |
27 | 0 | } |
28 | | |
29 | 0 | fn difference(self) -> Self::Output { |
30 | 0 | try_ordered_multi_op_owned::<_, _, DifferenceOp>( |
31 | 0 | self.into_iter().map(Ok::<_, core::convert::Infallible>), |
32 | | ) |
33 | 0 | .unwrap() |
34 | 0 | } |
35 | | |
36 | 0 | fn symmetric_difference(self) -> Self::Output { |
37 | 0 | try_simple_multi_op_owned::<_, _, SymmetricDifferenceOp>( |
38 | 0 | self.into_iter().map(Ok::<_, core::convert::Infallible>), |
39 | | ) |
40 | 0 | .unwrap() |
41 | 0 | } |
42 | | } |
43 | | |
44 | | impl<I, E> MultiOps<Result<RoaringTreemap, E>> for I |
45 | | where |
46 | | I: IntoIterator<Item = Result<RoaringTreemap, E>>, |
47 | | { |
48 | | type Output = Result<RoaringTreemap, E>; |
49 | | |
50 | 0 | fn union(self) -> Self::Output { |
51 | 0 | try_simple_multi_op_owned::<_, _, UnionOp>(self) |
52 | 0 | } |
53 | | |
54 | 0 | fn intersection(self) -> Self::Output { |
55 | 0 | try_ordered_multi_op_owned::<_, _, IntersectionOp>(self) |
56 | 0 | } |
57 | | |
58 | 0 | fn difference(self) -> Self::Output { |
59 | 0 | try_ordered_multi_op_owned::<_, _, DifferenceOp>(self) |
60 | 0 | } |
61 | | |
62 | 0 | fn symmetric_difference(self) -> Self::Output { |
63 | 0 | try_simple_multi_op_owned::<_, _, SymmetricDifferenceOp>(self) |
64 | 0 | } |
65 | | } |
66 | | |
67 | | #[inline] |
68 | 0 | fn try_simple_multi_op_owned<E, I, O: Op>(treemaps: I) -> Result<RoaringTreemap, E> |
69 | 0 | where |
70 | 0 | I: IntoIterator<Item = Result<RoaringTreemap, E>>, |
71 | | { |
72 | 0 | let treemaps = treemaps.into_iter().collect::<Result<Vec<_>, _>>()?; |
73 | | |
74 | 0 | let mut heap: BinaryHeap<_> = treemaps |
75 | 0 | .into_iter() |
76 | 0 | .filter_map(|treemap| { |
77 | 0 | let mut iter = treemap.map.into_iter(); |
78 | 0 | iter.next().map(|(key, bitmap)| PeekedRoaringBitmap { key, bitmap, iter }) |
79 | 0 | }) |
80 | 0 | .collect(); |
81 | | |
82 | 0 | let mut bitmaps = Vec::new(); |
83 | 0 | let mut map = BTreeMap::new(); |
84 | | |
85 | 0 | while let Some(mut peek) = heap.peek_mut() { |
86 | 0 | let (key, bitmap) = match peek.iter.next() { |
87 | 0 | Some((next_key, next_bitmap)) => { |
88 | 0 | let key = peek.key; |
89 | 0 | peek.key = next_key; |
90 | 0 | let bitmap = mem::replace(&mut peek.bitmap, next_bitmap); |
91 | 0 | (key, bitmap) |
92 | | } |
93 | | None => { |
94 | 0 | let popped = PeekMut::pop(peek); |
95 | 0 | (popped.key, popped.bitmap) |
96 | | } |
97 | | }; |
98 | | |
99 | 0 | if let Some((first_key, _)) = bitmaps.first() { |
100 | 0 | if *first_key != key { |
101 | 0 | let current_key = *first_key; |
102 | 0 | let computed_bitmap = O::op_owned(bitmaps.drain(..).map(|(_, rb)| rb)); |
103 | 0 | if !computed_bitmap.is_empty() { |
104 | 0 | map.insert(current_key, computed_bitmap); |
105 | 0 | } |
106 | 0 | } |
107 | 0 | } |
108 | | |
109 | 0 | bitmaps.push((key, bitmap)); |
110 | | } |
111 | | |
112 | 0 | if let Some((first_key, _)) = bitmaps.first() { |
113 | 0 | let current_key = *first_key; |
114 | 0 | let computed_bitmap = O::op_owned(bitmaps.drain(..).map(|(_, rb)| rb)); |
115 | 0 | if !computed_bitmap.is_empty() { |
116 | 0 | map.insert(current_key, computed_bitmap); |
117 | 0 | } |
118 | 0 | } |
119 | | |
120 | 0 | Ok(RoaringTreemap { map }) |
121 | 0 | } |
122 | | |
123 | | #[inline] |
124 | 0 | fn try_ordered_multi_op_owned<E, I, O: Op>(treemaps: I) -> Result<RoaringTreemap, E> |
125 | 0 | where |
126 | 0 | I: IntoIterator<Item = Result<RoaringTreemap, E>>, |
127 | | { |
128 | 0 | let mut treemaps = treemaps.into_iter(); |
129 | 0 | let mut treemap = match treemaps.next().transpose()? { |
130 | 0 | Some(treemap) => treemap, |
131 | 0 | None => return Ok(RoaringTreemap::new()), |
132 | | }; |
133 | 0 | let mut treemaps = treemaps.collect::<Result<Vec<_>, _>>()?; |
134 | | |
135 | | // for each key in the first treemap we're going to find and |
136 | | // accumulate all the corresponding bitmaps |
137 | 0 | let keys: Vec<_> = treemap.map.keys().copied().collect(); |
138 | 0 | for k in keys { |
139 | | // the unwrap is safe since we're iterating on our keys |
140 | 0 | let current_bitmap = treemap.map.remove(&k).unwrap(); |
141 | 0 | let new_bitmap = |
142 | 0 | O::op_owned(core::iter::once(current_bitmap).chain( |
143 | 0 | treemaps.iter_mut().map(|treemap| treemap.map.remove(&k).unwrap_or_default()), |
144 | | )); |
145 | 0 | if !new_bitmap.is_empty() { |
146 | 0 | treemap.map.insert(k, new_bitmap); |
147 | 0 | } |
148 | | } |
149 | | |
150 | 0 | Ok(treemap) |
151 | 0 | } |
152 | | |
153 | | #[inline] |
154 | 0 | fn try_ordered_multi_op_ref<'a, E: 'a, I, O: Op>(treemaps: I) -> Result<RoaringTreemap, E> |
155 | 0 | where |
156 | 0 | I: IntoIterator<Item = Result<&'a RoaringTreemap, E>>, |
157 | | { |
158 | 0 | let mut treemaps = treemaps.into_iter(); |
159 | 0 | let treemap = match treemaps.next().transpose()? { |
160 | 0 | Some(treemap) => treemap, |
161 | 0 | None => return Ok(RoaringTreemap::new()), |
162 | | }; |
163 | 0 | let treemaps = treemaps.collect::<Result<Vec<_>, _>>()?; |
164 | | |
165 | 0 | let mut ret = RoaringTreemap::new(); |
166 | | |
167 | | // for each keys in the first treemap we're going find and accumulate all the corresponding bitmaps |
168 | 0 | let keys: Vec<_> = treemap.map.keys().copied().collect(); |
169 | 0 | let empty_bitmap = RoaringBitmap::new(); |
170 | 0 | for k in keys { |
171 | | // the unwrap is safe since we're iterating on our keys |
172 | 0 | let current_bitmap = treemap.map.get(&k).unwrap(); |
173 | 0 | let new_bitmap = O::op_ref( |
174 | 0 | core::iter::once(current_bitmap) |
175 | 0 | .chain(treemaps.iter().map(|treemap| treemap.map.get(&k).unwrap_or(&empty_bitmap))), |
176 | | ); |
177 | 0 | if !new_bitmap.is_empty() { |
178 | 0 | ret.map.insert(k, new_bitmap); |
179 | 0 | } |
180 | | } |
181 | | |
182 | 0 | Ok(ret) |
183 | 0 | } |
184 | | |
185 | | #[inline] |
186 | 0 | fn try_simple_multi_op_ref<'a, E: 'a, I, O: Op>(treemaps: I) -> Result<RoaringTreemap, E> |
187 | 0 | where |
188 | 0 | I: IntoIterator<Item = Result<&'a RoaringTreemap, E>>, |
189 | | { |
190 | 0 | let treemaps = treemaps.into_iter().collect::<Result<Vec<_>, E>>()?; |
191 | | |
192 | 0 | let mut heap: BinaryHeap<_> = treemaps |
193 | 0 | .into_iter() |
194 | 0 | .filter_map(|treemap| { |
195 | 0 | let mut iter = treemap.map.iter(); |
196 | 0 | iter.next().map(|(&key, bitmap)| PeekedRoaringBitmap { key, bitmap, iter }) |
197 | 0 | }) |
198 | 0 | .collect(); |
199 | | |
200 | 0 | let mut bitmaps = Vec::new(); |
201 | 0 | let mut map = BTreeMap::new(); |
202 | | |
203 | 0 | while let Some(mut peek) = heap.peek_mut() { |
204 | 0 | let (key, bitmap) = match peek.iter.next() { |
205 | 0 | Some((&next_key, next_bitmap)) => { |
206 | 0 | let key = peek.key; |
207 | 0 | peek.key = next_key; |
208 | 0 | let bitmap = mem::replace(&mut peek.bitmap, next_bitmap); |
209 | 0 | (key, bitmap) |
210 | | } |
211 | | None => { |
212 | 0 | let popped = PeekMut::pop(peek); |
213 | 0 | (popped.key, popped.bitmap) |
214 | | } |
215 | | }; |
216 | | |
217 | 0 | if let Some((first_key, _)) = bitmaps.first() { |
218 | 0 | if *first_key != key { |
219 | 0 | let current_key = *first_key; |
220 | 0 | let computed_bitmap = O::op_ref(bitmaps.drain(..).map(|(_, rb)| rb)); |
221 | 0 | if !computed_bitmap.is_empty() { |
222 | 0 | map.insert(current_key, computed_bitmap); |
223 | 0 | } |
224 | 0 | } |
225 | 0 | } |
226 | | |
227 | 0 | bitmaps.push((key, bitmap)); |
228 | | } |
229 | | |
230 | 0 | if let Some((first_key, _)) = bitmaps.first() { |
231 | 0 | let current_key = *first_key; |
232 | 0 | let computed_bitmap = O::op_ref(bitmaps.drain(..).map(|(_, rb)| rb)); |
233 | 0 | if !computed_bitmap.is_empty() { |
234 | 0 | map.insert(current_key, computed_bitmap); |
235 | 0 | } |
236 | 0 | } |
237 | | |
238 | 0 | Ok(RoaringTreemap { map }) |
239 | 0 | } |
240 | | |
241 | | trait Op { |
242 | | fn op_owned<I: IntoIterator<Item = RoaringBitmap>>(iter: I) -> RoaringBitmap; |
243 | | fn op_ref<'a, I: IntoIterator<Item = &'a RoaringBitmap>>(iter: I) -> RoaringBitmap; |
244 | | } |
245 | | |
246 | | enum UnionOp {} |
247 | | |
248 | | impl Op for UnionOp { |
249 | 0 | fn op_owned<J: IntoIterator<Item = RoaringBitmap>>(iter: J) -> RoaringBitmap { |
250 | 0 | iter.union() |
251 | 0 | } |
252 | | |
253 | 0 | fn op_ref<'a, J: IntoIterator<Item = &'a RoaringBitmap>>(iter: J) -> RoaringBitmap { |
254 | 0 | iter.union() |
255 | 0 | } |
256 | | } |
257 | | |
258 | | enum IntersectionOp {} |
259 | | |
260 | | impl Op for IntersectionOp { |
261 | 0 | fn op_owned<J: IntoIterator<Item = RoaringBitmap>>(iter: J) -> RoaringBitmap { |
262 | 0 | iter.intersection() |
263 | 0 | } |
264 | | |
265 | 0 | fn op_ref<'a, J: IntoIterator<Item = &'a RoaringBitmap>>(iter: J) -> RoaringBitmap { |
266 | 0 | iter.intersection() |
267 | 0 | } |
268 | | } |
269 | | |
270 | | enum DifferenceOp {} |
271 | | |
272 | | impl Op for DifferenceOp { |
273 | 0 | fn op_owned<J: IntoIterator<Item = RoaringBitmap>>(iter: J) -> RoaringBitmap { |
274 | 0 | iter.difference() |
275 | 0 | } |
276 | | |
277 | 0 | fn op_ref<'a, J: IntoIterator<Item = &'a RoaringBitmap>>(iter: J) -> RoaringBitmap { |
278 | 0 | iter.difference() |
279 | 0 | } |
280 | | } |
281 | | |
282 | | enum SymmetricDifferenceOp {} |
283 | | |
284 | | impl Op for SymmetricDifferenceOp { |
285 | 0 | fn op_owned<J: IntoIterator<Item = RoaringBitmap>>(iter: J) -> RoaringBitmap { |
286 | 0 | iter.symmetric_difference() |
287 | 0 | } |
288 | | |
289 | 0 | fn op_ref<'a, J: IntoIterator<Item = &'a RoaringBitmap>>(iter: J) -> RoaringBitmap { |
290 | 0 | iter.symmetric_difference() |
291 | 0 | } |
292 | | } |
293 | | |
294 | | impl<'a, I> MultiOps<&'a RoaringTreemap> for I |
295 | | where |
296 | | I: IntoIterator<Item = &'a RoaringTreemap>, |
297 | | { |
298 | | type Output = RoaringTreemap; |
299 | | |
300 | 0 | fn union(self) -> Self::Output { |
301 | 0 | try_simple_multi_op_ref::<_, _, UnionOp>( |
302 | 0 | self.into_iter().map(Ok::<_, core::convert::Infallible>), |
303 | | ) |
304 | 0 | .unwrap() |
305 | 0 | } |
306 | | |
307 | 0 | fn intersection(self) -> Self::Output { |
308 | 0 | try_ordered_multi_op_ref::<_, _, IntersectionOp>( |
309 | 0 | self.into_iter().map(Ok::<_, core::convert::Infallible>), |
310 | | ) |
311 | 0 | .unwrap() |
312 | 0 | } |
313 | | |
314 | 0 | fn difference(self) -> Self::Output { |
315 | 0 | try_ordered_multi_op_ref::<_, _, DifferenceOp>( |
316 | 0 | self.into_iter().map(Ok::<_, core::convert::Infallible>), |
317 | | ) |
318 | 0 | .unwrap() |
319 | 0 | } |
320 | | |
321 | 0 | fn symmetric_difference(self) -> Self::Output { |
322 | 0 | try_simple_multi_op_ref::<_, _, SymmetricDifferenceOp>( |
323 | 0 | self.into_iter().map(Ok::<_, core::convert::Infallible>), |
324 | | ) |
325 | 0 | .unwrap() |
326 | 0 | } |
327 | | } |
328 | | |
329 | | impl<'a, I, E: 'a> MultiOps<Result<&'a RoaringTreemap, E>> for I |
330 | | where |
331 | | I: IntoIterator<Item = Result<&'a RoaringTreemap, E>>, |
332 | | { |
333 | | type Output = Result<RoaringTreemap, E>; |
334 | | |
335 | 0 | fn union(self) -> Self::Output { |
336 | 0 | try_simple_multi_op_ref::<_, _, UnionOp>(self) |
337 | 0 | } |
338 | | |
339 | 0 | fn intersection(self) -> Self::Output { |
340 | 0 | try_ordered_multi_op_ref::<_, _, IntersectionOp>(self) |
341 | 0 | } |
342 | | |
343 | 0 | fn difference(self) -> Self::Output { |
344 | 0 | try_ordered_multi_op_ref::<_, _, DifferenceOp>(self) |
345 | 0 | } |
346 | | |
347 | 0 | fn symmetric_difference(self) -> Self::Output { |
348 | 0 | try_simple_multi_op_ref::<_, _, SymmetricDifferenceOp>(self) |
349 | 0 | } |
350 | | } |
351 | | |
352 | | struct PeekedRoaringBitmap<R, I> { |
353 | | key: u32, |
354 | | bitmap: R, |
355 | | iter: I, |
356 | | } |
357 | | |
358 | | impl<R: Borrow<RoaringBitmap>, I> Ord for PeekedRoaringBitmap<R, I> { |
359 | 0 | fn cmp(&self, other: &Self) -> Ordering { |
360 | 0 | self.key.cmp(&other.key).reverse() |
361 | 0 | } |
362 | | } |
363 | | |
364 | | impl<R: Borrow<RoaringBitmap>, I> PartialOrd for PeekedRoaringBitmap<R, I> { |
365 | 0 | fn partial_cmp(&self, other: &Self) -> Option<Ordering> { |
366 | 0 | Some(self.cmp(other)) |
367 | 0 | } |
368 | | } |
369 | | |
370 | | impl<R: Borrow<RoaringBitmap>, I> Eq for PeekedRoaringBitmap<R, I> {} |
371 | | |
372 | | impl<R: Borrow<RoaringBitmap>, I> PartialEq for PeekedRoaringBitmap<R, I> { |
373 | 0 | fn eq(&self, other: &Self) -> bool { |
374 | 0 | self.key == other.key |
375 | 0 | } |
376 | | } |