diff --git a/src/iterators.rs b/src/iterators.rs new file mode 100644 index 0000000..b86717d --- /dev/null +++ b/src/iterators.rs @@ -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 { + blocks: [T; N], + index_front: usize, + index_back: usize, +} + +impl IntoIter { + pub(crate) fn new(blocks: [T; N]) -> Self { + Self { + blocks, + index_front: 0, + index_back: N - 1, + } + } +} + +impl Iterator for IntoIter { + type Item = usize; + + fn next(&mut self) -> Option { + 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) { + (0, Some(N * T::BITS)) + } +} + +impl DoubleEndedIterator for IntoIter { + fn next_back(&mut self) -> Option { + 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 FusedIterator for IntoIter {} + +#[cfg(test)] +mod tests { + use super::*; + + /// Iterator type for most tests + type TestIter = IntoIter; + + #[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::>()); + } + + #[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::>()); + } + + #[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()); + } +} diff --git a/src/lib.rs b/src/lib.rs index d0f9e90..e703840 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -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; @@ -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 @@ -151,6 +155,16 @@ impl TinyBitSet { self.blocks.iter().all(|&block| block == T::EMPTY) } + /// Iterates over the indices of set bits from lowest to highest. + pub fn iter(self) -> IntoIter { + IntoIter::new(self.blocks) + } + + /// Iterates over the indices of unset bits from lowest to highest. + pub fn iter_missing(self) -> IntoIter { + (!self).iter() + } + /// Set the given bit. /// /// # Panics @@ -280,6 +294,41 @@ impl Index for TinyBitSet { } } +impl IntoIterator for TinyBitSet { + type Item = usize; + + type IntoIter = IntoIter; + + fn into_iter(self) -> Self::IntoIter { + self.iter() + } +} + +impl IntoIterator for &TinyBitSet { + type Item = usize; + + type IntoIter = IntoIter; + + fn into_iter(self) -> Self::IntoIter { + self.iter() + } +} + +impl FromIterator for TinyBitSet { + /// 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>(iter: I) -> Self { + let mut bs = Self::EMPTY; + for i in iter { + bs.insert(i); + } + bs + } +} + impl Not for TinyBitSet { type Output = Self; @@ -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::>()); + } + + #[test] + fn iter_missing() { + let bs = TestBitSet::from([0b1101_0111, 0b1011_1101]); + assert_eq!(vec![3, 5, 9, 14], bs.iter_missing().collect::>()); + } + #[test] fn insert() { let mut bs = TestBitSet::EMPTY; @@ -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::>()); + } + + #[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::>()); + } + + #[test] + fn from_iterator() { + fn to_bs(indices: impl IntoIterator) -> 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);