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