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

use super::*;

pub trait FlatMapFn {
    type Input;
    type OutputList: ListFn;
    /// Map the given `input` item into a list.
    fn map(&self, input: Self::Input) -> Self::OutputList;
}

pub struct FlatMapList<I: ListFn, F: FlatMapFn<Input = I::Item>> {
    flat_map: F,
    input_list: I,
    output_list: Option<F::OutputList>,
}

impl<I: ListFn, F: FlatMapFn<Input = I::Item>> ListFn for FlatMapList<I, F> {
    type Item = <F::OutputList as ListFn>::Item;
    type End = I::End;
    fn next(mut self) -> ListState<Self> {
        loop {
            match self.output_list {
                Some(item_list) => match item_list.next() {
                    ListState::Some(some) => {
                        return ListState::some(
                            some.first,
                            FlatMapList {
                                flat_map: self.flat_map,
                                input_list: self.input_list,
                                output_list: Some(some.next),
                            },
                        )
                    }
                    ListState::End(..) => self.output_list = None,
                },
                None => match self.input_list.next() {
                    ListState::Some(some) => {
                        self.input_list = some.next;
                        self.output_list = Some(self.flat_map.map(some.first));
                    }
                    ListState::End(end) => return ListState::End(end),
                },
            }
        }
    }
}

pub trait FlatMap: ListFn {
    fn flat_map<F: FlatMapFn<Input = Self::Item>>(self, flat_map: F) -> FlatMapList<Self, F> {
        FlatMapList {
            input_list: self,
            flat_map,
            output_list: None,
        }
    }
}

impl<S: ListFn> FlatMap for S {}