iter_tree/tree/
iter.rs

1use std::iter::{once, Once};
2
3use super::Tree;
4
5pub enum IntoIter<T> {
6    Unique(Once<T>),
7    Multiple(NoneIntoIter<T>),
8}
9
10impl<T> From<Tree<T>> for IntoIter<T> {
11    fn from(value: Tree<T>) -> Self {
12        match value {
13            Tree::Leaf(value) => Self::Unique(once(value)),
14            Tree::Node(multiple) => Self::Multiple(NoneIntoIter {
15                state: NodeIterState::Normal,
16                iter: multiple.into_iter(),
17            }),
18        }
19    }
20}
21
22impl<T> Iterator for IntoIter<T> {
23    type Item = T;
24
25    fn next(&mut self) -> Option<Self::Item> {
26        match self {
27            Self::Unique(once) => once.next(),
28            Self::Multiple(multiple) => multiple.next(),
29        }
30    }
31}
32
33pub enum NodeIterState<T> {
34    Normal,
35    Recursion(Box<NoneIntoIter<T>>),
36}
37
38pub struct NoneIntoIter<T> {
39    state: NodeIterState<T>,
40    iter: std::vec::IntoIter<Tree<T>>,
41}
42
43impl<T> Iterator for NoneIntoIter<T> {
44    type Item = T;
45
46    fn next(&mut self) -> Option<Self::Item> {
47        match &mut self.state {
48            NodeIterState::Normal => {
49                let next = self.iter.next()?;
50
51                match next {
52                    Tree::Leaf(value) => Some(value),
53                    Tree::Node(trees) => {
54                        self.state = NodeIterState::Recursion(Box::new(Self {
55                            state: NodeIterState::Normal,
56                            iter: trees.into_iter(),
57                        }));
58                        self.next()
59                    }
60                }
61            }
62            NodeIterState::Recursion(rec) => {
63                let next = rec.next();
64
65                match next {
66                    None => {
67                        self.state = NodeIterState::Normal;
68                        self.next()
69                    }
70                    value => value,
71                }
72            }
73        }
74    }
75}
76
77impl<Item> IntoIterator for Tree<Item> {
78    type Item = Item;
79
80    type IntoIter = IntoIter<Item>;
81
82    fn into_iter(self) -> Self::IntoIter {
83        self.into()
84    }
85}