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
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
#[derive(PartialEq, Eq)]
pub struct TreeNode<Data, LeafData> {
    pub data: Data,
    pub open: bool,
    pub children: Vec<Tree<Data, LeafData>>,
}

#[derive(PartialEq, Eq)]
pub enum Tree<Data, LeafData> {
    Node(TreeNode<Data, LeafData>),
    Leaf(LeafData),
}

impl<Data, LeafData> Tree<Data, LeafData> {
    pub fn leaf(data: LeafData) -> Self {
        Self::Leaf(data)
    }
    pub fn node(data: Data, open: bool, children: Vec<Self>) -> Self {
        Self::Node(TreeNode {
            data,
            open,
            children,
        })
    }

    pub fn child_count(&self) -> usize {
        match self {
            Self::Leaf(_) => 0,
            Self::Node(TreeNode { children, .. }) => {
                children.len() + children.iter().map(Self::child_count).sum::<usize>()
            }
        }
    }

    pub fn leaf_data(&self) -> Option<&LeafData> {
        match self {
            Self::Node(_) => None,
            Self::Leaf(data) => Some(data),
        }
    }

    pub fn openable(&self) -> bool {
        matches!(self, Self::Node(node) if !node.open)
    }

    pub fn open(&mut self) {
        if let Self::Node(node) = self {
            node.open = true;
        }
    }
    pub fn close(&mut self) {
        if let Self::Node(node) = self {
            node.open = false;
        }
    }

    pub fn iter(&self) -> impl Iterator<Item = TreeIterItem<'_, Data, LeafData>> {
        TreeIter {
            tree: Some(self),
            iter_stack: Vec::new(),
            only_open: false,
        }
    }

    // iter open tree descending only into open nodes.
    // Note: This iterator skips the root node!
    pub fn iter_open(&self) -> impl Iterator<Item = TreeIterItem<'_, Data, LeafData>> {
        TreeIter {
            tree: Some(self),
            iter_stack: Vec::new(),
            only_open: true,
        }
        .skip(1)
    }

    pub fn nth_mut(&mut self, n: usize) -> Option<&mut Self> {
        let mut count = 0;
        let mut tree = Some(self);
        let mut iter_stack = Vec::new();
        loop {
            if count == n + 1 {
                return tree;
            }
            let item = tree?;
            if let Self::Node(node) = item {
                if node.open {
                    iter_stack.push(node.children.iter_mut());
                }
            }
            tree = next_from_iter_stack(&mut iter_stack);
            count += 1;
        }
    }
}

pub struct TreeIterItem<'a, Data, LeadData> {
    pub depth: usize,
    pub tree: &'a Tree<Data, LeadData>,
}

impl<'a, Data, LeafData> TreeIterItem<'a, Data, LeafData> {
    pub fn leaf_data(&self) -> Option<&LeafData> {
        self.tree.leaf_data()
    }
}

pub struct TreeIter<'a, Data, LeafData> {
    tree: Option<&'a Tree<Data, LeafData>>,
    iter_stack: Vec<std::slice::Iter<'a, Tree<Data, LeafData>>>,
    only_open: bool,
}

impl<'a, Data, LeafData> Iterator for TreeIter<'a, Data, LeafData> {
    type Item = TreeIterItem<'a, Data, LeafData>;
    fn next(&mut self) -> Option<Self::Item> {
        let item = self.tree?;
        let depth = self.iter_stack.len();
        if let Tree::Node(node) = item {
            if !self.only_open || node.open {
                self.iter_stack.push(node.children.iter());
            }
        }

        self.tree = next_from_iter_stack(&mut self.iter_stack);
        Some(TreeIterItem { depth, tree: item })
    }
}

// helper function to get next item from iteration stack when iterating over a Tree
fn next_from_iter_stack<T>(stack: &mut Vec<impl Iterator<Item = T>>) -> Option<T> {
    loop {
        match stack.pop() {
            None => {
                break None;
            }
            Some(mut iter) => {
                if let Some(next) = iter.next() {
                    stack.push(iter);
                    break Some(next);
                }
            }
        }
    }
}