polars_utils/
chunks.rs

1/// A copy of the [`std::slice::Chunks`] that exposes the inner `slice` and `chunk_size`.
2#[derive(Clone, Debug)]
3pub struct Chunks<'a, T> {
4    slice: &'a [T],
5    chunk_size: usize,
6}
7
8impl<'a, T> Iterator for Chunks<'a, T> {
9    type Item = &'a [T];
10
11    fn next(&mut self) -> Option<Self::Item> {
12        if self.slice.is_empty() {
13            return None;
14        }
15
16        let item;
17        (item, self.slice) = self.slice.split_at(self.chunk_size.min(self.slice.len()));
18
19        Some(item)
20    }
21
22    fn size_hint(&self) -> (usize, Option<usize>) {
23        let len = self.slice.len().div_ceil(self.chunk_size);
24        (len, Some(len))
25    }
26}
27
28impl<T> DoubleEndedIterator for Chunks<'_, T> {
29    fn next_back(&mut self) -> Option<Self::Item> {
30        if self.slice.is_empty() {
31            return None;
32        }
33
34        let rem = self.slice.len() % self.chunk_size;
35        let offset = if rem == 0 { self.chunk_size } else { rem };
36
37        let item;
38        (self.slice, item) = self.slice.split_at(self.slice.len() - offset);
39
40        Some(item)
41    }
42}
43
44impl<T> ExactSizeIterator for Chunks<'_, T> {}
45
46impl<'a, T> Chunks<'a, T> {
47    pub const fn new(slice: &'a [T], chunk_size: usize) -> Self {
48        Self { slice, chunk_size }
49    }
50
51    pub const fn as_slice(&self) -> &'a [T] {
52        self.slice
53    }
54
55    pub const fn chunk_size(&self) -> usize {
56        self.chunk_size
57    }
58
59    pub fn skip_in_place(&mut self, n: usize) {
60        let n = n * self.chunk_size;
61        self.slice = &self.slice[n.min(self.slice.len())..];
62    }
63}