1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150

use util::*;
use {BitSet, BitSetLike};

pub use self::drain::DrainBitIter;

#[cfg(feature = "parallel")]
pub use self::parallel::{BitParIter, BitProducer};

mod drain;
#[cfg(feature = "parallel")]
mod parallel;

/// An `Iterator` over a [`BitSetLike`] structure.
///
/// [`BitSetLike`]: ../trait.BitSetLike.html
#[derive(Debug, Clone)]
pub struct BitIter<T> {
    pub(crate) set: T,
    pub(crate) masks: [usize; LAYERS],
    pub(crate) prefix: [u32; LAYERS - 1],
}

impl<T> BitIter<T> {
    /// Creates a new `BitIter`. You usually don't call this function
    /// but just [`.iter()`] on a bit set.
    ///
    /// [`.iter()`]: ../trait.BitSetLike.html#method.iter
    pub fn new(set: T, masks: [usize; LAYERS], prefix: [u32; LAYERS - 1]) -> Self {
        BitIter {
            set: set,
            masks: masks,
            prefix: prefix,
        }
    }
}

impl<T: BitSetLike> BitIter<T> {
    /// Allows checking if set bit is contained in underlying bit set.
    pub fn contains(&self, i: Index) -> bool {
        self.set.contains(i)
    }
}

impl<'a> BitIter<&'a mut BitSet> {
    /// Clears the rest of the bitset starting from the next inner layer.
    pub(crate) fn clear(&mut self) {
        use self::State::Continue;
        while let Some(level) = (1..LAYERS).find(|&level| self.handle_level(level) == Continue) {
            let lower = level - 1;
            let idx = (self.prefix[lower] >> BITS) as usize;
            *self.set.layer_mut(lower, idx) = 0;
            if level == LAYERS - 1 {
                self.set.layer3 &= !((2 << idx) - 1);
            }
        }
    }
}

#[derive(PartialEq)]
pub(crate) enum State {
    Empty,
    Continue,
    Value(Index),
}

impl<T> Iterator for BitIter<T>
where
    T: BitSetLike,
{
    type Item = Index;

    fn next(&mut self) -> Option<Self::Item> {
        use self::State::*;
        'find: loop {
            for level in 0..LAYERS {
                match self.handle_level(level) {
                    Value(v) => return Some(v),
                    Continue => continue 'find,
                    Empty => {}
                }
            }
            // There is no set bits left
            return None;
        }
    }
}

impl<T: BitSetLike> BitIter<T> {
    pub(crate) fn handle_level(&mut self, level: usize) -> State {
        use self::State::*;
        if self.masks[level] == 0 {
            Empty
        } else {
            // Take the first bit that isn't zero
            let first_bit = self.masks[level].trailing_zeros();
            // Remove it from the mask
            self.masks[level] &= !(1 << first_bit);
            // Calculate the index of it
            let idx = self.prefix.get(level).cloned().unwrap_or(0) | first_bit;
            if level == 0 {
                // It's the lowest layer, so the `idx` is the next set bit
                Value(idx)
            } else {
                // Take the corresponding `usize` from the layer below
                self.masks[level - 1] = self.set.get_from_layer(level - 1, idx as usize);
                self.prefix[level - 1] = idx << BITS;
                Continue
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use {BitSet, BitSetLike};

    #[test]
    fn iterator_clear_empties() {
        use rand::prelude::*;

        let mut set = BitSet::new();
        let mut rng = thread_rng();
        let limit = 1_048_576;
        for _ in 0..(limit / 10) {
            set.add(rng.gen_range(0, limit));
        }
        (&mut set).iter().clear();
        assert_eq!(0, set.layer3);
        for &i in &set.layer2 {
            assert_eq!(0, i);
        }
        for &i in &set.layer1 {
            assert_eq!(0, i);
        }
        for &i in &set.layer0 {
            assert_eq!(0, i);
        }
    }

    #[test]
    fn iterator_clone() {
        let mut set = BitSet::new();
        set.add(1);
        set.add(3);
        let iter = set.iter().skip(1);
        for (a, b) in iter.clone().zip(iter) {
            assert_eq!(a, b);
        }
    }
}