hibitset/iter/
mod.rs

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