zhc_utils/iter/
interleave.rs

1/// Interleaves elements from this iterator with elements from a separator iterator.
2pub trait Interleave
3where
4    Self: Iterator + Sized,
5{
6    /// Creates an iterator that alternates between elements from `self` and `s`.
7    ///
8    /// The resulting iterator starts with an element from `self`, then takes one
9    /// from `s`, then back to `self`, and so on. If either iterator is exhausted,
10    /// the combined iterator terminates.
11    fn interleave_with<S: Iterator<Item = Self::Item>>(self, s: S) -> Interleaved<Self, S>;
12}
13
14impl<I: Iterator> Interleave for I {
15    fn interleave_with<S: Iterator<Item = Self::Item>>(self, s: S) -> Interleaved<Self, S> {
16        Interleaved::OnMain(self, s)
17    }
18}
19
20/// An iterator that alternates between elements from two source iterators.
21///
22/// Created by the `separate_with` method. This iterator yields elements by
23/// alternating between the main iterator and the separator iterator until
24/// either one is exhausted.
25pub enum Interleaved<A: Iterator, S: Iterator<Item = A::Item>> {
26    OnMain(A, S),
27    OnSep(S, A),
28    Finished,
29}
30
31impl<A: Iterator, S: Iterator<Item = A::Item>> Interleaved<A, S> {
32    fn to_sep(&mut self) {
33        let Interleaved::OnMain(a, s) = std::mem::replace(self, Interleaved::Finished) else {
34            panic!()
35        };
36        let _ = std::mem::replace(self, Interleaved::OnSep(s, a));
37    }
38
39    fn to_main(&mut self) {
40        let Interleaved::OnSep(s, a) = std::mem::replace(self, Interleaved::Finished) else {
41            panic!()
42        };
43        let _ = std::mem::replace(self, Interleaved::OnMain(a, s));
44    }
45
46    fn alternate(&mut self) {
47        match self {
48            Interleaved::OnMain(_, _) => self.to_sep(),
49            Interleaved::OnSep(_, _) => self.to_main(),
50            Interleaved::Finished => panic!(),
51        }
52    }
53
54    fn finish(&mut self) {
55        let _ = std::mem::replace(self, Interleaved::Finished);
56    }
57}
58
59impl<A: Iterator, S: Iterator<Item = A::Item>> Iterator for Interleaved<A, S> {
60    type Item = A::Item;
61
62    fn next(&mut self) -> Option<Self::Item> {
63        let out = match self {
64            Interleaved::OnMain(a, _) => a.next(),
65            Interleaved::OnSep(s, _) => s.next(),
66            Interleaved::Finished => None,
67        };
68        match out {
69            Some(_) => {
70                self.alternate();
71            }
72            None => {
73                self.finish();
74            }
75        };
76        out
77    }
78}
zhc_utils/iter/interleave.rs

zhc_utils/iter/
interleave.rs
