1/// Yields the bits of a `u64`.
  2pub struct BitIterator {
  3    /// The value that is being iterated bit-wise
  4    value: u64,
  5    /// True bits in the `mask` are the bits that have been visited.
  6    mask: u64,
  7}
  8
  9impl BitIterator {
 10    pub fn new(value: u64) -> BitIterator {
 11        BitIterator { value, mask: 0 }
 12    }
 13
 14    /// An efficient skip implementation.
 15    ///
 16    /// Note: The compiler is smart enough to translate a `skip(n)` into a single shift instruction
 17    /// if the code is inlined, however inlining does not always happen.
 18    pub fn skip_front(mut self, n: u32) -> Self {
 19        let mask = bitmask(n);
 20        let ones = self.mask.trailing_ones();
 21        let mask_position = ones;
 22        self.mask ^= mask << mask_position;
 23        self
 24    }
 25
 26    /// An efficient skip from the back.
 27    ///
 28    /// Note: The compiler is smart enough to translate a `skip(n)` into a single shift instruction
 29    /// if the code is inlined, however inlining does not always happen.
 30    pub fn skip_back(mut self, n: u32) -> Self {
 31        let mask = bitmask(n);
 32        let ones = self.mask.leading_ones();
 33        let mask_position = u64::BITS - ones - n;
 34        self.mask ^= mask << mask_position;
 35        self
 36    }
 37}
 38
 39impl Iterator for BitIterator {
 40    type Item = bool;
 41
 42    fn next(&mut self) -> Option<<Self as Iterator>::Item> {
 43        // trailing_ones is implemented with trailing_zeros, and the zeros are computed with the
 44        // intrinsic cttz. [Rust 1.67.0] x86 uses the `bsf` instruction. AArch64 uses the `rbit
 45        // clz` instructions.
 46        let ones = self.mask.trailing_ones();
 47
 48        if ones == u64::BITS {
 49            None
 50        } else {
 51            let bit_position = ones;
 52            let mask = 1 << bit_position;
 53            self.mask ^= mask;
 54            let bit = self.value & mask;
 55            Some(bit != 0)
 56        }
 57    }
 58}
 59
 60impl DoubleEndedIterator for BitIterator {
 61    fn next_back(&mut self) -> Option<<Self as Iterator>::Item> {
 62        // leading_ones is implemented with leading_zeros, and the zeros are computed with the
 63        // intrinsic ctlz. [Rust 1.67.0] x86 uses the `bsr` instruction. AArch64 uses the `clz`
 64        // instruction.
 65        let ones = self.mask.leading_ones();
 66
 67        if ones == u64::BITS {
 68            None
 69        } else {
 70            let bit_position = u64::BITS - ones - 1;
 71            let mask = 1 << bit_position;
 72            self.mask ^= mask;
 73            let bit = self.value & mask;
 74            Some(bit != 0)
 75        }
 76    }
 77}
 78
 79#[cfg(test)]
 80mod test {
 81    use super::BitIterator;
 82
 83    #[test]
 84    fn test_bit_iterator() {
 85        let v = 0b1;
 86        let mut it = BitIterator::new(v);
 87        assert!(it.next().unwrap(), "first bit is true");
 88        assert!(it.all(|v| v == false), "every other value is false");
 89
 90        let v = 0b10;
 91        let mut it = BitIterator::new(v);
 92        assert!(!it.next().unwrap(), "first bit is false");
 93        assert!(it.next().unwrap(), "first bit is true");
 94        assert!(it.all(|v| v == false), "every other value is false");
 95
 96        let v = 0b10;
 97        let mut it = BitIterator::new(v);
 98        assert!(!it.next_back().unwrap(), "last bit is false");
 99        assert!(!it.next().unwrap(), "first bit is false");
100        assert!(it.next().unwrap(), "first bit is true");
101        assert!(it.all(|v| v == false), "every other value is false");
102    }
103
104    #[test]
105    fn test_bit_iterator_skip() {
106        let v = 0b1;
107        let mut it = BitIterator::new(v).skip_front(1);
108        assert!(it.all(|v| v == false), "every other value is false");
109
110        let v = 0b10;
111        let mut it = BitIterator::new(v).skip_front(1);
112        assert!(it.next().unwrap(), "first bit is true");
113        assert!(it.all(|v| v == false), "every other value is false");
114
115        let high_bit = 0b1 << (u64::BITS - 1);
116        let mut it = BitIterator::new(high_bit).skip_back(1);
117        assert!(it.all(|v| v == false), "every other value is false");
118
119        let v = 0b10;
120        let mut it = BitIterator::new(v).skip_back(1);
121        assert!(!it.next_back().unwrap(), "last bit is false");
122        assert!(!it.next().unwrap(), "first bit is false");
123        assert!(it.next().unwrap(), "first bit is true");
124        assert!(it.all(|v| v == false), "every other value is false");
125    }
126
127    #[test]
128    fn test_skip_all() {
129        let v = 0b1;
130        let mut it = BitIterator::new(v).skip_front(u64::BITS);
131        assert!(it.next().is_none(), "iterator must be exhausted");
132
133        let v = 0b1;
134        let mut it = BitIterator::new(v).skip_back(u64::BITS);
135        assert!(it.next().is_none(), "iterator must be exhausted");
136    }
137
138    #[test]
139    fn test_bit_iterator_count_bits_after_skip() {
140        let any_value = 0b1;
141        for s in 0..u64::BITS {
142            let it = BitIterator::new(any_value).skip_front(s);
143            assert_eq!(it.count() as u32, u64::BITS - s)
144        }
145
146        let any_value = 0b1;
147        for s in 1..u64::BITS {
148            let it = BitIterator::new(any_value).skip_back(s);
149            assert_eq!(it.count() as u32, u64::BITS - s)
150        }
151    }
152
153    #[test]
154    fn test_bit_iterator_rev() {
155        let v = 0b1;
156        let mut it = BitIterator::new(v).rev();
157        assert!(it.nth(63).unwrap(), "the last value is true");
158    }
159}
160
161// UTILITIES
162// ===============================================================================================
163
164fn bitmask(s: u32) -> u64 {
165    match 1u64.checked_shl(s) {
166        Some(r) => r - 1,
167        None => u64::MAX,
168    }
169}