bit_set/
iter.rs

1use crate::{local_prelude::*, BitSet};
2
3#[derive(Clone)]
4struct BlockIter<T, B: BitBlock> {
5    head: B,
6    head_offset: usize,
7    tail: T,
8}
9
10impl<T, B: BitBlock> BlockIter<T, B>
11where
12    T: Iterator<Item = B>,
13{
14    fn from_blocks(mut blocks: T) -> Self {
15        let h = blocks.next().unwrap_or(B::zero());
16        BlockIter {
17            tail: blocks,
18            head: h,
19            head_offset: 0,
20        }
21    }
22}
23
24#[allow(clippy::multiple_inherent_impl)]
25impl<B: BitBlock> BitSet<B> {
26    /// Iterator over each usize stored in `self` union `other`.
27    /// See [`union_with`] for an efficient in-place version.
28    ///
29    /// # Examples
30    ///
31    /// ```
32    /// use bit_set::BitSet;
33    ///
34    /// let a = BitSet::from_bytes(&[0b01101000]);
35    /// let b = BitSet::from_bytes(&[0b10100000]);
36    ///
37    /// // Print 0, 1, 2, 4 in arbitrary order
38    /// for x in a.union(&b) {
39    ///     println!("{}", x);
40    /// }
41    /// ```
42    ///
43    /// [`union_with`]: Self::union_with
44    #[inline]
45    pub fn union<'a>(&'a self, other: &'a Self) -> Union<'a, B> {
46        fn or<B: BitBlock>(w1: B, w2: B) -> B {
47            w1 | w2
48        }
49
50        Union(BlockIter::from_blocks(TwoBitPositions {
51            set: self.bit_vec.blocks(),
52            other: other.bit_vec.blocks(),
53            merge: or,
54        }))
55    }
56
57    /// Iterator over each usize stored in `self` intersect `other`.
58    /// See [`intersect_with`] for an efficient in-place version.
59    ///
60    /// # Examples
61    ///
62    /// ```
63    /// use bit_set::BitSet;
64    ///
65    /// let a = BitSet::from_bytes(&[0b01101000]);
66    /// let b = BitSet::from_bytes(&[0b10100000]);
67    ///
68    /// // Print 2
69    /// for x in a.intersection(&b) {
70    ///     println!("{}", x);
71    /// }
72    /// ```
73    ///
74    /// [`intersect_with`]: Self::intersect_with
75    #[inline]
76    pub fn intersection<'a>(&'a self, other: &'a Self) -> Intersection<'a, B> {
77        fn bitand<B: BitBlock>(w1: B, w2: B) -> B {
78            w1 & w2
79        }
80        let min = cmp::min(self.bit_vec.len(), other.bit_vec.len());
81
82        Intersection {
83            iter: BlockIter::from_blocks(TwoBitPositions {
84                set: self.bit_vec.blocks(),
85                other: other.bit_vec.blocks(),
86                merge: bitand,
87            }),
88            n: min,
89        }
90    }
91
92    /// Iterator over each usize stored in the `self` setminus `other`.
93    /// See [`difference_with`] for an efficient in-place version.
94    ///
95    /// # Examples
96    ///
97    /// ```
98    /// use bit_set::BitSet;
99    ///
100    /// let a = BitSet::from_bytes(&[0b01101000]);
101    /// let b = BitSet::from_bytes(&[0b10100000]);
102    ///
103    /// // Print 1, 4 in arbitrary order
104    /// for x in a.difference(&b) {
105    ///     println!("{}", x);
106    /// }
107    ///
108    /// // Note that difference is not symmetric,
109    /// // and `b - a` means something else.
110    /// // This prints 0
111    /// for x in b.difference(&a) {
112    ///     println!("{}", x);
113    /// }
114    /// ```
115    ///
116    /// [`difference_with`]: Self::difference_with
117    #[inline]
118    pub fn difference<'a>(&'a self, other: &'a Self) -> Difference<'a, B> {
119        fn diff<B: BitBlock>(w1: B, w2: B) -> B {
120            w1 & !w2
121        }
122
123        Difference(BlockIter::from_blocks(TwoBitPositions {
124            set: self.bit_vec.blocks(),
125            other: other.bit_vec.blocks(),
126            merge: diff,
127        }))
128    }
129
130    /// Iterator over each usize stored in the symmetric difference of `self` and `other`.
131    /// See [`symmetric_difference_with`] for an efficient in-place version.
132    ///
133    /// # Examples
134    ///
135    /// ```
136    /// use bit_set::BitSet;
137    ///
138    /// let a = BitSet::from_bytes(&[0b01101000]);
139    /// let b = BitSet::from_bytes(&[0b10100000]);
140    ///
141    /// // Print 0, 1, 4 in arbitrary order
142    /// for x in a.symmetric_difference(&b) {
143    ///     println!("{}", x);
144    /// }
145    /// ```
146    ///
147    /// [`symmetric_difference_with`]: Self::symmetric_difference_with
148    #[inline]
149    pub fn symmetric_difference<'a>(&'a self, other: &'a Self) -> SymmetricDifference<'a, B> {
150        fn bitxor<B: BitBlock>(w1: B, w2: B) -> B {
151            w1 ^ w2
152        }
153
154        SymmetricDifference(BlockIter::from_blocks(TwoBitPositions {
155            set: self.bit_vec.blocks(),
156            other: other.bit_vec.blocks(),
157            merge: bitxor,
158        }))
159    }
160
161    /// Iterator over each usize stored in the `BitSet`.
162    ///
163    /// # Examples
164    ///
165    /// ```
166    /// use bit_set::BitSet;
167    ///
168    /// let s = BitSet::from_bytes(&[0b01001010]);
169    ///
170    /// // Print 1, 4, 6 in arbitrary order
171    /// for x in s.iter() {
172    ///     println!("{}", x);
173    /// }
174    /// ```
175    #[inline]
176    pub fn iter(&self) -> Iter<'_, B> {
177        Iter(BlockIter::from_blocks(self.bit_vec.blocks()))
178    }
179}
180
181/// An iterator combining two `BitSet` iterators.
182#[derive(Clone)]
183struct TwoBitPositions<'a, B: 'a + BitBlock> {
184    set: Blocks<'a, B>,
185    other: Blocks<'a, B>,
186    merge: fn(B, B) -> B,
187}
188
189/// An iterator for `BitSet`.
190#[derive(Clone)]
191pub struct Iter<'a, B: 'a + BitBlock>(BlockIter<Blocks<'a, B>, B>);
192#[derive(Clone)]
193pub struct Union<'a, B: 'a + BitBlock>(BlockIter<TwoBitPositions<'a, B>, B>);
194#[derive(Clone)]
195pub struct Intersection<'a, B: 'a + BitBlock> {
196    iter: BlockIter<TwoBitPositions<'a, B>, B>,
197    // as an optimization, we compute the maximum possible
198    // number of elements in the intersection, and count it
199    // down as we return elements. If we reach zero, we can
200    // stop.
201    n: usize,
202}
203#[derive(Clone)]
204pub struct Difference<'a, B: 'a + BitBlock>(BlockIter<TwoBitPositions<'a, B>, B>);
205#[derive(Clone)]
206pub struct SymmetricDifference<'a, B: 'a + BitBlock>(BlockIter<TwoBitPositions<'a, B>, B>);
207
208impl<T, B: BitBlock> Iterator for BlockIter<T, B>
209where
210    T: Iterator<Item = B>,
211{
212    type Item = usize;
213
214    fn next(&mut self) -> Option<Self::Item> {
215        while self.head == B::zero() {
216            match self.tail.next() {
217                Some(w) => self.head = w,
218                None => return None,
219            }
220            self.head_offset += B::bits();
221        }
222
223        // from the current block, isolate the
224        // LSB and subtract 1, producing k:
225        // a block with a number of set bits
226        // equal to the index of the LSB
227        let k = (self.head & (!self.head + B::one())) - B::one();
228        // update block, removing the LSB
229        self.head = self.head & (self.head - B::one());
230        // return offset + (index of LSB)
231        Some(self.head_offset + B::count_ones(k))
232    }
233
234    fn count(self) -> usize {
235        self.head.count_ones() + self.tail.map(|block| block.count_ones()).sum::<usize>()
236    }
237
238    #[inline]
239    fn size_hint(&self) -> (usize, Option<usize>) {
240        match self.tail.size_hint() {
241            (_, Some(h)) => (0, Some((1 + h) * B::bits())),
242            _ => (0, None),
243        }
244    }
245}
246
247impl<B: BitBlock> Iterator for TwoBitPositions<'_, B> {
248    type Item = B;
249
250    fn next(&mut self) -> Option<Self::Item> {
251        match (self.set.next(), self.other.next()) {
252            (Some(a), Some(b)) => Some((self.merge)(a, b)),
253            (Some(a), None) => Some((self.merge)(a, B::zero())),
254            (None, Some(b)) => Some((self.merge)(B::zero(), b)),
255            _ => None,
256        }
257    }
258
259    #[inline]
260    fn size_hint(&self) -> (usize, Option<usize>) {
261        let (first_lower_bound, first_upper_bound) = self.set.size_hint();
262        let (second_lower_bound, second_upper_bound) = self.other.size_hint();
263
264        let upper_bound = first_upper_bound.zip(second_upper_bound);
265
266        let get_max = |(a, b)| cmp::max(a, b);
267        (
268            cmp::max(first_lower_bound, second_lower_bound),
269            upper_bound.map(get_max),
270        )
271    }
272}
273
274impl<B: BitBlock> Iterator for Iter<'_, B> {
275    type Item = usize;
276
277    #[inline]
278    fn next(&mut self) -> Option<usize> {
279        self.0.next()
280    }
281    #[inline]
282    fn size_hint(&self) -> (usize, Option<usize>) {
283        self.0.size_hint()
284    }
285    #[inline]
286    fn count(self) -> usize {
287        self.0.count()
288    }
289}
290
291impl<B: BitBlock> Iterator for Union<'_, B> {
292    type Item = usize;
293
294    #[inline]
295    fn next(&mut self) -> Option<usize> {
296        self.0.next()
297    }
298    #[inline]
299    fn size_hint(&self) -> (usize, Option<usize>) {
300        self.0.size_hint()
301    }
302    #[inline]
303    fn count(self) -> usize {
304        self.0.count()
305    }
306}
307
308impl<B: BitBlock> Iterator for Intersection<'_, B> {
309    type Item = usize;
310
311    #[inline]
312    fn next(&mut self) -> Option<usize> {
313        if self.n != 0 {
314            self.n -= 1;
315            self.iter.next()
316        } else {
317            None
318        }
319    }
320    #[inline]
321    fn size_hint(&self) -> (usize, Option<usize>) {
322        // We could invoke self.iter.size_hint() and incorporate that into the hint.
323        // In practice, that does not seem worthwhile because the lower bound will
324        // always be zero and the upper bound could only possibly less then n in a
325        // partially iterated iterator. However, it makes little sense ask for size_hint
326        // in a partially iterated iterator, so it did not seem worthwhile.
327        (0, Some(self.n))
328    }
329    #[inline]
330    fn count(self) -> usize {
331        self.iter.count()
332    }
333}
334
335impl<B: BitBlock> Iterator for Difference<'_, B> {
336    type Item = usize;
337
338    #[inline]
339    fn next(&mut self) -> Option<usize> {
340        self.0.next()
341    }
342    #[inline]
343    fn size_hint(&self) -> (usize, Option<usize>) {
344        self.0.size_hint()
345    }
346    #[inline]
347    fn count(self) -> usize {
348        self.0.count()
349    }
350}
351
352impl<B: BitBlock> Iterator for SymmetricDifference<'_, B> {
353    type Item = usize;
354
355    #[inline]
356    fn next(&mut self) -> Option<usize> {
357        self.0.next()
358    }
359    #[inline]
360    fn size_hint(&self) -> (usize, Option<usize>) {
361        self.0.size_hint()
362    }
363    #[inline]
364    fn count(self) -> usize {
365        self.0.count()
366    }
367}
368
369impl<'a, B: BitBlock> IntoIterator for &'a BitSet<B> {
370    type Item = usize;
371    type IntoIter = Iter<'a, B>;
372
373    fn into_iter(self) -> Iter<'a, B> {
374        self.iter()
375    }
376}
bit_set/iter.rs

bit_set/
iter.rs
