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 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96
use crate::{Generator, GeneratorResult, SliceGenerator, ValueResult}; pub struct StoppingGen<'a, T> { stop_at: i32, stopped_data: Option<&'a T>, data: SliceGenerator<'a, T>, } impl<'a, T> StoppingGen<'a, T> { pub fn new(stop_at: i32, data: &'a [T]) -> Self { Self { stop_at, stopped_data: None, data: SliceGenerator::new(data), } } } impl<'a, T> Generator for StoppingGen<'a, T> { type Output = &'a T; fn run(&mut self, mut output: impl FnMut(Self::Output) -> ValueResult) -> GeneratorResult { if self.stop_at == 0 { self.stop_at -= 1; return GeneratorResult::Stopped; } if let Some(x) = self.stopped_data.take() { if output(x) == ValueResult::Stop { return GeneratorResult::Stopped; } } let stored_stop = &mut self.stopped_data; let stop_at = &mut self.stop_at; let result = self.data.run(|x| { let old_stop_at = *stop_at; *stop_at -= 1; if old_stop_at == 0 { *stored_stop = Some(x); ValueResult::Stop } else { output(x) } }); if result == GeneratorResult::Complete { *stop_at = -1; } result } } pub struct MultiStoppingGen<'a, T> { stop_at: &'a [usize], stop_at_index: usize, index: usize, data: &'a [T], } impl<'a, T> MultiStoppingGen<'a, T> { pub fn new(stop_at: &'a [usize], data: &'a [T]) -> Self { Self { stop_at, stop_at_index: 0, index: 0, data, } } } impl<'a, T> Generator for MultiStoppingGen<'a, T> { type Output = &'a T; fn run(&mut self, mut output: impl FnMut(Self::Output) -> ValueResult) -> GeneratorResult { while self.data.len() > 0 { if self.stop_at_index < self.stop_at.len() && self.index == self.stop_at[self.stop_at_index] { self.index += 1; self.stop_at_index += 1; return GeneratorResult::Stopped; } let data = &self.data[0]; let out_result = output(data); self.data = self.data.split_at(1).1; self.index += 1; if out_result == ValueResult::Stop { return GeneratorResult::Stopped; } } GeneratorResult::Complete } }