ra_ap_rustc_index/
bit_set.rs

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