1use std::iter::{FromIterator, IntoIterator};
2use std::ops::{BitAnd, BitAndAssign, BitOr, BitOrAssign, BitXor, BitXorAssign, Not};
3use std::usize;
4
5use util::*;
6
7use {AtomicBitSet, BitIter, BitSet, BitSetLike, DrainableBitSet};
8
9impl<'a, B> BitOrAssign<&'a B> for BitSet
10where
11 B: BitSetLike,
12{
13 fn bitor_assign(&mut self, lhs: &B) {
14 use iter::State::Continue;
15 let mut iter = lhs.iter();
16 while let Some(level) = (1..LAYERS).find(|&level| iter.handle_level(level) == Continue) {
17 let lower = level - 1;
18 let idx = iter.prefix[lower] as usize >> BITS;
19 *self.layer_mut(lower, idx) |= lhs.get_from_layer(lower, idx);
20 }
21 self.layer3 |= lhs.layer3();
22 }
23}
24
25impl<'a, B> BitAndAssign<&'a B> for BitSet
26where
27 B: BitSetLike,
28{
29 fn bitand_assign(&mut self, lhs: &B) {
30 use iter::State::*;
31 let mut iter = lhs.iter();
32 iter.masks[LAYERS - 1] &= self.layer3();
33 while let Some(level) = (1..LAYERS).find(|&level| iter.handle_level(level) == Continue) {
34 let lower = level - 1;
35 let idx = iter.prefix[lower] as usize >> BITS;
36 let our_layer = self.get_from_layer(lower, idx);
37 let their_layer = lhs.get_from_layer(lower, idx);
38
39 iter.masks[lower] &= our_layer;
40
41 let mut masks = [0; LAYERS];
42 masks[lower] = our_layer & !their_layer;
43 BitIter::new(&mut *self, masks, iter.prefix).clear();
44
45 *self.layer_mut(lower, idx) &= their_layer;
46 }
47 let mut masks = [0; LAYERS];
48 masks[LAYERS - 1] = self.layer3() & !lhs.layer3();
49 BitIter::new(&mut *self, masks, [0; LAYERS - 1]).clear();
50
51 self.layer3 &= lhs.layer3();
52 }
53}
54
55impl<'a, B> BitXorAssign<&'a B> for BitSet
56where
57 B: BitSetLike,
58{
59 fn bitxor_assign(&mut self, lhs: &B) {
60 use iter::State::*;
61 let mut iter = lhs.iter();
62 while let Some(level) = (1..LAYERS).find(|&level| iter.handle_level(level) == Continue) {
63 let lower = level - 1;
64 let idx = iter.prefix[lower] as usize >> BITS;
65
66 if lower == 0 {
67 *self.layer_mut(lower, idx) ^= lhs.get_from_layer(lower, idx);
68
69 let mut change_bit = |level| {
70 let lower = level - 1;
71 let h = iter.prefix.get(level).cloned().unwrap_or(0) as usize;
72 let l = iter.prefix[lower] as usize >> BITS;
73 let mask = 1 << (l & !h);
74
75 if self.get_from_layer(lower, l) == 0 {
76 *self.layer_mut(level, h >> BITS) &= !mask;
77 } else {
78 *self.layer_mut(level, h >> BITS) |= mask;
79 }
80 };
81
82 change_bit(level);
83 if iter.masks[level] == 0 {
84 (2..LAYERS).for_each(change_bit);
85 }
86 }
87 }
88 }
89}
90
91#[derive(Debug, Clone)]
97pub struct BitSetAnd<A: BitSetLike, B: BitSetLike>(pub A, pub B);
98
99impl<A: BitSetLike, B: BitSetLike> BitSetLike for BitSetAnd<A, B> {
100 #[inline]
101 fn layer3(&self) -> usize {
102 self.0.layer3() & self.1.layer3()
103 }
104 #[inline]
105 fn layer2(&self, i: usize) -> usize {
106 self.0.layer2(i) & self.1.layer2(i)
107 }
108 #[inline]
109 fn layer1(&self, i: usize) -> usize {
110 self.0.layer1(i) & self.1.layer1(i)
111 }
112 #[inline]
113 fn layer0(&self, i: usize) -> usize {
114 self.0.layer0(i) & self.1.layer0(i)
115 }
116 #[inline]
117 fn contains(&self, i: Index) -> bool {
118 self.0.contains(i) && self.1.contains(i)
119 }
120}
121
122impl<A: DrainableBitSet, B: DrainableBitSet> DrainableBitSet for BitSetAnd<A, B> {
123 #[inline]
124 fn remove(&mut self, i: Index) -> bool {
125 if self.contains(i) {
126 self.0.remove(i);
127 self.1.remove(i);
128 true
129 } else {
130 false
131 }
132 }
133}
134
135#[derive(Debug, Clone)]
141pub struct BitSetOr<A: BitSetLike, B: BitSetLike>(pub A, pub B);
142
143impl<A: BitSetLike, B: BitSetLike> BitSetLike for BitSetOr<A, B> {
144 #[inline]
145 fn layer3(&self) -> usize {
146 self.0.layer3() | self.1.layer3()
147 }
148 #[inline]
149 fn layer2(&self, i: usize) -> usize {
150 self.0.layer2(i) | self.1.layer2(i)
151 }
152 #[inline]
153 fn layer1(&self, i: usize) -> usize {
154 self.0.layer1(i) | self.1.layer1(i)
155 }
156 #[inline]
157 fn layer0(&self, i: usize) -> usize {
158 self.0.layer0(i) | self.1.layer0(i)
159 }
160 #[inline]
161 fn contains(&self, i: Index) -> bool {
162 self.0.contains(i) || self.1.contains(i)
163 }
164}
165
166impl<A: DrainableBitSet, B: DrainableBitSet> DrainableBitSet for BitSetOr<A, B> {
167 #[inline]
168 fn remove(&mut self, i: Index) -> bool {
169 if self.contains(i) {
170 self.0.remove(i);
171 self.1.remove(i);
172 true
173 } else {
174 false
175 }
176 }
177}
178
179#[derive(Debug, Clone)]
184pub struct BitSetNot<A: BitSetLike>(pub A);
185
186impl<A: BitSetLike> BitSetLike for BitSetNot<A> {
187 #[inline]
188 fn layer3(&self) -> usize {
189 !0
190 }
191 #[inline]
192 fn layer2(&self, _: usize) -> usize {
193 !0
194 }
195 #[inline]
196 fn layer1(&self, _: usize) -> usize {
197 !0
198 }
199 #[inline]
200 fn layer0(&self, i: usize) -> usize {
201 !self.0.layer0(i)
202 }
203 #[inline]
204 fn contains(&self, i: Index) -> bool {
205 !self.0.contains(i)
206 }
207}
208
209#[derive(Debug, Clone)]
215pub struct BitSetXor<A: BitSetLike, B: BitSetLike>(pub A, pub B);
216
217impl<A: BitSetLike, B: BitSetLike> BitSetLike for BitSetXor<A, B> {
218 #[inline]
219 fn layer3(&self) -> usize {
220 let xor = BitSetAnd(
221 BitSetOr(&self.0, &self.1),
222 BitSetNot(BitSetAnd(&self.0, &self.1)),
223 );
224 xor.layer3()
225 }
226 #[inline]
227 fn layer2(&self, id: usize) -> usize {
228 let xor = BitSetAnd(
229 BitSetOr(&self.0, &self.1),
230 BitSetNot(BitSetAnd(&self.0, &self.1)),
231 );
232 xor.layer2(id)
233 }
234 #[inline]
235 fn layer1(&self, id: usize) -> usize {
236 let xor = BitSetAnd(
237 BitSetOr(&self.0, &self.1),
238 BitSetNot(BitSetAnd(&self.0, &self.1)),
239 );
240 xor.layer1(id)
241 }
242 #[inline]
243 fn layer0(&self, id: usize) -> usize {
244 let xor = BitSetAnd(
245 BitSetOr(&self.0, &self.1),
246 BitSetNot(BitSetAnd(&self.0, &self.1)),
247 );
248 xor.layer0(id)
249 }
250 #[inline]
251 fn contains(&self, i: Index) -> bool {
252 BitSetAnd(
253 BitSetOr(&self.0, &self.1),
254 BitSetNot(BitSetAnd(&self.0, &self.1)),
255 )
256 .contains(i)
257 }
258}
259
260#[derive(Debug, Clone)]
263pub struct BitSetAll;
264impl BitSetLike for BitSetAll {
265 #[inline]
266 fn layer3(&self) -> usize {
267 usize::MAX
268 }
269 #[inline]
270 fn layer2(&self, _id: usize) -> usize {
271 usize::MAX
272 }
273 #[inline]
274 fn layer1(&self, _id: usize) -> usize {
275 usize::MAX
276 }
277 #[inline]
278 fn layer0(&self, _id: usize) -> usize {
279 usize::MAX
280 }
281 #[inline]
282 fn contains(&self, _i: Index) -> bool {
283 true
284 }
285}
286
287macro_rules! operator {
288 ( impl < ( $( $lifetime:tt )* ) ( $( $arg:ident ),* ) > for $bitset:ty ) => {
289 impl<$( $lifetime, )* $( $arg ),*> IntoIterator for $bitset
290 where $( $arg: BitSetLike ),*
291 {
292 type Item = <BitIter<Self> as Iterator>::Item;
293 type IntoIter = BitIter<Self>;
294 fn into_iter(self) -> Self::IntoIter {
295 self.iter()
296 }
297 }
298
299 impl<$( $lifetime, )* $( $arg ),*> Not for $bitset
300 where $( $arg: BitSetLike ),*
301 {
302 type Output = BitSetNot<Self>;
303 fn not(self) -> Self::Output {
304 BitSetNot(self)
305 }
306 }
307
308 impl<$( $lifetime, )* $( $arg, )* T> BitAnd<T> for $bitset
309 where T: BitSetLike,
310 $( $arg: BitSetLike ),*
311 {
312 type Output = BitSetAnd<Self, T>;
313 fn bitand(self, rhs: T) -> Self::Output {
314 BitSetAnd(self, rhs)
315 }
316 }
317
318 impl<$( $lifetime, )* $( $arg, )* T> BitOr<T> for $bitset
319 where T: BitSetLike,
320 $( $arg: BitSetLike ),*
321 {
322 type Output = BitSetOr<Self, T>;
323 fn bitor(self, rhs: T) -> Self::Output {
324 BitSetOr(self, rhs)
325 }
326 }
327
328 impl<$( $lifetime, )* $( $arg, )* T> BitXor<T> for $bitset
329 where T: BitSetLike,
330 $( $arg: BitSetLike ),*
331 {
332 type Output = BitSetXor<Self, T>;
333 fn bitxor(self, rhs: T) -> Self::Output {
334 BitSetXor(self, rhs)
335 }
336 }
337
338 }
339}
340
341operator!(impl<()()> for BitSet);
342operator!(impl<('a)()> for &'a BitSet);
343operator!(impl<()()> for AtomicBitSet);
344operator!(impl<('a)()> for &'a AtomicBitSet);
345operator!(impl<()(A)> for BitSetNot<A>);
346operator!(impl<('a)(A)> for &'a BitSetNot<A>);
347operator!(impl<()(A, B)> for BitSetAnd<A, B>);
348operator!(impl<('a)(A, B)> for &'a BitSetAnd<A, B>);
349operator!(impl<()(A, B)> for BitSetOr<A, B>);
350operator!(impl<('a)(A, B)> for &'a BitSetOr<A, B>);
351operator!(impl<()(A, B)> for BitSetXor<A, B>);
352operator!(impl<('a)(A, B)> for &'a BitSetXor<A, B>);
353operator!(impl<()()> for BitSetAll);
354operator!(impl<('a)()> for &'a BitSetAll);
355
356macro_rules! iterator {
357 ( $bitset:ident ) => {
358 impl FromIterator<Index> for $bitset {
359 fn from_iter<T>(iter: T) -> Self
360 where
361 T: IntoIterator<Item = Index>,
362 {
363 let mut bitset = $bitset::new();
364 for item in iter {
365 bitset.add(item);
366 }
367 bitset
368 }
369 }
370
371 impl<'a> FromIterator<&'a Index> for $bitset {
372 fn from_iter<T>(iter: T) -> Self
373 where
374 T: IntoIterator<Item = &'a Index>,
375 {
376 let mut bitset = $bitset::new();
377 for item in iter {
378 bitset.add(*item);
379 }
380 bitset
381 }
382 }
383
384 impl Extend<Index> for $bitset {
385 fn extend<T>(&mut self, iter: T)
386 where
387 T: IntoIterator<Item = Index>,
388 {
389 for item in iter {
390 self.add(item);
391 }
392 }
393 }
394
395 impl<'a> Extend<&'a Index> for $bitset {
396 fn extend<T>(&mut self, iter: T)
397 where
398 T: IntoIterator<Item = &'a Index>,
399 {
400 for item in iter {
401 self.add(*item);
402 }
403 }
404 }
405 };
406}
407
408iterator!(BitSet);
409iterator!(AtomicBitSet);
410
411#[cfg(test)]
412mod tests {
413 use {BitSet, BitSetLike, BitSetXor, Index};
414
415 #[test]
416 fn or_assign() {
417 use std::collections::HashSet;
418 use std::mem::size_of;
419
420 let usize_bits = size_of::<usize>() as u32 * 8;
421 let n = 10_000;
422 let f1 = &|n| 7 * usize_bits * n;
423 let f2 = &|n| 13 * usize_bits * n;
424
425 let mut c1: BitSet = (0..n).map(f1).collect();
426 let c2: BitSet = (0..n).map(f2).collect();
427
428 c1 |= &c2;
429
430 let h1: HashSet<_> = (0..n).map(f1).collect();
431 let h2: HashSet<_> = (0..n).map(f2).collect();
432 assert_eq!(c1.iter().collect::<HashSet<_>>(), &h1 | &h2);
433 }
434
435 #[test]
436 fn or_assign_random() {
437 use rand::prelude::*;
438
439 use std::collections::HashSet;
440 let limit = 1_048_576;
441 let mut rng = thread_rng();
442
443 let mut set1 = BitSet::new();
444 let mut check_set1 = HashSet::new();
445 for _ in 0..(limit / 100) {
446 let index = rng.gen_range(0, limit);
447 set1.add(index);
448 check_set1.insert(index);
449 }
450
451 let mut set2 = BitSet::new();
452 let mut check_set2 = HashSet::new();
453 for _ in 0..(limit / 100) {
454 let index = rng.gen_range(0, limit);
455 set2.add(index);
456 check_set2.insert(index);
457 }
458
459 let hs1 = (&set1).iter().collect::<HashSet<_>>();
460 let hs2 = (&set2).iter().collect::<HashSet<_>>();
461 let mut hs = (&hs1 | &hs2).iter().cloned().collect::<HashSet<_>>();
462
463 set1 |= &set2;
464
465 for _ in 0..(limit / 1000) {
466 let index = rng.gen_range(0, limit);
467 set1.add(index);
468 hs.insert(index);
469 }
470
471 assert_eq!(hs, set1.iter().collect());
472 }
473
474 #[test]
475 fn and_assign() {
476 use std::collections::HashSet;
477 use std::mem::size_of;
478
479 let usize_bits = size_of::<usize>() as u32 * 8;
480 let n = 10_000;
481 let f1 = &|n| 7 * usize_bits * n;
482 let f2 = &|n| 13 * usize_bits * n;
483
484 let mut c1: BitSet = (0..n).map(f1).collect();
485 let c2: BitSet = (0..n).map(f2).collect();
486
487 c1 &= &c2;
488
489 let h1: HashSet<_> = (0..n).map(f1).collect();
490 let h2: HashSet<_> = (0..n).map(f2).collect();
491 assert_eq!(c1.iter().collect::<HashSet<_>>(), &h1 & &h2);
492 }
493
494 #[test]
495 fn and_assign_specific() {
496 use util::BITS;
497
498 let mut c1 = BitSet::new();
499 c1.add(0);
500 let common = ((1 << BITS) << BITS) << BITS;
501 c1.add(common);
502 c1.add((((1 << BITS) << BITS) + 1) << BITS);
503
504 let mut c2: BitSet = BitSet::new();
505 c2.add(common);
506 c2.add((((1 << BITS) << BITS) + 2) << BITS);
507
508 c1 &= &c2;
509
510 assert_eq!(c1.iter().collect::<Vec<_>>(), [common]);
511 }
512
513 #[test]
514 fn and_assign_with_modification() {
515 use util::BITS;
516
517 let mut c1 = BitSet::new();
518 c1.add(0);
519 c1.add((1 << BITS) << BITS);
520
521 let mut c2: BitSet = BitSet::new();
522 c2.add(0);
523
524 c1 &= &c2;
525
526 let added = ((1 << BITS) + 1) << BITS;
527 c1.add(added);
528
529 assert_eq!(c1.iter().collect::<Vec<_>>(), [0, added]);
530 }
531
532 #[test]
533 fn and_assign_random() {
534 use rand::prelude::*;
535
536 use std::collections::HashSet;
537 let limit = 1_048_576;
538 let mut rng = thread_rng();
539
540 let mut set1 = BitSet::new();
541 let mut check_set1 = HashSet::new();
542 for _ in 0..(limit / 100) {
543 let index = rng.gen_range(0, limit);
544 set1.add(index);
545 check_set1.insert(index);
546 }
547
548 let mut set2 = BitSet::new();
549 let mut check_set2 = HashSet::new();
550 for _ in 0..(limit / 100) {
551 let index = rng.gen_range(0, limit);
552 set2.add(index);
553 check_set2.insert(index);
554 }
555
556 let hs1 = (&set1).iter().collect::<HashSet<_>>();
557 let hs2 = (&set2).iter().collect::<HashSet<_>>();
558 let mut hs = (&hs1 & &hs2).iter().cloned().collect::<HashSet<_>>();
559
560 set1 &= &set2;
561
562 for _ in 0..(limit / 1000) {
563 let index = rng.gen_range(0, limit);
564 set1.add(index);
565 hs.insert(index);
566 }
567
568 assert_eq!(hs, set1.iter().collect());
569 }
570
571 #[test]
572 fn xor_assign() {
573 use std::collections::HashSet;
574 use std::mem::size_of;
575
576 let usize_bits = size_of::<usize>() as u32 * 8;
577 let n = 10_000;
578 let f1 = &|n| 7 * usize_bits * n;
579 let f2 = &|n| 13 * usize_bits * n;
580
581 let mut c1: BitSet = (0..n).map(f1).collect();
582 let c2: BitSet = (0..n).map(f2).collect();
583 c1 ^= &c2;
584
585 let h1: HashSet<_> = (0..n).map(f1).collect();
586 let h2: HashSet<_> = (0..n).map(f2).collect();
587 assert_eq!(c1.iter().collect::<HashSet<_>>(), &h1 ^ &h2);
588 }
589
590 #[test]
591 fn xor_assign_specific() {
592 use util::BITS;
593
594 let mut c1 = BitSet::new();
595 c1.add(0);
596 let common = ((1 << BITS) << BITS) << BITS;
597 c1.add(common);
598 let a = (((1 << BITS) + 1) << BITS) << BITS;
599 c1.add(a);
600
601 let mut c2: BitSet = BitSet::new();
602 c2.add(common);
603 let b = (((1 << BITS) + 2) << BITS) << BITS;
604 c2.add(b);
605
606 c1 ^= &c2;
607
608 assert_eq!(c1.iter().collect::<Vec<_>>(), [0, a, b]);
609 }
610
611 #[test]
612 fn xor_assign_random() {
613 use rand::prelude::*;
614 use std::collections::HashSet;
615 let limit = 1_048_576;
616 let mut rng = thread_rng();
617
618 let mut set1 = BitSet::new();
619 let mut check_set1 = HashSet::new();
620 for _ in 0..(limit / 100) {
621 let index = rng.gen_range(0, limit);
622 set1.add(index);
623 check_set1.insert(index);
624 }
625
626 let mut set2 = BitSet::new();
627 let mut check_set2 = HashSet::new();
628 for _ in 0..(limit / 100) {
629 let index = rng.gen_range(0, limit);
630 set2.add(index);
631 check_set2.insert(index);
632 }
633
634 let hs1 = (&set1).iter().collect::<HashSet<_>>();
635 let hs2 = (&set2).iter().collect::<HashSet<_>>();
636 let mut hs = (&hs1 ^ &hs2).iter().cloned().collect::<HashSet<_>>();
637
638 set1 ^= &set2;
639
640 for _ in 0..(limit / 1000) {
641 let index = rng.gen_range(0, limit);
642 set1.add(index);
643 hs.insert(index);
644 }
645
646 assert_eq!(hs, set1.iter().collect());
647 }
648
649 #[test]
650 fn operators() {
651 let mut bitset = BitSet::new();
652 bitset.add(1);
653 bitset.add(3);
654 bitset.add(5);
655 bitset.add(15);
656 bitset.add(200);
657 bitset.add(50001);
658
659 let mut other = BitSet::new();
660 other.add(1);
661 other.add(3);
662 other.add(50000);
663 other.add(50001);
664
665 {
666 let not = &bitset & !&bitset;
667 assert_eq!(not.iter().count(), 0);
668 }
669
670 {
671 let either = &bitset | &other;
672 let collected = either.iter().collect::<Vec<Index>>();
673 assert_eq!(collected, vec![1, 3, 5, 15, 200, 50000, 50001]);
674
675 let either_sanity = bitset.clone() | other.clone();
676 assert_eq!(collected, either_sanity.iter().collect::<Vec<Index>>());
677 }
678
679 {
680 let same = &bitset & &other;
681 let collected = same.iter().collect::<Vec<Index>>();
682 assert_eq!(collected, vec![1, 3, 50001]);
683
684 let same_sanity = bitset.clone() & other.clone();
685 assert_eq!(collected, same_sanity.iter().collect::<Vec<Index>>());
686 }
687
688 {
689 let exclusive = &bitset ^ &other;
690 let collected = exclusive.iter().collect::<Vec<Index>>();
691 assert_eq!(collected, vec![5, 15, 200, 50000]);
692
693 let exclusive_sanity = bitset.clone() ^ other.clone();
694 assert_eq!(collected, exclusive_sanity.iter().collect::<Vec<Index>>());
695 }
696 }
697
698 #[test]
699 fn xor() {
700 let mut bitset = BitSet::new();
702 bitset.add(2);
703 bitset.add(3);
704 bitset.add(50000);
705
706 let mut other = BitSet::new();
708 other.add(1);
709 other.add(3);
710 other.add(50000);
711 other.add(50001);
712
713 {
714 let xor = BitSetXor(&bitset, &other);
716 let collected = xor.iter().collect::<Vec<Index>>();
717 assert_eq!(collected, vec![1, 2, 50001]);
718 }
719 }
720}