src/librustc_index/bit_set.rs

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