roaring/treemap/
multiops.rs

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

roaring/treemap/
multiops.rs
