deepmesa_collections/bitvec/
iter.rs

1/*
2   BitVector: A fast contiguous growable array of bits allocated
3   on the heap that allows storing and manipulating an arbitrary
4   number of bits. This collection is backed by a Vector<u8> which
5   manages the underlying memory.
6
7   Copyright 2021 "Rahul Singh <rsingh@arrsingh.com>"
8
9   Licensed under the Apache License, Version 2.0 (the "License");
10   you may not use this file except in compliance with the License.
11   You may obtain a copy of the License at
12
13       http://www.apache.org/licenses/LICENSE-2.0
14
15   Unless required by applicable law or agreed to in writing, software
16   distributed under the License is distributed on an "AS IS" BASIS,
17   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
18   See the License for the specific language governing permissions and
19   limitations under the License.
20*/
21
22use crate::bitvec::{
23    bitops, bitref::BitRefMut, bitslice::BitSlice, bitvec::BitVector, bytes,
24    traits::BitwiseClearAssign, BitCount, BitOrder,
25};
26
27macro_rules! iter_unsigned {
28    (
29        $(#[$outer:meta])*
30        $iter_name: ident, $i:ident, $max_bits: literal
31    ) => {
32        $(#[$outer])*
33        #[derive(Debug)]
34        pub struct $iter_name<'a> {
35            bits: &'a [u8],
36            cursor: usize,
37            bit_len: usize,
38            slice_offset: usize,
39        }
40
41        impl<'a> $iter_name<'a> {
42            pub(super) fn new(
43                bits: &'a [u8],
44                slice_offset: usize,
45                bit_len: usize,
46            ) -> $iter_name<'a> {
47                $iter_name {
48                    bits,
49                    cursor: 0,
50                    bit_len,
51                    slice_offset,
52                }
53            }
54        }
55
56        impl<'a> Iterator for $iter_name<'a> {
57            type Item = ($i, BitCount);
58            fn next(&mut self) -> Option<Self::Item> {
59                if self.cursor >= self.bit_len {
60                    return None;
61                }
62
63                let offset = self.slice_offset;
64                let len = self.bit_len + offset;
65                let st_index = self.cursor + offset;
66
67                let (val, bit_count) =
68                    bytes::read_bits(&self.bits, st_index, len, $max_bits, BitOrder::Lsb0);
69                self.cursor += bit_count;
70                Some((val as $i, bit_count))
71            }
72        }
73    };
74}
75
76iter_unsigned!(
77    /// An iterator that iterates over the bits of the
78    /// [`BitVector`](../struct.BitVector.html) 8 bits at a time.
79    ///
80    /// This struct is created by the
81    /// [`.iter_u8()`](BitVector#method.iter_u8)
82    /// method of [`BitVector`](../struct.BitVector.html) and
83    /// [`BitSlice`](BitSlice)
84    ///
85    /// # Examples
86    /// ```
87    /// use deepmesa_collections::BitVector;
88    /// use deepmesa_collections::bitvec::IterU8;
89    ///
90    /// let mut bv = BitVector::new();
91    /// bv.push_u16(0b0101_1101_0011_1010, Some(16));
92    ///
93    /// let mut iter: IterU8 = bv.iter_u8();
94    /// assert_eq!(iter.next(), Some((0b0101_1101, 8)));
95    /// assert_eq!(iter.next(), Some((0b0011_1010, 8)));
96    /// assert_eq!(iter.next(), None);
97    /// ```
98    IterU8,
99    u8,
100    8
101);
102iter_unsigned!(
103    /// An iterator that iterates over the bits of the
104    /// [`BitVector`](../struct.BitVector.html) 16 bits at a time.
105    ///
106    /// This struct is created by the
107    /// [`.iter_u16()`](BitVector#method.iter_u16)
108    /// method of [`BitVector`](../struct.BitVector.html) and
109    /// [`BitSlice`](BitSlice)
110    ///
111    /// # Examples
112    /// ```
113    /// use deepmesa_collections::BitVector;
114    /// use deepmesa_collections::bitvec::IterU16;
115    ///
116    /// let mut bv = BitVector::new();
117    /// bv.push_u16(0b0101_1101_0011_1010, Some(16));
118    ///
119    /// let mut iter:IterU16 = bv.iter_u16();
120    /// assert_eq!(iter.next(), Some((0b0101_1101_0011_1010, 16)));
121    /// assert_eq!(iter.next(), None);
122    /// ```
123    IterU16,
124    u16,
125    16
126);
127iter_unsigned!(
128    /// An iterator that iterates over the bits of the
129    /// [`BitVector`](../struct.BitVector.html) 32 bits at a time.
130    ///
131    /// This struct is created by the
132    /// [`.iter_u32()`](BitVector#method.iter_u32)
133    /// method of [`BitVector`](../struct.BitVector.html) and
134    /// [`BitSlice`](BitSlice)
135    ///
136    /// # Examples
137    /// ```
138    /// use deepmesa_collections::BitVector;
139    /// use deepmesa_collections::bitvec::IterU32;
140    ///
141    /// let mut bv = BitVector::new();
142    /// bv.push_u16(0b0101_1101_0011_1010, Some(16));
143    /// bv.push_u16(0b1111_0011_1100_0000, Some(16));
144    ///
145    /// let mut iter:IterU32 = bv.iter_u32();
146    /// assert_eq!(iter.next(), Some((0b0101_1101_0011_1010_1111_0011_1100_0000, 32)));
147    /// assert_eq!(iter.next(), None);
148    /// ```
149    IterU32,
150    u32,
151    32
152);
153iter_unsigned!(
154    /// An iterator that iterates over the bits of the
155    /// [`BitVector`](../struct.BitVector.html) 64 bits at a time.
156    ///
157    /// This struct is created by the
158    /// [`.iter_u64()`](BitVector#method.iter_u64)
159    /// method of [`BitVector`](../struct.BitVector.html) and
160    /// [`BitSlice`](BitSlice)
161    ///
162    /// # Examples
163    /// ```
164    /// use deepmesa_collections::BitVector;
165    /// use deepmesa_collections::bitvec::IterU64;
166    /// let mut bv = BitVector::new();
167    /// bv.push_u64(u64::MAX, Some(64));
168    ///
169    /// let mut iter:IterU64 = bv.iter_u64();
170    /// assert_eq!(iter.next(), Some((u64::MAX, 64)));
171    /// assert_eq!(iter.next(), None);
172    /// ```
173    IterU64,
174    u64,
175    64
176);
177iter_unsigned!(
178    /// An iterator that iterates over the bits of the
179    /// [`BitVector`](../struct.BitVector.html) 128 bits at a time.
180    ///
181    /// This struct is created by the
182    /// [`.iter_u128()`](BitVector#method.iter_u128)
183    /// method of [`BitVector`](../struct.BitVector.html) and
184    /// [`BitSlice`](BitSlice)
185    ///
186    /// # Examples
187    /// ```
188    /// use deepmesa_collections::BitVector;
189    /// use deepmesa_collections::bitvec::IterU128;
190    ///
191    /// let mut bv = BitVector::new();
192    /// bv.push_u64(u64::MAX, Some(64));
193    /// bv.push_u64(u64::MAX, Some(64));
194    ///
195    /// let mut iter = bv.iter_u128();
196    /// assert_eq!(iter.next(), Some((u128::MAX, 128)));
197    /// assert_eq!(iter.next(), None);
198    /// ```
199    IterU128,
200    u128,
201    128
202);
203
204/// An immutable iterator over the bits of the
205/// [`BitVector`](../struct.BitVector.html) or [`BitSlice`](BitSlice).
206///
207/// This struct is created by the [`.iter()`](BitVector#method.iter)
208/// method of [`BitVector`](../struct.BitVector.html) and
209/// [`BitSlice`](BitSlice)
210///
211/// # Examples
212/// ```
213/// use deepmesa_collections::BitVector;
214/// use deepmesa_collections::bitvec::Iter;
215///
216/// let mut bv = BitVector::new();
217/// bv.push_u8(0b101, None);
218///
219/// let mut iter:Iter = bv.iter();
220/// assert_eq!(iter.next(), Some(true));
221/// assert_eq!(iter.next(), Some(false));
222/// assert_eq!(iter.next(), Some(true));
223/// assert_eq!(iter.next(), None);
224/// ```
225#[derive(Debug)]
226pub struct Iter<'a> {
227    bits: &'a [u8],
228    cursor: usize,
229    bit_len: usize,
230    slice_offset: usize,
231}
232
233/// A mutable iterator over the bits of the
234/// [`BitVector`](../struct.BitVector.html) or [`BitSlice`](BitSlice).
235///
236/// This struct is created by the
237/// [`.iter_mut()`](BitVector#method.iter_mut) method of
238/// [`BitVector`](../struct.BitVector.html) and [`BitSlice`](BitSlice)
239///
240/// # Examples
241/// ```
242/// use deepmesa_collections::BitVector;
243/// use deepmesa_collections::bitvec::IterMut;
244///
245/// let mut bv = BitVector::with_capacity(20);
246/// bv.push_u8(0b1011_1100, Some(8));
247/// bv.push_u8(0b0011_1001, Some(8));
248/// let iter: IterMut = bv.iter_mut();
249/// for mut bit in iter {
250///     *bit = true;
251/// }
252/// assert_eq!(bv.read_u16(0), (0b1111_1111_1111_1111, 16));
253/// ```
254#[derive(Debug)]
255pub struct IterMut<'a> {
256    bits: &'a mut [u8],
257    cursor: usize,
258    bit_len: usize,
259    slice_offset: usize,
260}
261
262impl<'a> Iter<'a> {
263    pub(super) fn new(bits: &'a [u8], slice_offset: usize, bit_len: usize) -> Iter<'a> {
264        Iter {
265            bits,
266            cursor: 0,
267            bit_len,
268            slice_offset,
269        }
270    }
271}
272
273impl<'a> IterMut<'a> {
274    pub(super) fn new(bits: &'a mut [u8], slice_offset: usize, bit_len: usize) -> IterMut<'a> {
275        IterMut {
276            bits,
277            cursor: 0,
278            bit_len,
279            slice_offset,
280        }
281    }
282}
283
284macro_rules! iter_bits {
285    (
286        $(#[$outer:meta])*
287        $iter_name: ident
288    ) => {
289        $(#[$outer])*
290        #[derive(Debug)]
291        pub struct $iter_name<'a> {
292            bits: &'a [u8],
293            cur_byte: u8,
294            byte_idx: usize,
295            slice_offset: usize,
296            eb_idx: usize,
297            sb_idx: usize,
298            l_bit: usize,
299        }
300
301        impl<'a> $iter_name<'a> {
302            pub(super) fn new(
303                bits: &'a [u8],
304                slice_offset: usize,
305                bit_len: usize,
306            ) -> $iter_name<'a> {
307                let byte_idx = slice_offset / 8;
308                let eb_idx;
309                #[allow(unused_mut)]
310                let mut cur_byte: u8;
311                let l_bit = slice_offset + bit_len;
312                if bit_len == 0 {
313                    eb_idx = 0;
314                    cur_byte = 0;
315                } else {
316                    eb_idx = (l_bit - 1) / 8;
317                    cur_byte = bits[byte_idx];
318                    flip_bits!(cur_byte, $iter_name);
319                };
320                let sb_idx = byte_idx;
321                $iter_name {
322                    bits,
323                    slice_offset,
324                    byte_idx,
325                    cur_byte,
326                    eb_idx,
327                    sb_idx,
328                    l_bit,
329                }
330            }
331        }
332        impl<'a> Iterator for $iter_name<'a> {
333            type Item = usize;
334            fn next(&mut self) -> Option<Self::Item> {
335                if self.byte_idx == self.sb_idx {
336                    self.cur_byte
337                        .clear_msb_assign((self.slice_offset % 8) as u8);
338                }
339
340                if self.cur_byte == 0 {
341                    self.byte_idx += 1;
342                    if self.byte_idx > self.eb_idx {
343                        return None;
344                    }
345                    self.cur_byte = self.bits[self.byte_idx];
346                    flip_bits!(self.cur_byte, $iter_name);
347                }
348
349                let bit_idx = self.cur_byte.leading_zeros() as usize;
350                let idx = bit_idx + (self.byte_idx * 8);
351                if idx >= self.l_bit {
352                    return None;
353                }
354                self.cur_byte.clear_msb_assign((bit_idx + 1) as u8);
355
356                return Some(idx - self.slice_offset);
357            }
358        }
359    };
360}
361
362iter_bits!(
363    /// An iterator that iterates over the `1` bits of the
364    /// [`BitVector`](../struct.BitVector.html).
365    ///
366    /// This struct is created by the
367    /// [`.iter_ones()`](BitVector#method.iter_ones)
368    /// method of [`BitVector`](../struct.BitVector.html) and
369    /// [`BitSlice`](BitSlice)
370    ///
371    /// # Examples
372    /// ```
373    /// use deepmesa_collections::BitVector;
374    /// use deepmesa_collections::bitvec::IterOnes;
375    ///
376    /// let mut bv = BitVector::with_capacity(20);
377    /// bv.push_u8(0b0010_0101, Some(8));
378    ///
379    /// let mut iter:IterOnes = bv.iter_ones();
380    /// assert_eq!(iter.next(), Some(2));
381    /// assert_eq!(iter.next(), Some(5));
382    /// assert_eq!(iter.next(), Some(7));
383    /// assert_eq!(iter.next(), None);
384    /// ```
385    IterOnes
386);
387
388iter_bits!(
389    /// An iterator that iterates over the `0` bits of the
390    /// [`BitVector`](../struct.BitVector.html).
391    ///
392    /// This struct is created by the
393    /// [`.iter_zeros()`](BitVector#method.iter_zeros)
394    /// method of [`BitVector`](../struct.BitVector.html) and
395    /// [`BitSlice`](BitSlice)
396    ///
397    /// # Examples
398    /// ```
399    /// use deepmesa_collections::BitVector;
400    /// use deepmesa_collections::bitvec::IterZeros;
401    ///
402    /// let mut bv = BitVector::with_capacity(20);
403    /// bv.push_u8(0b1101_1010, Some(8));
404    ///
405    /// let mut iter:IterZeros = bv.iter_zeros();
406    /// assert_eq!(iter.next(), Some(2));
407    /// assert_eq!(iter.next(), Some(5));
408    /// assert_eq!(iter.next(), Some(7));
409    /// assert_eq!(iter.next(), None);
410    /// ```
411    IterZeros
412);
413
414impl<'a> Iterator for Iter<'a> {
415    type Item = bool;
416    fn next(&mut self) -> Option<Self::Item> {
417        if self.cursor >= self.bit_len {
418            return None;
419        }
420
421        let index = self.cursor + self.slice_offset;
422        self.cursor += 1;
423        return Some(bit_at_unchecked!(index, self.bits));
424    }
425}
426
427impl<'a> Iterator for IterMut<'a> {
428    type Item = BitRefMut<'a, bool>;
429    fn next(&mut self) -> Option<Self::Item> {
430        if self.cursor >= self.bit_len {
431            return None;
432        }
433
434        let slice_index = self.cursor / 8;
435        let byte_ptr = self.bits[slice_index..slice_index].as_mut_ptr();
436        let index = self.cursor + self.slice_offset;
437
438        unsafe {
439            let byte = *byte_ptr;
440            let bit = bitops::is_msb_nset(byte, (index % 8) as u8);
441
442            self.cursor += 1;
443            return Some(BitRefMut::<bool>::new(bit, byte_ptr, index));
444        }
445    }
446}
447
448impl<'a> IntoIterator for &'a BitSlice {
449    type Item = bool;
450    type IntoIter = Iter<'a>;
451    fn into_iter(self) -> Self::IntoIter {
452        self.iter()
453    }
454}
455
456impl<'a> IntoIterator for &'a BitVector {
457    type Item = bool;
458    type IntoIter = Iter<'a>;
459    fn into_iter(self) -> Self::IntoIter {
460        self.iter()
461    }
462}
463
464#[cfg(test)]
465mod tests {
466    use super::*;
467    use crate::bitvec::bitvec::BitVector;
468
469    #[test]
470    fn test_iter_slice() {
471        let mut bv = BitVector::with_capacity(128);
472        bv.push_u8(0b1100_1010, None);
473
474        let slice = &bv[2..6];
475        let mut iter = slice.iter();
476        assert_eq!(&iter.next(), &Some(false));
477        assert_eq!(&iter.next(), &Some(false));
478        assert_eq!(&iter.next(), &Some(true));
479        assert_eq!(&iter.next(), &Some(false));
480        assert_eq!(&iter.next(), &None);
481
482        let mut i = 0;
483        for v in slice {
484            assert_eq!(v, slice.get(i).unwrap());
485            i += 1;
486        }
487    }
488
489    #[test]
490    fn test_iter_bitvec() {
491        let mut bv = BitVector::with_capacity(128);
492        bv.push(false);
493        bv.push(true);
494        bv.push(false);
495        bv.push(false);
496
497        let mut iter = bv.iter();
498        assert_eq!(&iter.next(), &Some(false));
499        assert_eq!(&iter.next(), &Some(true));
500        assert_eq!(&iter.next(), &Some(false));
501        assert_eq!(&iter.next(), &Some(false));
502        assert_eq!(&iter.next(), &None);
503
504        let mut i = 0;
505        for v in &bv {
506            assert_eq!(v, bv.get(i).unwrap());
507            i += 1;
508        }
509    }
510
511    #[test]
512    fn test_iter_u16_bitvec() {
513        let mut bv = BitVector::with_capacity(128);
514        bv.push_u16(0b1100_1010_0011_0101, None);
515        bv.push_u16(0b1000_1100_0011_1111, None);
516        bv.push_u8(0b0000_0011, None);
517        assert_eq!(bv.len(), 34);
518        let mut iter = bv.iter_u16();
519        assert_eq!(iter.next(), Some((0b1100_1010_0011_0101, 16)));
520        assert_eq!(iter.next(), Some((0b1000_1100_0011_1111, 16)));
521        assert_eq!(iter.next(), Some((0b0000_0000_0000_0011, 2)));
522        assert_eq!(&iter.next(), &None);
523    }
524
525    #[test]
526    fn test_iter_ones() {
527        let mut bv = BitVector::with_capacity(128);
528        let mut iter = IterOnes::new(&bv.bits, 0, bv.len());
529        assert_eq!(iter.next(), None);
530
531        bv.push_u16(0b1100_1010_0011_0101, None);
532        bv.push_u16(0b1000_1100_0011_1111, None);
533
534        let mut iter = IterOnes::new(&bv.bits, 0, bv.len());
535        assert_eq!(iter.next(), Some(0));
536        assert_eq!(iter.next(), Some(1));
537        assert_eq!(iter.next(), Some(4));
538        assert_eq!(iter.next(), Some(6));
539        assert_eq!(iter.next(), Some(10));
540        assert_eq!(iter.next(), Some(11));
541        assert_eq!(iter.next(), Some(13));
542        assert_eq!(iter.next(), Some(15));
543        assert_eq!(iter.next(), Some(16));
544        assert_eq!(iter.next(), Some(20));
545        assert_eq!(iter.next(), Some(21));
546        assert_eq!(iter.next(), Some(26));
547        assert_eq!(iter.next(), Some(27));
548        assert_eq!(iter.next(), Some(28));
549        assert_eq!(iter.next(), Some(29));
550        assert_eq!(iter.next(), Some(30));
551        assert_eq!(iter.next(), Some(31));
552        assert_eq!(iter.next(), None);
553    }
554
555    #[test]
556    fn test_iter_zeros() {
557        let mut bv = BitVector::with_capacity(128);
558        let mut iter = IterZeros::new(&bv.bits, 0, bv.len());
559        assert_eq!(iter.next(), None);
560
561        bv.push_u16(0b0011_0101_1100_1010, Some(16));
562        bv.push_u16(0b0111_0011_1100_0000, Some(16));
563        let mut iter = IterZeros::new(&bv.bits, 0, bv.len());
564        assert_eq!(iter.next(), Some(0));
565        assert_eq!(iter.next(), Some(1));
566        assert_eq!(iter.next(), Some(4));
567        assert_eq!(iter.next(), Some(6));
568        assert_eq!(iter.next(), Some(10));
569        assert_eq!(iter.next(), Some(11));
570        assert_eq!(iter.next(), Some(13));
571        assert_eq!(iter.next(), Some(15));
572        assert_eq!(iter.next(), Some(16));
573        assert_eq!(iter.next(), Some(20));
574        assert_eq!(iter.next(), Some(21));
575        assert_eq!(iter.next(), Some(26));
576        assert_eq!(iter.next(), Some(27));
577        assert_eq!(iter.next(), Some(28));
578        assert_eq!(iter.next(), Some(29));
579        assert_eq!(iter.next(), Some(30));
580        assert_eq!(iter.next(), Some(31));
581        assert_eq!(iter.next(), None);
582    }
583}