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,
62{
63    chunk: core::iter::Enumerate<I>,
64    begin_idx: usize,
65}
66
67impl<I> Iterator for EnumeratedChunk<I>
68where
69    I: ExactSizeIterator,
70{
71    type Item = (usize, I::Item);
72
73    fn next(&mut self) -> Option<Self::Item> {
74        self.chunk.next().map(|(i, x)| (self.begin_idx + i, x))
75    }
76
77    fn size_hint(&self) -> (usize, Option<usize>) {
78        let len = self.chunk.len();
79        (len, Some(len))
80    }
81}
82
83impl<I> ExactSizeIterator for EnumeratedChunk<I>
84where
85    I: ExactSizeIterator,
86{
87    fn len(&self) -> usize {
88        self.chunk.len()
89    }
90}