Add element iterator for ChunkedBitSet

This commit is contained in:
Tomasz Miąsko 2022-04-30 00:00:00 +00:00
parent 76d4862fdd
commit cdfdb99c9e
2 changed files with 82 additions and 0 deletions

View File

@ -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) {

View File

@ -342,6 +342,40 @@ fn chunked_bitset() {
b10000b.assert_valid();
}
#[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);