roaring/treemap/
ops.rs

1use alloc::collections::btree_map::Entry;
2use core::mem;
3use core::ops::{BitAnd, BitAndAssign, BitOr, BitOrAssign, BitXor, BitXorAssign, Sub, SubAssign};
4
5use crate::RoaringTreemap;
6
7#[cfg(not(feature = "std"))]
8use alloc::vec::Vec;
9
10impl RoaringTreemap {
11    /// Computes the len of the union with the specified other treemap without creating a new
12    /// treemap.
13    ///
14    /// This is faster and more space efficient when you're only interested in the cardinality of
15    /// the union.
16    ///
17    /// # Examples
18    ///
19    /// ```rust
20    /// use roaring::RoaringTreemap;
21    ///
22    /// let rb1: RoaringTreemap = (1..4).collect();
23    /// let rb2: RoaringTreemap = (3..5).collect();
24    ///
25    ///
26    /// assert_eq!(rb1.union_len(&rb2), (rb1 | rb2).len());
27    /// ```
28    pub fn union_len(&self, other: &RoaringTreemap) -> u64 {
29        self.len().wrapping_add(other.len()).wrapping_sub(self.intersection_len(other))
30    }
31
32    /// Computes the len of the intersection with the specified other treemap without creating a
33    /// new treemap.
34    ///
35    /// This is faster and more space efficient when you're only interested in the cardinality of
36    /// the intersection.
37    ///
38    /// # Examples
39    ///
40    /// ```rust
41    /// use roaring::RoaringTreemap;
42    ///
43    /// let rb1: RoaringTreemap = (1..4).collect();
44    /// let rb2: RoaringTreemap = (3..5).collect();
45    ///
46    ///
47    /// assert_eq!(rb1.intersection_len(&rb2), (rb1 & rb2).len());
48    /// ```
49    pub fn intersection_len(&self, other: &RoaringTreemap) -> u64 {
50        self.pairs(other)
51            .map(|pair| match pair {
52                (Some(..), None) => 0,
53                (None, Some(..)) => 0,
54                (Some(lhs), Some(rhs)) => lhs.intersection_len(rhs),
55                (None, None) => 0,
56            })
57            .sum()
58    }
59
60    /// Computes the len of the difference with the specified other treemap without creating a new
61    /// treemap.
62    ///
63    /// This is faster and more space efficient when you're only interested in the cardinality of
64    /// the difference.
65    ///
66    /// # Examples
67    ///
68    /// ```rust
69    /// use roaring::RoaringTreemap;
70    ///
71    /// let rb1: RoaringTreemap = (1..4).collect();
72    /// let rb2: RoaringTreemap = (3..5).collect();
73    ///
74    ///
75    /// assert_eq!(rb1.difference_len(&rb2), (rb1 - rb2).len());
76    /// ```
77    pub fn difference_len(&self, other: &RoaringTreemap) -> u64 {
78        self.len() - self.intersection_len(other)
79    }
80
81    /// Computes the len of the symmetric difference with the specified other treemap without
82    /// creating a new bitmap.
83    ///
84    /// This is faster and more space efficient when you're only interested in the cardinality of
85    /// the symmetric difference.
86    ///
87    /// # Examples
88    ///
89    /// ```rust
90    /// use roaring::RoaringTreemap;
91    ///
92    /// let rb1: RoaringTreemap = (1..4).collect();
93    /// let rb2: RoaringTreemap = (3..5).collect();
94    ///
95    ///
96    /// assert_eq!(rb1.symmetric_difference_len(&rb2), (rb1 ^ rb2).len());
97    /// ```
98    pub fn symmetric_difference_len(&self, other: &RoaringTreemap) -> u64 {
99        let intersection_len = self.intersection_len(other);
100
101        self.len()
102            .wrapping_add(other.len())
103            .wrapping_sub(intersection_len)
104            .wrapping_sub(intersection_len)
105    }
106}
107
108impl BitOr<RoaringTreemap> for RoaringTreemap {
109    type Output = RoaringTreemap;
110
111    /// An `union` between two sets.
112    fn bitor(mut self, rhs: RoaringTreemap) -> RoaringTreemap {
113        BitOrAssign::bitor_assign(&mut self, rhs);
114        self
115    }
116}
117
118impl BitOr<&RoaringTreemap> for RoaringTreemap {
119    type Output = RoaringTreemap;
120
121    /// An `union` between two sets.
122    fn bitor(mut self, rhs: &RoaringTreemap) -> RoaringTreemap {
123        BitOrAssign::bitor_assign(&mut self, rhs);
124        self
125    }
126}
127
128impl BitOr<RoaringTreemap> for &RoaringTreemap {
129    type Output = RoaringTreemap;
130
131    /// An `union` between two sets.
132    fn bitor(self, rhs: RoaringTreemap) -> RoaringTreemap {
133        BitOr::bitor(rhs, self)
134    }
135}
136
137impl BitOr<&RoaringTreemap> for &RoaringTreemap {
138    type Output = RoaringTreemap;
139
140    /// An `union` between two sets.
141    fn bitor(self, rhs: &RoaringTreemap) -> RoaringTreemap {
142        if self.len() <= rhs.len() {
143            BitOr::bitor(rhs.clone(), self)
144        } else {
145            BitOr::bitor(self.clone(), rhs)
146        }
147    }
148}
149
150impl BitOrAssign<RoaringTreemap> for RoaringTreemap {
151    /// An `union` between two sets.
152    fn bitor_assign(&mut self, mut rhs: RoaringTreemap) {
153        // We make sure that we apply the union operation on the biggest map.
154        if self.len() < rhs.len() {
155            mem::swap(self, &mut rhs);
156        }
157
158        for (key, other_rb) in rhs.map {
159            match self.map.entry(key) {
160                Entry::Vacant(ent) => {
161                    ent.insert(other_rb);
162                }
163                Entry::Occupied(mut ent) => {
164                    BitOrAssign::bitor_assign(ent.get_mut(), other_rb);
165                }
166            }
167        }
168    }
169}
170
171impl BitOrAssign<&RoaringTreemap> for RoaringTreemap {
172    /// An `union` between two sets.
173    fn bitor_assign(&mut self, rhs: &RoaringTreemap) {
174        for (key, other_rb) in &rhs.map {
175            match self.map.entry(*key) {
176                Entry::Vacant(ent) => {
177                    ent.insert(other_rb.clone());
178                }
179                Entry::Occupied(mut ent) => {
180                    BitOrAssign::bitor_assign(ent.get_mut(), other_rb);
181                }
182            }
183        }
184    }
185}
186
187impl BitAnd<RoaringTreemap> for RoaringTreemap {
188    type Output = RoaringTreemap;
189
190    /// An `intersection` between two sets.
191    fn bitand(mut self, rhs: RoaringTreemap) -> RoaringTreemap {
192        BitAndAssign::bitand_assign(&mut self, rhs);
193        self
194    }
195}
196
197impl BitAnd<&RoaringTreemap> for RoaringTreemap {
198    type Output = RoaringTreemap;
199
200    /// An `intersection` between two sets.
201    fn bitand(mut self, rhs: &RoaringTreemap) -> RoaringTreemap {
202        BitAndAssign::bitand_assign(&mut self, rhs);
203        self
204    }
205}
206
207impl BitAnd<RoaringTreemap> for &RoaringTreemap {
208    type Output = RoaringTreemap;
209
210    /// An `intersection` between two sets.
211    fn bitand(self, rhs: RoaringTreemap) -> RoaringTreemap {
212        BitAnd::bitand(rhs, self)
213    }
214}
215
216impl BitAnd<&RoaringTreemap> for &RoaringTreemap {
217    type Output = RoaringTreemap;
218
219    /// An `intersection` between two sets.
220    fn bitand(self, rhs: &RoaringTreemap) -> RoaringTreemap {
221        if rhs.len() < self.len() {
222            BitAnd::bitand(self.clone(), rhs)
223        } else {
224            BitAnd::bitand(rhs.clone(), self)
225        }
226    }
227}
228
229impl BitAndAssign<RoaringTreemap> for RoaringTreemap {
230    /// An `intersection` between two sets.
231    fn bitand_assign(&mut self, mut rhs: RoaringTreemap) {
232        // We make sure that we apply the intersection operation on the smallest map.
233        if rhs.len() < self.len() {
234            mem::swap(self, &mut rhs);
235        }
236
237        BitAndAssign::bitand_assign(self, &rhs)
238    }
239}
240
241impl BitAndAssign<&RoaringTreemap> for RoaringTreemap {
242    /// An `intersection` between two sets.
243    fn bitand_assign(&mut self, rhs: &RoaringTreemap) {
244        let mut keys_to_remove: Vec<u32> = Vec::new();
245        for (key, self_rb) in &mut self.map {
246            match rhs.map.get(key) {
247                Some(other_rb) => {
248                    BitAndAssign::bitand_assign(self_rb, other_rb);
249                    if self_rb.is_empty() {
250                        keys_to_remove.push(*key);
251                    }
252                }
253                None => keys_to_remove.push(*key),
254            }
255        }
256
257        for key in keys_to_remove {
258            self.map.remove(&key);
259        }
260    }
261}
262
263impl Sub<RoaringTreemap> for RoaringTreemap {
264    type Output = RoaringTreemap;
265
266    /// A `difference` between two sets.
267    fn sub(mut self, rhs: RoaringTreemap) -> RoaringTreemap {
268        SubAssign::sub_assign(&mut self, rhs);
269        self
270    }
271}
272
273impl Sub<&RoaringTreemap> for RoaringTreemap {
274    type Output = RoaringTreemap;
275
276    /// A `difference` between two sets.
277    fn sub(mut self, rhs: &RoaringTreemap) -> RoaringTreemap {
278        SubAssign::sub_assign(&mut self, rhs);
279        self
280    }
281}
282
283impl Sub<RoaringTreemap> for &RoaringTreemap {
284    type Output = RoaringTreemap;
285
286    /// A `difference` between two sets.
287    fn sub(self, rhs: RoaringTreemap) -> RoaringTreemap {
288        Sub::sub(self.clone(), rhs)
289    }
290}
291
292impl Sub<&RoaringTreemap> for &RoaringTreemap {
293    type Output = RoaringTreemap;
294
295    /// A `difference` between two sets.
296    fn sub(self, rhs: &RoaringTreemap) -> RoaringTreemap {
297        Sub::sub(self.clone(), rhs)
298    }
299}
300
301impl SubAssign<RoaringTreemap> for RoaringTreemap {
302    /// A `difference` between two sets.
303    fn sub_assign(&mut self, rhs: RoaringTreemap) {
304        SubAssign::sub_assign(self, &rhs)
305    }
306}
307
308impl SubAssign<&RoaringTreemap> for RoaringTreemap {
309    /// A `difference` between two sets.
310    fn sub_assign(&mut self, rhs: &RoaringTreemap) {
311        for (key, rhs_rb) in &rhs.map {
312            match self.map.entry(*key) {
313                Entry::Vacant(_entry) => (),
314                Entry::Occupied(mut entry) => {
315                    SubAssign::sub_assign(entry.get_mut(), rhs_rb);
316                    if entry.get().is_empty() {
317                        entry.remove_entry();
318                    }
319                }
320            }
321        }
322    }
323}
324
325impl BitXor<RoaringTreemap> for RoaringTreemap {
326    type Output = RoaringTreemap;
327
328    /// A `symmetric difference` between two sets.
329    fn bitxor(mut self, rhs: RoaringTreemap) -> RoaringTreemap {
330        BitXorAssign::bitxor_assign(&mut self, rhs);
331        self
332    }
333}
334
335impl BitXor<&RoaringTreemap> for RoaringTreemap {
336    type Output = RoaringTreemap;
337
338    /// A `symmetric difference` between two sets.
339    fn bitxor(mut self, rhs: &RoaringTreemap) -> RoaringTreemap {
340        BitXorAssign::bitxor_assign(&mut self, rhs);
341        self
342    }
343}
344
345impl BitXor<RoaringTreemap> for &RoaringTreemap {
346    type Output = RoaringTreemap;
347
348    /// A `symmetric difference` between two sets.
349    fn bitxor(self, rhs: RoaringTreemap) -> RoaringTreemap {
350        BitXor::bitxor(rhs, self)
351    }
352}
353
354impl BitXor<&RoaringTreemap> for &RoaringTreemap {
355    type Output = RoaringTreemap;
356
357    /// A `symmetric difference` between two sets.
358    fn bitxor(self, rhs: &RoaringTreemap) -> RoaringTreemap {
359        if self.len() < rhs.len() {
360            BitXor::bitxor(self, rhs.clone())
361        } else {
362            BitXor::bitxor(self.clone(), rhs)
363        }
364    }
365}
366
367impl BitXorAssign<RoaringTreemap> for RoaringTreemap {
368    /// A `symmetric difference` between two sets.
369    fn bitxor_assign(&mut self, rhs: RoaringTreemap) {
370        for (key, other_rb) in rhs.map {
371            match self.map.entry(key) {
372                Entry::Vacant(entry) => {
373                    entry.insert(other_rb);
374                }
375                Entry::Occupied(mut entry) => {
376                    BitXorAssign::bitxor_assign(entry.get_mut(), other_rb);
377                    if entry.get().is_empty() {
378                        entry.remove_entry();
379                    }
380                }
381            }
382        }
383    }
384}
385
386impl BitXorAssign<&RoaringTreemap> for RoaringTreemap {
387    /// A `symmetric difference` between two sets.
388    fn bitxor_assign(&mut self, rhs: &RoaringTreemap) {
389        for (key, other_rb) in &rhs.map {
390            match self.map.entry(*key) {
391                Entry::Vacant(entry) => {
392                    entry.insert(other_rb.clone());
393                }
394                Entry::Occupied(mut entry) => {
395                    BitXorAssign::bitxor_assign(entry.get_mut(), other_rb);
396                    if entry.get().is_empty() {
397                        entry.remove_entry();
398                    }
399                }
400            }
401        }
402    }
403}
404
405#[cfg(test)]
406mod test {
407    use crate::{MultiOps, RoaringTreemap};
408    use proptest::prelude::*;
409
410    // fast count tests
411    proptest! {
412        #[test]
413        fn union_len_eq_len_of_materialized_union(
414            a in RoaringTreemap::arbitrary(),
415            b in RoaringTreemap::arbitrary()
416        ) {
417            prop_assert_eq!(a.union_len(&b), (a | b).len());
418        }
419
420        #[test]
421        fn intersection_len_eq_len_of_materialized_intersection(
422            a in RoaringTreemap::arbitrary(),
423            b in RoaringTreemap::arbitrary()
424        ) {
425            prop_assert_eq!(a.intersection_len(&b), (a & b).len());
426        }
427
428        #[test]
429        fn difference_len_eq_len_of_materialized_difference(
430            a in RoaringTreemap::arbitrary(),
431            b in RoaringTreemap::arbitrary()
432        ) {
433            prop_assert_eq!(a.difference_len(&b), (a - b).len());
434        }
435
436        #[test]
437        fn symmetric_difference_len_eq_len_of_materialized_symmetric_difference(
438            a in RoaringTreemap::arbitrary(),
439            b in RoaringTreemap::arbitrary()
440        ) {
441            prop_assert_eq!(a.symmetric_difference_len(&b), (a ^ b).len());
442        }
443
444        #[test]
445        fn all_union_give_the_same_result(
446            a in RoaringTreemap::arbitrary(),
447            b in RoaringTreemap::arbitrary(),
448            c in RoaringTreemap::arbitrary()
449        ) {
450            let mut ref_assign = a.clone();
451            ref_assign |= &b;
452            ref_assign |= &c;
453
454            let mut own_assign = a.clone();
455            own_assign |= b.clone();
456            own_assign |= c.clone();
457
458            let ref_inline = &a | &b | &c;
459            let own_inline = a.clone() | b.clone() | c.clone();
460
461            let ref_multiop = [&a, &b, &c].union();
462            let own_multiop = [a, b.clone(), c.clone()].union();
463
464            for roar in &[own_assign, ref_inline, own_inline, ref_multiop, own_multiop] {
465                prop_assert_eq!(&ref_assign, roar);
466            }
467        }
468
469        #[test]
470        fn all_intersection_give_the_same_result(
471            a in RoaringTreemap::arbitrary(),
472            b in RoaringTreemap::arbitrary(),
473            c in RoaringTreemap::arbitrary()
474        ) {
475            let mut ref_assign = a.clone();
476            ref_assign &= &b;
477            ref_assign &= &c;
478
479            let mut own_assign = a.clone();
480            own_assign &= b.clone();
481            own_assign &= c.clone();
482
483            let ref_inline = &a & &b & &c;
484            let own_inline = a.clone() & b.clone() & c.clone();
485
486            let ref_multiop = [&a, &b, &c].intersection();
487            let own_multiop = [a, b.clone(), c.clone()].intersection();
488
489            for roar in &[own_assign, ref_inline, own_inline, ref_multiop, own_multiop] {
490                prop_assert_eq!(&ref_assign, roar);
491            }
492        }
493
494        #[test]
495        fn all_difference_give_the_same_result(
496            a in RoaringTreemap::arbitrary(),
497            b in RoaringTreemap::arbitrary(),
498            c in RoaringTreemap::arbitrary()
499        ) {
500            let mut ref_assign = a.clone();
501            ref_assign -= &b;
502            ref_assign -= &c;
503
504            let mut own_assign = a.clone();
505            own_assign -= b.clone();
506            own_assign -= c.clone();
507
508            let ref_inline = &a - &b - &c;
509            let own_inline = a.clone() - b.clone() - c.clone();
510
511            let ref_multiop = [&a, &b, &c].difference();
512            let own_multiop = [a, b.clone(), c.clone()].difference();
513
514            for roar in &[own_assign, ref_inline, own_inline, ref_multiop, own_multiop] {
515                prop_assert_eq!(&ref_assign, roar);
516            }
517        }
518
519        #[test]
520        fn all_symmetric_difference_give_the_same_result(
521            a in RoaringTreemap::arbitrary(),
522            b in RoaringTreemap::arbitrary(),
523            c in RoaringTreemap::arbitrary()
524        ) {
525            let mut ref_assign = a.clone();
526            ref_assign ^= &b;
527            ref_assign ^= &c;
528
529            let mut own_assign = a.clone();
530            own_assign ^= b.clone();
531            own_assign ^= c.clone();
532
533            let ref_inline = &a ^ &b ^ &c;
534            let own_inline = a.clone() ^ b.clone() ^ c.clone();
535
536            let ref_multiop = [&a, &b, &c].symmetric_difference();
537            let own_multiop = [a, b.clone(), c.clone()].symmetric_difference();
538
539            for roar in &[own_assign, ref_inline, own_inline, ref_multiop, own_multiop] {
540                prop_assert_eq!(&ref_assign, roar);
541            }
542        }
543    }
544}
roaring/treemap/ops.rs

roaring/treemap/
ops.rs
