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}