1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56
use serde::{Deserialize, Serialize}; use super::{Consumer, DistributedIterator}; #[must_use] pub struct Chain<A, B> { a: A, b: B, } impl<A, B> Chain<A, B> { pub(super) fn new(a: A, b: B) -> Self { Self { a, b } } } impl<A: DistributedIterator, B: DistributedIterator<Item = A::Item>> DistributedIterator for Chain<A, B> { type Item = A::Item; type Task = ChainConsumer<A::Task, B::Task>; fn size_hint(&self) -> (usize, Option<usize>) { let (a_lower, a_upper) = self.a.size_hint(); let (b_lower, b_upper) = self.b.size_hint(); ( a_lower + b_lower, if let (Some(a), Some(b)) = (a_upper, b_upper) { Some(a + b) } else { None }, ) } fn next_task(&mut self) -> Option<Self::Task> { self.a .next_task() .map(ChainConsumer::A) .or_else(|| self.b.next_task().map(ChainConsumer::B)) } } #[derive(Serialize, Deserialize)] pub enum ChainConsumer<A, B> { A(A), B(B), } impl<A: Consumer, B: Consumer<Item = A::Item>> Consumer for ChainConsumer<A, B> { type Item = A::Item; fn run(self, i: &mut impl FnMut(Self::Item) -> bool) -> bool { match self { Self::A(a) => a.run(i), Self::B(b) => b.run(i), } } }