diff --git a/src/lib.rs b/src/lib.rs index c4a7cc2..ddf0fb3 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,7 +1,7 @@ extern crate unreachable; use std::ops; -use std::iter::FromIterator; +use std::iter::{Iterator, IntoIterator, FromIterator}; use std::default::Default; use std::collections::HashMap; @@ -132,6 +132,103 @@ impl BitSet { pub fn is_superset(&self, other: &BitSet) -> bool { other.is_subset(self) } + + pub fn iter(&self) -> Iter { + Iter { iter: self.blocks.iter(), block: 0, bits: 0, bit: BITS } + } +} + + +pub struct Iter<'a> { + iter: std::collections::hash_map::Iter<'a, u64, u64>, + block: u64, + bits: u64, + bit: u64, +} + +impl<'a> Iterator for Iter<'a> { + type Item = u64; + + #[inline] + fn next(&mut self) -> Option { + loop { + if self.bits == 0 || self.bit == BITS { + match self.iter.next() { + Some((block, bits)) => { + self.block = *block; + self.bits = *bits; + self.bit = 0; + } + None => return None, + } + } + + for i in self.bit..BITS { + if self.bits & (1 << i) != 0 { + self.bit = i + 1; + + return Some((self.block * BITS) + i); + } + } + + self.bit = BITS; + } + } +} + +pub struct IntoIter { + iter: std::collections::hash_map::IntoIter, + block: u64, + bits: u64, + bit: u64, +} + +impl Iterator for IntoIter { + type Item = u64; + + #[inline] + fn next(&mut self) -> Option { + loop { + if self.bits == 0 || self.bit == BITS { + match self.iter.next() { + Some((block, bits)) => { + self.block = block; + self.bits = bits; + self.bit = 0; + } + None => return None, + } + } + + for i in self.bit..BITS { + if self.bits & (1 << i) != 0 { + self.bit = i + 1; + + return Some((self.block * BITS) + i); + } + } + + self.bit = BITS; + } + } +} + +impl<'a> IntoIterator for &'a BitSet { + type Item = u64; + type IntoIter = Iter<'a>; + + fn into_iter(self) -> Iter<'a> { + self.iter() + } +} + +impl IntoIterator for BitSet { + type Item = u64; + type IntoIter = IntoIter; + + fn into_iter(self) -> IntoIter { + IntoIter { iter: self.blocks.into_iter(), block: 0, bits: 0, bit: BITS } + } } impl Default for BitSet { @@ -284,4 +381,28 @@ mod tests { set.insert(4); assert_eq!(set.is_subset(&sup), false); } + + #[test] + fn iter() { + let set = [1, 2, 3].iter().cloned().collect::(); + let mut iter = set.iter(); + + assert_eq!(Some(1), iter.next()); + assert_eq!(Some(2), iter.next()); + assert_eq!(Some(3), iter.next()); + assert_eq!(None, iter.next()); + assert_eq!(None, iter.next()); + } + + #[test] + fn into_iter() { + let set = [1, 2, 3].iter().cloned().collect::(); + let mut iter = set.into_iter(); + + assert_eq!(Some(1), iter.next()); + assert_eq!(Some(2), iter.next()); + assert_eq!(Some(3), iter.next()); + assert_eq!(None, iter.next()); + assert_eq!(None, iter.next()); + } }