orx_concurrent_iter/enumerate/
chunk_puller.rs

1use crate::ChunkPuller;
2
3/// Chunk puller of an enumerated concurrent iterator; i.e., [`Enumerate`]
4///
5/// [`Enumerate`]: crate::enumerate::Enumerate
6pub struct EnumeratedChunkPuller<P>
7where
8    P: ChunkPuller,
9{
10    puller: P,
11}
12
13impl<P> EnumeratedChunkPuller<P>
14where
15    P: ChunkPuller,
16{
17    pub(crate) fn new(puller: P) -> Self {
18        Self { puller }
19    }
20}
21
22impl<P> ChunkPuller for EnumeratedChunkPuller<P>
23where
24    P: ChunkPuller,
25{
26    type ChunkItem = (usize, P::ChunkItem);
27
28    type Chunk<'c>
29        = EnumeratedChunk<P::Chunk<'c>>
30    where
31        Self: 'c;
32
33    fn chunk_size(&self) -> usize {
34        self.puller.chunk_size()
35    }
36
37    fn pull(&mut self) -> Option<Self::Chunk<'_>> {
38        self.puller
39            .pull_with_idx()
40            .map(|(begin_idx, x)| EnumeratedChunk {
41                begin_idx,
42                chunk: x.enumerate(),
43            })
44    }
45
46    fn pull_with_idx(&mut self) -> Option<(usize, Self::Chunk<'_>)> {
47        self.puller.pull_with_idx().map(|(begin_idx, x)| {
48            (
49                begin_idx,
50                EnumeratedChunk {
51                    begin_idx,
52                    chunk: x.enumerate(),
53                },
54            )
55        })
56    }
57}
58
59pub struct EnumeratedChunk<I>
60where
61    I: ExactSizeIterator + Default,
62{
63    chunk: core::iter::Enumerate<I>,
64    begin_idx: usize,
65}
66
67impl<I> Default for EnumeratedChunk<I>
68where
69    I: ExactSizeIterator + Default,
70{
71    fn default() -> Self {
72        Self {
73            chunk: Default::default(),
74            begin_idx: 0,
75        }
76    }
77}
78
79impl<I> Iterator for EnumeratedChunk<I>
80where
81    I: ExactSizeIterator + Default,
82{
83    type Item = (usize, I::Item);
84
85    fn next(&mut self) -> Option<Self::Item> {
86        self.chunk.next().map(|(i, x)| (self.begin_idx + i, x))
87    }
88
89    fn size_hint(&self) -> (usize, Option<usize>) {
90        let len = self.chunk.len();
91        (len, Some(len))
92    }
93}
94
95impl<I> ExactSizeIterator for EnumeratedChunk<I>
96where
97    I: ExactSizeIterator + Default,
98{
99    fn len(&self) -> usize {
100        self.chunk.len()
101    }
102}