Skip to content

Instantly share code, notes, and snippets.

@herabit
Created March 22, 2024 04:12
Show Gist options
  • Save herabit/e691a1a643f9b614e781d0c304283da2 to your computer and use it in GitHub Desktop.
Save herabit/e691a1a643f9b614e781d0c304283da2 to your computer and use it in GitHub Desktop.
I think this should work
use std::{
iter::FusedIterator,
ops::Range,
};
use bitvec::{
order::{BitOrder, Lsb0},
store::BitStore,
};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum BitRange {
Ones { start: usize, end: usize },
Zeroes { start: usize, end: usize },
}
impl BitRange {
#[inline]
pub const fn into_range(self) -> Range<usize> {
let (BitRange::Ones { start, end } | BitRange::Zeroes { start, end }) = self;
start..end
}
#[inline]
pub const fn ones(range: Range<usize>) -> BitRange {
BitRange::Ones {
start: range.start,
end: range.end,
}
}
#[inline]
pub const fn zeroes(range: Range<usize>) -> BitRange {
BitRange::Zeroes {
start: range.start,
end: range.end,
}
}
}
#[derive(Debug, Clone)]
pub struct BitRangeIter<'a, T: BitStore = usize, O: BitOrder = Lsb0> {
bits: &'a bitvec::slice::BitSlice<T, O>,
offset: usize,
}
impl<'a, T: BitStore, O: BitOrder> BitRangeIter<'a, T, O> {
#[inline]
pub const fn new(bits: &'a bitvec::slice::BitSlice<T, O>) -> Self {
Self { bits, offset: 0 }
}
#[inline]
pub fn remaining_bits(&self) -> &'a bitvec::slice::BitSlice<T, O> {
self.bits
}
}
impl<'a, T: BitStore, O: BitOrder> Iterator for BitRangeIter<'a, T, O> {
type Item = BitRange;
fn next(&mut self) -> Option<Self::Item> {
let first_bit = *self.bits.first()?;
let bits = match first_bit {
true => self.bits.leading_ones(),
false => self.bits.leading_zeros(),
};
let start = self.offset;
self.offset += bits;
self.bits = self.bits.split_at(self.offset).0;
Some(match first_bit {
true => BitRange::ones(start..self.offset),
false => BitRange::zeroes(start..self.offset),
})
}
}
impl<'a, T: BitStore, O: BitOrder> DoubleEndedIterator for BitRangeIter<'a, T, O> {
fn next_back(&mut self) -> Option<Self::Item> {
let last_bit = *self.bits.last()?;
let bits = match last_bit {
true => self.bits.trailing_ones(),
false => self.bits.trailing_zeros(),
};
let end = self.bits.len();
let start = end - bits;
self.bits = self.bits.split_at(start).0;
Some(match last_bit {
true => BitRange::ones(start..end),
false => BitRange::zeroes(start..end),
})
}
}
impl<'a, T: BitStore, O: BitOrder> FusedIterator for BitRangeIter<'a, T, O> {}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment