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 ChunkedBitSet { } } + #[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 Clone for ChunkedBitSet { } } +pub struct ChunkedBitIter<'a, T: Idx> { + index: usize, + bitset: &'a ChunkedBitSet, +} + +impl<'a, T: Idx> ChunkedBitIter<'a, T> { + #[inline] + fn new(bitset: &'a ChunkedBitSet) -> ChunkedBitIter<'a, T> { + ChunkedBitIter { index: 0, bitset } + } +} + +impl<'a, T: Idx> Iterator for ChunkedBitIter<'a, T> { + type Item = T; + fn next(&mut self) -> Option { + 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 @@ -342,6 +342,40 @@ fn chunked_bitset() { b10000b.assert_valid(); } +#[test] +fn chunked_bitset_iter() { + fn with_elements(elements: &[usize], domain_size: usize) -> ChunkedBitSet { + let mut s = ChunkedBitSet::new_empty(domain_size); + for &e in elements { + s.insert(e); + } + s + } + + // Empty + let vec: Vec = Vec::new(); + let bit = with_elements(&vec, 9000); + assert_eq!(vec, bit.iter().collect::>()); + + // Filled + let n = 10000; + let vec: Vec = (0..n).collect(); + let bit = with_elements(&vec, n); + assert_eq!(vec, bit.iter().collect::>()); + + // Filled with trailing zeros + let n = 10000; + let vec: Vec = (0..n).collect(); + let bit = with_elements(&vec, 2 * n); + assert_eq!(vec, bit.iter().collect::>()); + + // Mixed + let n = 12345; + let vec: Vec = vec![0, 1, 2, 2010, 2047, 2099, 6000, 6002, 6004]; + let bit = with_elements(&vec, n); + assert_eq!(vec, bit.iter().collect::>()); +} + #[test] fn grow() { let mut set: GrowableBitSet = GrowableBitSet::with_capacity(65);