swc_common/util/
iter.rs

1pub trait IteratorExt: Iterator {
2    /// Copied from https://stackoverflow.com/a/49456265/6193633
3    fn chain_with<F, I>(self, f: F) -> ChainWith<Self, F, I::IntoIter>
4    where
5        Self: Sized,
6        F: FnOnce() -> I,
7        I: IntoIterator<Item = Self::Item>,
8    {
9        ChainWith {
10            base: self,
11            factory: Some(f),
12            iterator: None,
13        }
14    }
15}
16
17impl<I: Iterator> IteratorExt for I {}
18
19pub struct ChainWith<B, F, I> {
20    base: B,
21    factory: Option<F>,
22    iterator: Option<I>,
23}
24
25impl<B, F, I> Iterator for ChainWith<B, F, I::IntoIter>
26where
27    B: Iterator,
28    F: FnOnce() -> I,
29    I: IntoIterator<Item = B::Item>,
30{
31    type Item = I::Item;
32
33    fn next(&mut self) -> Option<Self::Item> {
34        if let Some(b) = self.base.next() {
35            return Some(b);
36        }
37
38        // Exhausted the first, generate the second
39
40        if let Some(f) = self.factory.take() {
41            self.iterator = Some(f().into_iter());
42        }
43
44        self.iterator
45            .as_mut()
46            .expect("There must be an iterator")
47            .next()
48    }
49}