primal_bit/
iter.rs

1use std::iter;
2use std::mem;
3use std::ops::Range;
4use std::slice;
5use std::vec;
6
7use crate::BitVec;
8use crate::BITS;
9
10impl BitVec {
11    /// Returns an iterator over the elements of the vector in order.
12    ///
13    /// # Examples
14    ///
15    /// ```ignore
16    /// use std::collections::BitVec;
17    ///
18    /// let bv = BitVec::from_bytes(&[0b01110100, 0b10010010]);
19    /// assert_eq!(bv.iter().filter(|x| *x).count(), 7);
20    /// ```
21    #[inline]
22    pub fn iter(&self) -> Iter<'_> {
23        Iter {
24            bit_vec: self,
25            idx: 0..self.len(),
26        }
27    }
28}
29
30/// An iterator for `BitVec`.
31#[derive(Clone)]
32pub struct Iter<'a> {
33    bit_vec: &'a BitVec,
34    idx: Range<usize>,
35}
36
37impl<'a> Iterator for Iter<'a> {
38    type Item = bool;
39
40    #[inline]
41    fn next(&mut self) -> Option<bool> {
42        self.bit_vec.get(self.idx.next()?)
43    }
44
45    #[inline]
46    fn size_hint(&self) -> (usize, Option<usize>) {
47        self.idx.size_hint()
48    }
49}
50
51impl<'a> DoubleEndedIterator for Iter<'a> {
52    #[inline]
53    fn next_back(&mut self) -> Option<bool> {
54        self.bit_vec.get(self.idx.next_back()?)
55    }
56}
57
58impl<'a> ExactSizeIterator for Iter<'a> {}
59
60impl<'a> IntoIterator for &'a BitVec {
61    type Item = bool;
62    type IntoIter = Iter<'a>;
63
64    #[inline]
65    fn into_iter(self) -> Iter<'a> {
66        self.iter()
67    }
68}
69
70impl BitVec {
71    #[inline]
72    pub fn ones_from(&self, from: usize) -> Ones<'_> {
73        let (byte, bit) = (from / BITS, from % BITS);
74        let mask = (!0) << bit;
75        let mut iter = self.as_bytes()[byte..].iter().copied();
76        let (current, _) = usize_from_bytes(&mut iter);
77        InnerOnes {
78            base: byte * BITS,
79            current: current & mask,
80            iter,
81        }
82    }
83
84    #[inline]
85    pub fn into_ones(self) -> IntoOnes {
86        let mut iter = self.into_bytes().into_iter();
87        let (current, _) = usize_from_bytes(&mut iter);
88        InnerOnes {
89            base: 0,
90            current,
91            iter,
92        }
93    }
94}
95
96pub type Ones<'a> = InnerOnes<iter::Copied<slice::Iter<'a, u8>>>;
97pub type IntoOnes = InnerOnes<vec::IntoIter<u8>>;
98
99#[derive(Clone)]
100pub struct InnerOnes<I> {
101    base: usize,
102    current: usize,
103    iter: I,
104}
105
106impl<I: Iterator<Item = u8>> Iterator for InnerOnes<I> {
107    type Item = usize;
108
109    #[inline]
110    fn next(&mut self) -> Option<usize> {
111        let mut c = self.current;
112        while c == 0 {
113            let (next, done) = usize_from_bytes(&mut self.iter);
114            if done {
115                return None;
116            }
117            self.base += BITS * mem::size_of::<usize>();
118            c = next;
119        }
120
121        let lsb = c.trailing_zeros();
122        self.current = c & (c - 1);
123        Some(self.base + lsb as usize)
124    }
125}
126
127/// Extract a `usize` from an iterator over bytes.
128///
129/// It helps our iterator performance to pull bits from a native word at a time.
130#[inline]
131fn usize_from_bytes(iter: &mut impl Iterator<Item = u8>) -> (usize, bool) {
132    let mut n = 0;
133    for i in 0..mem::size_of::<usize>() {
134        match iter.next() {
135            Some(b) => n |= usize::from(b) << (i * BITS),
136            None => return (n, i == 0),
137        }
138    }
139    (n, false)
140}
141
142#[test]
143fn iter_ones_from() {
144    let len = 5000;
145
146    let zeros = BitVec::from_elem(len, false);
147    assert_eq!(zeros.ones_from(0).next(), None);
148
149    let ones = BitVec::from_elem(len, true);
150    let mut iter = ones.ones_from(0);
151    for i in 0..len {
152        assert_eq!(iter.next(), Some(i));
153        assert_eq!(ones.ones_from(i).next(), Some(i));
154    }
155
156    let mut halves = BitVec::from_elem(len * 2, false);
157    for i in 0..len {
158        halves.set(i * 2, true);
159    }
160    let mut iter = halves.ones_from(0);
161    for i in 0..len {
162        assert_eq!(iter.next(), Some(i * 2));
163        assert_eq!(halves.ones_from(i * 2).next(), Some(i * 2));
164        if i > 0 {
165            assert_eq!(halves.ones_from(i * 2 - 1).next(), Some(i * 2));
166        }
167        assert_eq!(halves.ones_from(i).next(), Some(i + (i % 2)));
168    }
169
170    let mut sparse = BitVec::from_elem(len * 101, false);
171    for i in 0..len {
172        sparse.set(i * 101, true);
173    }
174    let mut iter = sparse.ones_from(0);
175    for i in 0..len {
176        let i101 = i * 101;
177        assert_eq!(iter.next(), Some(i101));
178        if i > 0 {
179            for j in i101 - 100..=i101 {
180                assert_eq!(sparse.ones_from(j).next(), Some(i101));
181            }
182        }
183    }
184}
185
186#[test]
187fn iter_into_ones() {
188    let len = 5000;
189
190    let zeros = BitVec::from_elem(len, false);
191    assert_eq!(zeros.into_ones().next(), None);
192
193    let ones = BitVec::from_elem(len, true);
194    let mut iter = ones.into_ones();
195    for i in 0..len {
196        assert_eq!(iter.next(), Some(i));
197    }
198
199    let mut halves = BitVec::from_elem(len * 2, false);
200    for i in 0..len {
201        halves.set(i * 2, true);
202    }
203    let mut iter = halves.into_ones();
204    for i in 0..len {
205        assert_eq!(iter.next(), Some(i * 2));
206    }
207
208    let mut sparse = BitVec::from_elem(len * 101, false);
209    for i in 0..len {
210        sparse.set(i * 101, true);
211    }
212    let mut iter = sparse.into_ones();
213    for i in 0..len {
214        assert_eq!(iter.next(), Some(i * 101));
215    }
216}
primal_bit/iter.rs

primal_bit/
iter.rs
