roaring/bitmap/
ops.rs

1use core::mem;
2use core::ops::{BitAnd, BitAndAssign, BitOr, BitOrAssign, BitXor, BitXorAssign, Sub, SubAssign};
3
4use crate::bitmap::container::Container;
5use crate::bitmap::Pairs;
6use crate::RoaringBitmap;
7
8#[cfg(not(feature = "std"))]
9use alloc::vec::Vec;
10
11impl RoaringBitmap {
12    /// Computes the len of the intersection with the specified other bitmap without creating a
13    /// new bitmap.
14    ///
15    /// This is faster and more space efficient when you're only interested in the cardinality of
16    /// the intersection.
17    ///
18    /// # Examples
19    ///
20    /// ```rust
21    /// use roaring::RoaringBitmap;
22    ///
23    /// let rb1: RoaringBitmap = (1..4).collect();
24    /// let rb2: RoaringBitmap = (3..5).collect();
25    ///
26    ///
27    /// assert_eq!(rb1.intersection_len(&rb2), (rb1 & rb2).len());
28    /// ```
29    pub fn intersection_len(&self, other: &RoaringBitmap) -> u64 {
30        Pairs::new(&self.containers, &other.containers)
31            .map(|pair| match pair {
32                (Some(..), None) => 0,
33                (None, Some(..)) => 0,
34                (Some(lhs), Some(rhs)) => lhs.intersection_len(rhs),
35                (None, None) => 0,
36            })
37            .sum()
38    }
39
40    /// Computes the len of the union with the specified other bitmap without creating a new bitmap.
41    ///
42    /// This is faster and more space efficient when you're only interested in the cardinality of
43    /// the union.
44    ///
45    /// # Examples
46    ///
47    /// ```rust
48    /// use roaring::RoaringBitmap;
49    ///
50    /// let rb1: RoaringBitmap = (1..4).collect();
51    /// let rb2: RoaringBitmap = (3..5).collect();
52    ///
53    ///
54    /// assert_eq!(rb1.union_len(&rb2), (rb1 | rb2).len());
55    /// ```
56    pub fn union_len(&self, other: &RoaringBitmap) -> u64 {
57        self.len().wrapping_add(other.len()).wrapping_sub(self.intersection_len(other))
58    }
59
60    /// Computes the len of the difference with the specified other bitmap without creating a new
61    /// bitmap.
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::RoaringBitmap;
70    ///
71    /// let rb1: RoaringBitmap = (1..4).collect();
72    /// let rb2: RoaringBitmap = (3..5).collect();
73    ///
74    ///
75    /// assert_eq!(rb1.difference_len(&rb2), (rb1 - rb2).len());
76    /// ```
77    pub fn difference_len(&self, other: &RoaringBitmap) -> u64 {
78        self.len() - self.intersection_len(other)
79    }
80
81    /// Computes the len of the symmetric difference with the specified other bitmap 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::RoaringBitmap;
91    ///
92    /// let rb1: RoaringBitmap = (1..4).collect();
93    /// let rb2: RoaringBitmap = (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: &RoaringBitmap) -> u64 {
99        let intersection_len = self.intersection_len(other);
100        self.len()
101            .wrapping_add(other.len())
102            .wrapping_sub(intersection_len)
103            .wrapping_sub(intersection_len)
104    }
105}
106
107impl BitOr<RoaringBitmap> for RoaringBitmap {
108    type Output = RoaringBitmap;
109
110    /// An `union` between two sets.
111    fn bitor(mut self, rhs: RoaringBitmap) -> RoaringBitmap {
112        BitOrAssign::bitor_assign(&mut self, rhs);
113        self
114    }
115}
116
117impl BitOr<&RoaringBitmap> for RoaringBitmap {
118    type Output = RoaringBitmap;
119
120    /// An `union` between two sets.
121    fn bitor(mut self, rhs: &RoaringBitmap) -> RoaringBitmap {
122        BitOrAssign::bitor_assign(&mut self, rhs);
123        self
124    }
125}
126
127impl BitOr<RoaringBitmap> for &RoaringBitmap {
128    type Output = RoaringBitmap;
129
130    /// An `union` between two sets.
131    fn bitor(self, rhs: RoaringBitmap) -> RoaringBitmap {
132        BitOr::bitor(rhs, self)
133    }
134}
135
136impl BitOr<&RoaringBitmap> for &RoaringBitmap {
137    type Output = RoaringBitmap;
138
139    /// An `union` between two sets.
140    fn bitor(self, rhs: &RoaringBitmap) -> RoaringBitmap {
141        let mut containers = Vec::new();
142
143        for pair in Pairs::new(&self.containers, &rhs.containers) {
144            match pair {
145                (Some(lhs), None) => containers.push(lhs.clone()),
146                (None, Some(rhs)) => containers.push(rhs.clone()),
147                (Some(lhs), Some(rhs)) => containers.push(BitOr::bitor(lhs, rhs)),
148                (None, None) => break,
149            }
150        }
151
152        RoaringBitmap { containers }
153    }
154}
155
156impl BitOrAssign<RoaringBitmap> for RoaringBitmap {
157    /// An `union` between two sets.
158    fn bitor_assign(&mut self, mut rhs: RoaringBitmap) {
159        // We make sure that we apply the union operation on the biggest map.
160        if self.len() < rhs.len() {
161            mem::swap(self, &mut rhs);
162        }
163
164        for container in rhs.containers {
165            let key = container.key;
166            match self.containers.binary_search_by_key(&key, |c| c.key) {
167                Err(loc) => self.containers.insert(loc, container),
168                Ok(loc) => BitOrAssign::bitor_assign(&mut self.containers[loc], container),
169            }
170        }
171    }
172}
173
174impl BitOrAssign<&RoaringBitmap> for RoaringBitmap {
175    /// An `union` between two sets.
176    fn bitor_assign(&mut self, rhs: &RoaringBitmap) {
177        for container in &rhs.containers {
178            let key = container.key;
179            match self.containers.binary_search_by_key(&key, |c| c.key) {
180                Err(loc) => self.containers.insert(loc, container.clone()),
181                Ok(loc) => BitOrAssign::bitor_assign(&mut self.containers[loc], container),
182            }
183        }
184    }
185}
186
187impl BitAnd<RoaringBitmap> for RoaringBitmap {
188    type Output = RoaringBitmap;
189
190    /// An `intersection` between two sets.
191    fn bitand(mut self, rhs: RoaringBitmap) -> RoaringBitmap {
192        BitAndAssign::bitand_assign(&mut self, rhs);
193        self
194    }
195}
196
197impl BitAnd<&RoaringBitmap> for RoaringBitmap {
198    type Output = RoaringBitmap;
199
200    /// An `intersection` between two sets.
201    fn bitand(mut self, rhs: &RoaringBitmap) -> RoaringBitmap {
202        BitAndAssign::bitand_assign(&mut self, rhs);
203        self
204    }
205}
206
207impl BitAnd<RoaringBitmap> for &RoaringBitmap {
208    type Output = RoaringBitmap;
209
210    /// An `intersection` between two sets.
211    fn bitand(self, rhs: RoaringBitmap) -> RoaringBitmap {
212        BitAnd::bitand(rhs, self)
213    }
214}
215
216impl BitAnd<&RoaringBitmap> for &RoaringBitmap {
217    type Output = RoaringBitmap;
218
219    /// An `intersection` between two sets.
220    fn bitand(self, rhs: &RoaringBitmap) -> RoaringBitmap {
221        let mut containers = Vec::new();
222
223        for pair in Pairs::new(&self.containers, &rhs.containers) {
224            if let (Some(lhs), Some(rhs)) = pair {
225                let container = BitAnd::bitand(lhs, rhs);
226                if !container.is_empty() {
227                    containers.push(container);
228                }
229            }
230        }
231
232        RoaringBitmap { containers }
233    }
234}
235
236impl BitAndAssign<RoaringBitmap> for RoaringBitmap {
237    /// An `intersection` between two sets.
238    fn bitand_assign(&mut self, mut rhs: RoaringBitmap) {
239        // We make sure that we apply the intersection operation on the smallest map.
240        if rhs.containers.len() < self.containers.len() {
241            mem::swap(self, &mut rhs);
242        }
243
244        self.containers.retain_mut(|cont| {
245            let key = cont.key;
246            match rhs.containers.binary_search_by_key(&key, |c| c.key) {
247                Ok(loc) => {
248                    let rhs_cont = &mut rhs.containers[loc];
249                    let rhs_cont = mem::replace(rhs_cont, Container::new(rhs_cont.key));
250                    BitAndAssign::bitand_assign(cont, rhs_cont);
251                    !cont.is_empty()
252                }
253                Err(_) => false,
254            }
255        })
256    }
257}
258
259impl BitAndAssign<&RoaringBitmap> for RoaringBitmap {
260    /// An `intersection` between two sets.
261    fn bitand_assign(&mut self, rhs: &RoaringBitmap) {
262        self.containers.retain_mut(|cont| {
263            let key = cont.key;
264            match rhs.containers.binary_search_by_key(&key, |c| c.key) {
265                Ok(loc) => {
266                    BitAndAssign::bitand_assign(cont, &rhs.containers[loc]);
267                    !cont.is_empty()
268                }
269                Err(_) => false,
270            }
271        })
272    }
273}
274
275impl Sub<RoaringBitmap> for RoaringBitmap {
276    type Output = RoaringBitmap;
277
278    /// A `difference` between two sets.
279    fn sub(mut self, rhs: RoaringBitmap) -> RoaringBitmap {
280        SubAssign::sub_assign(&mut self, &rhs);
281        self
282    }
283}
284
285impl Sub<&RoaringBitmap> for RoaringBitmap {
286    type Output = RoaringBitmap;
287
288    /// A `difference` between two sets.
289    fn sub(mut self, rhs: &RoaringBitmap) -> RoaringBitmap {
290        SubAssign::sub_assign(&mut self, rhs);
291        self
292    }
293}
294
295impl Sub<RoaringBitmap> for &RoaringBitmap {
296    type Output = RoaringBitmap;
297
298    /// A `difference` between two sets.
299    fn sub(self, rhs: RoaringBitmap) -> RoaringBitmap {
300        Sub::sub(self, &rhs)
301    }
302}
303
304impl Sub<&RoaringBitmap> for &RoaringBitmap {
305    type Output = RoaringBitmap;
306
307    /// A `difference` between two sets.
308    fn sub(self, rhs: &RoaringBitmap) -> RoaringBitmap {
309        let mut containers = Vec::new();
310
311        for pair in Pairs::new(&self.containers, &rhs.containers) {
312            match pair {
313                (Some(lhs), None) => containers.push(lhs.clone()),
314                (None, Some(_)) => (),
315                (Some(lhs), Some(rhs)) => {
316                    let container = Sub::sub(lhs, rhs);
317                    if !container.is_empty() {
318                        containers.push(container);
319                    }
320                }
321                (None, None) => break,
322            }
323        }
324
325        RoaringBitmap { containers }
326    }
327}
328
329impl SubAssign<RoaringBitmap> for RoaringBitmap {
330    /// A `difference` between two sets.
331    fn sub_assign(&mut self, rhs: RoaringBitmap) {
332        SubAssign::sub_assign(self, &rhs)
333    }
334}
335
336impl SubAssign<&RoaringBitmap> for RoaringBitmap {
337    /// A `difference` between two sets.
338    fn sub_assign(&mut self, rhs: &RoaringBitmap) {
339        self.containers.retain_mut(|cont| {
340            match rhs.containers.binary_search_by_key(&cont.key, |c| c.key) {
341                Ok(loc) => {
342                    SubAssign::sub_assign(cont, &rhs.containers[loc]);
343                    !cont.is_empty()
344                }
345                Err(_) => true,
346            }
347        })
348    }
349}
350
351impl BitXor<RoaringBitmap> for RoaringBitmap {
352    type Output = RoaringBitmap;
353
354    /// A `symmetric difference` between two sets.
355    fn bitxor(mut self, rhs: RoaringBitmap) -> RoaringBitmap {
356        BitXorAssign::bitxor_assign(&mut self, rhs);
357        self
358    }
359}
360
361impl BitXor<&RoaringBitmap> for RoaringBitmap {
362    type Output = RoaringBitmap;
363
364    /// A `symmetric difference` between two sets.
365    fn bitxor(mut self, rhs: &RoaringBitmap) -> RoaringBitmap {
366        BitXorAssign::bitxor_assign(&mut self, rhs);
367        self
368    }
369}
370
371impl BitXor<RoaringBitmap> for &RoaringBitmap {
372    type Output = RoaringBitmap;
373
374    /// A `symmetric difference` between two sets.
375    fn bitxor(self, rhs: RoaringBitmap) -> RoaringBitmap {
376        BitXor::bitxor(rhs, self)
377    }
378}
379
380impl BitXor<&RoaringBitmap> for &RoaringBitmap {
381    type Output = RoaringBitmap;
382
383    /// A `symmetric difference` between two sets.
384    fn bitxor(self, rhs: &RoaringBitmap) -> RoaringBitmap {
385        let mut containers = Vec::new();
386
387        for pair in Pairs::new(&self.containers, &rhs.containers) {
388            match pair {
389                (Some(lhs), None) => containers.push(lhs.clone()),
390                (None, Some(rhs)) => containers.push(rhs.clone()),
391                (Some(lhs), Some(rhs)) => {
392                    let container = BitXor::bitxor(lhs, rhs);
393                    if !container.is_empty() {
394                        containers.push(container);
395                    }
396                }
397                (None, None) => break,
398            }
399        }
400
401        RoaringBitmap { containers }
402    }
403}
404
405impl BitXorAssign<RoaringBitmap> for RoaringBitmap {
406    /// A `symmetric difference` between two sets.
407    fn bitxor_assign(&mut self, rhs: RoaringBitmap) {
408        for pair in Pairs::new(mem::take(&mut self.containers), rhs.containers) {
409            match pair {
410                (Some(mut lhs), Some(rhs)) => {
411                    BitXorAssign::bitxor_assign(&mut lhs, rhs);
412                    if !lhs.is_empty() {
413                        self.containers.push(lhs);
414                    }
415                }
416                (Some(lhs), None) => self.containers.push(lhs),
417                (None, Some(rhs)) => self.containers.push(rhs),
418                (None, None) => break,
419            }
420        }
421    }
422}
423
424impl BitXorAssign<&RoaringBitmap> for RoaringBitmap {
425    /// A `symmetric difference` between two sets.
426    fn bitxor_assign(&mut self, rhs: &RoaringBitmap) {
427        for pair in Pairs::new(mem::take(&mut self.containers), &rhs.containers) {
428            match pair {
429                (Some(mut lhs), Some(rhs)) => {
430                    BitXorAssign::bitxor_assign(&mut lhs, rhs);
431                    if !lhs.is_empty() {
432                        self.containers.push(lhs);
433                    }
434                }
435                (Some(lhs), None) => self.containers.push(lhs),
436                (None, Some(rhs)) => self.containers.push(rhs.clone()),
437                (None, None) => break,
438            }
439        }
440    }
441}
442
443#[cfg(test)]
444mod test {
445    use crate::{MultiOps, RoaringBitmap};
446    use core::convert::Infallible;
447    use proptest::prelude::*;
448
449    // fast count tests
450    proptest! {
451        #[test]
452        fn union_len_eq_len_of_materialized_union(
453            a in RoaringBitmap::arbitrary(),
454            b in RoaringBitmap::arbitrary()
455        ) {
456            prop_assert_eq!(a.union_len(&b), (a | b).len());
457        }
458
459        #[test]
460        fn intersection_len_eq_len_of_materialized_intersection(
461            a in RoaringBitmap::arbitrary(),
462            b in RoaringBitmap::arbitrary()
463        ) {
464            prop_assert_eq!(a.intersection_len(&b), (a & b).len());
465        }
466
467        #[test]
468        fn difference_len_eq_len_of_materialized_difference(
469            a in RoaringBitmap::arbitrary(),
470            b in RoaringBitmap::arbitrary()
471        ) {
472            prop_assert_eq!(a.difference_len(&b), (a - b).len());
473        }
474
475        #[test]
476        fn symmetric_difference_len_eq_len_of_materialized_symmetric_difference(
477            a in RoaringBitmap::arbitrary(),
478            b in RoaringBitmap::arbitrary()
479        ) {
480            prop_assert_eq!(a.symmetric_difference_len(&b), (a ^ b).len());
481        }
482
483        #[test]
484        fn all_union_give_the_same_result(
485            a in RoaringBitmap::arbitrary(),
486            b in RoaringBitmap::arbitrary(),
487            c in RoaringBitmap::arbitrary()
488        ) {
489            let mut ref_assign = a.clone();
490            ref_assign |= &b;
491            ref_assign |= &c;
492
493            let mut own_assign = a.clone();
494            own_assign |= b.clone();
495            own_assign |= c.clone();
496
497            let ref_inline = &a | &b | &c;
498            let own_inline = a.clone() | b.clone() | c.clone();
499
500            let ref_multiop = [&a, &b, &c].union();
501            let own_multiop = [a.clone(), b.clone(), c.clone()].union();
502
503            let ref_multiop_try = [&a, &b, &c].map(Ok::<_, Infallible>).union().unwrap();
504            let own_multiop_try = [a, b, c].map(Ok::<_, Infallible>).union().unwrap();
505
506            for roar in &[
507                own_assign,
508                ref_inline,
509                own_inline,
510                ref_multiop,
511                own_multiop,
512                ref_multiop_try,
513                own_multiop_try,
514            ] {
515                prop_assert_eq!(&ref_assign, roar);
516            }
517        }
518
519        #[test]
520        fn all_intersection_give_the_same_result(
521            a in RoaringBitmap::arbitrary(),
522            b in RoaringBitmap::arbitrary(),
523            c in RoaringBitmap::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].intersection();
537            let own_multiop = [a.clone(), b.clone(), c.clone()].intersection();
538
539            let ref_multiop_try = [&a, &b, &c].map(Ok::<_, Infallible>).intersection().unwrap();
540            let own_multiop_try = [a, b, c].map(Ok::<_, Infallible>).intersection().unwrap();
541
542            for roar in &[
543                own_assign,
544                ref_inline,
545                own_inline,
546                ref_multiop,
547                own_multiop,
548                ref_multiop_try,
549                own_multiop_try,
550            ] {
551                prop_assert_eq!(&ref_assign, roar);
552            }
553        }
554
555        #[test]
556        fn all_difference_give_the_same_result(
557            a in RoaringBitmap::arbitrary(),
558            b in RoaringBitmap::arbitrary(),
559            c in RoaringBitmap::arbitrary()
560        ) {
561            let mut ref_assign = a.clone();
562            ref_assign -= &b;
563            ref_assign -= &c;
564
565            let mut own_assign = a.clone();
566            own_assign -= b.clone();
567            own_assign -= c.clone();
568
569            let ref_inline = &a - &b - &c;
570            let own_inline = a.clone() - b.clone() - c.clone();
571
572            let ref_multiop = [&a, &b, &c].difference();
573            let own_multiop = [a.clone(), b.clone(), c.clone()].difference();
574
575            let ref_multiop_try = [&a, &b, &c].map(Ok::<_, Infallible>).difference().unwrap();
576            let own_multiop_try = [a, b, c].map(Ok::<_, Infallible>).difference().unwrap();
577
578            for roar in &[
579                own_assign,
580                ref_inline,
581                own_inline,
582                ref_multiop,
583                own_multiop,
584                ref_multiop_try,
585                own_multiop_try,
586            ] {
587                prop_assert_eq!(&ref_assign, roar);
588            }
589        }
590
591        #[test]
592        fn all_symmetric_difference_give_the_same_result(
593            a in RoaringBitmap::arbitrary(),
594            b in RoaringBitmap::arbitrary(),
595            c in RoaringBitmap::arbitrary()
596        ) {
597            let mut ref_assign = a.clone();
598            ref_assign ^= &b;
599            ref_assign ^= &c;
600
601            let mut own_assign = a.clone();
602            own_assign ^= b.clone();
603            own_assign ^= c.clone();
604
605            let ref_inline = &a ^ &b ^ &c;
606            let own_inline = a.clone() ^ b.clone() ^ c.clone();
607
608            let ref_multiop = [&a, &b, &c].symmetric_difference();
609            let own_multiop = [a.clone(), b.clone(), c.clone()].symmetric_difference();
610
611            let ref_multiop_try = [&a, &b, &c]
612                .map(Ok::<_, Infallible>)
613                .symmetric_difference()
614                .unwrap();
615            let own_multiop_try = [a, b, c]
616                .map(Ok::<_, Infallible>)
617                .symmetric_difference()
618                .unwrap();
619
620            for roar in &[
621                own_assign,
622                ref_inline,
623                own_inline,
624                ref_multiop,
625                own_multiop,
626                ref_multiop_try,
627                own_multiop_try,
628            ] {
629                prop_assert_eq!(&ref_assign, roar);
630            }
631        }
632    }
633}