rustic_rs/commands/tui/
tree.rs

1#[derive(PartialEq, Eq)]
2pub struct TreeNode<Data, LeafData> {
3    pub data: Data,
4    pub open: bool,
5    pub children: Vec<Tree<Data, LeafData>>,
6}
7
8#[derive(PartialEq, Eq)]
9pub enum Tree<Data, LeafData> {
10    Node(TreeNode<Data, LeafData>),
11    Leaf(LeafData),
12}
13
14impl<Data, LeafData> Tree<Data, LeafData> {
15    pub fn leaf(data: LeafData) -> Self {
16        Self::Leaf(data)
17    }
18    pub fn node(data: Data, open: bool, children: Vec<Self>) -> Self {
19        Self::Node(TreeNode {
20            data,
21            open,
22            children,
23        })
24    }
25
26    pub fn child_count(&self) -> usize {
27        match self {
28            Self::Leaf(_) => 0,
29            Self::Node(TreeNode { children, .. }) => {
30                children.len() + children.iter().map(Self::child_count).sum::<usize>()
31            }
32        }
33    }
34
35    pub fn leaf_data(&self) -> Option<&LeafData> {
36        match self {
37            Self::Node(_) => None,
38            Self::Leaf(data) => Some(data),
39        }
40    }
41
42    pub fn openable(&self) -> bool {
43        matches!(self, Self::Node(node) if !node.open)
44    }
45
46    pub fn open(&mut self) {
47        if let Self::Node(node) = self {
48            node.open = true;
49        }
50    }
51    pub fn close(&mut self) {
52        if let Self::Node(node) = self {
53            node.open = false;
54        }
55    }
56
57    pub fn iter(&self) -> impl Iterator<Item = TreeIterItem<'_, Data, LeafData>> {
58        TreeIter {
59            tree: Some(self),
60            iter_stack: Vec::new(),
61            only_open: false,
62        }
63    }
64
65    // iter open tree descending only into open nodes.
66    // Note: This iterator skips the root node!
67    pub fn iter_open(&self) -> impl Iterator<Item = TreeIterItem<'_, Data, LeafData>> {
68        TreeIter {
69            tree: Some(self),
70            iter_stack: Vec::new(),
71            only_open: true,
72        }
73        .skip(1)
74    }
75
76    pub fn nth_mut(&mut self, n: usize) -> Option<&mut Self> {
77        let mut count = 0;
78        let mut tree = Some(self);
79        let mut iter_stack = Vec::new();
80        loop {
81            if count == n + 1 {
82                return tree;
83            }
84            let item = tree?;
85            if let Self::Node(node) = item {
86                if node.open {
87                    iter_stack.push(node.children.iter_mut());
88                }
89            }
90            tree = next_from_iter_stack(&mut iter_stack);
91            count += 1;
92        }
93    }
94}
95
96pub struct TreeIterItem<'a, Data, LeadData> {
97    pub depth: usize,
98    pub tree: &'a Tree<Data, LeadData>,
99}
100
101impl<Data, LeafData> TreeIterItem<'_, Data, LeafData> {
102    pub fn leaf_data(&self) -> Option<&LeafData> {
103        self.tree.leaf_data()
104    }
105}
106
107pub struct TreeIter<'a, Data, LeafData> {
108    tree: Option<&'a Tree<Data, LeafData>>,
109    iter_stack: Vec<std::slice::Iter<'a, Tree<Data, LeafData>>>,
110    only_open: bool,
111}
112
113impl<'a, Data, LeafData> Iterator for TreeIter<'a, Data, LeafData> {
114    type Item = TreeIterItem<'a, Data, LeafData>;
115    fn next(&mut self) -> Option<Self::Item> {
116        let item = self.tree?;
117        let depth = self.iter_stack.len();
118        if let Tree::Node(node) = item {
119            if !self.only_open || node.open {
120                self.iter_stack.push(node.children.iter());
121            }
122        }
123
124        self.tree = next_from_iter_stack(&mut self.iter_stack);
125        Some(TreeIterItem { depth, tree: item })
126    }
127}
128
129// helper function to get next item from iteration stack when iterating over a Tree
130fn next_from_iter_stack<T>(stack: &mut Vec<impl Iterator<Item = T>>) -> Option<T> {
131    loop {
132        match stack.pop() {
133            None => {
134                break None;
135            }
136            Some(mut iter) => {
137                if let Some(next) = iter.next() {
138                    stack.push(iter);
139                    break Some(next);
140                }
141            }
142        }
143    }
144}
rustic_rs/commands/tui/tree.rs

rustic_rs/commands/tui/
tree.rs
