grass/dev/iterator/
iter.rs

1use std::iter::repeat_with;
2
3use tracing::error;
4
5pub struct MarkEndIterator<'a, T: Iterator> {
6    source: &'a mut T,
7    result: Vec<Option<T::Item>>,
8}
9
10pub enum TakeUntilIteratorItem<T> {
11    Start(T),
12    End(T),
13}
14
15impl<'a, T: Iterator<Item = U>, U: Clone> Iterator for MarkEndIterator<'a, T> {
16    type Item = TakeUntilIteratorItem<T::Item>;
17
18    fn next(&mut self) -> Option<Self::Item> {
19        loop {
20            let next = if let Some(next) = self.source.next() {
21                next
22            } else {
23                for value in 0..self.result.len() {
24                    let result = match &self.result[value] {
25                        Some(value) => Some(TakeUntilIteratorItem::End(value.clone())),
26                        None => continue,
27                    };
28
29                    self.result[value] = None;
30                    return result;
31                }
32                break;
33            };
34
35            let value = self.result.first().map(|value| (*value).clone());
36            match value {
37                Some(Some(value)) => {
38                    for i in 0..self.result.len() - 1 {
39                        self.result[i] = self.result[i + 1].clone();
40                    }
41                    let length = self.result.len();
42                    self.result[length - 1] = Some(next);
43                    let result = value.clone();
44                    return Some(TakeUntilIteratorItem::Start(result));
45                }
46                Some(None) => {
47                    for i in 0..self.result.len() - 1 {
48                        self.result[i] = self.result[i + 1].clone();
49                    }
50                    let length = self.result.len();
51                    self.result[length - 1] = Some(next);
52                }
53                None => {
54                    error!("TakeUntilIterator called with size 0");
55                    break;
56                }
57            };
58        }
59        None
60    }
61}
62
63pub trait IterExtensions: Iterator + Sized {
64    fn mark_end(&mut self, n: usize) -> MarkEndIterator<Self>;
65}
66
67impl<T: Iterator<Item = U> + Sized, U: Default> IterExtensions for T {
68    fn mark_end(&mut self, n: usize) -> MarkEndIterator<T> {
69        MarkEndIterator {
70            source: self,
71            result: Vec::from_iter(repeat_with(Default::default).take(n)),
72        }
73    }
74}