]>
Commit | Line | Data |
---|---|---|
1b1a35ee XL |
1 | // Copyright 2018 Google LLC |
2 | // | |
3 | // Use of this source code is governed by an MIT-style | |
4 | // license that can be found in the LICENSE file or at | |
5 | // https://opensource.org/licenses/MIT. | |
6 | ||
7 | //! A Vec-based container for a tree structure. | |
8 | ||
9 | use std::num::NonZeroUsize; | |
10 | use std::ops::{Add, Sub}; | |
11 | ||
12 | #[derive(Debug, Eq, PartialEq, Copy, Clone, PartialOrd)] | |
13 | pub struct TreeIndex(NonZeroUsize); | |
14 | ||
15 | impl TreeIndex { | |
16 | fn new(i: usize) -> Self { | |
17 | TreeIndex(NonZeroUsize::new(i).unwrap()) | |
18 | } | |
19 | ||
20 | pub fn get(self) -> usize { | |
21 | self.0.get() | |
22 | } | |
23 | } | |
24 | ||
25 | impl Add<usize> for TreeIndex { | |
26 | type Output = TreeIndex; | |
27 | ||
28 | fn add(self, rhs: usize) -> Self { | |
29 | let inner = self.0.get() + rhs; | |
30 | TreeIndex::new(inner) | |
31 | } | |
32 | } | |
33 | ||
34 | impl Sub<usize> for TreeIndex { | |
35 | type Output = TreeIndex; | |
36 | ||
37 | fn sub(self, rhs: usize) -> Self { | |
38 | let inner = self.0.get().checked_sub(rhs).unwrap(); | |
39 | TreeIndex::new(inner) | |
40 | } | |
41 | } | |
42 | ||
43 | #[derive(Debug, Clone, Copy)] | |
44 | pub struct Node<T> { | |
45 | pub child: Option<TreeIndex>, | |
46 | pub next: Option<TreeIndex>, | |
47 | pub item: T, | |
48 | } | |
49 | ||
50 | /// A tree abstraction, intended for fast building as a preorder traversal. | |
51 | #[derive(Clone)] | |
52 | pub struct Tree<T> { | |
53 | nodes: Vec<Node<T>>, | |
54 | spine: Vec<TreeIndex>, // indices of nodes on path to current node | |
55 | cur: Option<TreeIndex>, | |
56 | } | |
57 | ||
58 | impl<T: Default> Tree<T> { | |
59 | // Indices start at one, so we place a dummy value at index zero. | |
60 | // The alternative would be subtracting one from every TreeIndex | |
61 | // every time we convert it to usize to index our nodes. | |
62 | pub fn with_capacity(cap: usize) -> Tree<T> { | |
63 | let mut nodes = Vec::with_capacity(cap); | |
64 | nodes.push(Node { | |
65 | child: None, | |
66 | next: None, | |
67 | item: <T as Default>::default(), | |
68 | }); | |
69 | Tree { | |
70 | nodes, | |
71 | spine: Vec::new(), | |
72 | cur: None, | |
73 | } | |
74 | } | |
75 | ||
76 | /// Returns the index of the element currently in focus. | |
77 | pub fn cur(&self) -> Option<TreeIndex> { | |
78 | self.cur | |
79 | } | |
80 | ||
81 | /// Append one item to the current position in the tree. | |
82 | pub fn append(&mut self, item: T) -> TreeIndex { | |
83 | let ix = self.create_node(item); | |
84 | let this = Some(ix); | |
85 | ||
86 | if let Some(ix) = self.cur { | |
87 | self[ix].next = this; | |
88 | } else if let Some(&parent) = self.spine.last() { | |
89 | self[parent].child = this; | |
90 | } | |
91 | self.cur = this; | |
92 | ix | |
93 | } | |
94 | ||
95 | /// Create an isolated node. | |
96 | pub fn create_node(&mut self, item: T) -> TreeIndex { | |
97 | let this = self.nodes.len(); | |
98 | self.nodes.push(Node { | |
99 | child: None, | |
100 | next: None, | |
101 | item, | |
102 | }); | |
103 | TreeIndex::new(this) | |
104 | } | |
105 | ||
106 | /// Push down one level, so that new items become children of the current node. | |
107 | /// The new focus index is returned. | |
108 | pub fn push(&mut self) -> TreeIndex { | |
109 | let cur_ix = self.cur.unwrap(); | |
110 | self.spine.push(cur_ix); | |
111 | self.cur = self[cur_ix].child; | |
112 | cur_ix | |
113 | } | |
114 | ||
115 | /// Pop back up a level. | |
116 | pub fn pop(&mut self) -> Option<TreeIndex> { | |
117 | let ix = Some(self.spine.pop()?); | |
118 | self.cur = ix; | |
119 | ix | |
120 | } | |
121 | ||
122 | /// Look at the parent node. | |
123 | pub fn peek_up(&self) -> Option<TreeIndex> { | |
124 | self.spine.last().copied() | |
125 | } | |
126 | ||
127 | /// Look at grandparent node. | |
128 | pub fn peek_grandparent(&self) -> Option<TreeIndex> { | |
129 | if self.spine.len() >= 2 { | |
130 | Some(self.spine[self.spine.len() - 2]) | |
131 | } else { | |
132 | None | |
133 | } | |
134 | } | |
135 | ||
136 | /// Returns true when there are no nodes in the tree, false otherwise. | |
137 | pub fn is_empty(&self) -> bool { | |
138 | self.nodes.len() <= 1 | |
139 | } | |
140 | ||
141 | /// Returns the length of the spine. | |
142 | pub fn spine_len(&self) -> usize { | |
143 | self.spine.len() | |
144 | } | |
145 | ||
146 | /// Resets the focus to the first node added to the tree, if it exists. | |
147 | pub fn reset(&mut self) { | |
148 | self.cur = if self.is_empty() { | |
149 | None | |
150 | } else { | |
151 | Some(TreeIndex::new(1)) | |
152 | }; | |
153 | self.spine.clear(); | |
154 | } | |
155 | ||
156 | /// Walks the spine from a root node up to, but not including, the current node. | |
157 | pub fn walk_spine(&self) -> impl std::iter::DoubleEndedIterator<Item = &TreeIndex> { | |
158 | self.spine.iter() | |
159 | } | |
160 | ||
161 | /// Moves focus to the next sibling of the given node. | |
162 | pub fn next_sibling(&mut self, cur_ix: TreeIndex) -> Option<TreeIndex> { | |
163 | self.cur = self[cur_ix].next; | |
164 | self.cur | |
165 | } | |
166 | } | |
167 | ||
168 | impl<T> std::fmt::Debug for Tree<T> | |
169 | where | |
170 | T: std::fmt::Debug, | |
171 | { | |
172 | fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { | |
173 | fn debug_tree<T>( | |
174 | tree: &Tree<T>, | |
175 | cur: TreeIndex, | |
176 | indent: usize, | |
177 | f: &mut std::fmt::Formatter, | |
178 | ) -> std::fmt::Result | |
179 | where | |
180 | T: std::fmt::Debug, | |
181 | { | |
182 | for _ in 0..indent { | |
183 | write!(f, " ")?; | |
184 | } | |
185 | writeln!(f, "{:?}", &tree[cur].item)?; | |
186 | if let Some(child_ix) = tree[cur].child { | |
187 | debug_tree(tree, child_ix, indent + 1, f)?; | |
188 | } | |
189 | if let Some(next_ix) = tree[cur].next { | |
190 | debug_tree(tree, next_ix, indent, f)?; | |
191 | } | |
192 | Ok(()) | |
193 | } | |
194 | ||
195 | if self.nodes.len() > 1 { | |
196 | let cur = TreeIndex(NonZeroUsize::new(1).unwrap()); | |
197 | debug_tree(self, cur, 0, f) | |
198 | } else { | |
199 | write!(f, "Empty tree") | |
200 | } | |
201 | } | |
202 | } | |
203 | ||
204 | impl<T> std::ops::Index<TreeIndex> for Tree<T> { | |
205 | type Output = Node<T>; | |
206 | ||
207 | fn index(&self, ix: TreeIndex) -> &Self::Output { | |
208 | self.nodes.index(ix.get()) | |
209 | } | |
210 | } | |
211 | ||
212 | impl<T> std::ops::IndexMut<TreeIndex> for Tree<T> { | |
213 | fn index_mut(&mut self, ix: TreeIndex) -> &mut Node<T> { | |
214 | self.nodes.index_mut(ix.get()) | |
215 | } | |
216 | } |