arrow2/bitmap/utils/
iterator.rs

1use crate::trusted_len::TrustedLen;
2
3use super::get_bit_unchecked;
4
5/// An iterator over bits according to the [LSB](https://en.wikipedia.org/wiki/Bit_numbering#Least_significant_bit),
6/// i.e. the bytes `[4u8, 128u8]` correspond to `[false, false, true, false, ..., true]`.
7#[derive(Debug, Clone)]
8pub struct BitmapIter<'a> {
9    bytes: &'a [u8],
10    index: usize,
11    end: usize,
12}
13
14impl<'a> BitmapIter<'a> {
15    /// Creates a new [`BitmapIter`].
16    pub fn new(slice: &'a [u8], offset: usize, len: usize) -> Self {
17        // example:
18        // slice.len() = 4
19        // offset = 9
20        // len = 23
21        // result:
22        let bytes = &slice[offset / 8..];
23        // bytes.len() = 3
24        let index = offset % 8;
25        // index = 9 % 8 = 1
26        let end = len + index;
27        // end = 23 + 1 = 24
28        assert!(end <= bytes.len() * 8);
29        // maximum read before UB in bits: bytes.len() * 8 = 24
30        // the first read from the end is `end - 1`, thus, end = 24 is ok
31
32        Self { bytes, index, end }
33    }
34}
35
36impl<'a> Iterator for BitmapIter<'a> {
37    type Item = bool;
38
39    #[inline]
40    fn next(&mut self) -> Option<Self::Item> {
41        if self.index == self.end {
42            return None;
43        }
44        let old = self.index;
45        self.index += 1;
46        // See comment in `new`
47        Some(unsafe { get_bit_unchecked(self.bytes, old) })
48    }
49
50    #[inline]
51    fn size_hint(&self) -> (usize, Option<usize>) {
52        let exact = self.end - self.index;
53        (exact, Some(exact))
54    }
55
56    #[inline]
57    fn nth(&mut self, n: usize) -> Option<Self::Item> {
58        let new_index = self.index + n;
59        if new_index > self.end {
60            self.index = self.end;
61            None
62        } else {
63            self.index = new_index;
64            self.next()
65        }
66    }
67}
68
69impl<'a> DoubleEndedIterator for BitmapIter<'a> {
70    #[inline]
71    fn next_back(&mut self) -> Option<bool> {
72        if self.index == self.end {
73            None
74        } else {
75            self.end -= 1;
76            // See comment in `new`; end was first decreased
77            Some(unsafe { get_bit_unchecked(self.bytes, self.end) })
78        }
79    }
80}
81
82unsafe impl TrustedLen for BitmapIter<'_> {}
83impl ExactSizeIterator for BitmapIter<'_> {}
arrow2/bitmap/utils/iterator.rs

arrow2/bitmap/utils/
iterator.rs
