orx_iterable/transformations/
flat_mapped.rs

1use crate::Iterable;
2
3/// Wraps an `Iterable` and creates a new `Iterable` which maps-and-flattens the elements
4/// of the original iterable.
5pub struct FlatMapped<I, M, U>
6where
7    I: Iterable,
8    U: IntoIterator,
9    M: Fn(I::Item) -> U + Copy,
10{
11    pub(crate) it: I,
12    pub(crate) flat_map: M,
13}
14
15impl<I, M, U> Iterable for FlatMapped<I, M, U>
16where
17    I: Iterable,
18    U: IntoIterator,
19    M: Fn(I::Item) -> U + Copy,
20{
21    type Item = U::Item;
22
23    type Iter = FlatMappedIter<I, M, U>;
24
25    fn iter(&self) -> Self::Iter {
26        let mut iter1 = self.it.iter();
27        let iterable2: Option<U> = iter1.next().map(self.flat_map);
28        let iter2: Option<U::IntoIter> = iterable2.map(|x| x.into_iter());
29
30        FlatMappedIter {
31            flat_map: self.flat_map,
32            iter1,
33            iter2,
34        }
35    }
36}
37
38/// Flat mapped iterator for iterables.
39pub struct FlatMappedIter<I, M, U>
40where
41    I: Iterable,
42    U: IntoIterator,
43    M: Fn(I::Item) -> U + Copy,
44{
45    pub(crate) iter1: I::Iter,
46    pub(crate) iter2: Option<U::IntoIter>,
47    pub(crate) flat_map: M,
48}
49
50impl<I, M, U> Iterator for FlatMappedIter<I, M, U>
51where
52    I: Iterable,
53    U: IntoIterator,
54    M: Fn(I::Item) -> U + Copy,
55{
56    type Item = U::Item;
57
58    fn next(&mut self) -> Option<Self::Item> {
59        loop {
60            let iter2 = self.iter2.as_mut()?;
61            let value = iter2.next();
62
63            match value.is_some() {
64                true => return value,
65                false => {
66                    let x = self.iter1.next()?;
67                    let iterable2: U = (self.flat_map)(x);
68                    let iter2: U::IntoIter = iterable2.into_iter();
69                    self.iter2 = Some(iter2);
70                }
71            }
72        }
73    }
74}