grass/dev/iterator/
iter.rs1use 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}