trailtree/implement/update.rs
1use crate::{TrailTree, Node, Cell};
2
3use core::cmp::Ordering;
4use alloc::{vec::Vec, vec};
5
6impl<K,T> TrailTree<K,T>
7where K: Clone + Ord
8{
9 /// Inserts an index path into the tree, returning the ID of the last node in the path.
10 ///
11 /// # Examples
12 ///
13 /// ```
14 /// use trailtree::TrailTree;
15 ///
16 /// let mut con = TrailTree::<u32,()>::new();
17 ///
18 /// let a = con.insert(&[5, 10]);
19 ///
20 /// assert_eq!(con.search(&[5, 10]), Some(a));
21 /// ```
22 ///
23 /// See also [`TrailTree::extend`].
24 pub fn insert(&mut self, path:&[K]) -> usize
25 {
26 self.modify(self.root, path)
27 }
28
29 /// Extends an existing path in the tree with a new path, returning the ID of the last node in the extended path.
30 ///
31 /// # Examples
32 ///
33 /// ```
34 /// use trailtree::TrailTree;
35 ///
36 /// let mut con = TrailTree::<u32,()>::new();
37 ///
38 /// let a = con.insert(&[5]);
39 /// let b = con.extend(a, &[10]);
40 ///
41 /// assert_eq!(con.search(&[5, 10]), Some(b));
42 /// assert_eq!(con.parent(b), Some(a));
43 /// ```
44 ///
45 /// See also [`TrailTree::insert`].
46 pub fn extend(&mut self, mut id:usize, path:&[K]) -> usize
47 {
48 let root = self.rel(id).child[1];
49 if root == usize::MAX {
50 let mut first = None;
51 self.fill_path(&mut first, &mut id, path);
52 first.unwrap_or(id)
53 } else {
54 self.modify(root, path)
55 }
56 }
57
58 fn modify(&mut self, id:usize, path:&[K]) -> usize
59 {
60 let mut current = id;
61
62 // Create root node if empty.
63 if self.root == usize::MAX {
64 let mut first = None;
65 self.fill_path(&mut first, &mut current, path);
66 self.root = first.unwrap_or(usize::MAX);
67 } else {
68 let mut path_index = 0;
69 while path_index < path.len() {
70 let mut traverse = true;
71
72 while traverse {
73 let node = self.node(current).unwrap();
74 let rel = node.rel;
75
76 let compare = path[path_index].cmp(&node.key);
77 match compare {
78 Ordering::Less | Ordering::Greater => {
79 let index = (compare as i8 + 1) as usize;
80 if node.rel.child[index] == usize::MAX {
81 let (next_id, next) = self.acquire(path[path_index].clone());
82
83 next.rel.parent = current;
84 next.rel.root = rel.root;
85
86 self.rel_mut(current).child[index] = next_id;
87 current = next_id;
88
89 self.balance_tree(current);
90
91 self.fill_path(&mut None, &mut current, &path[path_index+1..]);
92
93 traverse = false;
94 path_index = path.len();
95 } else {
96 current = node.rel.child[index];
97 }
98 }
99
100 Ordering::Equal => if node.rel.child[1] == usize::MAX {
101 self.fill_path(&mut None, &mut current, &path[path_index+1..]);
102
103 traverse = false;
104 path_index = path.len();
105 } else {
106 current = if path_index + 1 < path.len() { node.rel.child[1] } else { current };
107 traverse = false;
108 }
109 };
110 }
111
112 path_index += 1;
113 }
114 }
115
116 current
117 }
118
119 fn fill_path(&mut self, first:&mut Option<usize>, current:&mut usize, path:&[K])
120 {
121 for key in path {
122 let next = {
123 let (next_id, next_node) = self.acquire(key.clone());
124 next_node.rel.root = *current;
125 next_id
126 };
127
128 first.get_or_insert(next);
129
130 if *current != usize::MAX {
131 self.rel_mut(*current).child[1] = next;
132 }
133
134 *current = next;
135 }
136 }
137
138 /// Removes the subtree starting at the specified element.
139 ///
140 /// # Examples
141 ///
142 /// ```
143 /// use trailtree::TrailTree;
144 ///
145 /// let mut con = TrailTree::<u32,()>::new();
146 ///
147 /// let a = con.insert(&[10]);
148 /// con.extend(a, &[20, 30]);
149 /// let b = con.insert(&[5, 2]);
150 /// con.insert(&[20]);
151 /// con.insert(&[15]);
152 ///
153 /// con.remove(a);
154 ///
155 /// assert_eq!(con.search(&[10]), None);
156 /// assert_eq!(con.search(&[5, 2]), Some(b));
157 /// ```
158 pub fn remove(&mut self, id:usize)
159 {
160 if self.node(id).is_some() {
161 let rel = self.rel(id);
162
163 // Remove child subtree.
164 self.remove_subtree(rel.child[1]);
165
166 let lrel = self.rel(rel.child[0]);
167 let rrel = self.rel(rel.child[2]);
168
169 let mut replacement = usize::MAX;
170 let mut balance = rel.parent;
171
172 let mask = (((rel.child[0] != usize::MAX) as u32) << 1)
173 | ((rel.child[2] != usize::MAX) as u32);
174 match mask {
175 2 => {
176 self.rel_mut(rel.child[0]).parent = rel.parent;
177
178 let prel = self.rel_mut(rel.parent);
179 if prel.child[0] == id {
180 prel.child[0] = rel.child[0];
181 } else {
182 prel.child[2] = rel.child[0];
183 }
184
185 replacement = rel.child[0];
186 balance = rel.parent;
187 }
188 1 => {
189 self.rel_mut(rel.child[2]).parent = rel.parent;
190
191 let prel = self.rel_mut(rel.parent);
192 if prel.child[0] == id {
193 prel.child[0] = rel.child[0];
194 } else {
195 prel.child[2] = rel.child[0];
196 }
197
198 replacement = rel.child[2];
199 balance = rel.parent;
200 }
201 3 => {
202 let sel = if rrel.height > lrel.height { 2 } else { 0 };
203 let isel = if sel == 0 { 2 } else { 0 };
204
205 // Move nearest descendant to top.
206 let mut current = rel.child[sel];
207 while let Some(node) = self.node(current) && node.rel.child[isel] != usize::MAX {
208 current = node.rel.child[isel];
209 }
210
211 // Update parent and child nodes.
212 let srel = self.rel(current);
213 self.rel_mut(srel.parent).child[isel] = srel.child[sel];
214
215 if srel.child[sel] != usize::MAX {
216 self.rel_mut(srel.child[sel]).parent = srel.parent;
217 }
218
219 // Update moved node.
220 {
221 let srel = self.rel_mut(current);
222 srel.parent = rel.parent;
223 srel.child[0] = rel.child[0];
224 srel.child[2] = rel.child[2];
225 replacement = current;
226 }
227
228 self.rel_mut(rel.child[0]).parent = current;
229 self.rel_mut(rel.child[2]).parent = current;
230
231 balance = srel.parent;
232 }
233 _ => { }
234 }
235
236 if balance != usize::MAX {
237 self.balance_tree(balance);
238 }
239
240 // Update node references.
241 if rel.parent != usize::MAX {
242 let prel = self.rel_mut(rel.parent);
243 if prel.child[0] == id {
244 prel.child[0] = replacement;
245 } else {
246 prel.child[2] = replacement;
247 }
248 } else if rel.root != usize::MAX {
249 self.rel_mut(rel.root).child[1] = replacement;
250 } else {
251 self.root = replacement;
252 }
253
254 // Replace cell with next pointer.
255 self.cells[id] = Cell::pointer(self.next);
256 self.next = id;
257 }
258 }
259
260 pub(crate) fn remove_subtree(&mut self, id:usize)
261 {
262 if id != usize::MAX {
263 let mut stack = vec![(id, 0)];
264
265 while let Some((top, index)) = stack.last_mut() {
266 let rel = self.rel(*top);
267
268 if *index < 3 {
269 // Push next child to stack.
270 let i = *index;
271 *index += 1;
272
273 if rel.child[i] != usize::MAX {
274 stack.push((rel.child[i], 0));
275 }
276 } else {
277 // Update node references.
278 if rel.parent != usize::MAX {
279 let prel = self.rel_mut(rel.parent);
280 if prel.child[0] == *top {
281 prel.child[0] = usize::MAX;
282 } else {
283 prel.child[2] = usize::MAX;
284 }
285 } else if rel.root != usize::MAX {
286 self.rel_mut(rel.root).child[1] = usize::MAX;
287 } else {
288 self.root = usize::MAX;
289 }
290
291 // Replace cell with next pointer.
292 self.cells[*top] = Cell::pointer(self.next);
293 self.next = *top;
294
295 stack.pop();
296 }
297 }
298 }
299 }
300
301 /// Removes all items from the tree, but does not release memory.
302 ///
303 /// # Examples
304 ///
305 /// ```
306 /// use trailtree::TrailTree;
307 ///
308 /// let mut con = TrailTree::<u32,()>::new();
309 ///
310 /// con.insert(&[5, 10, 15]);
311 /// con.clear();
312 ///
313 /// assert_eq!(con.first(None), None);
314 /// ```
315 pub fn clear(&mut self)
316 {
317 self.cells.clear();
318 self.root = usize::MAX;
319 self.next = 0;
320 }
321
322 /// Releases all memory used by the tree.
323 pub fn release(&mut self)
324 {
325 self.cells = Vec::new();
326 self.root = usize::MAX;
327 self.next = 0;
328 }
329
330 /// Replaces data associated with the specified element and returns the old value.
331 ///
332 /// # Examples
333 ///
334 /// ```
335 /// use trailtree::TrailTree;
336 ///
337 /// let mut con = TrailTree::<u32,u32>::new();
338 ///
339 /// let a = con.insert(&[5, 10]);
340 ///
341 /// assert_eq!(con.set(a, 42), None);
342 /// assert_eq!(con.set(a, 43), Some(42));
343 /// ```
344 pub fn set(&mut self, id:usize, data:T) -> Option<T>
345 {
346 if id < self.cells.len() && let Cell::Node(node) = &mut self.cells[id] {
347 let old = node.data.take();
348 node.data = Some(data);
349 old
350 } else {
351 None
352 }
353 }
354
355 /// Removes data associated with the specified element and returns the old value.
356 ///
357 /// # Examples
358 ///
359 /// ```
360 /// use trailtree::TrailTree;
361 ///
362 /// let mut con = TrailTree::<u32,u32>::new();
363 ///
364 /// let a = con.insert(&[5, 10]);
365 ///
366 /// con.set(a, 24);
367 ///
368 /// assert_eq!(con.unset(a), Some(24));
369 /// assert_eq!(con.unset(a), None);
370 /// ```
371 pub fn unset(&mut self, id:usize) -> Option<T>
372 {
373 if id < self.cells.len() && let Cell::Node(node) = &mut self.cells[id] {
374 node.data.take()
375 } else {
376 None
377 }
378 }
379
380 pub(crate) fn acquire(&mut self, key:K) -> (usize, &mut Node<K,T>)
381 {
382 let id = self.next;
383
384 // Acquire block if next exceeds current allocation.
385 if id >= self.cells.len() {
386 let next = self.cells.len() + 1;
387 self.cells.push(Cell::pointer(next));
388 }
389
390 if let Cell::Pointer(addr) = self.cells[id] {
391 self.next = addr.root;
392 }
393 self.cells[id] = Cell::Node(Node::new(key));
394
395 if let Cell::Node(node) = &mut self.cells[id] {
396 (id, node)
397 } else {
398 unreachable!()
399 }
400 }
401}