arrow2/bitmap/utils/chunk_iterator/
chunks_exact.rs

1use std::{convert::TryInto, slice::ChunksExact};
2
3use crate::trusted_len::TrustedLen;
4
5use super::{BitChunk, BitChunkIterExact};
6
7/// An iterator over a slice of bytes in [`BitChunk`]s.
8#[derive(Debug)]
9pub struct BitChunksExact<'a, T: BitChunk> {
10    iter: ChunksExact<'a, u8>,
11    remainder: &'a [u8],
12    remainder_len: usize,
13    phantom: std::marker::PhantomData<T>,
14}
15
16impl<'a, T: BitChunk> BitChunksExact<'a, T> {
17    /// Creates a new [`BitChunksExact`].
18    #[inline]
19    pub fn new(bitmap: &'a [u8], length: usize) -> Self {
20        assert!(length <= bitmap.len() * 8);
21        let size_of = std::mem::size_of::<T>();
22
23        let bitmap = &bitmap[..length.saturating_add(7) / 8];
24
25        let split = (length / 8 / size_of) * size_of;
26        let (chunks, remainder) = bitmap.split_at(split);
27        let remainder_len = length - chunks.len() * 8;
28        let iter = chunks.chunks_exact(size_of);
29
30        Self {
31            iter,
32            remainder,
33            remainder_len,
34            phantom: std::marker::PhantomData,
35        }
36    }
37
38    /// Returns the number of chunks of this iterator
39    #[inline]
40    pub fn len(&self) -> usize {
41        self.iter.len()
42    }
43
44    /// Returns whether there are still elements in this iterator
45    #[inline]
46    pub fn is_empty(&self) -> bool {
47        self.len() == 0
48    }
49
50    /// Returns the remaining [`BitChunk`]. It is zero iff `len / 8 == 0`.
51    #[inline]
52    pub fn remainder(&self) -> T {
53        let remainder_bytes = self.remainder;
54        if remainder_bytes.is_empty() {
55            return T::zero();
56        }
57        let remainder = match remainder_bytes.try_into() {
58            Ok(a) => a,
59            Err(_) => {
60                let mut remainder = T::zero().to_ne_bytes();
61                remainder_bytes
62                    .iter()
63                    .enumerate()
64                    .for_each(|(index, b)| remainder[index] = *b);
65                remainder
66            }
67        };
68        T::from_ne_bytes(remainder)
69    }
70}
71
72impl<T: BitChunk> Iterator for BitChunksExact<'_, T> {
73    type Item = T;
74
75    #[inline]
76    fn next(&mut self) -> Option<Self::Item> {
77        self.iter.next().map(|x| match x.try_into() {
78            Ok(a) => T::from_ne_bytes(a),
79            Err(_) => unreachable!(),
80        })
81    }
82
83    #[inline]
84    fn size_hint(&self) -> (usize, Option<usize>) {
85        self.iter.size_hint()
86    }
87}
88
89unsafe impl<T: BitChunk> TrustedLen for BitChunksExact<'_, T> {}
90
91impl<T: BitChunk> BitChunkIterExact<T> for BitChunksExact<'_, T> {
92    #[inline]
93    fn remainder(&self) -> T {
94        self.remainder()
95    }
96
97    #[inline]
98    fn remainder_len(&self) -> usize {
99        self.remainder_len
100    }
101}
arrow2/bitmap/utils/chunk_iterator/chunks_exact.rs

arrow2/bitmap/utils/chunk_iterator/
chunks_exact.rs
