moirai_iter/parallel/adapters/
blocks.rs

1use super::super::{Consumer, ParallelIterator};
2use std::ops::ControlFlow;
3
4#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
5struct ExponentialBlockPolicy;
6
7#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
8struct UniformBlockPolicy;
9
10#[repr(transparent)]
11#[derive(Clone, Copy, Debug, Eq, PartialEq)]
12struct BlockSize(usize);
13
14impl BlockSize {
15    fn new(value: usize) -> Self {
16        assert!(value != 0, "block size must be non-zero");
17        Self(value)
18    }
19}
20
21/// Indexed source adapter for exponential logical block scheduling.
22pub struct ExponentialBlocks<I> {
23    base: I,
24    policy: ExponentialBlockPolicy,
25}
26
27impl<I> ExponentialBlocks<I> {
28    pub(in crate::parallel) fn new(base: I) -> Self {
29        Self {
30            base,
31            policy: ExponentialBlockPolicy,
32        }
33    }
34}
35
36impl<I> ParallelIterator for ExponentialBlocks<I>
37where
38    I: ParallelIterator,
39    I::Item: Sync + 'static,
40{
41    type Item = I::Item;
42
43    fn seq_items(self) -> Vec<Self::Item> {
44        let _policy = self.policy;
45        self.base.seq_items()
46    }
47
48    fn seq_try_fold<Acc, B, FoldFn>(self, init: Acc, fold_fn: FoldFn) -> ControlFlow<B, Acc>
49    where
50        FoldFn: FnMut(Acc, Self::Item) -> ControlFlow<B, Acc>,
51    {
52        let _policy = self.policy;
53        self.base.seq_try_fold(init, fold_fn)
54    }
55
56    fn drive<C, R>(self, consumer: C) -> R
57    where
58        C: Consumer<Self::Item, Result = R> + Send + Sync,
59        R: Send,
60    {
61        // The block policy selects a logical scheduling shape and leaves the
62        // item stream identical to the base's, so this adapter drives the base
63        // directly rather than collecting it. Collecting cost the source's
64        // shards to express an identity.
65        let _policy = self.policy;
66        self.base.drive(consumer)
67    }
68}
69
70/// Indexed source adapter for uniform logical block scheduling.
71pub struct UniformBlocks<I> {
72    base: I,
73    block_size: BlockSize,
74    policy: UniformBlockPolicy,
75}
76
77impl<I> UniformBlocks<I> {
78    pub(in crate::parallel) fn new(base: I, block_size: usize) -> Self {
79        Self {
80            base,
81            block_size: BlockSize::new(block_size),
82            policy: UniformBlockPolicy,
83        }
84    }
85}
86
87impl<I> ParallelIterator for UniformBlocks<I>
88where
89    I: ParallelIterator,
90    I::Item: Sync + 'static,
91{
92    type Item = I::Item;
93
94    fn seq_items(self) -> Vec<Self::Item> {
95        let _block_size = self.block_size;
96        let _policy = self.policy;
97        self.base.seq_items()
98    }
99
100    fn seq_try_fold<Acc, B, FoldFn>(self, init: Acc, fold_fn: FoldFn) -> ControlFlow<B, Acc>
101    where
102        FoldFn: FnMut(Acc, Self::Item) -> ControlFlow<B, Acc>,
103    {
104        let _block_size = self.block_size;
105        let _policy = self.policy;
106        self.base.seq_try_fold(init, fold_fn)
107    }
108
109    fn drive<C, R>(self, consumer: C) -> R
110    where
111        C: Consumer<Self::Item, Result = R> + Send + Sync,
112        R: Send,
113    {
114        // Identity item stream, as on [`ExponentialBlocks`].
115        let _block_size = self.block_size;
116        let _policy = self.policy;
117        self.base.drive(consumer)
118    }
119}
120
121#[cfg(test)]
122mod tests {
123    use super::*;
124
125    #[test]
126    fn block_policy_markers_are_zero_sized() {
127        assert_eq!(std::mem::size_of::<ExponentialBlockPolicy>(), 0);
128        assert_eq!(std::mem::size_of::<UniformBlockPolicy>(), 0);
129    }
130}
moirai_iter/parallel/adapters/blocks.rs

moirai_iter/parallel/adapters/
blocks.rs
