Bitset conversion to and from iterators (#5)

* Add iterators

* Iterator methods and traits for `TinyBitSet`

* Implement `FromIterator<usize>` for `TinyBitSet`

* Add `size_hint()`

Co-authored-by: Rinde van Lon <[email protected]>

---------

Co-authored-by: Rinde van Lon <[email protected]>

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<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());
+    }
+}

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