Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Bitset conversion to and from iterators #5

Merged
merged 6 commits into from
Dec 15, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
101 changes: 101 additions & 0 deletions src/iterators.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,101 @@
use std::iter::FusedIterator;

use crate::BitBlock;

/// Iterator over the set bits in a bitset.
///
/// Yields the indices of the set bits in ascending order.
#[derive(Debug, Clone)]
pub struct IntoIter<T: BitBlock, const N: usize> {
blocks: [T; N],
index_front: usize,
index_back: usize,
}

impl<T: BitBlock, const N: usize> IntoIter<T, N> {
pub(crate) fn new(blocks: [T; N]) -> Self {
Self {
blocks,
index_front: 0,
index_back: N - 1,
}
}
}

impl<T: BitBlock, const N: usize> Iterator for IntoIter<T, N> {
type Item = usize;

fn next(&mut self) -> Option<Self::Item> {
while self.index_front <= self.index_back {
let block = &mut self.blocks[self.index_front];
if *block == T::EMPTY {
self.index_front += 1;
continue;
}

let idx_in_block = block.trailing_zeros() as usize;
*block ^= T::LSB << idx_in_block;
return Some(idx_in_block + self.index_front * T::BITS);
}

None
}

fn size_hint(&self) -> (usize, Option<usize>) {
(0, Some(N * T::BITS))
}
}

impl<T: BitBlock, const N: usize> DoubleEndedIterator for IntoIter<T, N> {
fn next_back(&mut self) -> Option<Self::Item> {
while self.index_front <= self.index_back {
let block = &mut self.blocks[self.index_back];
if *block == T::EMPTY {
let Some(new_index_back) = self.index_back.checked_sub(1) else {
break;
};
self.index_back = new_index_back;
continue;
}

let idx_in_block = T::BITS - 1 - block.leading_zeros() as usize;
*block ^= T::LSB << idx_in_block;
return Some(idx_in_block + self.index_back * T::BITS);
}

None
}
}

impl<T: BitBlock, const N: usize> FusedIterator for IntoIter<T, N> {}

#[cfg(test)]
mod tests {
use super::*;

/// Iterator type for most tests
type TestIter = IntoIter<u8, 2>;

#[test]
fn forward_iteration() {
let iter = TestIter::new([0b1000_1001, 0b1000_0100]);
let expected = vec![0, 3, 7, 10, 15];
assert_eq!(expected, iter.collect::<Vec<_>>());
}

#[test]
fn backward_iteration() {
let iter = TestIter::new([0b0100_0100, 0b0110_0001]);
let expected = vec![14, 13, 8, 6, 2];
assert_eq!(expected, iter.rev().collect::<Vec<_>>());
}

#[test]
fn mixed_forward_backward_iteration() {
let mut iter = TestIter::new([0b0100_0000, 0b0110_0001]);
assert_eq!(Some(14), iter.next_back());
assert_eq!(Some(6), iter.next());
assert_eq!(Some(8), iter.next());
assert_eq!(Some(13), iter.next_back());
}
}
89 changes: 89 additions & 0 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,8 @@
//! depending on the use-case.
//!
//! [fixedbitset]: https://github.com/petgraph/fixedbitset
mod iterators;

use std::array;
use std::fmt;
use std::fmt::Binary;
Expand All @@ -41,6 +43,8 @@ use std::ops::Not;

use num_traits::PrimInt;

pub use iterators::IntoIter;

/// Integer that can be used as a block of bits in a bitset.
pub trait BitBlock:
PrimInt + BitAndAssign + BitOrAssign + BitXorAssign + Binary + LowerHex + UpperHex + 'static
Expand Down Expand Up @@ -151,6 +155,16 @@ impl<T: BitBlock, const N: usize> TinyBitSet<T, N> {
self.blocks.iter().all(|&block| block == T::EMPTY)
}

/// Iterates over the indices of set bits from lowest to highest.
pub fn iter(self) -> IntoIter<T, N> {
IntoIter::new(self.blocks)
}

/// Iterates over the indices of unset bits from lowest to highest.
pub fn iter_missing(self) -> IntoIter<T, N> {
(!self).iter()
}

/// Set the given bit.
///
/// # Panics
Expand Down Expand Up @@ -280,6 +294,41 @@ impl<T: BitBlock, const N: usize> Index<usize> for TinyBitSet<T, N> {
}
}

impl<T: BitBlock, const N: usize> IntoIterator for TinyBitSet<T, N> {
type Item = usize;

type IntoIter = IntoIter<T, N>;

fn into_iter(self) -> Self::IntoIter {
self.iter()
}
}

impl<T: BitBlock, const N: usize> IntoIterator for &TinyBitSet<T, N> {
type Item = usize;

type IntoIter = IntoIter<T, N>;

fn into_iter(self) -> Self::IntoIter {
self.iter()
}
}

impl<T: BitBlock, const N: usize> FromIterator<usize> for TinyBitSet<T, N> {
/// Creates a bitset with an iterator of indices of set bits.
///
/// # Panics
///
/// Panics if any of the indices are out of range.
fn from_iter<I: IntoIterator<Item = usize>>(iter: I) -> Self {
let mut bs = Self::EMPTY;
for i in iter {
bs.insert(i);
}
bs
}
}

impl<T: BitBlock, const N: usize> Not for TinyBitSet<T, N> {
type Output = Self;

Expand Down Expand Up @@ -475,6 +524,18 @@ mod tests {
assert!(TestBitSet::from([0b0000_0000, 0b0000_0000]).is_empty());
}

#[test]
fn iter() {
let bs = TestBitSet::from([0b1000_0001, 0b0011_1100]);
assert_eq!(vec![0, 7, 10, 11, 12, 13], bs.iter().collect::<Vec<_>>());
}

#[test]
fn iter_missing() {
let bs = TestBitSet::from([0b1101_0111, 0b1011_1101]);
assert_eq!(vec![3, 5, 9, 14], bs.iter_missing().collect::<Vec<_>>());
}

#[test]
fn insert() {
let mut bs = TestBitSet::EMPTY;
Expand Down Expand Up @@ -639,6 +700,34 @@ mod tests {
assert!(!bs[9]);
}

#[test]
fn into_iterator() {
let bs = TestBitSet::from([0b0010_1000, 0b0100_1100]);
assert_eq!(vec![3, 5, 10, 11, 14], bs.into_iter().collect::<Vec<_>>());
}

#[test]
fn ref_into_iterator() {
let bs = TestBitSet::from([0b0000_0010, 0b0001_0110]);
let iter = (&bs).into_iter();
assert_eq!(vec![1, 9, 10, 12], iter.collect::<Vec<_>>());
}

#[test]
fn from_iterator() {
fn to_bs(indices: impl IntoIterator<Item = usize>) -> TestBitSet {
indices.into_iter().collect()
}

assert_eq!(TestBitSet::EMPTY, to_bs([]));
assert_eq!(TestBitSet::singleton(5), to_bs([5]));
assert_eq!(TestBitSet::singleton(6), to_bs([6, 6, 6]));
assert_eq!(
TestBitSet::singleton(6) | TestBitSet::singleton(11),
to_bs([11, 6])
);
}

#[test]
fn not() {
assert_eq!(TestBitSet::ALL, !TestBitSet::EMPTY);
Expand Down