rustc_index/
bit_set.rs

1use std::marker::PhantomData;
2#[cfg(not(feature = "nightly"))]
3use std::mem;
4use std::ops::{BitAnd, BitAndAssign, BitOrAssign, Bound, Not, Range, RangeBounds, Shl};
5use std::rc::Rc;
6use std::{fmt, iter, slice};
7
8use Chunk::*;
9#[cfg(feature = "nightly")]
10use rustc_macros::{Decodable_NoContext, Encodable_NoContext};
11use smallvec::{SmallVec, smallvec};
12
13use crate::{Idx, IndexVec};
14
15#[cfg(test)]
16mod tests;
17
18type Word = u64;
19const WORD_BYTES: usize = size_of::<Word>();
20const WORD_BITS: usize = WORD_BYTES * 8;
21
22// The choice of chunk size has some trade-offs.
23//
24// A big chunk size tends to favour cases where many large `ChunkedBitSet`s are
25// present, because they require fewer `Chunk`s, reducing the number of
26// allocations and reducing peak memory usage. Also, fewer chunk operations are
27// required, though more of them might be `Mixed`.
28//
29// A small chunk size tends to favour cases where many small `ChunkedBitSet`s
30// are present, because less space is wasted at the end of the final chunk (if
31// it's not full).
32const CHUNK_WORDS: usize = 32;
33const CHUNK_BITS: usize = CHUNK_WORDS * WORD_BITS; // 2048 bits
34
35/// ChunkSize is small to keep `Chunk` small. The static assertion ensures it's
36/// not too small.
37type ChunkSize = u16;
38const _: () = assert!(CHUNK_BITS <= ChunkSize::MAX as usize);
39
40pub trait BitRelations<Rhs> {
41    fn union(&mut self, other: &Rhs) -> bool;
42    fn subtract(&mut self, other: &Rhs) -> bool;
43    fn intersect(&mut self, other: &Rhs) -> bool;
44}
45
46#[inline]
47fn inclusive_start_end<T: Idx>(
48    range: impl RangeBounds<T>,
49    domain: usize,
50) -> Option<(usize, usize)> {
51    // Both start and end are inclusive.
52    let start = match range.start_bound().cloned() {
53        Bound::Included(start) => start.index(),
54        Bound::Excluded(start) => start.index() + 1,
55        Bound::Unbounded => 0,
56    };
57    let end = match range.end_bound().cloned() {
58        Bound::Included(end) => end.index(),
59        Bound::Excluded(end) => end.index().checked_sub(1)?,
60        Bound::Unbounded => domain - 1,
61    };
62    assert!(end < domain);
63    if start > end {
64        return None;
65    }
66    Some((start, end))
67}
68
69macro_rules! bit_relations_inherent_impls {
70    () => {
71        /// Sets `self = self | other` and returns `true` if `self` changed
72        /// (i.e., if new bits were added).
73        pub fn union<Rhs>(&mut self, other: &Rhs) -> bool
74        where
75            Self: BitRelations<Rhs>,
76        {
77            <Self as BitRelations<Rhs>>::union(self, other)
78        }
79
80        /// Sets `self = self - other` and returns `true` if `self` changed.
81        /// (i.e., if any bits were removed).
82        pub fn subtract<Rhs>(&mut self, other: &Rhs) -> bool
83        where
84            Self: BitRelations<Rhs>,
85        {
86            <Self as BitRelations<Rhs>>::subtract(self, other)
87        }
88
89        /// Sets `self = self & other` and return `true` if `self` changed.
90        /// (i.e., if any bits were removed).
91        pub fn intersect<Rhs>(&mut self, other: &Rhs) -> bool
92        where
93            Self: BitRelations<Rhs>,
94        {
95            <Self as BitRelations<Rhs>>::intersect(self, other)
96        }
97    };
98}
99
100/// A fixed-size bitset type with a dense representation.
101///
102/// Note 1: Since this bitset is dense, if your domain is big, and/or relatively
103/// homogeneous (for example, with long runs of bits set or unset), then it may
104/// be preferable to instead use a [MixedBitSet], or an
105/// [IntervalSet](crate::interval::IntervalSet). They should be more suited to
106/// sparse, or highly-compressible, domains.
107///
108/// Note 2: Use [`GrowableBitSet`] if you need support for resizing after creation.
109///
110/// `T` is an index type, typically a newtyped `usize` wrapper, but it can also
111/// just be `usize`.
112///
113/// All operations that involve an element will panic if the element is equal
114/// to or greater than the domain size. All operations that involve two bitsets
115/// will panic if the bitsets have differing domain sizes.
116///
117#[cfg_attr(feature = "nightly", derive(Decodable_NoContext, Encodable_NoContext))]
118#[derive(Eq, PartialEq, Hash)]
119pub struct DenseBitSet<T> {
120    domain_size: usize,
121    words: SmallVec<[Word; 2]>,
122    marker: PhantomData<T>,
123}
124
125impl<T> DenseBitSet<T> {
126    /// Gets the domain size.
127    pub fn domain_size(&self) -> usize {
128        self.domain_size
129    }
130}
131
132impl<T: Idx> DenseBitSet<T> {
133    /// Creates a new, empty bitset with a given `domain_size`.
134    #[inline]
135    pub fn new_empty(domain_size: usize) -> DenseBitSet<T> {
136        let num_words = num_words(domain_size);
137        DenseBitSet { domain_size, words: smallvec![0; num_words], marker: PhantomData }
138    }
139
140    /// Creates a new, filled bitset with a given `domain_size`.
141    #[inline]
142    pub fn new_filled(domain_size: usize) -> DenseBitSet<T> {
143        let num_words = num_words(domain_size);
144        let mut result =
145            DenseBitSet { domain_size, words: smallvec![!0; num_words], marker: PhantomData };
146        result.clear_excess_bits();
147        result
148    }
149
150    /// Clear all elements.
151    #[inline]
152    pub fn clear(&mut self) {
153        self.words.fill(0);
154    }
155
156    /// Clear excess bits in the final word.
157    fn clear_excess_bits(&mut self) {
158        clear_excess_bits_in_final_word(self.domain_size, &mut self.words);
159    }
160
161    /// Count the number of set bits in the set.
162    pub fn count(&self) -> usize {
163        self.words.iter().map(|e| e.count_ones() as usize).sum()
164    }
165
166    /// Returns `true` if `self` contains `elem`.
167    #[inline]
168    pub fn contains(&self, elem: T) -> bool {
169        assert!(elem.index() < self.domain_size);
170        let (word_index, mask) = word_index_and_mask(elem);
171        (self.words[word_index] & mask) != 0
172    }
173
174    /// Is `self` is a (non-strict) superset of `other`?
175    #[inline]
176    pub fn superset(&self, other: &DenseBitSet<T>) -> bool {
177        assert_eq!(self.domain_size, other.domain_size);
178        self.words.iter().zip(&other.words).all(|(a, b)| (a & b) == *b)
179    }
180
181    /// Is the set empty?
182    #[inline]
183    pub fn is_empty(&self) -> bool {
184        self.words.iter().all(|a| *a == 0)
185    }
186
187    /// Insert `elem`. Returns whether the set has changed.
188    #[inline]
189    pub fn insert(&mut self, elem: T) -> bool {
190        assert!(
191            elem.index() < self.domain_size,
192            "inserting element at index {} but domain size is {}",
193            elem.index(),
194            self.domain_size,
195        );
196        let (word_index, mask) = word_index_and_mask(elem);
197        let word_ref = &mut self.words[word_index];
198        let word = *word_ref;
199        let new_word = word | mask;
200        *word_ref = new_word;
201        new_word != word
202    }
203
204    #[inline]
205    pub fn insert_range(&mut self, elems: impl RangeBounds<T>) {
206        let Some((start, end)) = inclusive_start_end(elems, self.domain_size) else {
207            return;
208        };
209
210        let (start_word_index, start_mask) = word_index_and_mask(start);
211        let (end_word_index, end_mask) = word_index_and_mask(end);
212
213        // Set all words in between start and end (exclusively of both).
214        for word_index in (start_word_index + 1)..end_word_index {
215            self.words[word_index] = !0;
216        }
217
218        if start_word_index != end_word_index {
219            // Start and end are in different words, so we handle each in turn.
220            //
221            // We set all leading bits. This includes the start_mask bit.
222            self.words[start_word_index] |= !(start_mask - 1);
223            // And all trailing bits (i.e. from 0..=end) in the end word,
224            // including the end.
225            self.words[end_word_index] |= end_mask | (end_mask - 1);
226        } else {
227            self.words[start_word_index] |= end_mask | (end_mask - start_mask);
228        }
229    }
230
231    /// Sets all bits to true.
232    pub fn insert_all(&mut self) {
233        self.words.fill(!0);
234        self.clear_excess_bits();
235    }
236
237    /// Checks whether any bit in the given range is a 1.
238    #[inline]
239    pub fn contains_any(&self, elems: impl RangeBounds<T>) -> bool {
240        let Some((start, end)) = inclusive_start_end(elems, self.domain_size) else {
241            return false;
242        };
243        let (start_word_index, start_mask) = word_index_and_mask(start);
244        let (end_word_index, end_mask) = word_index_and_mask(end);
245
246        if start_word_index == end_word_index {
247            self.words[start_word_index] & (end_mask | (end_mask - start_mask)) != 0
248        } else {
249            if self.words[start_word_index] & !(start_mask - 1) != 0 {
250                return true;
251            }
252
253            let remaining = start_word_index + 1..end_word_index;
254            if remaining.start <= remaining.end {
255                self.words[remaining].iter().any(|&w| w != 0)
256                    || self.words[end_word_index] & (end_mask | (end_mask - 1)) != 0
257            } else {
258                false
259            }
260        }
261    }
262
263    /// Returns `true` if the set has changed.
264    #[inline]
265    pub fn remove(&mut self, elem: T) -> bool {
266        assert!(elem.index() < self.domain_size);
267        let (word_index, mask) = word_index_and_mask(elem);
268        let word_ref = &mut self.words[word_index];
269        let word = *word_ref;
270        let new_word = word & !mask;
271        *word_ref = new_word;
272        new_word != word
273    }
274
275    /// Iterates over the indices of set bits in a sorted order.
276    #[inline]
277    pub fn iter(&self) -> BitIter<'_, T> {
278        BitIter::new(&self.words)
279    }
280
281    pub fn last_set_in(&self, range: impl RangeBounds<T>) -> Option<T> {
282        let (start, end) = inclusive_start_end(range, self.domain_size)?;
283        let (start_word_index, _) = word_index_and_mask(start);
284        let (end_word_index, end_mask) = word_index_and_mask(end);
285
286        let end_word = self.words[end_word_index] & (end_mask | (end_mask - 1));
287        if end_word != 0 {
288            let pos = max_bit(end_word) + WORD_BITS * end_word_index;
289            if start <= pos {
290                return Some(T::new(pos));
291            }
292        }
293
294        // We exclude end_word_index from the range here, because we don't want
295        // to limit ourselves to *just* the last word: the bits set it in may be
296        // after `end`, so it may not work out.
297        if let Some(offset) =
298            self.words[start_word_index..end_word_index].iter().rposition(|&w| w != 0)
299        {
300            let word_idx = start_word_index + offset;
301            let start_word = self.words[word_idx];
302            let pos = max_bit(start_word) + WORD_BITS * word_idx;
303            if start <= pos {
304                return Some(T::new(pos));
305            }
306        }
307
308        None
309    }
310
311    bit_relations_inherent_impls! {}
312
313    /// Sets `self = self | !other`.
314    ///
315    /// FIXME: Incorporate this into [`BitRelations`] and fill out
316    /// implementations for other bitset types, if needed.
317    pub fn union_not(&mut self, other: &DenseBitSet<T>) {
318        assert_eq!(self.domain_size, other.domain_size);
319
320        // FIXME(Zalathar): If we were to forcibly _set_ all excess bits before
321        // the bitwise update, and then clear them again afterwards, we could
322        // quickly and accurately detect whether the update changed anything.
323        // But that's only worth doing if there's an actual use-case.
324
325        bitwise(&mut self.words, &other.words, |a, b| a | !b);
326        // The bitwise update `a | !b` can result in the last word containing
327        // out-of-domain bits, so we need to clear them.
328        self.clear_excess_bits();
329    }
330}
331
332// dense REL dense
333impl<T: Idx> BitRelations<DenseBitSet<T>> for DenseBitSet<T> {
334    fn union(&mut self, other: &DenseBitSet<T>) -> bool {
335        assert_eq!(self.domain_size, other.domain_size);
336        bitwise(&mut self.words, &other.words, |a, b| a | b)
337    }
338
339    fn subtract(&mut self, other: &DenseBitSet<T>) -> bool {
340        assert_eq!(self.domain_size, other.domain_size);
341        bitwise(&mut self.words, &other.words, |a, b| a & !b)
342    }
343
344    fn intersect(&mut self, other: &DenseBitSet<T>) -> bool {
345        assert_eq!(self.domain_size, other.domain_size);
346        bitwise(&mut self.words, &other.words, |a, b| a & b)
347    }
348}
349
350impl<T: Idx> From<GrowableBitSet<T>> for DenseBitSet<T> {
351    fn from(bit_set: GrowableBitSet<T>) -> Self {
352        bit_set.bit_set
353    }
354}
355
356impl<T> Clone for DenseBitSet<T> {
357    fn clone(&self) -> Self {
358        DenseBitSet {
359            domain_size: self.domain_size,
360            words: self.words.clone(),
361            marker: PhantomData,
362        }
363    }
364
365    fn clone_from(&mut self, from: &Self) {
366        self.domain_size = from.domain_size;
367        self.words.clone_from(&from.words);
368    }
369}
370
371impl<T: Idx> fmt::Debug for DenseBitSet<T> {
372    fn fmt(&self, w: &mut fmt::Formatter<'_>) -> fmt::Result {
373        w.debug_list().entries(self.iter()).finish()
374    }
375}
376
377impl<T: Idx> ToString for DenseBitSet<T> {
378    fn to_string(&self) -> String {
379        let mut result = String::new();
380        let mut sep = '[';
381
382        // Note: this is a little endian printout of bytes.
383
384        // i tracks how many bits we have printed so far.
385        let mut i = 0;
386        for word in &self.words {
387            let mut word = *word;
388            for _ in 0..WORD_BYTES {
389                // for each byte in `word`:
390                let remain = self.domain_size - i;
391                // If less than a byte remains, then mask just that many bits.
392                let mask = if remain <= 8 { (1 << remain) - 1 } else { 0xFF };
393                assert!(mask <= 0xFF);
394                let byte = word & mask;
395
396                result.push_str(&format!("{sep}{byte:02x}"));
397
398                if remain <= 8 {
399                    break;
400                }
401                word >>= 8;
402                i += 8;
403                sep = '-';
404            }
405            sep = '|';
406        }
407        result.push(']');
408
409        result
410    }
411}
412
413pub struct BitIter<'a, T: Idx> {
414    /// A copy of the current word, but with any already-visited bits cleared.
415    /// (This lets us use `trailing_zeros()` to find the next set bit.) When it
416    /// is reduced to 0, we move onto the next word.
417    word: Word,
418
419    /// The offset (measured in bits) of the current word.
420    offset: usize,
421
422    /// Underlying iterator over the words.
423    iter: slice::Iter<'a, Word>,
424
425    marker: PhantomData<T>,
426}
427
428impl<'a, T: Idx> BitIter<'a, T> {
429    #[inline]
430    fn new(words: &'a [Word]) -> BitIter<'a, T> {
431        // We initialize `word` and `offset` to degenerate values. On the first
432        // call to `next()` we will fall through to getting the first word from
433        // `iter`, which sets `word` to the first word (if there is one) and
434        // `offset` to 0. Doing it this way saves us from having to maintain
435        // additional state about whether we have started.
436        BitIter {
437            word: 0,
438            offset: usize::MAX - (WORD_BITS - 1),
439            iter: words.iter(),
440            marker: PhantomData,
441        }
442    }
443}
444
445impl<'a, T: Idx> Iterator for BitIter<'a, T> {
446    type Item = T;
447    fn next(&mut self) -> Option<T> {
448        loop {
449            if self.word != 0 {
450                // Get the position of the next set bit in the current word,
451                // then clear the bit.
452                let bit_pos = self.word.trailing_zeros() as usize;
453                self.word ^= 1 << bit_pos;
454                return Some(T::new(bit_pos + self.offset));
455            }
456
457            // Move onto the next word. `wrapping_add()` is needed to handle
458            // the degenerate initial value given to `offset` in `new()`.
459            self.word = *self.iter.next()?;
460            self.offset = self.offset.wrapping_add(WORD_BITS);
461        }
462    }
463}
464
465/// A fixed-size bitset type with a partially dense, partially sparse
466/// representation. The bitset is broken into chunks, and chunks that are all
467/// zeros or all ones are represented and handled very efficiently.
468///
469/// This type is especially efficient for sets that typically have a large
470/// `domain_size` with significant stretches of all zeros or all ones, and also
471/// some stretches with lots of 0s and 1s mixed in a way that causes trouble
472/// for `IntervalSet`.
473///
474/// Best used via `MixedBitSet`, rather than directly, because `MixedBitSet`
475/// has better performance for small bitsets.
476///
477/// `T` is an index type, typically a newtyped `usize` wrapper, but it can also
478/// just be `usize`.
479///
480/// All operations that involve an element will panic if the element is equal
481/// to or greater than the domain size. All operations that involve two bitsets
482/// will panic if the bitsets have differing domain sizes.
483#[derive(PartialEq, Eq)]
484pub struct ChunkedBitSet<T> {
485    domain_size: usize,
486
487    /// The chunks. Each one contains exactly CHUNK_BITS values, except the
488    /// last one which contains 1..=CHUNK_BITS values.
489    chunks: Box<[Chunk]>,
490
491    marker: PhantomData<T>,
492}
493
494// Note: the chunk domain size is duplicated in each variant. This is a bit
495// inconvenient, but it allows the type size to be smaller than if we had an
496// outer struct containing a chunk domain size plus the `Chunk`, because the
497// compiler can place the chunk domain size after the tag.
498#[derive(Clone, Debug, PartialEq, Eq)]
499enum Chunk {
500    /// A chunk that is all zeros; we don't represent the zeros explicitly.
501    /// The `ChunkSize` is always non-zero.
502    Zeros(ChunkSize),
503
504    /// A chunk that is all ones; we don't represent the ones explicitly.
505    /// `ChunkSize` is always non-zero.
506    Ones(ChunkSize),
507
508    /// A chunk that has a mix of zeros and ones, which are represented
509    /// explicitly and densely. It never has all zeros or all ones.
510    ///
511    /// If this is the final chunk there may be excess, unused words. This
512    /// turns out to be both simpler and have better performance than
513    /// allocating the minimum number of words, largely because we avoid having
514    /// to store the length, which would make this type larger. These excess
515    /// words are always zero, as are any excess bits in the final in-use word.
516    ///
517    /// The first `ChunkSize` field is always non-zero.
518    ///
519    /// The second `ChunkSize` field is the count of 1s set in the chunk, and
520    /// must satisfy `0 < count < chunk_domain_size`.
521    ///
522    /// The words are within an `Rc` because it's surprisingly common to
523    /// duplicate an entire chunk, e.g. in `ChunkedBitSet::clone_from()`, or
524    /// when a `Mixed` chunk is union'd into a `Zeros` chunk. When we do need
525    /// to modify a chunk we use `Rc::make_mut`.
526    Mixed(ChunkSize, ChunkSize, Rc<[Word; CHUNK_WORDS]>),
527}
528
529// This type is used a lot. Make sure it doesn't unintentionally get bigger.
530#[cfg(target_pointer_width = "64")]
531crate::static_assert_size!(Chunk, 16);
532
533impl<T> ChunkedBitSet<T> {
534    pub fn domain_size(&self) -> usize {
535        self.domain_size
536    }
537
538    #[cfg(test)]
539    fn assert_valid(&self) {
540        if self.domain_size == 0 {
541            assert!(self.chunks.is_empty());
542            return;
543        }
544
545        assert!((self.chunks.len() - 1) * CHUNK_BITS <= self.domain_size);
546        assert!(self.chunks.len() * CHUNK_BITS >= self.domain_size);
547        for chunk in self.chunks.iter() {
548            chunk.assert_valid();
549        }
550    }
551}
552
553impl<T: Idx> ChunkedBitSet<T> {
554    /// Creates a new bitset with a given `domain_size` and chunk kind.
555    fn new(domain_size: usize, is_empty: bool) -> Self {
556        let chunks = if domain_size == 0 {
557            Box::new([])
558        } else {
559            // All the chunks have a chunk_domain_size of `CHUNK_BITS` except
560            // the final one.
561            let final_chunk_domain_size = {
562                let n = domain_size % CHUNK_BITS;
563                if n == 0 { CHUNK_BITS } else { n }
564            };
565            let mut chunks =
566                vec![Chunk::new(CHUNK_BITS, is_empty); num_chunks(domain_size)].into_boxed_slice();
567            *chunks.last_mut().unwrap() = Chunk::new(final_chunk_domain_size, is_empty);
568            chunks
569        };
570        ChunkedBitSet { domain_size, chunks, marker: PhantomData }
571    }
572
573    /// Creates a new, empty bitset with a given `domain_size`.
574    #[inline]
575    pub fn new_empty(domain_size: usize) -> Self {
576        ChunkedBitSet::new(domain_size, /* is_empty */ true)
577    }
578
579    /// Creates a new, filled bitset with a given `domain_size`.
580    #[inline]
581    pub fn new_filled(domain_size: usize) -> Self {
582        ChunkedBitSet::new(domain_size, /* is_empty */ false)
583    }
584
585    pub fn clear(&mut self) {
586        let domain_size = self.domain_size();
587        *self = ChunkedBitSet::new_empty(domain_size);
588    }
589
590    #[cfg(test)]
591    fn chunks(&self) -> &[Chunk] {
592        &self.chunks
593    }
594
595    /// Count the number of bits in the set.
596    pub fn count(&self) -> usize {
597        self.chunks.iter().map(|chunk| chunk.count()).sum()
598    }
599
600    pub fn is_empty(&self) -> bool {
601        self.chunks.iter().all(|chunk| matches!(chunk, Zeros(..)))
602    }
603
604    /// Returns `true` if `self` contains `elem`.
605    #[inline]
606    pub fn contains(&self, elem: T) -> bool {
607        assert!(elem.index() < self.domain_size);
608        let chunk = &self.chunks[chunk_index(elem)];
609        match &chunk {
610            Zeros(_) => false,
611            Ones(_) => true,
612            Mixed(_, _, words) => {
613                let (word_index, mask) = chunk_word_index_and_mask(elem);
614                (words[word_index] & mask) != 0
615            }
616        }
617    }
618
619    #[inline]
620    pub fn iter(&self) -> ChunkedBitIter<'_, T> {
621        ChunkedBitIter::new(self)
622    }
623
624    /// Insert `elem`. Returns whether the set has changed.
625    pub fn insert(&mut self, elem: T) -> bool {
626        assert!(elem.index() < self.domain_size);
627        let chunk_index = chunk_index(elem);
628        let chunk = &mut self.chunks[chunk_index];
629        match *chunk {
630            Zeros(chunk_domain_size) => {
631                if chunk_domain_size > 1 {
632                    #[cfg(feature = "nightly")]
633                    let mut words = {
634                        // We take some effort to avoid copying the words.
635                        let words = Rc::<[Word; CHUNK_WORDS]>::new_zeroed();
636                        // SAFETY: `words` can safely be all zeroes.
637                        unsafe { words.assume_init() }
638                    };
639                    #[cfg(not(feature = "nightly"))]
640                    let mut words = {
641                        // FIXME: unconditionally use `Rc::new_zeroed` once it is stable (#63291).
642                        let words = mem::MaybeUninit::<[Word; CHUNK_WORDS]>::zeroed();
643                        // SAFETY: `words` can safely be all zeroes.
644                        let words = unsafe { words.assume_init() };
645                        // Unfortunate possibly-large copy
646                        Rc::new(words)
647                    };
648                    let words_ref = Rc::get_mut(&mut words).unwrap();
649
650                    let (word_index, mask) = chunk_word_index_and_mask(elem);
651                    words_ref[word_index] |= mask;
652                    *chunk = Mixed(chunk_domain_size, 1, words);
653                } else {
654                    *chunk = Ones(chunk_domain_size);
655                }
656                true
657            }
658            Ones(_) => false,
659            Mixed(chunk_domain_size, ref mut count, ref mut words) => {
660                // We skip all the work if the bit is already set.
661                let (word_index, mask) = chunk_word_index_and_mask(elem);
662                if (words[word_index] & mask) == 0 {
663                    *count += 1;
664                    if *count < chunk_domain_size {
665                        let words = Rc::make_mut(words);
666                        words[word_index] |= mask;
667                    } else {
668                        *chunk = Ones(chunk_domain_size);
669                    }
670                    true
671                } else {
672                    false
673                }
674            }
675        }
676    }
677
678    /// Sets all bits to true.
679    pub fn insert_all(&mut self) {
680        for chunk in self.chunks.iter_mut() {
681            *chunk = match *chunk {
682                Zeros(chunk_domain_size)
683                | Ones(chunk_domain_size)
684                | Mixed(chunk_domain_size, ..) => Ones(chunk_domain_size),
685            }
686        }
687    }
688
689    /// Returns `true` if the set has changed.
690    pub fn remove(&mut self, elem: T) -> bool {
691        assert!(elem.index() < self.domain_size);
692        let chunk_index = chunk_index(elem);
693        let chunk = &mut self.chunks[chunk_index];
694        match *chunk {
695            Zeros(_) => false,
696            Ones(chunk_domain_size) => {
697                if chunk_domain_size > 1 {
698                    #[cfg(feature = "nightly")]
699                    let mut words = {
700                        // We take some effort to avoid copying the words.
701                        let words = Rc::<[Word; CHUNK_WORDS]>::new_zeroed();
702                        // SAFETY: `words` can safely be all zeroes.
703                        unsafe { words.assume_init() }
704                    };
705                    #[cfg(not(feature = "nightly"))]
706                    let mut words = {
707                        // FIXME: unconditionally use `Rc::new_zeroed` once it is stable (#63291).
708                        let words = mem::MaybeUninit::<[Word; CHUNK_WORDS]>::zeroed();
709                        // SAFETY: `words` can safely be all zeroes.
710                        let words = unsafe { words.assume_init() };
711                        // Unfortunate possibly-large copy
712                        Rc::new(words)
713                    };
714                    let words_ref = Rc::get_mut(&mut words).unwrap();
715
716                    // Set only the bits in use.
717                    let num_words = num_words(chunk_domain_size as usize);
718                    words_ref[..num_words].fill(!0);
719                    clear_excess_bits_in_final_word(
720                        chunk_domain_size as usize,
721                        &mut words_ref[..num_words],
722                    );
723                    let (word_index, mask) = chunk_word_index_and_mask(elem);
724                    words_ref[word_index] &= !mask;
725                    *chunk = Mixed(chunk_domain_size, chunk_domain_size - 1, words);
726                } else {
727                    *chunk = Zeros(chunk_domain_size);
728                }
729                true
730            }
731            Mixed(chunk_domain_size, ref mut count, ref mut words) => {
732                // We skip all the work if the bit is already clear.
733                let (word_index, mask) = chunk_word_index_and_mask(elem);
734                if (words[word_index] & mask) != 0 {
735                    *count -= 1;
736                    if *count > 0 {
737                        let words = Rc::make_mut(words);
738                        words[word_index] &= !mask;
739                    } else {
740                        *chunk = Zeros(chunk_domain_size);
741                    }
742                    true
743                } else {
744                    false
745                }
746            }
747        }
748    }
749
750    fn chunk_iter(&self, chunk_index: usize) -> ChunkIter<'_> {
751        match self.chunks.get(chunk_index) {
752            Some(Zeros(_chunk_domain_size)) => ChunkIter::Zeros,
753            Some(Ones(chunk_domain_size)) => ChunkIter::Ones(0..*chunk_domain_size as usize),
754            Some(Mixed(chunk_domain_size, _, words)) => {
755                let num_words = num_words(*chunk_domain_size as usize);
756                ChunkIter::Mixed(BitIter::new(&words[0..num_words]))
757            }
758            None => ChunkIter::Finished,
759        }
760    }
761
762    bit_relations_inherent_impls! {}
763}
764
765impl<T: Idx> BitRelations<ChunkedBitSet<T>> for ChunkedBitSet<T> {
766    fn union(&mut self, other: &ChunkedBitSet<T>) -> bool {
767        assert_eq!(self.domain_size, other.domain_size);
768        debug_assert_eq!(self.chunks.len(), other.chunks.len());
769
770        let mut changed = false;
771        for (mut self_chunk, other_chunk) in self.chunks.iter_mut().zip(other.chunks.iter()) {
772            match (&mut self_chunk, &other_chunk) {
773                (_, Zeros(_)) | (Ones(_), _) => {}
774                (Zeros(self_chunk_domain_size), Ones(other_chunk_domain_size))
775                | (Mixed(self_chunk_domain_size, ..), Ones(other_chunk_domain_size))
776                | (Zeros(self_chunk_domain_size), Mixed(other_chunk_domain_size, ..)) => {
777                    // `other_chunk` fully overwrites `self_chunk`
778                    debug_assert_eq!(self_chunk_domain_size, other_chunk_domain_size);
779                    *self_chunk = other_chunk.clone();
780                    changed = true;
781                }
782                (
783                    Mixed(self_chunk_domain_size, self_chunk_count, self_chunk_words),
784                    Mixed(_other_chunk_domain_size, _other_chunk_count, other_chunk_words),
785                ) => {
786                    // First check if the operation would change
787                    // `self_chunk.words`. If not, we can avoid allocating some
788                    // words, and this happens often enough that it's a
789                    // performance win. Also, we only need to operate on the
790                    // in-use words, hence the slicing.
791                    let op = |a, b| a | b;
792                    let num_words = num_words(*self_chunk_domain_size as usize);
793                    if bitwise_changes(
794                        &self_chunk_words[0..num_words],
795                        &other_chunk_words[0..num_words],
796                        op,
797                    ) {
798                        let self_chunk_words = Rc::make_mut(self_chunk_words);
799                        let has_changed = bitwise(
800                            &mut self_chunk_words[0..num_words],
801                            &other_chunk_words[0..num_words],
802                            op,
803                        );
804                        debug_assert!(has_changed);
805                        *self_chunk_count = self_chunk_words[0..num_words]
806                            .iter()
807                            .map(|w| w.count_ones() as ChunkSize)
808                            .sum();
809                        if *self_chunk_count == *self_chunk_domain_size {
810                            *self_chunk = Ones(*self_chunk_domain_size);
811                        }
812                        changed = true;
813                    }
814                }
815            }
816        }
817        changed
818    }
819
820    fn subtract(&mut self, other: &ChunkedBitSet<T>) -> bool {
821        assert_eq!(self.domain_size, other.domain_size);
822        debug_assert_eq!(self.chunks.len(), other.chunks.len());
823
824        let mut changed = false;
825        for (mut self_chunk, other_chunk) in self.chunks.iter_mut().zip(other.chunks.iter()) {
826            match (&mut self_chunk, &other_chunk) {
827                (Zeros(..), _) | (_, Zeros(..)) => {}
828                (
829                    Ones(self_chunk_domain_size) | Mixed(self_chunk_domain_size, _, _),
830                    Ones(other_chunk_domain_size),
831                ) => {
832                    debug_assert_eq!(self_chunk_domain_size, other_chunk_domain_size);
833                    changed = true;
834                    *self_chunk = Zeros(*self_chunk_domain_size);
835                }
836                (
837                    Ones(self_chunk_domain_size),
838                    Mixed(other_chunk_domain_size, other_chunk_count, other_chunk_words),
839                ) => {
840                    debug_assert_eq!(self_chunk_domain_size, other_chunk_domain_size);
841                    changed = true;
842                    let num_words = num_words(*self_chunk_domain_size as usize);
843                    debug_assert!(num_words > 0 && num_words <= CHUNK_WORDS);
844                    let mut tail_mask =
845                        1 << (*other_chunk_domain_size - ((num_words - 1) * WORD_BITS) as u16) - 1;
846                    let mut self_chunk_words = **other_chunk_words;
847                    for word in self_chunk_words[0..num_words].iter_mut().rev() {
848                        *word = !*word & tail_mask;
849                        tail_mask = u64::MAX;
850                    }
851                    let self_chunk_count = *self_chunk_domain_size - *other_chunk_count;
852                    debug_assert_eq!(
853                        self_chunk_count,
854                        self_chunk_words[0..num_words]
855                            .iter()
856                            .map(|w| w.count_ones() as ChunkSize)
857                            .sum()
858                    );
859                    *self_chunk =
860                        Mixed(*self_chunk_domain_size, self_chunk_count, Rc::new(self_chunk_words));
861                }
862                (
863                    Mixed(self_chunk_domain_size, self_chunk_count, self_chunk_words),
864                    Mixed(_other_chunk_domain_size, _other_chunk_count, other_chunk_words),
865                ) => {
866                    // See [`<Self as BitRelations<ChunkedBitSet<T>>>::union`] for the explanation
867                    let op = |a: u64, b: u64| a & !b;
868                    let num_words = num_words(*self_chunk_domain_size as usize);
869                    if bitwise_changes(
870                        &self_chunk_words[0..num_words],
871                        &other_chunk_words[0..num_words],
872                        op,
873                    ) {
874                        let self_chunk_words = Rc::make_mut(self_chunk_words);
875                        let has_changed = bitwise(
876                            &mut self_chunk_words[0..num_words],
877                            &other_chunk_words[0..num_words],
878                            op,
879                        );
880                        debug_assert!(has_changed);
881                        *self_chunk_count = self_chunk_words[0..num_words]
882                            .iter()
883                            .map(|w| w.count_ones() as ChunkSize)
884                            .sum();
885                        if *self_chunk_count == 0 {
886                            *self_chunk = Zeros(*self_chunk_domain_size);
887                        }
888                        changed = true;
889                    }
890                }
891            }
892        }
893        changed
894    }
895
896    fn intersect(&mut self, other: &ChunkedBitSet<T>) -> bool {
897        assert_eq!(self.domain_size, other.domain_size);
898        debug_assert_eq!(self.chunks.len(), other.chunks.len());
899
900        let mut changed = false;
901        for (mut self_chunk, other_chunk) in self.chunks.iter_mut().zip(other.chunks.iter()) {
902            match (&mut self_chunk, &other_chunk) {
903                (Zeros(..), _) | (_, Ones(..)) => {}
904                (
905                    Ones(self_chunk_domain_size),
906                    Zeros(other_chunk_domain_size) | Mixed(other_chunk_domain_size, ..),
907                )
908                | (Mixed(self_chunk_domain_size, ..), Zeros(other_chunk_domain_size)) => {
909                    debug_assert_eq!(self_chunk_domain_size, other_chunk_domain_size);
910                    changed = true;
911                    *self_chunk = other_chunk.clone();
912                }
913                (
914                    Mixed(self_chunk_domain_size, self_chunk_count, self_chunk_words),
915                    Mixed(_other_chunk_domain_size, _other_chunk_count, other_chunk_words),
916                ) => {
917                    // See [`<Self as BitRelations<ChunkedBitSet<T>>>::union`] for the explanation
918                    let op = |a, b| a & b;
919                    let num_words = num_words(*self_chunk_domain_size as usize);
920                    if bitwise_changes(
921                        &self_chunk_words[0..num_words],
922                        &other_chunk_words[0..num_words],
923                        op,
924                    ) {
925                        let self_chunk_words = Rc::make_mut(self_chunk_words);
926                        let has_changed = bitwise(
927                            &mut self_chunk_words[0..num_words],
928                            &other_chunk_words[0..num_words],
929                            op,
930                        );
931                        debug_assert!(has_changed);
932                        *self_chunk_count = self_chunk_words[0..num_words]
933                            .iter()
934                            .map(|w| w.count_ones() as ChunkSize)
935                            .sum();
936                        if *self_chunk_count == 0 {
937                            *self_chunk = Zeros(*self_chunk_domain_size);
938                        }
939                        changed = true;
940                    }
941                }
942            }
943        }
944
945        changed
946    }
947}
948
949impl<T: Idx> BitRelations<ChunkedBitSet<T>> for DenseBitSet<T> {
950    fn union(&mut self, other: &ChunkedBitSet<T>) -> bool {
951        sequential_update(|elem| self.insert(elem), other.iter())
952    }
953
954    fn subtract(&mut self, _other: &ChunkedBitSet<T>) -> bool {
955        unimplemented!("implement if/when necessary");
956    }
957
958    fn intersect(&mut self, other: &ChunkedBitSet<T>) -> bool {
959        assert_eq!(self.domain_size(), other.domain_size);
960        let mut changed = false;
961        for (i, chunk) in other.chunks.iter().enumerate() {
962            let mut words = &mut self.words[i * CHUNK_WORDS..];
963            if words.len() > CHUNK_WORDS {
964                words = &mut words[..CHUNK_WORDS];
965            }
966            match chunk {
967                Zeros(..) => {
968                    for word in words {
969                        if *word != 0 {
970                            changed = true;
971                            *word = 0;
972                        }
973                    }
974                }
975                Ones(..) => (),
976                Mixed(_, _, data) => {
977                    for (i, word) in words.iter_mut().enumerate() {
978                        let new_val = *word & data[i];
979                        if new_val != *word {
980                            changed = true;
981                            *word = new_val;
982                        }
983                    }
984                }
985            }
986        }
987        changed
988    }
989}
990
991impl<T> Clone for ChunkedBitSet<T> {
992    fn clone(&self) -> Self {
993        ChunkedBitSet {
994            domain_size: self.domain_size,
995            chunks: self.chunks.clone(),
996            marker: PhantomData,
997        }
998    }
999
1000    /// WARNING: this implementation of clone_from will panic if the two
1001    /// bitsets have different domain sizes. This constraint is not inherent to
1002    /// `clone_from`, but it works with the existing call sites and allows a
1003    /// faster implementation, which is important because this function is hot.
1004    fn clone_from(&mut self, from: &Self) {
1005        assert_eq!(self.domain_size, from.domain_size);
1006        debug_assert_eq!(self.chunks.len(), from.chunks.len());
1007
1008        self.chunks.clone_from(&from.chunks)
1009    }
1010}
1011
1012pub struct ChunkedBitIter<'a, T: Idx> {
1013    bit_set: &'a ChunkedBitSet<T>,
1014
1015    // The index of the current chunk.
1016    chunk_index: usize,
1017
1018    // The sub-iterator for the current chunk.
1019    chunk_iter: ChunkIter<'a>,
1020}
1021
1022impl<'a, T: Idx> ChunkedBitIter<'a, T> {
1023    #[inline]
1024    fn new(bit_set: &'a ChunkedBitSet<T>) -> ChunkedBitIter<'a, T> {
1025        ChunkedBitIter { bit_set, chunk_index: 0, chunk_iter: bit_set.chunk_iter(0) }
1026    }
1027}
1028
1029impl<'a, T: Idx> Iterator for ChunkedBitIter<'a, T> {
1030    type Item = T;
1031
1032    fn next(&mut self) -> Option<T> {
1033        loop {
1034            match &mut self.chunk_iter {
1035                ChunkIter::Zeros => {}
1036                ChunkIter::Ones(iter) => {
1037                    if let Some(next) = iter.next() {
1038                        return Some(T::new(next + self.chunk_index * CHUNK_BITS));
1039                    }
1040                }
1041                ChunkIter::Mixed(iter) => {
1042                    if let Some(next) = iter.next() {
1043                        return Some(T::new(next + self.chunk_index * CHUNK_BITS));
1044                    }
1045                }
1046                ChunkIter::Finished => return None,
1047            }
1048            self.chunk_index += 1;
1049            self.chunk_iter = self.bit_set.chunk_iter(self.chunk_index);
1050        }
1051    }
1052}
1053
1054impl Chunk {
1055    #[cfg(test)]
1056    fn assert_valid(&self) {
1057        match *self {
1058            Zeros(chunk_domain_size) | Ones(chunk_domain_size) => {
1059                assert!(chunk_domain_size as usize <= CHUNK_BITS);
1060            }
1061            Mixed(chunk_domain_size, count, ref words) => {
1062                assert!(chunk_domain_size as usize <= CHUNK_BITS);
1063                assert!(0 < count && count < chunk_domain_size);
1064
1065                // Check the number of set bits matches `count`.
1066                assert_eq!(
1067                    words.iter().map(|w| w.count_ones() as ChunkSize).sum::<ChunkSize>(),
1068                    count
1069                );
1070
1071                // Check the not-in-use words are all zeroed.
1072                let num_words = num_words(chunk_domain_size as usize);
1073                if num_words < CHUNK_WORDS {
1074                    assert_eq!(
1075                        words[num_words..]
1076                            .iter()
1077                            .map(|w| w.count_ones() as ChunkSize)
1078                            .sum::<ChunkSize>(),
1079                        0
1080                    );
1081                }
1082            }
1083        }
1084    }
1085
1086    fn new(chunk_domain_size: usize, is_empty: bool) -> Self {
1087        debug_assert!(0 < chunk_domain_size && chunk_domain_size <= CHUNK_BITS);
1088        let chunk_domain_size = chunk_domain_size as ChunkSize;
1089        if is_empty { Zeros(chunk_domain_size) } else { Ones(chunk_domain_size) }
1090    }
1091
1092    /// Count the number of 1s in the chunk.
1093    fn count(&self) -> usize {
1094        match *self {
1095            Zeros(_) => 0,
1096            Ones(chunk_domain_size) => chunk_domain_size as usize,
1097            Mixed(_, count, _) => count as usize,
1098        }
1099    }
1100}
1101
1102enum ChunkIter<'a> {
1103    Zeros,
1104    Ones(Range<usize>),
1105    Mixed(BitIter<'a, usize>),
1106    Finished,
1107}
1108
1109// Applies a function to mutate a bitset, and returns true if any
1110// of the applications return true
1111fn sequential_update<T: Idx>(
1112    mut self_update: impl FnMut(T) -> bool,
1113    it: impl Iterator<Item = T>,
1114) -> bool {
1115    it.fold(false, |changed, elem| self_update(elem) | changed)
1116}
1117
1118impl<T: Idx> fmt::Debug for ChunkedBitSet<T> {
1119    fn fmt(&self, w: &mut fmt::Formatter<'_>) -> fmt::Result {
1120        w.debug_list().entries(self.iter()).finish()
1121    }
1122}
1123
1124/// Sets `out_vec[i] = op(out_vec[i], in_vec[i])` for each index `i` in both
1125/// slices. The slices must have the same length.
1126///
1127/// Returns true if at least one bit in `out_vec` was changed.
1128///
1129/// ## Warning
1130/// Some bitwise operations (e.g. union-not, xor) can set output bits that were
1131/// unset in in both inputs. If this happens in the last word/chunk of a bitset,
1132/// it can cause the bitset to contain out-of-domain values, which need to
1133/// be cleared with `clear_excess_bits_in_final_word`. This also makes the
1134/// "changed" return value unreliable, because the change might have only
1135/// affected excess bits.
1136#[inline]
1137fn bitwise<Op>(out_vec: &mut [Word], in_vec: &[Word], op: Op) -> bool
1138where
1139    Op: Fn(Word, Word) -> Word,
1140{
1141    assert_eq!(out_vec.len(), in_vec.len());
1142    let mut changed = 0;
1143    for (out_elem, in_elem) in iter::zip(out_vec, in_vec) {
1144        let old_val = *out_elem;
1145        let new_val = op(old_val, *in_elem);
1146        *out_elem = new_val;
1147        // This is essentially equivalent to a != with changed being a bool, but
1148        // in practice this code gets auto-vectorized by the compiler for most
1149        // operators. Using != here causes us to generate quite poor code as the
1150        // compiler tries to go back to a boolean on each loop iteration.
1151        changed |= old_val ^ new_val;
1152    }
1153    changed != 0
1154}
1155
1156/// Does this bitwise operation change `out_vec`?
1157#[inline]
1158fn bitwise_changes<Op>(out_vec: &[Word], in_vec: &[Word], op: Op) -> bool
1159where
1160    Op: Fn(Word, Word) -> Word,
1161{
1162    assert_eq!(out_vec.len(), in_vec.len());
1163    for (out_elem, in_elem) in iter::zip(out_vec, in_vec) {
1164        let old_val = *out_elem;
1165        let new_val = op(old_val, *in_elem);
1166        if old_val != new_val {
1167            return true;
1168        }
1169    }
1170    false
1171}
1172
1173/// A bitset with a mixed representation, using `DenseBitSet` for small and
1174/// medium bitsets, and `ChunkedBitSet` for large bitsets, i.e. those with
1175/// enough bits for at least two chunks. This is a good choice for many bitsets
1176/// that can have large domain sizes (e.g. 5000+).
1177///
1178/// `T` is an index type, typically a newtyped `usize` wrapper, but it can also
1179/// just be `usize`.
1180///
1181/// All operations that involve an element will panic if the element is equal
1182/// to or greater than the domain size. All operations that involve two bitsets
1183/// will panic if the bitsets have differing domain sizes.
1184#[derive(PartialEq, Eq)]
1185pub enum MixedBitSet<T> {
1186    Small(DenseBitSet<T>),
1187    Large(ChunkedBitSet<T>),
1188}
1189
1190impl<T> MixedBitSet<T> {
1191    pub fn domain_size(&self) -> usize {
1192        match self {
1193            MixedBitSet::Small(set) => set.domain_size(),
1194            MixedBitSet::Large(set) => set.domain_size(),
1195        }
1196    }
1197}
1198
1199impl<T: Idx> MixedBitSet<T> {
1200    #[inline]
1201    pub fn new_empty(domain_size: usize) -> MixedBitSet<T> {
1202        if domain_size <= CHUNK_BITS {
1203            MixedBitSet::Small(DenseBitSet::new_empty(domain_size))
1204        } else {
1205            MixedBitSet::Large(ChunkedBitSet::new_empty(domain_size))
1206        }
1207    }
1208
1209    #[inline]
1210    pub fn is_empty(&self) -> bool {
1211        match self {
1212            MixedBitSet::Small(set) => set.is_empty(),
1213            MixedBitSet::Large(set) => set.is_empty(),
1214        }
1215    }
1216
1217    #[inline]
1218    pub fn contains(&self, elem: T) -> bool {
1219        match self {
1220            MixedBitSet::Small(set) => set.contains(elem),
1221            MixedBitSet::Large(set) => set.contains(elem),
1222        }
1223    }
1224
1225    #[inline]
1226    pub fn insert(&mut self, elem: T) -> bool {
1227        match self {
1228            MixedBitSet::Small(set) => set.insert(elem),
1229            MixedBitSet::Large(set) => set.insert(elem),
1230        }
1231    }
1232
1233    pub fn insert_all(&mut self) {
1234        match self {
1235            MixedBitSet::Small(set) => set.insert_all(),
1236            MixedBitSet::Large(set) => set.insert_all(),
1237        }
1238    }
1239
1240    #[inline]
1241    pub fn remove(&mut self, elem: T) -> bool {
1242        match self {
1243            MixedBitSet::Small(set) => set.remove(elem),
1244            MixedBitSet::Large(set) => set.remove(elem),
1245        }
1246    }
1247
1248    pub fn iter(&self) -> MixedBitIter<'_, T> {
1249        match self {
1250            MixedBitSet::Small(set) => MixedBitIter::Small(set.iter()),
1251            MixedBitSet::Large(set) => MixedBitIter::Large(set.iter()),
1252        }
1253    }
1254
1255    #[inline]
1256    pub fn clear(&mut self) {
1257        match self {
1258            MixedBitSet::Small(set) => set.clear(),
1259            MixedBitSet::Large(set) => set.clear(),
1260        }
1261    }
1262
1263    bit_relations_inherent_impls! {}
1264}
1265
1266impl<T> Clone for MixedBitSet<T> {
1267    fn clone(&self) -> Self {
1268        match self {
1269            MixedBitSet::Small(set) => MixedBitSet::Small(set.clone()),
1270            MixedBitSet::Large(set) => MixedBitSet::Large(set.clone()),
1271        }
1272    }
1273
1274    /// WARNING: this implementation of clone_from may panic if the two
1275    /// bitsets have different domain sizes. This constraint is not inherent to
1276    /// `clone_from`, but it works with the existing call sites and allows a
1277    /// faster implementation, which is important because this function is hot.
1278    fn clone_from(&mut self, from: &Self) {
1279        match (self, from) {
1280            (MixedBitSet::Small(set), MixedBitSet::Small(from)) => set.clone_from(from),
1281            (MixedBitSet::Large(set), MixedBitSet::Large(from)) => set.clone_from(from),
1282            _ => panic!("MixedBitSet size mismatch"),
1283        }
1284    }
1285}
1286
1287impl<T: Idx> BitRelations<MixedBitSet<T>> for MixedBitSet<T> {
1288    fn union(&mut self, other: &MixedBitSet<T>) -> bool {
1289        match (self, other) {
1290            (MixedBitSet::Small(set), MixedBitSet::Small(other)) => set.union(other),
1291            (MixedBitSet::Large(set), MixedBitSet::Large(other)) => set.union(other),
1292            _ => panic!("MixedBitSet size mismatch"),
1293        }
1294    }
1295
1296    fn subtract(&mut self, other: &MixedBitSet<T>) -> bool {
1297        match (self, other) {
1298            (MixedBitSet::Small(set), MixedBitSet::Small(other)) => set.subtract(other),
1299            (MixedBitSet::Large(set), MixedBitSet::Large(other)) => set.subtract(other),
1300            _ => panic!("MixedBitSet size mismatch"),
1301        }
1302    }
1303
1304    fn intersect(&mut self, _other: &MixedBitSet<T>) -> bool {
1305        unimplemented!("implement if/when necessary");
1306    }
1307}
1308
1309impl<T: Idx> fmt::Debug for MixedBitSet<T> {
1310    fn fmt(&self, w: &mut fmt::Formatter<'_>) -> fmt::Result {
1311        match self {
1312            MixedBitSet::Small(set) => set.fmt(w),
1313            MixedBitSet::Large(set) => set.fmt(w),
1314        }
1315    }
1316}
1317
1318pub enum MixedBitIter<'a, T: Idx> {
1319    Small(BitIter<'a, T>),
1320    Large(ChunkedBitIter<'a, T>),
1321}
1322
1323impl<'a, T: Idx> Iterator for MixedBitIter<'a, T> {
1324    type Item = T;
1325    fn next(&mut self) -> Option<T> {
1326        match self {
1327            MixedBitIter::Small(iter) => iter.next(),
1328            MixedBitIter::Large(iter) => iter.next(),
1329        }
1330    }
1331}
1332
1333/// A resizable bitset type with a dense representation.
1334///
1335/// `T` is an index type, typically a newtyped `usize` wrapper, but it can also
1336/// just be `usize`.
1337///
1338/// All operations that involve an element will panic if the element is equal
1339/// to or greater than the domain size.
1340#[derive(Clone, Debug, PartialEq)]
1341pub struct GrowableBitSet<T: Idx> {
1342    bit_set: DenseBitSet<T>,
1343}
1344
1345impl<T: Idx> Default for GrowableBitSet<T> {
1346    fn default() -> Self {
1347        GrowableBitSet::new_empty()
1348    }
1349}
1350
1351impl<T: Idx> GrowableBitSet<T> {
1352    /// Ensure that the set can hold at least `min_domain_size` elements.
1353    pub fn ensure(&mut self, min_domain_size: usize) {
1354        if self.bit_set.domain_size < min_domain_size {
1355            self.bit_set.domain_size = min_domain_size;
1356        }
1357
1358        let min_num_words = num_words(min_domain_size);
1359        if self.bit_set.words.len() < min_num_words {
1360            self.bit_set.words.resize(min_num_words, 0)
1361        }
1362    }
1363
1364    pub fn new_empty() -> GrowableBitSet<T> {
1365        GrowableBitSet { bit_set: DenseBitSet::new_empty(0) }
1366    }
1367
1368    pub fn with_capacity(capacity: usize) -> GrowableBitSet<T> {
1369        GrowableBitSet { bit_set: DenseBitSet::new_empty(capacity) }
1370    }
1371
1372    /// Returns `true` if the set has changed.
1373    #[inline]
1374    pub fn insert(&mut self, elem: T) -> bool {
1375        self.ensure(elem.index() + 1);
1376        self.bit_set.insert(elem)
1377    }
1378
1379    /// Returns `true` if the set has changed.
1380    #[inline]
1381    pub fn remove(&mut self, elem: T) -> bool {
1382        self.ensure(elem.index() + 1);
1383        self.bit_set.remove(elem)
1384    }
1385
1386    #[inline]
1387    pub fn is_empty(&self) -> bool {
1388        self.bit_set.is_empty()
1389    }
1390
1391    #[inline]
1392    pub fn contains(&self, elem: T) -> bool {
1393        let (word_index, mask) = word_index_and_mask(elem);
1394        self.bit_set.words.get(word_index).is_some_and(|word| (word & mask) != 0)
1395    }
1396
1397    #[inline]
1398    pub fn iter(&self) -> BitIter<'_, T> {
1399        self.bit_set.iter()
1400    }
1401
1402    #[inline]
1403    pub fn len(&self) -> usize {
1404        self.bit_set.count()
1405    }
1406}
1407
1408impl<T: Idx> From<DenseBitSet<T>> for GrowableBitSet<T> {
1409    fn from(bit_set: DenseBitSet<T>) -> Self {
1410        Self { bit_set }
1411    }
1412}
1413
1414/// A fixed-size 2D bit matrix type with a dense representation.
1415///
1416/// `R` and `C` are index types used to identify rows and columns respectively;
1417/// typically newtyped `usize` wrappers, but they can also just be `usize`.
1418///
1419/// All operations that involve a row and/or column index will panic if the
1420/// index exceeds the relevant bound.
1421#[cfg_attr(feature = "nightly", derive(Decodable_NoContext, Encodable_NoContext))]
1422#[derive(Clone, Eq, PartialEq, Hash)]
1423pub struct BitMatrix<R: Idx, C: Idx> {
1424    num_rows: usize,
1425    num_columns: usize,
1426    words: SmallVec<[Word; 2]>,
1427    marker: PhantomData<(R, C)>,
1428}
1429
1430impl<R: Idx, C: Idx> BitMatrix<R, C> {
1431    /// Creates a new `rows x columns` matrix, initially empty.
1432    pub fn new(num_rows: usize, num_columns: usize) -> BitMatrix<R, C> {
1433        // For every element, we need one bit for every other
1434        // element. Round up to an even number of words.
1435        let words_per_row = num_words(num_columns);
1436        BitMatrix {
1437            num_rows,
1438            num_columns,
1439            words: smallvec![0; num_rows * words_per_row],
1440            marker: PhantomData,
1441        }
1442    }
1443
1444    /// Creates a new matrix, with `row` used as the value for every row.
1445    pub fn from_row_n(row: &DenseBitSet<C>, num_rows: usize) -> BitMatrix<R, C> {
1446        let num_columns = row.domain_size();
1447        let words_per_row = num_words(num_columns);
1448        assert_eq!(words_per_row, row.words.len());
1449        BitMatrix {
1450            num_rows,
1451            num_columns,
1452            words: iter::repeat(&row.words).take(num_rows).flatten().cloned().collect(),
1453            marker: PhantomData,
1454        }
1455    }
1456
1457    pub fn rows(&self) -> impl Iterator<Item = R> {
1458        (0..self.num_rows).map(R::new)
1459    }
1460
1461    /// The range of bits for a given row.
1462    fn range(&self, row: R) -> (usize, usize) {
1463        let words_per_row = num_words(self.num_columns);
1464        let start = row.index() * words_per_row;
1465        (start, start + words_per_row)
1466    }
1467
1468    /// Sets the cell at `(row, column)` to true. Put another way, insert
1469    /// `column` to the bitset for `row`.
1470    ///
1471    /// Returns `true` if this changed the matrix.
1472    pub fn insert(&mut self, row: R, column: C) -> bool {
1473        assert!(row.index() < self.num_rows && column.index() < self.num_columns);
1474        let (start, _) = self.range(row);
1475        let (word_index, mask) = word_index_and_mask(column);
1476        let words = &mut self.words[..];
1477        let word = words[start + word_index];
1478        let new_word = word | mask;
1479        words[start + word_index] = new_word;
1480        word != new_word
1481    }
1482
1483    /// Do the bits from `row` contain `column`? Put another way, is
1484    /// the matrix cell at `(row, column)` true?  Put yet another way,
1485    /// if the matrix represents (transitive) reachability, can
1486    /// `row` reach `column`?
1487    pub fn contains(&self, row: R, column: C) -> bool {
1488        assert!(row.index() < self.num_rows && column.index() < self.num_columns);
1489        let (start, _) = self.range(row);
1490        let (word_index, mask) = word_index_and_mask(column);
1491        (self.words[start + word_index] & mask) != 0
1492    }
1493
1494    /// Returns those indices that are true in rows `a` and `b`. This
1495    /// is an *O*(*n*) operation where *n* is the number of elements
1496    /// (somewhat independent from the actual size of the
1497    /// intersection, in particular).
1498    pub fn intersect_rows(&self, row1: R, row2: R) -> Vec<C> {
1499        assert!(row1.index() < self.num_rows && row2.index() < self.num_rows);
1500        let (row1_start, row1_end) = self.range(row1);
1501        let (row2_start, row2_end) = self.range(row2);
1502        let mut result = Vec::with_capacity(self.num_columns);
1503        for (base, (i, j)) in (row1_start..row1_end).zip(row2_start..row2_end).enumerate() {
1504            let mut v = self.words[i] & self.words[j];
1505            for bit in 0..WORD_BITS {
1506                if v == 0 {
1507                    break;
1508                }
1509                if v & 0x1 != 0 {
1510                    result.push(C::new(base * WORD_BITS + bit));
1511                }
1512                v >>= 1;
1513            }
1514        }
1515        result
1516    }
1517
1518    /// Adds the bits from row `read` to the bits from row `write`, and
1519    /// returns `true` if anything changed.
1520    ///
1521    /// This is used when computing transitive reachability because if
1522    /// you have an edge `write -> read`, because in that case
1523    /// `write` can reach everything that `read` can (and
1524    /// potentially more).
1525    pub fn union_rows(&mut self, read: R, write: R) -> bool {
1526        assert!(read.index() < self.num_rows && write.index() < self.num_rows);
1527        let (read_start, read_end) = self.range(read);
1528        let (write_start, write_end) = self.range(write);
1529        let words = &mut self.words[..];
1530        let mut changed = 0;
1531        for (read_index, write_index) in iter::zip(read_start..read_end, write_start..write_end) {
1532            let word = words[write_index];
1533            let new_word = word | words[read_index];
1534            words[write_index] = new_word;
1535            // See `bitwise` for the rationale.
1536            changed |= word ^ new_word;
1537        }
1538        changed != 0
1539    }
1540
1541    /// Adds the bits from `with` to the bits from row `write`, and
1542    /// returns `true` if anything changed.
1543    pub fn union_row_with(&mut self, with: &DenseBitSet<C>, write: R) -> bool {
1544        assert!(write.index() < self.num_rows);
1545        assert_eq!(with.domain_size(), self.num_columns);
1546        let (write_start, write_end) = self.range(write);
1547        bitwise(&mut self.words[write_start..write_end], &with.words, |a, b| a | b)
1548    }
1549
1550    /// Sets every cell in `row` to true.
1551    pub fn insert_all_into_row(&mut self, row: R) {
1552        assert!(row.index() < self.num_rows);
1553        let (start, end) = self.range(row);
1554        let words = &mut self.words[..];
1555        for index in start..end {
1556            words[index] = !0;
1557        }
1558        clear_excess_bits_in_final_word(self.num_columns, &mut self.words[..end]);
1559    }
1560
1561    /// Gets a slice of the underlying words.
1562    pub fn words(&self) -> &[Word] {
1563        &self.words
1564    }
1565
1566    /// Iterates through all the columns set to true in a given row of
1567    /// the matrix.
1568    pub fn iter(&self, row: R) -> BitIter<'_, C> {
1569        assert!(row.index() < self.num_rows);
1570        let (start, end) = self.range(row);
1571        BitIter::new(&self.words[start..end])
1572    }
1573
1574    /// Returns the number of elements in `row`.
1575    pub fn count(&self, row: R) -> usize {
1576        let (start, end) = self.range(row);
1577        self.words[start..end].iter().map(|e| e.count_ones() as usize).sum()
1578    }
1579}
1580
1581impl<R: Idx, C: Idx> fmt::Debug for BitMatrix<R, C> {
1582    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
1583        /// Forces its contents to print in regular mode instead of alternate mode.
1584        struct OneLinePrinter<T>(T);
1585        impl<T: fmt::Debug> fmt::Debug for OneLinePrinter<T> {
1586            fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
1587                write!(fmt, "{:?}", self.0)
1588            }
1589        }
1590
1591        write!(fmt, "BitMatrix({}x{}) ", self.num_rows, self.num_columns)?;
1592        let items = self.rows().flat_map(|r| self.iter(r).map(move |c| (r, c)));
1593        fmt.debug_set().entries(items.map(OneLinePrinter)).finish()
1594    }
1595}
1596
1597/// A fixed-column-size, variable-row-size 2D bit matrix with a moderately
1598/// sparse representation.
1599///
1600/// Initially, every row has no explicit representation. If any bit within a row
1601/// is set, the entire row is instantiated as `Some(<DenseBitSet>)`.
1602/// Furthermore, any previously uninstantiated rows prior to it will be
1603/// instantiated as `None`. Those prior rows may themselves become fully
1604/// instantiated later on if any of their bits are set.
1605///
1606/// `R` and `C` are index types used to identify rows and columns respectively;
1607/// typically newtyped `usize` wrappers, but they can also just be `usize`.
1608#[derive(Clone, Debug)]
1609pub struct SparseBitMatrix<R, C>
1610where
1611    R: Idx,
1612    C: Idx,
1613{
1614    num_columns: usize,
1615    rows: IndexVec<R, Option<DenseBitSet<C>>>,
1616}
1617
1618impl<R: Idx, C: Idx> SparseBitMatrix<R, C> {
1619    /// Creates a new empty sparse bit matrix with no rows or columns.
1620    pub fn new(num_columns: usize) -> Self {
1621        Self { num_columns, rows: IndexVec::new() }
1622    }
1623
1624    fn ensure_row(&mut self, row: R) -> &mut DenseBitSet<C> {
1625        // Instantiate any missing rows up to and including row `row` with an empty `DenseBitSet`.
1626        // Then replace row `row` with a full `DenseBitSet` if necessary.
1627        self.rows.get_or_insert_with(row, || DenseBitSet::new_empty(self.num_columns))
1628    }
1629
1630    /// Sets the cell at `(row, column)` to true. Put another way, insert
1631    /// `column` to the bitset for `row`.
1632    ///
1633    /// Returns `true` if this changed the matrix.
1634    pub fn insert(&mut self, row: R, column: C) -> bool {
1635        self.ensure_row(row).insert(column)
1636    }
1637
1638    /// Sets the cell at `(row, column)` to false. Put another way, delete
1639    /// `column` from the bitset for `row`. Has no effect if `row` does not
1640    /// exist.
1641    ///
1642    /// Returns `true` if this changed the matrix.
1643    pub fn remove(&mut self, row: R, column: C) -> bool {
1644        match self.rows.get_mut(row) {
1645            Some(Some(row)) => row.remove(column),
1646            _ => false,
1647        }
1648    }
1649
1650    /// Sets all columns at `row` to false. Has no effect if `row` does
1651    /// not exist.
1652    pub fn clear(&mut self, row: R) {
1653        if let Some(Some(row)) = self.rows.get_mut(row) {
1654            row.clear();
1655        }
1656    }
1657
1658    /// Do the bits from `row` contain `column`? Put another way, is
1659    /// the matrix cell at `(row, column)` true?  Put yet another way,
1660    /// if the matrix represents (transitive) reachability, can
1661    /// `row` reach `column`?
1662    pub fn contains(&self, row: R, column: C) -> bool {
1663        self.row(row).is_some_and(|r| r.contains(column))
1664    }
1665
1666    /// Adds the bits from row `read` to the bits from row `write`, and
1667    /// returns `true` if anything changed.
1668    ///
1669    /// This is used when computing transitive reachability because if
1670    /// you have an edge `write -> read`, because in that case
1671    /// `write` can reach everything that `read` can (and
1672    /// potentially more).
1673    pub fn union_rows(&mut self, read: R, write: R) -> bool {
1674        if read == write || self.row(read).is_none() {
1675            return false;
1676        }
1677
1678        self.ensure_row(write);
1679        if let (Some(read_row), Some(write_row)) = self.rows.pick2_mut(read, write) {
1680            write_row.union(read_row)
1681        } else {
1682            unreachable!()
1683        }
1684    }
1685
1686    /// Insert all bits in the given row.
1687    pub fn insert_all_into_row(&mut self, row: R) {
1688        self.ensure_row(row).insert_all();
1689    }
1690
1691    pub fn rows(&self) -> impl Iterator<Item = R> {
1692        self.rows.indices()
1693    }
1694
1695    /// Iterates through all the columns set to true in a given row of
1696    /// the matrix.
1697    pub fn iter(&self, row: R) -> impl Iterator<Item = C> {
1698        self.row(row).into_iter().flat_map(|r| r.iter())
1699    }
1700
1701    pub fn row(&self, row: R) -> Option<&DenseBitSet<C>> {
1702        self.rows.get(row)?.as_ref()
1703    }
1704
1705    /// Intersects `row` with `set`. `set` can be either `DenseBitSet` or
1706    /// `ChunkedBitSet`. Has no effect if `row` does not exist.
1707    ///
1708    /// Returns true if the row was changed.
1709    pub fn intersect_row<Set>(&mut self, row: R, set: &Set) -> bool
1710    where
1711        DenseBitSet<C>: BitRelations<Set>,
1712    {
1713        match self.rows.get_mut(row) {
1714            Some(Some(row)) => row.intersect(set),
1715            _ => false,
1716        }
1717    }
1718
1719    /// Subtracts `set` from `row`. `set` can be either `DenseBitSet` or
1720    /// `ChunkedBitSet`. Has no effect if `row` does not exist.
1721    ///
1722    /// Returns true if the row was changed.
1723    pub fn subtract_row<Set>(&mut self, row: R, set: &Set) -> bool
1724    where
1725        DenseBitSet<C>: BitRelations<Set>,
1726    {
1727        match self.rows.get_mut(row) {
1728            Some(Some(row)) => row.subtract(set),
1729            _ => false,
1730        }
1731    }
1732
1733    /// Unions `row` with `set`. `set` can be either `DenseBitSet` or
1734    /// `ChunkedBitSet`.
1735    ///
1736    /// Returns true if the row was changed.
1737    pub fn union_row<Set>(&mut self, row: R, set: &Set) -> bool
1738    where
1739        DenseBitSet<C>: BitRelations<Set>,
1740    {
1741        self.ensure_row(row).union(set)
1742    }
1743}
1744
1745#[inline]
1746fn num_words<T: Idx>(domain_size: T) -> usize {
1747    (domain_size.index() + WORD_BITS - 1) / WORD_BITS
1748}
1749
1750#[inline]
1751fn num_chunks<T: Idx>(domain_size: T) -> usize {
1752    assert!(domain_size.index() > 0);
1753    (domain_size.index() + CHUNK_BITS - 1) / CHUNK_BITS
1754}
1755
1756#[inline]
1757fn word_index_and_mask<T: Idx>(elem: T) -> (usize, Word) {
1758    let elem = elem.index();
1759    let word_index = elem / WORD_BITS;
1760    let mask = 1 << (elem % WORD_BITS);
1761    (word_index, mask)
1762}
1763
1764#[inline]
1765fn chunk_index<T: Idx>(elem: T) -> usize {
1766    elem.index() / CHUNK_BITS
1767}
1768
1769#[inline]
1770fn chunk_word_index_and_mask<T: Idx>(elem: T) -> (usize, Word) {
1771    let chunk_elem = elem.index() % CHUNK_BITS;
1772    word_index_and_mask(chunk_elem)
1773}
1774
1775fn clear_excess_bits_in_final_word(domain_size: usize, words: &mut [Word]) {
1776    let num_bits_in_final_word = domain_size % WORD_BITS;
1777    if num_bits_in_final_word > 0 {
1778        let mask = (1 << num_bits_in_final_word) - 1;
1779        words[words.len() - 1] &= mask;
1780    }
1781}
1782
1783#[inline]
1784fn max_bit(word: Word) -> usize {
1785    WORD_BITS - 1 - word.leading_zeros() as usize
1786}
1787
1788/// Integral type used to represent the bit set.
1789pub trait FiniteBitSetTy:
1790    BitAnd<Output = Self>
1791    + BitAndAssign
1792    + BitOrAssign
1793    + Clone
1794    + Copy
1795    + Shl
1796    + Not<Output = Self>
1797    + PartialEq
1798    + Sized
1799{
1800    /// Size of the domain representable by this type, e.g. 64 for `u64`.
1801    const DOMAIN_SIZE: u32;
1802
1803    /// Value which represents the `FiniteBitSet` having every bit set.
1804    const FILLED: Self;
1805    /// Value which represents the `FiniteBitSet` having no bits set.
1806    const EMPTY: Self;
1807
1808    /// Value for one as the integral type.
1809    const ONE: Self;
1810    /// Value for zero as the integral type.
1811    const ZERO: Self;
1812
1813    /// Perform a checked left shift on the integral type.
1814    fn checked_shl(self, rhs: u32) -> Option<Self>;
1815    /// Perform a checked right shift on the integral type.
1816    fn checked_shr(self, rhs: u32) -> Option<Self>;
1817}
1818
1819impl FiniteBitSetTy for u32 {
1820    const DOMAIN_SIZE: u32 = 32;
1821
1822    const FILLED: Self = Self::MAX;
1823    const EMPTY: Self = Self::MIN;
1824
1825    const ONE: Self = 1u32;
1826    const ZERO: Self = 0u32;
1827
1828    fn checked_shl(self, rhs: u32) -> Option<Self> {
1829        self.checked_shl(rhs)
1830    }
1831
1832    fn checked_shr(self, rhs: u32) -> Option<Self> {
1833        self.checked_shr(rhs)
1834    }
1835}
1836
1837impl std::fmt::Debug for FiniteBitSet<u32> {
1838    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1839        write!(f, "{:032b}", self.0)
1840    }
1841}
1842
1843/// A fixed-sized bitset type represented by an integer type. Indices outwith than the range
1844/// representable by `T` are considered set.
1845#[cfg_attr(feature = "nightly", derive(Decodable_NoContext, Encodable_NoContext))]
1846#[derive(Copy, Clone, Eq, PartialEq)]
1847pub struct FiniteBitSet<T: FiniteBitSetTy>(pub T);
1848
1849impl<T: FiniteBitSetTy> FiniteBitSet<T> {
1850    /// Creates a new, empty bitset.
1851    pub fn new_empty() -> Self {
1852        Self(T::EMPTY)
1853    }
1854
1855    /// Sets the `index`th bit.
1856    pub fn set(&mut self, index: u32) {
1857        self.0 |= T::ONE.checked_shl(index).unwrap_or(T::ZERO);
1858    }
1859
1860    /// Unsets the `index`th bit.
1861    pub fn clear(&mut self, index: u32) {
1862        self.0 &= !T::ONE.checked_shl(index).unwrap_or(T::ZERO);
1863    }
1864
1865    /// Sets the `i`th to `j`th bits.
1866    pub fn set_range(&mut self, range: Range<u32>) {
1867        let bits = T::FILLED
1868            .checked_shl(range.end - range.start)
1869            .unwrap_or(T::ZERO)
1870            .not()
1871            .checked_shl(range.start)
1872            .unwrap_or(T::ZERO);
1873        self.0 |= bits;
1874    }
1875
1876    /// Is the set empty?
1877    pub fn is_empty(&self) -> bool {
1878        self.0 == T::EMPTY
1879    }
1880
1881    /// Returns the domain size of the bitset.
1882    pub fn within_domain(&self, index: u32) -> bool {
1883        index < T::DOMAIN_SIZE
1884    }
1885
1886    /// Returns if the `index`th bit is set.
1887    pub fn contains(&self, index: u32) -> Option<bool> {
1888        self.within_domain(index)
1889            .then(|| ((self.0.checked_shr(index).unwrap_or(T::ONE)) & T::ONE) == T::ONE)
1890    }
1891}
1892
1893impl<T: FiniteBitSetTy> Default for FiniteBitSet<T> {
1894    fn default() -> Self {
1895        Self::new_empty()
1896    }
1897}