Skip to main content

solar_data_structures/
bit_set.rs

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