From 2918a7d5a9e65721a31598b50e975ae0882feac3 Mon Sep 17 00:00:00 2001
From: Nicholas Nethercote <nnethercote@mozilla.com>
Date: Tue, 15 Oct 2019 10:02:55 +1100
Subject: [PATCH 1/2] Add `BitIter::new()`.

This factors out some duplicated code.
---
 src/librustc_index/bit_set.rs | 23 +++++++++++++----------
 1 file changed, 13 insertions(+), 10 deletions(-)

diff --git a/src/librustc_index/bit_set.rs b/src/librustc_index/bit_set.rs
index 8c49e0dde0dcd..bfaeb847d80b8 100644
--- a/src/librustc_index/bit_set.rs
+++ b/src/librustc_index/bit_set.rs
@@ -168,11 +168,7 @@ impl<T: Idx> BitSet<T> {
     /// Iterates over the indices of set bits in a sorted order.
     #[inline]
     pub fn iter(&self) -> BitIter<'_, T> {
-        BitIter {
-            cur: None,
-            iter: self.words.iter().enumerate(),
-            marker: PhantomData,
-        }
+        BitIter::new(&self.words)
     }
 
     /// Duplicates the set as a hybrid set.
@@ -296,6 +292,17 @@ pub struct BitIter<'a, T: Idx> {
     marker: PhantomData<T>
 }
 
+impl<'a, T: Idx> BitIter<'a, T> {
+    #[inline]
+    fn new(words: &'a [Word]) -> BitIter<'a, T> {
+        BitIter {
+            cur: None,
+            iter: words.iter().enumerate(),
+            marker: PhantomData,
+        }
+    }
+}
+
 impl<'a, T: Idx> Iterator for BitIter<'a, T> {
     type Item = T;
     fn next(&mut self) -> Option<T> {
@@ -851,11 +858,7 @@ impl<R: Idx, C: Idx> BitMatrix<R, C> {
     pub fn iter(&self, row: R) -> BitIter<'_, C> {
         assert!(row.index() < self.num_rows);
         let (start, end) = self.range(row);
-        BitIter {
-            cur: None,
-            iter: self.words[start..end].iter().enumerate(),
-            marker: PhantomData,
-        }
+        BitIter::new(&self.words[start..end])
     }
 
     /// Returns the number of elements in `row`.

From 60851b08e523e1a3ab4defc8b6049751b6bf1ad5 Mon Sep 17 00:00:00 2001
From: Nicholas Nethercote <nnethercote@mozilla.com>
Date: Tue, 15 Oct 2019 10:11:30 +1100
Subject: [PATCH 2/2] 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.
---
 src/librustc_index/bit_set.rs | 44 ++++++++++++++++++++++++-----------
 1 file changed, 31 insertions(+), 13 deletions(-)

diff --git a/src/librustc_index/bit_set.rs b/src/librustc_index/bit_set.rs
index bfaeb847d80b8..d8c6e4c33e2fa 100644
--- a/src/librustc_index/bit_set.rs
+++ b/src/librustc_index/bit_set.rs
@@ -287,17 +287,32 @@ impl<T: Idx> ToString for BitSet<T> {
 }
 
 pub struct BitIter<'a, T: Idx> {
-    cur: Option<(Word, usize)>,
-    iter: iter::Enumerate<slice::Iter<'a, Word>>,
+    /// 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>,
+
     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 {
-            cur: None,
-            iter: words.iter().enumerate(),
+            word: 0,
+            offset: std::usize::MAX - (WORD_BITS - 1),
+            iter: words.iter(),
             marker: PhantomData,
         }
     }
@@ -307,17 +322,20 @@ impl<'a, T: Idx> Iterator for BitIter<'a, T> {
     type Item = T;
     fn next(&mut self) -> Option<T> {
         loop {
-            if let Some((ref mut word, offset)) = self.cur {
-                let bit_pos = word.trailing_zeros() as usize;
-                if bit_pos != WORD_BITS {
-                    let bit = 1 << bit_pos;
-                    *word ^= bit;
-                    return Some(T::new(bit_pos + offset))
-                }
+            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;
+                let bit = 1 << bit_pos;
+                self.word ^= bit;
+                return Some(T::new(bit_pos + self.offset))
             }
 
-            let (i, word) = self.iter.next()?;
-            self.cur = Some((*word, WORD_BITS * i));
+            // Move onto the next word. `wrapping_add()` is needed to handle
+            // the degenerate initial value given to `offset` in `new()`.
+            let word = self.iter.next()?;
+            self.word = *word;
+            self.offset = self.offset.wrapping_add(WORD_BITS);
         }
     }
 }