use super::WORD_LEN;
use crate::bit_vectors::BitVector;
use crate::bit_vectors::NumBits;
use crate::broadword;
pub struct UnaryIter<'a> {
bv: &'a BitVector,
pos: usize,
buf: usize,
}
impl<'a> UnaryIter<'a> {
pub fn new(bv: &'a BitVector, pos: usize) -> Self {
let buf = bv.words()[pos / WORD_LEN] & (usize::MAX.wrapping_shl((pos % WORD_LEN) as u32));
Self { bv, pos, buf }
}
#[inline(always)]
pub const fn position(&self) -> usize {
self.pos
}
#[inline(always)]
pub fn skip1(&mut self, k: usize) -> Option<usize> {
let mut skipped = 0;
let mut buf = self.buf;
loop {
let w = broadword::popcount(buf);
if skipped + w > k {
break;
}
skipped += w;
self.pos += WORD_LEN;
let word_pos = self.pos / WORD_LEN;
if self.bv.num_words() <= word_pos {
return None;
}
buf = self.bv.words()[word_pos];
}
debug_assert!(buf != 0);
let pos_in_word = broadword::select_in_word(buf, k - skipped).unwrap();
self.buf = buf & usize::MAX.wrapping_shl(pos_in_word as u32);
self.pos = (self.pos & !(WORD_LEN - 1)) + pos_in_word;
Some(self.pos)
}
#[inline(always)]
pub fn skip0(&mut self, k: usize) -> Option<usize> {
let mut skipped = 0;
let pos_in_word = self.pos % WORD_LEN;
let mut buf = !self.buf & usize::MAX.wrapping_shl(pos_in_word as u32);
loop {
let w = broadword::popcount(buf);
if skipped + w > k {
break;
}
skipped += w;
self.pos += WORD_LEN;
let word_pos = self.pos / WORD_LEN;
if self.bv.num_words() <= word_pos {
return None;
}
buf = !self.bv.words()[word_pos];
}
debug_assert!(buf != 0);
let pos_in_word = broadword::select_in_word(buf, k - skipped).unwrap();
self.buf = !buf & usize::MAX.wrapping_shl(pos_in_word as u32);
self.pos = (self.pos & !(WORD_LEN - 1)) + pos_in_word;
Some(self.pos).filter(|&x| x < self.bv.num_bits())
}
}
impl Iterator for UnaryIter<'_> {
type Item = usize;
#[inline(always)]
fn next(&mut self) -> Option<Self::Item> {
let mut buf = self.buf;
while buf == 0 {
self.pos += WORD_LEN;
let word_pos = self.pos / WORD_LEN;
if self.bv.num_words() <= word_pos {
return None;
}
buf = self.bv.words()[word_pos];
}
let pos_in_word = broadword::lsb(buf).unwrap();
self.buf = buf & (buf - 1); self.pos = (self.pos & !(WORD_LEN - 1)) + pos_in_word;
Some(self.pos)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_next_all_zeros() {
let bv = BitVector::from_bit(false, 100);
let mut it = bv.unary_iter(0);
assert_eq!(it.next(), None);
}
#[test]
fn test_skip1_all_zeros() {
let bv = BitVector::from_bit(false, 100);
let mut it = bv.unary_iter(0);
assert_eq!(it.skip1(0), None);
}
#[test]
fn test_skip0_all_ones() {
let bv = BitVector::from_bit(true, 100);
let mut it = bv.unary_iter(0);
assert_eq!(it.skip0(0), None);
}
}