polars_arrow/bitmap/utils/
chunks_exact_mut.rs

1use super::BitChunk;
2
3/// An iterator over mutable slices of bytes of exact size.
4///
5/// # Safety
6/// The slices returned by this iterator are guaranteed to have length equal to
7/// `size_of::<T>()`.
8#[derive(Debug)]
9pub struct BitChunksExactMut<'a, T: BitChunk> {
10    chunks: std::slice::ChunksExactMut<'a, u8>,
11    remainder: &'a mut [u8],
12    remainder_len: usize,
13    marker: std::marker::PhantomData<T>,
14}
15
16impl<'a, T: BitChunk> BitChunksExactMut<'a, T> {
17    /// Returns a new [`BitChunksExactMut`]
18    #[inline]
19    pub fn new(bitmap: &'a mut [u8], length: usize) -> Self {
20        assert!(length <= bitmap.len() * 8);
21        let size_of = size_of::<T>();
22
23        let bitmap = &mut bitmap[..length.saturating_add(7) / 8];
24
25        let split = (length / 8 / size_of) * size_of;
26        let (chunks, remainder) = bitmap.split_at_mut(split);
27        let remainder_len = length - chunks.len() * 8;
28
29        let chunks = chunks.chunks_exact_mut(size_of);
30        Self {
31            chunks,
32            remainder,
33            remainder_len,
34            marker: std::marker::PhantomData,
35        }
36    }
37
38    /// The remainder slice
39    #[inline]
40    pub fn remainder(&mut self) -> &mut [u8] {
41        self.remainder
42    }
43
44    /// The length of the remainder slice in bits.
45    #[inline]
46    pub fn remainder_len(&mut self) -> usize {
47        self.remainder_len
48    }
49}
50
51impl<'a, T: BitChunk> Iterator for BitChunksExactMut<'a, T> {
52    type Item = &'a mut [u8];
53
54    #[inline]
55    fn next(&mut self) -> Option<Self::Item> {
56        self.chunks.next()
57    }
58
59    #[inline]
60    fn size_hint(&self) -> (usize, Option<usize>) {
61        self.chunks.size_hint()
62    }
63}