Skip to content

Commit f5bfffd

Browse files
committed
Optimize BitSet iteration.
This commit removes an `Option` check in `BitIter::next()`, avoids calling `trailing_zeros()` when it's not necessary, and avoids the need for `enumerate()`. This gives a tiny (0.2%) instruction count win on a couple of benchmarks. The commit also adds some comments, which is good because this iteration code is moderately complex.
1 parent ad0dbf8 commit f5bfffd

File tree

1 file changed

+30
-13
lines changed

1 file changed

+30
-13
lines changed

src/librustc_index/bit_set.rs

+30-13
Original file line numberDiff line numberDiff line change
@@ -287,17 +287,32 @@ impl<T: Idx> ToString for BitSet<T> {
287287
}
288288

289289
pub struct BitIter<'a, T: Idx> {
290-
cur: Option<(Word, usize)>,
291-
iter: iter::Enumerate<slice::Iter<'a, Word>>,
290+
/// A copy of the current word, but with any already-visited bits cleared.
291+
/// (This lets us use `trailing_zeros()` to find the next set bit.) When it
292+
/// is reduced to 0, we move onto the next word.
293+
word: Word,
294+
295+
/// The offset (measured in bits) of the current word.
296+
offset: usize,
297+
298+
/// Underlying iterator over the words.
299+
iter: slice::Iter<'a, Word>,
300+
292301
marker: PhantomData<T>
293302
}
294303

295304
impl<'a, T: Idx> BitIter<'a, T> {
296305
#[inline]
297306
fn new(words: &'a [Word]) -> BitIter<'a, T> {
307+
// We initialize `word` and `offset` to degenerate values. On the first
308+
// call to `next()` we will fall through to getting the first word from
309+
// `iter`, which sets `word` to the first word (if there is one) and
310+
// `offset` to 0. Doing it this way saves us from having to maintain
311+
// additional state about whether we have started.
298312
BitIter {
299-
cur: None,
300-
iter: words.iter().enumerate(),
313+
word: 0,
314+
offset: std::usize::MAX - (WORD_BITS - 1),
315+
iter: words.iter(),
301316
marker: PhantomData,
302317
}
303318
}
@@ -307,17 +322,19 @@ impl<'a, T: Idx> Iterator for BitIter<'a, T> {
307322
type Item = T;
308323
fn next(&mut self) -> Option<T> {
309324
loop {
310-
if let Some((ref mut word, offset)) = self.cur {
311-
let bit_pos = word.trailing_zeros() as usize;
312-
if bit_pos != WORD_BITS {
313-
let bit = 1 << bit_pos;
314-
*word ^= bit;
315-
return Some(T::new(bit_pos + offset))
316-
}
325+
if self.word != 0 {
326+
// Get the position of the next set bit in the current word,
327+
// then clear the bit.
328+
let bit_pos = self.word.trailing_zeros() as usize;
329+
let bit = 1 << bit_pos;
330+
self.word ^= bit;
331+
return Some(T::new(bit_pos + self.offset))
317332
}
318333

319-
let (i, word) = self.iter.next()?;
320-
self.cur = Some((*word, WORD_BITS * i));
334+
// Move onto the next word.
335+
let word = self.iter.next()?;
336+
self.word = *word;
337+
self.offset += WORD_BITS;
321338
}
322339
}
323340
}

0 commit comments

Comments
 (0)