[EXPERIMENT] RawBitSet experiments

compiler/rustc_index/src/bit_set.rs

@@ -1,14 +1,16 @@
 use std::marker::PhantomData;
 use std::ops::{Bound, Range, RangeBounds};
 use std::rc::Rc;
-use std::{fmt, iter, slice};
+use std::{fmt, iter};
 
 use Chunk::*;
 #[cfg(feature = "nightly")]
 use rustc_macros::{Decodable_NoContext, Encodable_NoContext};
 
+use crate::bit_set::raw::RawBitIter;
 use crate::{Idx, IndexVec};
 
+mod raw;
 #[cfg(test)]
 mod tests;
 
@@ -162,30 +164,10 @@ impl<T: Idx> DenseBitSet<T> {
     }
 
     #[inline]
-    pub fn insert_range(&mut self, elems: impl RangeBounds<T>) {
-        let Some((start, end)) = inclusive_start_end(elems, self.domain_size) else {
-            return;
-        };
-
-        let (start_word_index, start_mask) = word_index_and_mask(start);
-        let (end_word_index, end_mask) = word_index_and_mask(end);
-
-        // Set all words in between start and end (exclusively of both).
-        for word_index in (start_word_index + 1)..end_word_index {
-            self.words[word_index] = !0;
-        }
-
-        if start_word_index != end_word_index {
-            // Start and end are in different words, so we handle each in turn.
-            //
-            // We set all leading bits. This includes the start_mask bit.
-            self.words[start_word_index] |= !(start_mask - 1);
-            // And all trailing bits (i.e. from 0..=end) in the end word,
-            // including the end.
-            self.words[end_word_index] |= end_mask | (end_mask - 1);
-        } else {
-            self.words[start_word_index] |= end_mask | (end_mask - start_mask);
-        }
+    pub fn insert_range(&mut self, range: impl RangeBounds<T>) {
+        let start = range.start_bound().map(|i| i.index());
+        let end = range.end_bound().map(|i| i.index());
+        raw::insert_range(self.domain_size, &mut self.words, (start, end));
     }
 
     /// Sets all bits to true.
@@ -196,28 +178,10 @@ impl<T: Idx> DenseBitSet<T> {
 
     /// Checks whether any bit in the given range is a 1.
     #[inline]
-    pub fn contains_any(&self, elems: impl RangeBounds<T>) -> bool {
-        let Some((start, end)) = inclusive_start_end(elems, self.domain_size) else {
-            return false;
-        };
-        let (start_word_index, start_mask) = word_index_and_mask(start);
-        let (end_word_index, end_mask) = word_index_and_mask(end);
-
-        if start_word_index == end_word_index {
-            self.words[start_word_index] & (end_mask | (end_mask - start_mask)) != 0
-        } else {
-            if self.words[start_word_index] & !(start_mask - 1) != 0 {
-                return true;
-            }
-
-            let remaining = start_word_index + 1..end_word_index;
-            if remaining.start <= remaining.end {
-                self.words[remaining].iter().any(|&w| w != 0)
-                    || self.words[end_word_index] & (end_mask | (end_mask - 1)) != 0
-            } else {
-                false
-            }
-        }
+    pub fn contains_any(&self, range: impl RangeBounds<T>) -> bool {
+        let start = range.start_bound().map(|i| i.index());
+        let end = range.end_bound().map(|i| i.index());
+        raw::contains_any(self.domain_size, &self.words, (start, end))
     }
 
     /// Returns `true` if the set has changed.
@@ -254,33 +218,9 @@ impl<T: Idx> DenseBitSet<T> {
     }
 
     pub fn last_set_in(&self, range: impl RangeBounds<T>) -> Option<T> {
-        let (start, end) = inclusive_start_end(range, self.domain_size)?;
-        let (start_word_index, _) = word_index_and_mask(start);
-        let (end_word_index, end_mask) = word_index_and_mask(end);
-
-        let end_word = self.words[end_word_index] & (end_mask | (end_mask - 1));
-        if end_word != 0 {
-            let pos = max_bit(end_word) + WORD_BITS * end_word_index;
-            if start <= pos {
-                return Some(T::new(pos));
-            }
-        }
-
-        // We exclude end_word_index from the range here, because we don't want
-        // to limit ourselves to *just* the last word: the bits set it in may be
-        // after `end`, so it may not work out.
-        if let Some(offset) =
-            self.words[start_word_index..end_word_index].iter().rposition(|&w| w != 0)
-        {
-            let word_idx = start_word_index + offset;
-            let start_word = self.words[word_idx];
-            let pos = max_bit(start_word) + WORD_BITS * word_idx;
-            if start <= pos {
-                return Some(T::new(pos));
-            }
-        }
-
-        None
+        let start = range.start_bound().map(|i| i.index());
+        let end = range.end_bound().map(|i| i.index());
+        raw::last_set_in(self.domain_size, &self.words, (start, end)).map(T::new)
     }
 
     /// Sets `self = self | !other`.
@@ -381,54 +321,22 @@ impl<T: Idx> ToString for DenseBitSet<T> {
 }
 
 pub struct BitIter<'a, T: Idx> {
-    /// A copy of the current word, but with any already-visited bits cleared.
-    /// (This lets us use `trailing_zeros()` to find the next set bit.) When it
-    /// is reduced to 0, we move onto the next word.
-    word: Word,
-
-    /// The offset (measured in bits) of the current word.
-    offset: usize,
-
-    /// Underlying iterator over the words.
-    iter: slice::Iter<'a, Word>,
-
+    raw: RawBitIter<'a>,
     marker: PhantomData<T>,
 }
 
 impl<'a, T: Idx> BitIter<'a, T> {
-    #[inline]
-    fn new(words: &'a [Word]) -> BitIter<'a, T> {
-        // We initialize `word` and `offset` to degenerate values. On the first
-        // call to `next()` we will fall through to getting the first word from
-        // `iter`, which sets `word` to the first word (if there is one) and
-        // `offset` to 0. Doing it this way saves us from having to maintain
-        // additional state about whether we have started.
-        BitIter {
-            word: 0,
-            offset: usize::MAX - (WORD_BITS - 1),
-            iter: words.iter(),
-            marker: PhantomData,
-        }
+    #[inline(always)]
+    fn new(words: &'a [Word]) -> Self {
+        BitIter { raw: RawBitIter::new(words), marker: PhantomData }
     }
 }
 
 impl<'a, T: Idx> Iterator for BitIter<'a, T> {
     type Item = T;
-    fn next(&mut self) -> Option<T> {
-        loop {
-            if self.word != 0 {
-                // Get the position of the next set bit in the current word,
-                // then clear the bit.
-                let bit_pos = self.word.trailing_zeros() as usize;
-                self.word ^= 1 << bit_pos;
-                return Some(T::new(bit_pos + self.offset));
-            }
 
-            // Move onto the next word. `wrapping_add()` is needed to handle
-            // the degenerate initial value given to `offset` in `new()`.
-            self.word = *self.iter.next()?;
-            self.offset = self.offset.wrapping_add(WORD_BITS);
-        }
+    fn next(&mut self) -> Option<Self::Item> {
+        self.raw.next().map(T::new)
     }
 }
 
@@ -713,7 +621,7 @@ impl<T: Idx> ChunkedBitSet<T> {
             Some(Ones { chunk_domain_size }) => ChunkIter::Ones(0..*chunk_domain_size as usize),
             Some(Mixed { chunk_domain_size, words, .. }) => {
                 let num_words = num_words(*chunk_domain_size as usize);
-                ChunkIter::Mixed(BitIter::new(&words[0..num_words]))
+                ChunkIter::Mixed(RawBitIter::new(&words[0..num_words]))
             }
             None => ChunkIter::Finished,
         }
@@ -963,7 +871,7 @@ impl Chunk {
 enum ChunkIter<'a> {
     Zeros,
     Ones(Range<usize>),
-    Mixed(BitIter<'a, usize>),
+    Mixed(RawBitIter<'a>),
     Finished,
 }
 

compiler/rustc_index/src/bit_set/raw.rs

@@ -0,0 +1,151 @@
+use std::ops::Bound;
+use std::slice;
+
+use crate::bit_set::{WORD_BITS, Word, inclusive_start_end, max_bit, word_index_and_mask};
+
+#[inline]
+pub(crate) fn contains_any(
+    domain_size: usize,
+    words: &[Word],
+    range: (Bound<usize>, Bound<usize>),
+) -> bool {
+    let Some((start, end)) = inclusive_start_end(range, domain_size) else {
+        return false;
+    };
+
+    let (start_word_index, start_mask) = word_index_and_mask(start);
+    let (end_word_index, end_mask) = word_index_and_mask(end);
+
+    if start_word_index == end_word_index {
+        words[start_word_index] & (end_mask | (end_mask - start_mask)) != 0
+    } else {
+        if words[start_word_index] & !(start_mask - 1) != 0 {
+            return true;
+        }
+
+        let remaining = start_word_index + 1..end_word_index;
+        if remaining.start <= remaining.end {
+            words[remaining].iter().any(|&w| w != 0)
+                || words[end_word_index] & (end_mask | (end_mask - 1)) != 0
+        } else {
+            false
+        }
+    }
+}
+
+#[inline]
+pub(crate) fn last_set_in(
+    domain_size: usize,
+    words: &[Word],
+    range: (Bound<usize>, Bound<usize>),
+) -> Option<usize> {
+    let (start, end) = inclusive_start_end(range, domain_size)?;
+
+    let (start_word_index, _) = word_index_and_mask(start);
+    let (end_word_index, end_mask) = word_index_and_mask(end);
+
+    let end_word = words[end_word_index] & (end_mask | (end_mask - 1));
+    if end_word != 0 {
+        let pos = max_bit(end_word) + WORD_BITS * end_word_index;
+        if start <= pos {
+            return Some(pos);
+        }
+    }
+
+    // We exclude end_word_index from the range here, because we don't want
+    // to limit ourselves to *just* the last word: the bits set it in may be
+    // after `end`, so it may not work out.
+    if let Some(offset) = words[start_word_index..end_word_index].iter().rposition(|&w| w != 0) {
+        let word_idx = start_word_index + offset;
+        let start_word = words[word_idx];
+        let pos = max_bit(start_word) + WORD_BITS * word_idx;
+        if start <= pos {
+            return Some(pos);
+        }
+    }
+
+    None
+}
+
+#[inline]
+pub(crate) fn insert_range(
+    domain_size: usize,
+    words: &mut [Word],
+    range: (Bound<usize>, Bound<usize>),
+) {
+    let Some((start, end)) = inclusive_start_end(range, domain_size) else {
+        return;
+    };
+
+    let (start_word_index, start_mask) = word_index_and_mask(start);
+    let (end_word_index, end_mask) = word_index_and_mask(end);
+
+    // Set all words in between start and end (exclusively of both).
+    for word_index in (start_word_index + 1)..end_word_index {
+        words[word_index] = !0;
+    }
+
+    if start_word_index != end_word_index {
+        // Start and end are in different words, so we handle each in turn.
+        //
+        // We set all leading bits. This includes the start_mask bit.
+        words[start_word_index] |= !(start_mask - 1);
+        // And all trailing bits (i.e. from 0..=end) in the end word,
+        // including the end.
+        words[end_word_index] |= end_mask | (end_mask - 1);
+    } else {
+        words[start_word_index] |= end_mask | (end_mask - start_mask);
+    }
+}
+
+pub(crate) struct RawBitIter<'a> {
+    /// A copy of the current word, but with any already-visited bits cleared.
+    /// (This lets us use `trailing_zeros()` to find the next set bit.) When it
+    /// is reduced to 0, we move onto the next word.
+    word: Word,
+
+    /// The offset (measured in bits) of the current word.
+    offset: usize,
+
+    /// Underlying iterator over the words.
+    iter: slice::Iter<'a, Word>,
+}
+
+impl<'a> RawBitIter<'a> {
+    #[inline(always)]
+    pub(crate) fn new(words: &'a [Word]) -> Self {
+        // Initialize `offset` to `0 - WORD_BITS`, so that the first iteration
+        // will see `word == 0` and increase the offset to its starting value of 0.
+        //
+        // This avoids having to explicitly track whether the iterator has started.
+        RawBitIter {
+            word: 0,
+            offset: const { (0usize).wrapping_sub(WORD_BITS) },
+            iter: words.iter(),
+        }
+    }
+}
+
+impl<'a> Iterator for RawBitIter<'a> {
+    type Item = usize;
+
+    #[inline]
+    fn next(&mut self) -> Option<Self::Item> {
+        // Keep looping until we find a non-empty word, or run out of words.
+        loop {
+            if self.word != 0 {
+                // Get the position of the next set bit in the current word,
+                // then clear the bit.
+                let bit_pos = self.word.trailing_zeros() as usize;
+                self.word ^= 1 << bit_pos;
+                return Some(bit_pos + self.offset);
+            }
+
+            // Move onto the next word, or stop if there isn't one.
+            self.word = self.iter.next().copied()?;
+            // This needs to be a wrapping add so that the first iteration will
+            // correctly overflow to a starting offset of 0.
+            self.offset = self.offset.wrapping_add(WORD_BITS);
+        }
+    }
+}