]> git.lizzy.rs Git - rust.git/blob - src/librustc_data_structures/bit_set.rs
Merge Variant and Variant_
[rust.git] / src / librustc_data_structures / bit_set.rs
1 use crate::indexed_vec::{Idx, IndexVec};
2 use smallvec::SmallVec;
3 use std::fmt;
4 use std::iter;
5 use std::marker::PhantomData;
6 use std::mem;
7 use std::slice;
8
9 #[cfg(test)]
10 mod tests;
11
12 pub type Word = u64;
13 pub const WORD_BYTES: usize = mem::size_of::<Word>();
14 pub const WORD_BITS: usize = WORD_BYTES * 8;
15
16 /// A fixed-size bitset type with a dense representation. It does not support
17 /// resizing after creation; use `GrowableBitSet` for that.
18 ///
19 /// `T` is an index type, typically a newtyped `usize` wrapper, but it can also
20 /// just be `usize`.
21 ///
22 /// All operations that involve an element will panic if the element is equal
23 /// to or greater than the domain size. All operations that involve two bitsets
24 /// will panic if the bitsets have differing domain sizes.
25 #[derive(Clone, Eq, PartialEq, RustcDecodable, RustcEncodable)]
26 pub struct BitSet<T: Idx> {
27     domain_size: usize,
28     words: Vec<Word>,
29     marker: PhantomData<T>,
30 }
31
32 impl<T: Idx> BitSet<T> {
33     /// Creates a new, empty bitset with a given `domain_size`.
34     #[inline]
35     pub fn new_empty(domain_size: usize) -> BitSet<T> {
36         let num_words = num_words(domain_size);
37         BitSet {
38             domain_size,
39             words: vec![0; num_words],
40             marker: PhantomData,
41         }
42     }
43
44     /// Creates a new, filled bitset with a given `domain_size`.
45     #[inline]
46     pub fn new_filled(domain_size: usize) -> BitSet<T> {
47         let num_words = num_words(domain_size);
48         let mut result = BitSet {
49             domain_size,
50             words: vec![!0; num_words],
51             marker: PhantomData,
52         };
53         result.clear_excess_bits();
54         result
55     }
56
57     /// Gets the domain size.
58     pub fn domain_size(&self) -> usize {
59         self.domain_size
60     }
61
62     /// Clear all elements.
63     #[inline]
64     pub fn clear(&mut self) {
65         for word in &mut self.words {
66             *word = 0;
67         }
68     }
69
70     /// Clear excess bits in the final word.
71     fn clear_excess_bits(&mut self) {
72         let num_bits_in_final_word = self.domain_size % WORD_BITS;
73         if num_bits_in_final_word > 0 {
74             let mask = (1 << num_bits_in_final_word) - 1;
75             let final_word_idx = self.words.len() - 1;
76             self.words[final_word_idx] &= mask;
77         }
78     }
79
80     /// Efficiently overwrite `self` with `other`.
81     pub fn overwrite(&mut self, other: &BitSet<T>) {
82         assert!(self.domain_size == other.domain_size);
83         self.words.clone_from_slice(&other.words);
84     }
85
86     /// Count the number of set bits in the set.
87     pub fn count(&self) -> usize {
88         self.words.iter().map(|e| e.count_ones() as usize).sum()
89     }
90
91     /// Returns `true` if `self` contains `elem`.
92     #[inline]
93     pub fn contains(&self, elem: T) -> bool {
94         assert!(elem.index() < self.domain_size);
95         let (word_index, mask) = word_index_and_mask(elem);
96         (self.words[word_index] & mask) != 0
97     }
98
99     /// Is `self` is a (non-strict) superset of `other`?
100     #[inline]
101     pub fn superset(&self, other: &BitSet<T>) -> bool {
102         assert_eq!(self.domain_size, other.domain_size);
103         self.words.iter().zip(&other.words).all(|(a, b)| (a & b) == *b)
104     }
105
106     /// Is the set empty?
107     #[inline]
108     pub fn is_empty(&self) -> bool {
109         self.words.iter().all(|a| *a == 0)
110     }
111
112     /// Insert `elem`. Returns whether the set has changed.
113     #[inline]
114     pub fn insert(&mut self, elem: T) -> bool {
115         assert!(elem.index() < self.domain_size);
116         let (word_index, mask) = word_index_and_mask(elem);
117         let word_ref = &mut self.words[word_index];
118         let word = *word_ref;
119         let new_word = word | mask;
120         *word_ref = new_word;
121         new_word != word
122     }
123
124     /// Sets all bits to true.
125     pub fn insert_all(&mut self) {
126         for word in &mut self.words {
127             *word = !0;
128         }
129         self.clear_excess_bits();
130     }
131
132     /// Returns `true` if the set has changed.
133     #[inline]
134     pub fn remove(&mut self, elem: T) -> bool {
135         assert!(elem.index() < self.domain_size);
136         let (word_index, mask) = word_index_and_mask(elem);
137         let word_ref = &mut self.words[word_index];
138         let word = *word_ref;
139         let new_word = word & !mask;
140         *word_ref = new_word;
141         new_word != word
142     }
143
144     /// Sets `self = self | other` and returns `true` if `self` changed
145     /// (i.e., if new bits were added).
146     pub fn union(&mut self, other: &impl UnionIntoBitSet<T>) -> bool {
147         other.union_into(self)
148     }
149
150     /// Sets `self = self - other` and returns `true` if `self` changed.
151     /// (i.e., if any bits were removed).
152     pub fn subtract(&mut self, other: &impl SubtractFromBitSet<T>) -> bool {
153         other.subtract_from(self)
154     }
155
156     /// Sets `self = self & other` and return `true` if `self` changed.
157     /// (i.e., if any bits were removed).
158     pub fn intersect(&mut self, other: &BitSet<T>) -> bool {
159         assert_eq!(self.domain_size, other.domain_size);
160         bitwise(&mut self.words, &other.words, |a, b| { a & b })
161     }
162
163     /// Gets a slice of the underlying words.
164     pub fn words(&self) -> &[Word] {
165         &self.words
166     }
167
168     /// Iterates over the indices of set bits in a sorted order.
169     #[inline]
170     pub fn iter(&self) -> BitIter<'_, T> {
171         BitIter {
172             cur: None,
173             iter: self.words.iter().enumerate(),
174             marker: PhantomData,
175         }
176     }
177
178     /// Duplicates the set as a hybrid set.
179     pub fn to_hybrid(&self) -> HybridBitSet<T> {
180         // Note: we currently don't bother trying to make a Sparse set.
181         HybridBitSet::Dense(self.to_owned())
182     }
183
184     /// Set `self = self | other`. In contrast to `union` returns `true` if the set contains at
185     /// least one bit that is not in `other` (i.e. `other` is not a superset of `self`).
186     ///
187     /// This is an optimization for union of a hybrid bitset.
188     fn reverse_union_sparse(&mut self, sparse: &SparseBitSet<T>) -> bool {
189         assert!(sparse.domain_size == self.domain_size);
190         self.clear_excess_bits();
191
192         let mut not_already = false;
193         // Index of the current word not yet merged.
194         let mut current_index = 0;
195         // Mask of bits that came from the sparse set in the current word.
196         let mut new_bit_mask = 0;
197         for (word_index, mask) in sparse.iter().map(|x| word_index_and_mask(*x)) {
198             // Next bit is in a word not inspected yet.
199             if word_index > current_index {
200                 self.words[current_index] |= new_bit_mask;
201                 // Were there any bits in the old word that did not occur in the sparse set?
202                 not_already |= (self.words[current_index] ^ new_bit_mask) != 0;
203                 // Check all words we skipped for any set bit.
204                 not_already |= self.words[current_index+1..word_index].iter().any(|&x| x != 0);
205                 // Update next word.
206                 current_index = word_index;
207                 // Reset bit mask, no bits have been merged yet.
208                 new_bit_mask = 0;
209             }
210             // Add bit and mark it as coming from the sparse set.
211             // self.words[word_index] |= mask;
212             new_bit_mask |= mask;
213         }
214         self.words[current_index] |= new_bit_mask;
215         // Any bits in the last inspected word that were not in the sparse set?
216         not_already |= (self.words[current_index] ^ new_bit_mask) != 0;
217         // Any bits in the tail? Note `clear_excess_bits` before.
218         not_already |= self.words[current_index+1..].iter().any(|&x| x != 0);
219
220         not_already
221     }
222 }
223
224 /// This is implemented by all the bitsets so that BitSet::union() can be
225 /// passed any type of bitset.
226 pub trait UnionIntoBitSet<T: Idx> {
227     // Performs `other = other | self`.
228     fn union_into(&self, other: &mut BitSet<T>) -> bool;
229 }
230
231 /// This is implemented by all the bitsets so that BitSet::subtract() can be
232 /// passed any type of bitset.
233 pub trait SubtractFromBitSet<T: Idx> {
234     // Performs `other = other - self`.
235     fn subtract_from(&self, other: &mut BitSet<T>) -> bool;
236 }
237
238 impl<T: Idx> UnionIntoBitSet<T> for BitSet<T> {
239     fn union_into(&self, other: &mut BitSet<T>) -> bool {
240         assert_eq!(self.domain_size, other.domain_size);
241         bitwise(&mut other.words, &self.words, |a, b| { a | b })
242     }
243 }
244
245 impl<T: Idx> SubtractFromBitSet<T> for BitSet<T> {
246     fn subtract_from(&self, other: &mut BitSet<T>) -> bool {
247         assert_eq!(self.domain_size, other.domain_size);
248         bitwise(&mut other.words, &self.words, |a, b| { a & !b })
249     }
250 }
251
252 impl<T: Idx> fmt::Debug for BitSet<T> {
253     fn fmt(&self, w: &mut fmt::Formatter<'_>) -> fmt::Result {
254         w.debug_list()
255          .entries(self.iter())
256          .finish()
257     }
258 }
259
260 impl<T: Idx> ToString for BitSet<T> {
261     fn to_string(&self) -> String {
262         let mut result = String::new();
263         let mut sep = '[';
264
265         // Note: this is a little endian printout of bytes.
266
267         // i tracks how many bits we have printed so far.
268         let mut i = 0;
269         for word in &self.words {
270             let mut word = *word;
271             for _ in 0..WORD_BYTES { // for each byte in `word`:
272                 let remain = self.domain_size - i;
273                 // If less than a byte remains, then mask just that many bits.
274                 let mask = if remain <= 8 { (1 << remain) - 1 } else { 0xFF };
275                 assert!(mask <= 0xFF);
276                 let byte = word & mask;
277
278                 result.push_str(&format!("{}{:02x}", sep, byte));
279
280                 if remain <= 8 { break; }
281                 word >>= 8;
282                 i += 8;
283                 sep = '-';
284             }
285             sep = '|';
286         }
287         result.push(']');
288
289         result
290     }
291 }
292
293 pub struct BitIter<'a, T: Idx> {
294     cur: Option<(Word, usize)>,
295     iter: iter::Enumerate<slice::Iter<'a, Word>>,
296     marker: PhantomData<T>
297 }
298
299 impl<'a, T: Idx> Iterator for BitIter<'a, T> {
300     type Item = T;
301     fn next(&mut self) -> Option<T> {
302         loop {
303             if let Some((ref mut word, offset)) = self.cur {
304                 let bit_pos = word.trailing_zeros() as usize;
305                 if bit_pos != WORD_BITS {
306                     let bit = 1 << bit_pos;
307                     *word ^= bit;
308                     return Some(T::new(bit_pos + offset))
309                 }
310             }
311
312             let (i, word) = self.iter.next()?;
313             self.cur = Some((*word, WORD_BITS * i));
314         }
315     }
316 }
317
318 #[inline]
319 fn bitwise<Op>(out_vec: &mut [Word], in_vec: &[Word], op: Op) -> bool
320     where Op: Fn(Word, Word) -> Word
321 {
322     assert_eq!(out_vec.len(), in_vec.len());
323     let mut changed = false;
324     for (out_elem, in_elem) in out_vec.iter_mut().zip(in_vec.iter()) {
325         let old_val = *out_elem;
326         let new_val = op(old_val, *in_elem);
327         *out_elem = new_val;
328         changed |= old_val != new_val;
329     }
330     changed
331 }
332
333 const SPARSE_MAX: usize = 8;
334
335 /// A fixed-size bitset type with a sparse representation and a maximum of
336 /// `SPARSE_MAX` elements. The elements are stored as a sorted `SmallVec` with
337 /// no duplicates; although `SmallVec` can spill its elements to the heap, that
338 /// never happens within this type because of the `SPARSE_MAX` limit.
339 ///
340 /// This type is used by `HybridBitSet`; do not use directly.
341 #[derive(Clone, Debug)]
342 pub struct SparseBitSet<T: Idx> {
343     domain_size: usize,
344     elems: SmallVec<[T; SPARSE_MAX]>,
345 }
346
347 impl<T: Idx> SparseBitSet<T> {
348     fn new_empty(domain_size: usize) -> Self {
349         SparseBitSet {
350             domain_size,
351             elems: SmallVec::new()
352         }
353     }
354
355     fn len(&self) -> usize {
356         self.elems.len()
357     }
358
359     fn is_empty(&self) -> bool {
360         self.elems.len() == 0
361     }
362
363     fn contains(&self, elem: T) -> bool {
364         assert!(elem.index() < self.domain_size);
365         self.elems.contains(&elem)
366     }
367
368     fn insert(&mut self, elem: T) -> bool {
369         assert!(elem.index() < self.domain_size);
370         let changed = if let Some(i) = self.elems.iter().position(|&e| e >= elem) {
371             if self.elems[i] == elem {
372                 // `elem` is already in the set.
373                 false
374             } else {
375                 // `elem` is smaller than one or more existing elements.
376                 self.elems.insert(i, elem);
377                 true
378             }
379         } else {
380             // `elem` is larger than all existing elements.
381             self.elems.push(elem);
382             true
383         };
384         assert!(self.len() <= SPARSE_MAX);
385         changed
386     }
387
388     fn remove(&mut self, elem: T) -> bool {
389         assert!(elem.index() < self.domain_size);
390         if let Some(i) = self.elems.iter().position(|&e| e == elem) {
391             self.elems.remove(i);
392             true
393         } else {
394             false
395         }
396     }
397
398     fn to_dense(&self) -> BitSet<T> {
399         let mut dense = BitSet::new_empty(self.domain_size);
400         for elem in self.elems.iter() {
401             dense.insert(*elem);
402         }
403         dense
404     }
405
406     fn iter(&self) -> slice::Iter<'_, T> {
407         self.elems.iter()
408     }
409 }
410
411 impl<T: Idx> UnionIntoBitSet<T> for SparseBitSet<T> {
412     fn union_into(&self, other: &mut BitSet<T>) -> bool {
413         assert_eq!(self.domain_size, other.domain_size);
414         let mut changed = false;
415         for elem in self.iter() {
416             changed |= other.insert(*elem);
417         }
418         changed
419     }
420 }
421
422 impl<T: Idx> SubtractFromBitSet<T> for SparseBitSet<T> {
423     fn subtract_from(&self, other: &mut BitSet<T>) -> bool {
424         assert_eq!(self.domain_size, other.domain_size);
425         let mut changed = false;
426         for elem in self.iter() {
427             changed |= other.remove(*elem);
428         }
429         changed
430     }
431 }
432
433 /// A fixed-size bitset type with a hybrid representation: sparse when there
434 /// are up to a `SPARSE_MAX` elements in the set, but dense when there are more
435 /// than `SPARSE_MAX`.
436 ///
437 /// This type is especially efficient for sets that typically have a small
438 /// number of elements, but a large `domain_size`, and are cleared frequently.
439 ///
440 /// `T` is an index type, typically a newtyped `usize` wrapper, but it can also
441 /// just be `usize`.
442 ///
443 /// All operations that involve an element will panic if the element is equal
444 /// to or greater than the domain size. All operations that involve two bitsets
445 /// will panic if the bitsets have differing domain sizes.
446 #[derive(Clone, Debug)]
447 pub enum HybridBitSet<T: Idx> {
448     Sparse(SparseBitSet<T>),
449     Dense(BitSet<T>),
450 }
451
452 impl<T: Idx> HybridBitSet<T> {
453     pub fn new_empty(domain_size: usize) -> Self {
454         HybridBitSet::Sparse(SparseBitSet::new_empty(domain_size))
455     }
456
457     fn domain_size(&self) -> usize {
458         match self {
459             HybridBitSet::Sparse(sparse) => sparse.domain_size,
460             HybridBitSet::Dense(dense) => dense.domain_size,
461         }
462     }
463
464     pub fn clear(&mut self) {
465         let domain_size = self.domain_size();
466         *self = HybridBitSet::new_empty(domain_size);
467     }
468
469     pub fn contains(&self, elem: T) -> bool {
470         match self {
471             HybridBitSet::Sparse(sparse) => sparse.contains(elem),
472             HybridBitSet::Dense(dense) => dense.contains(elem),
473         }
474     }
475
476     pub fn superset(&self, other: &HybridBitSet<T>) -> bool {
477         match (self, other) {
478             (HybridBitSet::Dense(self_dense), HybridBitSet::Dense(other_dense)) => {
479                 self_dense.superset(other_dense)
480             }
481             _ => {
482                 assert!(self.domain_size() == other.domain_size());
483                 other.iter().all(|elem| self.contains(elem))
484             }
485         }
486     }
487
488     pub fn is_empty(&self) -> bool {
489         match self {
490             HybridBitSet::Sparse(sparse) => sparse.is_empty(),
491             HybridBitSet::Dense(dense) => dense.is_empty(),
492         }
493     }
494
495     pub fn insert(&mut self, elem: T) -> bool {
496         // No need to check `elem` against `self.domain_size` here because all
497         // the match cases check it, one way or another.
498         match self {
499             HybridBitSet::Sparse(sparse) if sparse.len() < SPARSE_MAX => {
500                 // The set is sparse and has space for `elem`.
501                 sparse.insert(elem)
502             }
503             HybridBitSet::Sparse(sparse) if sparse.contains(elem) => {
504                 // The set is sparse and does not have space for `elem`, but
505                 // that doesn't matter because `elem` is already present.
506                 false
507             }
508             HybridBitSet::Sparse(sparse) => {
509                 // The set is sparse and full. Convert to a dense set.
510                 let mut dense = sparse.to_dense();
511                 let changed = dense.insert(elem);
512                 assert!(changed);
513                 *self = HybridBitSet::Dense(dense);
514                 changed
515             }
516             HybridBitSet::Dense(dense) => dense.insert(elem),
517         }
518     }
519
520     pub fn insert_all(&mut self) {
521         let domain_size = self.domain_size();
522         match self {
523             HybridBitSet::Sparse(_) => {
524                 *self = HybridBitSet::Dense(BitSet::new_filled(domain_size));
525             }
526             HybridBitSet::Dense(dense) => dense.insert_all(),
527         }
528     }
529
530     pub fn remove(&mut self, elem: T) -> bool {
531         // Note: we currently don't bother going from Dense back to Sparse.
532         match self {
533             HybridBitSet::Sparse(sparse) => sparse.remove(elem),
534             HybridBitSet::Dense(dense) => dense.remove(elem),
535         }
536     }
537
538     pub fn union(&mut self, other: &HybridBitSet<T>) -> bool {
539         match self {
540             HybridBitSet::Sparse(self_sparse) => {
541                 match other {
542                     HybridBitSet::Sparse(other_sparse) => {
543                         // Both sets are sparse. Add the elements in
544                         // `other_sparse` to `self` one at a time. This
545                         // may or may not cause `self` to be densified.
546                         assert_eq!(self.domain_size(), other.domain_size());
547                         let mut changed = false;
548                         for elem in other_sparse.iter() {
549                             changed |= self.insert(*elem);
550                         }
551                         changed
552                     }
553                     HybridBitSet::Dense(other_dense) => {
554                         // `self` is sparse and `other` is dense. To
555                         // merge them, we have two available strategies:
556                         // * Densify `self` then merge other
557                         // * Clone other then integrate bits from `self`
558                         // The second strategy requires dedicated method
559                         // since the usual `union` returns the wrong
560                         // result. In the dedicated case the computation
561                         // is slightly faster if the bits of the sparse
562                         // bitset map to only few words of the dense
563                         // representation, i.e. indices are near each
564                         // other.
565                         //
566                         // Benchmarking seems to suggest that the second
567                         // option is worth it.
568                         let mut new_dense = other_dense.clone();
569                         let changed = new_dense.reverse_union_sparse(self_sparse);
570                         *self = HybridBitSet::Dense(new_dense);
571                         changed
572                     }
573                 }
574             }
575
576             HybridBitSet::Dense(self_dense) => self_dense.union(other),
577         }
578     }
579
580     /// Converts to a dense set, consuming itself in the process.
581     pub fn to_dense(self) -> BitSet<T> {
582         match self {
583             HybridBitSet::Sparse(sparse) => sparse.to_dense(),
584             HybridBitSet::Dense(dense) => dense,
585         }
586     }
587
588     pub fn iter(&self) -> HybridIter<'_, T> {
589         match self {
590             HybridBitSet::Sparse(sparse) => HybridIter::Sparse(sparse.iter()),
591             HybridBitSet::Dense(dense) => HybridIter::Dense(dense.iter()),
592         }
593     }
594 }
595
596 impl<T: Idx> UnionIntoBitSet<T> for HybridBitSet<T> {
597     fn union_into(&self, other: &mut BitSet<T>) -> bool {
598         match self {
599             HybridBitSet::Sparse(sparse) => sparse.union_into(other),
600             HybridBitSet::Dense(dense) => dense.union_into(other),
601         }
602     }
603 }
604
605 impl<T: Idx> SubtractFromBitSet<T> for HybridBitSet<T> {
606     fn subtract_from(&self, other: &mut BitSet<T>) -> bool {
607         match self {
608             HybridBitSet::Sparse(sparse) => sparse.subtract_from(other),
609             HybridBitSet::Dense(dense) => dense.subtract_from(other),
610         }
611     }
612 }
613
614 pub enum HybridIter<'a, T: Idx> {
615     Sparse(slice::Iter<'a, T>),
616     Dense(BitIter<'a, T>),
617 }
618
619 impl<'a, T: Idx> Iterator for HybridIter<'a, T> {
620     type Item = T;
621
622     fn next(&mut self) -> Option<T> {
623         match self {
624             HybridIter::Sparse(sparse) => sparse.next().map(|e| *e),
625             HybridIter::Dense(dense) => dense.next(),
626         }
627     }
628 }
629
630 /// A resizable bitset type with a dense representation.
631 ///
632 /// `T` is an index type, typically a newtyped `usize` wrapper, but it can also
633 /// just be `usize`.
634 ///
635 /// All operations that involve an element will panic if the element is equal
636 /// to or greater than the domain size.
637 #[derive(Clone, Debug, PartialEq)]
638 pub struct GrowableBitSet<T: Idx> {
639     bit_set: BitSet<T>,
640 }
641
642 impl<T: Idx> GrowableBitSet<T> {
643     /// Ensure that the set can hold at least `min_domain_size` elements.
644     pub fn ensure(&mut self, min_domain_size: usize) {
645         if self.bit_set.domain_size < min_domain_size {
646             self.bit_set.domain_size = min_domain_size;
647         }
648
649         let min_num_words = num_words(min_domain_size);
650         if self.bit_set.words.len() < min_num_words {
651             self.bit_set.words.resize(min_num_words, 0)
652         }
653     }
654
655     pub fn new_empty() -> GrowableBitSet<T> {
656         GrowableBitSet { bit_set: BitSet::new_empty(0) }
657     }
658
659     pub fn with_capacity(capacity: usize) -> GrowableBitSet<T> {
660         GrowableBitSet { bit_set: BitSet::new_empty(capacity) }
661     }
662
663     /// Returns `true` if the set has changed.
664     #[inline]
665     pub fn insert(&mut self, elem: T) -> bool {
666         self.ensure(elem.index() + 1);
667         self.bit_set.insert(elem)
668     }
669
670     #[inline]
671     pub fn contains(&self, elem: T) -> bool {
672         let (word_index, mask) = word_index_and_mask(elem);
673         if let Some(word) = self.bit_set.words.get(word_index) {
674             (word & mask) != 0
675         } else {
676             false
677         }
678     }
679 }
680
681 /// A fixed-size 2D bit matrix type with a dense representation.
682 ///
683 /// `R` and `C` are index types used to identify rows and columns respectively;
684 /// typically newtyped `usize` wrappers, but they can also just be `usize`.
685 ///
686 /// All operations that involve a row and/or column index will panic if the
687 /// index exceeds the relevant bound.
688 #[derive(Clone, Debug, Eq, PartialEq, RustcDecodable, RustcEncodable)]
689 pub struct BitMatrix<R: Idx, C: Idx> {
690     num_rows: usize,
691     num_columns: usize,
692     words: Vec<Word>,
693     marker: PhantomData<(R, C)>,
694 }
695
696 impl<R: Idx, C: Idx> BitMatrix<R, C> {
697     /// Creates a new `rows x columns` matrix, initially empty.
698     pub fn new(num_rows: usize, num_columns: usize) -> BitMatrix<R, C> {
699         // For every element, we need one bit for every other
700         // element. Round up to an even number of words.
701         let words_per_row = num_words(num_columns);
702         BitMatrix {
703             num_rows,
704             num_columns,
705             words: vec![0; num_rows * words_per_row],
706             marker: PhantomData,
707         }
708     }
709
710     /// Creates a new matrix, with `row` used as the value for every row.
711     pub fn from_row_n(row: &BitSet<C>, num_rows: usize) -> BitMatrix<R, C> {
712         let num_columns = row.domain_size();
713         let words_per_row = num_words(num_columns);
714         assert_eq!(words_per_row, row.words().len());
715         BitMatrix {
716             num_rows,
717             num_columns,
718             words: iter::repeat(row.words()).take(num_rows).flatten().cloned().collect(),
719             marker: PhantomData,
720         }
721     }
722
723     pub fn rows(&self) -> impl Iterator<Item = R> {
724         (0..self.num_rows).map(R::new)
725     }
726
727     /// The range of bits for a given row.
728     fn range(&self, row: R) -> (usize, usize) {
729         let words_per_row = num_words(self.num_columns);
730         let start = row.index() * words_per_row;
731         (start, start + words_per_row)
732     }
733
734     /// Sets the cell at `(row, column)` to true. Put another way, insert
735     /// `column` to the bitset for `row`.
736     ///
737     /// Returns `true` if this changed the matrix.
738     pub fn insert(&mut self, row: R, column: C) -> bool {
739         assert!(row.index() < self.num_rows && column.index() < self.num_columns);
740         let (start, _) = self.range(row);
741         let (word_index, mask) = word_index_and_mask(column);
742         let words = &mut self.words[..];
743         let word = words[start + word_index];
744         let new_word = word | mask;
745         words[start + word_index] = new_word;
746         word != new_word
747     }
748
749     /// Do the bits from `row` contain `column`? Put another way, is
750     /// the matrix cell at `(row, column)` true?  Put yet another way,
751     /// if the matrix represents (transitive) reachability, can
752     /// `row` reach `column`?
753     pub fn contains(&self, row: R, column: C) -> bool {
754         assert!(row.index() < self.num_rows && column.index() < self.num_columns);
755         let (start, _) = self.range(row);
756         let (word_index, mask) = word_index_and_mask(column);
757         (self.words[start + word_index] & mask) != 0
758     }
759
760     /// Returns those indices that are true in rows `a` and `b`. This
761     /// is an O(n) operation where `n` is the number of elements
762     /// (somewhat independent from the actual size of the
763     /// intersection, in particular).
764     pub fn intersect_rows(&self, row1: R, row2: R) -> Vec<C> {
765         assert!(row1.index() < self.num_rows && row2.index() < self.num_rows);
766         let (row1_start, row1_end) = self.range(row1);
767         let (row2_start, row2_end) = self.range(row2);
768         let mut result = Vec::with_capacity(self.num_columns);
769         for (base, (i, j)) in (row1_start..row1_end).zip(row2_start..row2_end).enumerate() {
770             let mut v = self.words[i] & self.words[j];
771             for bit in 0..WORD_BITS {
772                 if v == 0 {
773                     break;
774                 }
775                 if v & 0x1 != 0 {
776                     result.push(C::new(base * WORD_BITS + bit));
777                 }
778                 v >>= 1;
779             }
780         }
781         result
782     }
783
784     /// Adds the bits from row `read` to the bits from row `write`, and
785     /// returns `true` if anything changed.
786     ///
787     /// This is used when computing transitive reachability because if
788     /// you have an edge `write -> read`, because in that case
789     /// `write` can reach everything that `read` can (and
790     /// potentially more).
791     pub fn union_rows(&mut self, read: R, write: R) -> bool {
792         assert!(read.index() < self.num_rows && write.index() < self.num_rows);
793         let (read_start, read_end) = self.range(read);
794         let (write_start, write_end) = self.range(write);
795         let words = &mut self.words[..];
796         let mut changed = false;
797         for (read_index, write_index) in (read_start..read_end).zip(write_start..write_end) {
798             let word = words[write_index];
799             let new_word = word | words[read_index];
800             words[write_index] = new_word;
801             changed |= word != new_word;
802         }
803         changed
804     }
805
806     /// Adds the bits from `with` to the bits from row `write`, and
807     /// returns `true` if anything changed.
808     pub fn union_row_with(&mut self, with: &BitSet<C>, write: R) -> bool {
809         assert!(write.index() < self.num_rows);
810         assert_eq!(with.domain_size(), self.num_columns);
811         let (write_start, write_end) = self.range(write);
812         let mut changed = false;
813         for (read_index, write_index) in (0..with.words().len()).zip(write_start..write_end) {
814             let word = self.words[write_index];
815             let new_word = word | with.words()[read_index];
816             self.words[write_index] = new_word;
817             changed |= word != new_word;
818         }
819         changed
820     }
821
822     /// Sets every cell in `row` to true.
823     pub fn insert_all_into_row(&mut self, row: R) {
824         assert!(row.index() < self.num_rows);
825         let (start, end) = self.range(row);
826         let words = &mut self.words[..];
827         for index in start..end {
828             words[index] = !0;
829         }
830         self.clear_excess_bits(row);
831     }
832
833     /// Clear excess bits in the final word of the row.
834     fn clear_excess_bits(&mut self, row: R) {
835         let num_bits_in_final_word = self.num_columns % WORD_BITS;
836         if num_bits_in_final_word > 0 {
837             let mask = (1 << num_bits_in_final_word) - 1;
838             let (_, end) = self.range(row);
839             let final_word_idx = end - 1;
840             self.words[final_word_idx] &= mask;
841         }
842     }
843
844     /// Gets a slice of the underlying words.
845     pub fn words(&self) -> &[Word] {
846         &self.words
847     }
848
849     /// Iterates through all the columns set to true in a given row of
850     /// the matrix.
851     pub fn iter(&self, row: R) -> BitIter<'_, C> {
852         assert!(row.index() < self.num_rows);
853         let (start, end) = self.range(row);
854         BitIter {
855             cur: None,
856             iter: self.words[start..end].iter().enumerate(),
857             marker: PhantomData,
858         }
859     }
860
861     /// Returns the number of elements in `row`.
862     pub fn count(&self, row: R) -> usize {
863         let (start, end) = self.range(row);
864         self.words[start..end].iter().map(|e| e.count_ones() as usize).sum()
865     }
866 }
867
868 /// A fixed-column-size, variable-row-size 2D bit matrix with a moderately
869 /// sparse representation.
870 ///
871 /// Initially, every row has no explicit representation. If any bit within a
872 /// row is set, the entire row is instantiated as `Some(<HybridBitSet>)`.
873 /// Furthermore, any previously uninstantiated rows prior to it will be
874 /// instantiated as `None`. Those prior rows may themselves become fully
875 /// instantiated later on if any of their bits are set.
876 ///
877 /// `R` and `C` are index types used to identify rows and columns respectively;
878 /// typically newtyped `usize` wrappers, but they can also just be `usize`.
879 #[derive(Clone, Debug)]
880 pub struct SparseBitMatrix<R, C>
881 where
882     R: Idx,
883     C: Idx,
884 {
885     num_columns: usize,
886     rows: IndexVec<R, Option<HybridBitSet<C>>>,
887 }
888
889 impl<R: Idx, C: Idx> SparseBitMatrix<R, C> {
890     /// Creates a new empty sparse bit matrix with no rows or columns.
891     pub fn new(num_columns: usize) -> Self {
892         Self {
893             num_columns,
894             rows: IndexVec::new(),
895         }
896     }
897
898     fn ensure_row(&mut self, row: R) -> &mut HybridBitSet<C> {
899         // Instantiate any missing rows up to and including row `row` with an
900         // empty HybridBitSet.
901         self.rows.ensure_contains_elem(row, || None);
902
903         // Then replace row `row` with a full HybridBitSet if necessary.
904         let num_columns = self.num_columns;
905         self.rows[row].get_or_insert_with(|| HybridBitSet::new_empty(num_columns))
906     }
907
908     /// Sets the cell at `(row, column)` to true. Put another way, insert
909     /// `column` to the bitset for `row`.
910     ///
911     /// Returns `true` if this changed the matrix.
912     pub fn insert(&mut self, row: R, column: C) -> bool {
913         self.ensure_row(row).insert(column)
914     }
915
916     /// Do the bits from `row` contain `column`? Put another way, is
917     /// the matrix cell at `(row, column)` true?  Put yet another way,
918     /// if the matrix represents (transitive) reachability, can
919     /// `row` reach `column`?
920     pub fn contains(&self, row: R, column: C) -> bool {
921         self.row(row).map_or(false, |r| r.contains(column))
922     }
923
924     /// Adds the bits from row `read` to the bits from row `write`, and
925     /// returns `true` if anything changed.
926     ///
927     /// This is used when computing transitive reachability because if
928     /// you have an edge `write -> read`, because in that case
929     /// `write` can reach everything that `read` can (and
930     /// potentially more).
931     pub fn union_rows(&mut self, read: R, write: R) -> bool {
932         if read == write || self.row(read).is_none() {
933             return false;
934         }
935
936         self.ensure_row(write);
937         if let (Some(read_row), Some(write_row)) = self.rows.pick2_mut(read, write) {
938             write_row.union(read_row)
939         } else {
940             unreachable!()
941         }
942     }
943
944     /// Union a row, `from`, into the `into` row.
945     pub fn union_into_row(&mut self, into: R, from: &HybridBitSet<C>) -> bool {
946         self.ensure_row(into).union(from)
947     }
948
949     /// Insert all bits in the given row.
950     pub fn insert_all_into_row(&mut self, row: R) {
951         self.ensure_row(row).insert_all();
952     }
953
954     pub fn rows(&self) -> impl Iterator<Item = R> {
955         self.rows.indices()
956     }
957
958     /// Iterates through all the columns set to true in a given row of
959     /// the matrix.
960     pub fn iter<'a>(&'a self, row: R) -> impl Iterator<Item = C> + 'a {
961         self.row(row).into_iter().flat_map(|r| r.iter())
962     }
963
964     pub fn row(&self, row: R) -> Option<&HybridBitSet<C>> {
965         if let Some(Some(row)) = self.rows.get(row) {
966             Some(row)
967         } else {
968             None
969         }
970     }
971 }
972
973 #[inline]
974 fn num_words<T: Idx>(domain_size: T) -> usize {
975     (domain_size.index() + WORD_BITS - 1) / WORD_BITS
976 }
977
978 #[inline]
979 fn word_index_and_mask<T: Idx>(elem: T) -> (usize, Word) {
980     let elem = elem.index();
981     let word_index = elem / WORD_BITS;
982     let mask = 1 << (elem % WORD_BITS);
983     (word_index, mask)
984 }