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