orx_iterable/transformations/
flat_mapped.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
use crate::Iterable;

/// Wraps an `Iterable` and creates a new `Iterable` which maps-and-flattens the elements
/// of the original iterable.
pub struct FlatMapped<I, M, U>
where
    I: Iterable,
    U: IntoIterator,
    M: Fn(I::Item) -> U + Copy,
{
    pub(crate) it: I,
    pub(crate) flat_map: M,
}

impl<I, M, U> Iterable for FlatMapped<I, M, U>
where
    I: Iterable,
    U: IntoIterator,
    M: Fn(I::Item) -> U + Copy,
{
    type Item = U::Item;

    type Iter = FlatMappedIter<I, M, U>;

    fn iter(&self) -> Self::Iter {
        let mut iter1 = self.it.iter();
        let iterable2: Option<U> = iter1.next().map(self.flat_map);
        let iter2: Option<U::IntoIter> = iterable2.map(|x| x.into_iter());

        FlatMappedIter {
            flat_map: self.flat_map,
            iter1,
            iter2,
        }
    }
}

/// Flat mapped iterator for iterables.
pub struct FlatMappedIter<I, M, U>
where
    I: Iterable,
    U: IntoIterator,
    M: Fn(I::Item) -> U + Copy,
{
    pub(crate) iter1: I::Iter,
    pub(crate) iter2: Option<U::IntoIter>,
    pub(crate) flat_map: M,
}

impl<I, M, U> Iterator for FlatMappedIter<I, M, U>
where
    I: Iterable,
    U: IntoIterator,
    M: Fn(I::Item) -> U + Copy,
{
    type Item = U::Item;

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            let iter2 = self.iter2.as_mut()?;
            let value = iter2.next();

            match value.is_some() {
                true => return value,
                false => {
                    let x = self.iter1.next()?;
                    let iterable2: U = (self.flat_map)(x);
                    let iter2: U::IntoIter = iterable2.into_iter();
                    self.iter2 = Some(iter2);
                }
            }
        }
    }
}