tree_layout/tree/
mod.rs

1use std::vec::IntoIter;
2
3/// A vector backed tree implementation used to essentially cache the user's tree.
4use crate::NodeInfo;
5
6#[cfg(test)]
7mod test;
8
9#[derive(Clone, Debug, Eq, PartialEq)]
10pub struct TreeNode<D> {
11    /// Data needed by the actual algorithm.
12    pub data: D,
13
14    /// The position of this node among it's siblings.
15    ///
16    /// Can also be thought of as the number of left-siblings this node has.
17    pub order: usize,
18    /// The depth of this node.
19    ///
20    /// Can be thought of as the number of edges between this node and the root node.
21    pub depth: usize,
22
23    /// The index into `Tree` of this node's parent.
24    pub parent: Option<usize>,
25    /// The indices into `Tree` of this node's children.
26    pub children: Vec<usize>,
27}
28
29impl<D> TreeNode<D> {
30    pub fn is_leaf(&self) -> bool {
31        self.children.is_empty()
32    }
33
34    pub fn is_root(&self) -> bool {
35        self.parent.is_none()
36    }
37}
38
39#[derive(Clone, Debug, Eq, PartialEq)]
40pub struct TreeLayout<D> {
41    arena: Vec<TreeNode<D>>,
42}
43
44impl<D> Default for TreeLayout<D> {
45    fn default() -> Self {
46        Self { arena: vec![] }
47    }
48}
49
50impl<D> TreeLayout<D> {
51    pub fn new<F, N, T>(user_tree: &T, root: N, data: F) -> Self
52    where
53        F: Fn(&T, N) -> D,
54        N: Clone,
55        T: NodeInfo<N>,
56    {
57        let mut tree = Vec::new();
58        tree.push(TreeNode { data: data(user_tree, root.clone()), order: 0, depth: 0, parent: None, children: Vec::new() });
59
60        let mut queue = std::collections::VecDeque::new();
61        queue.push_back((0, root));
62
63        while let Some((parent, node)) = queue.pop_front() {
64            let index = tree.len();
65
66            for (i, child) in user_tree.children(node).into_iter().enumerate() {
67                let index = index + i;
68                let depth = tree[parent].depth + 1;
69
70                tree[parent].children.push(index);
71
72                tree.push(TreeNode {
73                    data: data(user_tree, child.clone()),
74
75                    order: i,
76                    depth,
77
78                    parent: Some(parent),
79                    children: Vec::new(),
80                });
81
82                queue.push_back((index, child));
83            }
84        }
85
86        TreeLayout { arena: tree }
87    }
88
89    pub fn root(&self) -> Option<usize> {
90        if self.arena.is_empty() { None } else { Some(0) }
91    }
92
93    pub fn breadth_first(&self, node: usize) -> Vec<usize> {
94        let mut breadth_first = vec![node];
95        let mut index = 0;
96
97        while index < breadth_first.len() {
98            let node = breadth_first[index];
99            breadth_first.extend_from_slice(&self[node].children);
100            index += 1;
101        }
102
103        breadth_first
104    }
105
106    pub fn post_order(&self, node: usize) -> Vec<usize> {
107        let mut breadth_first = vec![node];
108        let mut post_order = Vec::new();
109
110        while let Some(node) = breadth_first.pop() {
111            breadth_first.extend_from_slice(&self[node].children);
112            post_order.push(node);
113        }
114
115        post_order.reverse();
116        post_order
117    }
118
119    pub fn left_siblings(&self, node: usize) -> Vec<usize> {
120        let order = self[node].order;
121
122        if let Some(parent) = self[node].parent { self[parent].children[0..order].into() } else { Vec::new() }
123    }
124
125    pub fn siblings_between(&self, left: usize, right: usize) -> Vec<usize> {
126        let left_order = self[left].order;
127        let right_order = self[right].order;
128
129        if self[left].is_root() || self[right].is_root() {
130            assert!(self[left].is_root(), "If one node is the root then both nodes must be.");
131            assert!(self[right].is_root(), "If one node is the root then both nodes must be.");
132
133            return Vec::new();
134        }
135
136        let left_parent = self[left].parent.expect("`is_none` has already been checked.");
137
138        let right_parent = self[right].parent.expect("`is_none` has already been checked.");
139
140        assert_eq!(left_parent, right_parent, "Nodes must actually be siblings.");
141
142        let parent = left_parent;
143
144        self[parent].children[left_order + 1..right_order].into()
145    }
146
147    pub fn previous_sibling(&self, node: usize) -> Option<usize> {
148        let order = self[node].order;
149        if order == 0 {
150            return None;
151        }
152
153        let parent = self[node].parent.expect("Nodes where `order != 0` always have parents.");
154
155        Some(self[parent].children[order - 1])
156    }
157
158    pub fn find_parent(&self, node: &TreeNode<D>) -> Option<&TreeNode<D>> {
159        self.arena.get(node.parent?)
160    }
161}
162
163impl<D> IntoIterator for TreeLayout<D> {
164    type Item = TreeNode<D>;
165    type IntoIter = IntoIter<TreeNode<D>>;
166
167    fn into_iter(self) -> Self::IntoIter {
168        self.arena.into_iter()
169    }
170}
171
172impl<D> std::ops::Index<usize> for TreeLayout<D> {
173    type Output = TreeNode<D>;
174
175    fn index(&self, index: usize) -> &Self::Output {
176        &self.arena[index]
177    }
178}
179
180impl<D> std::ops::IndexMut<usize> for TreeLayout<D> {
181    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
182        &mut self.arena[index]
183    }
184}
tree_layout/tree/mod.rs

tree_layout/tree/
mod.rs
