hibitset/iter/
mod.rs

1use util::*;
2use {BitSet, BitSetLike};
3
4pub use self::drain::DrainBitIter;
5
6#[cfg(feature = "parallel")]
7pub use self::parallel::{BitParIter, BitProducer};
8
9mod drain;
10#[cfg(feature = "parallel")]
11mod parallel;
12
13/// An `Iterator` over a [`BitSetLike`] structure.
14///
15/// [`BitSetLike`]: ../trait.BitSetLike.html
16#[derive(Debug, Clone)]
17pub struct BitIter<T> {
18    pub(crate) set: T,
19    pub(crate) masks: [usize; LAYERS],
20    pub(crate) prefix: [u32; LAYERS - 1],
21}
22
23impl<T> BitIter<T> {
24    /// Creates a new `BitIter`. You usually don't call this function
25    /// but just [`.iter()`] on a bit set.
26    ///
27    /// [`.iter()`]: ../trait.BitSetLike.html#method.iter
28    pub fn new(set: T, masks: [usize; LAYERS], prefix: [u32; LAYERS - 1]) -> Self {
29        BitIter {
30            set: set,
31            masks: masks,
32            prefix: prefix,
33        }
34    }
35}
36
37impl<T: BitSetLike> BitIter<T> {
38    /// Allows checking if set bit is contained in underlying bit set.
39    pub fn contains(&self, i: Index) -> bool {
40        self.set.contains(i)
41    }
42}
43
44impl<'a> BitIter<&'a mut BitSet> {
45    /// Clears the rest of the bitset starting from the next inner layer.
46    pub(crate) fn clear(&mut self) {
47        use self::State::Continue;
48        while let Some(level) = (1..LAYERS).find(|&level| self.handle_level(level) == Continue) {
49            let lower = level - 1;
50            let idx = (self.prefix[lower] >> BITS) as usize;
51            *self.set.layer_mut(lower, idx) = 0;
52            if level == LAYERS - 1 {
53                self.set.layer3 &= !((2 << idx) - 1);
54            }
55        }
56    }
57}
58
59#[derive(PartialEq)]
60pub(crate) enum State {
61    Empty,
62    Continue,
63    Value(Index),
64}
65
66impl<T> Iterator for BitIter<T>
67where
68    T: BitSetLike,
69{
70    type Item = Index;
71
72    fn next(&mut self) -> Option<Self::Item> {
73        use self::State::*;
74        'find: loop {
75            for level in 0..LAYERS {
76                match self.handle_level(level) {
77                    Value(v) => return Some(v),
78                    Continue => continue 'find,
79                    Empty => {}
80                }
81            }
82            // There is no set bits left
83            return None;
84        }
85    }
86}
87
88impl<T: BitSetLike> BitIter<T> {
89    pub(crate) fn handle_level(&mut self, level: usize) -> State {
90        use self::State::*;
91        if self.masks[level] == 0 {
92            Empty
93        } else {
94            // Take the first bit that isn't zero
95            let first_bit = self.masks[level].trailing_zeros();
96            // Remove it from the mask
97            self.masks[level] &= !(1 << first_bit);
98            // Calculate the index of it
99            let idx = self.prefix.get(level).cloned().unwrap_or(0) | first_bit;
100            if level == 0 {
101                // It's the lowest layer, so the `idx` is the next set bit
102                Value(idx)
103            } else {
104                // Take the corresponding `usize` from the layer below
105                self.masks[level - 1] = self.set.get_from_layer(level - 1, idx as usize);
106                self.prefix[level - 1] = idx << BITS;
107                Continue
108            }
109        }
110    }
111}
112
113#[cfg(test)]
114mod tests {
115    use {BitSet, BitSetLike};
116
117    #[test]
118    fn iterator_clear_empties() {
119        use rand::prelude::*;
120
121        let mut set = BitSet::new();
122        let mut rng = thread_rng();
123        let limit = 1_048_576;
124        for _ in 0..(limit / 10) {
125            set.add(rng.gen_range(0, limit));
126        }
127        (&mut set).iter().clear();
128        assert_eq!(0, set.layer3);
129        for &i in &set.layer2 {
130            assert_eq!(0, i);
131        }
132        for &i in &set.layer1 {
133            assert_eq!(0, i);
134        }
135        for &i in &set.layer0 {
136            assert_eq!(0, i);
137        }
138    }
139
140    #[test]
141    fn iterator_clone() {
142        let mut set = BitSet::new();
143        set.add(1);
144        set.add(3);
145        let iter = set.iter().skip(1);
146        for (a, b) in iter.clone().zip(iter) {
147            assert_eq!(a, b);
148        }
149    }
150}