Add element iterator for ChunkedBitSet

2 files changed, 82 insertions, 0 deletions

diff --git a/compiler/rustc_index/src/bit_set.rs b/compiler/rustc_index/src/bit_set.rs
index 33a8d6c11ff..059755a743b 100644
--- a/compiler/rustc_index/src/bit_set.rs
+++ b/compiler/rustc_index/src/bit_set.rs
@@ -479,6 +479,11 @@ impl<T: Idx> ChunkedBitSet<T> {
         }
     }
 
+    #[inline]
+    pub fn iter(&self) -> ChunkedBitIter<'_, T> {
+        ChunkedBitIter::new(self)
+    }
+
     /// Insert `elem`. Returns whether the set has changed.
     pub fn insert(&mut self, elem: T) -> bool {
         assert!(elem.index() < self.domain_size);
@@ -697,6 +702,49 @@ impl<T> Clone for ChunkedBitSet<T> {
     }
 }
 
+pub struct ChunkedBitIter<'a, T: Idx> {
+    index: usize,
+    bitset: &'a ChunkedBitSet<T>,
+}
+
+impl<'a, T: Idx> ChunkedBitIter<'a, T> {
+    #[inline]
+    fn new(bitset: &'a ChunkedBitSet<T>) -> ChunkedBitIter<'a, T> {
+        ChunkedBitIter { index: 0, bitset }
+    }
+}
+
+impl<'a, T: Idx> Iterator for ChunkedBitIter<'a, T> {
+    type Item = T;
+    fn next(&mut self) -> Option<T> {
+        while self.index < self.bitset.domain_size() {
+            let elem = T::new(self.index);
+            let chunk = &self.bitset.chunks[chunk_index(elem)];
+            match &chunk {
+                Zeros(chunk_domain_size) => {
+                    self.index += *chunk_domain_size as usize;
+                }
+                Ones(_chunk_domain_size) => {
+                    self.index += 1;
+                    return Some(elem);
+                }
+                Mixed(_chunk_domain_size, _, words) => loop {
+                    let elem = T::new(self.index);
+                    self.index += 1;
+                    let (word_index, mask) = chunk_word_index_and_mask(elem);
+                    if (words[word_index] & mask) != 0 {
+                        return Some(elem);
+                    }
+                    if self.index % CHUNK_BITS == 0 {
+                        break;
+                    }
+                },
+            }
+        }
+        None
+    }
+}
+
 impl Chunk {
     #[cfg(test)]
     fn assert_valid(&self) {
diff --git a/compiler/rustc_index/src/bit_set/tests.rs b/compiler/rustc_index/src/bit_set/tests.rs
index eec7dab5189..cfc891e97a3 100644
--- a/compiler/rustc_index/src/bit_set/tests.rs
+++ b/compiler/rustc_index/src/bit_set/tests.rs
@@ -343,6 +343,40 @@ fn chunked_bitset() {
 }
 
 #[test]
+fn chunked_bitset_iter() {
+    fn with_elements(elements: &[usize], domain_size: usize) -> ChunkedBitSet<usize> {
+        let mut s = ChunkedBitSet::new_empty(domain_size);
+        for &e in elements {
+            s.insert(e);
+        }
+        s
+    }
+
+    // Empty
+    let vec: Vec<usize> = Vec::new();
+    let bit = with_elements(&vec, 9000);
+    assert_eq!(vec, bit.iter().collect::<Vec<_>>());
+
+    // Filled
+    let n = 10000;
+    let vec: Vec<usize> = (0..n).collect();
+    let bit = with_elements(&vec, n);
+    assert_eq!(vec, bit.iter().collect::<Vec<_>>());
+
+    // Filled with trailing zeros
+    let n = 10000;
+    let vec: Vec<usize> = (0..n).collect();
+    let bit = with_elements(&vec, 2 * n);
+    assert_eq!(vec, bit.iter().collect::<Vec<_>>());
+
+    // Mixed
+    let n = 12345;
+    let vec: Vec<usize> = vec![0, 1, 2, 2010, 2047, 2099, 6000, 6002, 6004];
+    let bit = with_elements(&vec, n);
+    assert_eq!(vec, bit.iter().collect::<Vec<_>>());
+}
+
+#[test]
 fn grow() {
     let mut set: GrowableBitSet<usize> = GrowableBitSet::with_capacity(65);
     for index in 0..65 {