oakwood/
lib.rs

1//! Vector-backed trees
2
3#![no_std]
4
5extern crate alloc;
6
7use alloc::vec::Vec;
8use core::{fmt, marker::PhantomData, ops::{Index, IndexMut}, hash::Hash};
9
10pub trait Cookie: fmt::Debug + Default + Copy + Eq {
11    fn next(&mut self) -> Self;
12    fn invalid() -> Self;
13}
14
15/// 64-bit safety cookie
16#[derive(Debug, Default, Copy, Clone, Eq, Ord, PartialEq, PartialOrd, Hash)]
17#[repr(transparent)]
18pub struct Cookie64(u64);
19
20impl Cookie for Cookie64 {
21    fn next(&mut self) -> Self {
22        let bck = self.0;
23        self.0 += 1;
24        Self(bck)
25    }
26
27    fn invalid() -> Self {
28        Self(u64::MAX)
29    }
30}
31
32/// 32-bit safety cookie
33#[derive(Debug, Default, Copy, Clone, Eq, Ord, PartialEq, PartialOrd, Hash)]
34#[repr(transparent)]
35pub struct Cookie32(u32);
36
37impl Cookie for Cookie32 {
38    fn next(&mut self) -> Self {
39        let bck = self.0;
40        self.0 += 1;
41        Self(bck)
42    }
43
44    fn invalid() -> Self {
45        Self(u32::MAX)
46    }
47}
48
49/// 0-bit safety cookie
50#[derive(Debug, Default, Copy, Clone, Eq, Ord, PartialEq, PartialOrd, Hash)]
51#[repr(transparent)]
52pub struct NoCookie;
53
54impl Cookie for NoCookie {
55    fn next(&mut self) -> Self { Self }
56    fn invalid() -> Self { Self }
57}
58
59pub trait NodeIndex: Default + Copy + fmt::Debug + Eq + Ord + Into<usize> + From<usize> {}
60pub trait OptionalNodeIndex<I: NodeIndex>: Default + Copy + Eq + Into<Option<I>> + From<Option<I>> {}
61
62pub trait NodeKey<C: Cookie, I: NodeIndex>: fmt::Debug + Copy + Eq + Ord + Hash {
63    fn new(index: I, cookie: C) -> Self;
64    fn index(&self) -> I;
65    fn cookie(&self) -> C;
66}
67
68/// Create [`NodeIndex`], [`OptionalNodeIndex`] implementors automatically
69#[macro_export]
70macro_rules! index {
71    ($index:ident, $option:ident) => { $crate::index!($index, $option, u32); };
72    ($index:ident, $option:ident, $typ:ty) => {
73        #[derive(Copy, Clone, Debug, Eq, Ord, PartialEq, PartialOrd, Hash)]
74        #[repr(transparent)]
75        pub struct $index($typ);
76
77        impl Default for $index {
78            fn default() -> Self {
79                Self(<$typ>::MAX)
80            }
81        }
82
83        impl From<usize> for $index {
84            fn from(primitive: usize) -> Self {
85                $index(primitive as $typ)
86            }
87        }
88
89        impl From<$index> for usize {
90            fn from(index: $index) -> Self {
91                index.0 as usize
92            }
93        }
94
95        impl $crate::NodeIndex for $index {}
96
97        impl ::core::fmt::Display for $index {
98            fn fmt(&self, f: &mut ::core::fmt::Formatter<'_>) -> ::core::fmt::Result {
99                write!(f, "{}", self.0)
100            }
101        }
102
103        #[derive(Copy, Clone, PartialEq, Eq, Hash)]
104        #[repr(transparent)]
105        pub struct $option($typ);
106
107        impl $option {
108            pub fn get(self) -> Option<$index> {
109                self.into()
110            }
111        }
112
113        impl Default for $option {
114            fn default() -> Self {
115                None.into()
116            }
117        }
118
119        impl From<$option> for Option<$index> {
120            fn from(opt: $option) -> Self {
121                match opt.0 {
122                    <$typ>::MAX => None,
123                    value => Some($index(value)),
124                }
125            }
126        }
127
128        impl From<Option<$index>> for $option {
129            fn from(opt: Option<$index>) -> Self {
130                match opt.map(|v| v.0) {
131                    Some(<$typ>::MAX) => panic!("Reserved value!"),
132                    Some(value) => $option(value),
133                    None => $option(<$typ>::MAX),
134                }
135            }
136        }
137
138        impl $crate::OptionalNodeIndex<$index> for $option {}
139
140        impl ::core::fmt::Debug for $option {
141            fn fmt(&self, f: &mut ::core::fmt::Formatter<'_>) -> ::core::fmt::Result {
142                write!(f, "{:?}", Option::<$index>::from(*self))
143            }
144        }
145    };
146}
147
148/// Create [`NodeKey`], [`NodeIndex`], [`OptionalNodeIndex`] implementors automatically
149#[macro_export]
150macro_rules! tree_params {
151    ($key:ident, $index:ident, $option:ident, $cookie:ty) => {
152        $crate::index!($index, $option);
153
154        #[derive(Copy, Clone, Debug, Eq, Ord, PartialEq, PartialOrd, Default, Hash)]
155        pub struct $key {
156            index: $index,
157            cookie: $cookie,
158        }
159
160        impl $crate::NodeKey<$cookie, $index> for $key {
161            fn new(index: $index, cookie: $cookie) -> Self { Self { index, cookie } }
162            fn index(&self) -> $index { self.index }
163            fn cookie(&self) -> $cookie { self.cookie }
164        }
165    };
166}
167
168#[macro_export]
169macro_rules! tree {
170    ($tree:ident, $node:ty, $key:ident, $index:ident, $option:ident, $cookie:ty) => {
171        $crate::tree_params!($key, $index, $option, $cookie);
172        pub type $tree = $crate::Tree<$cookie, $index, $option, $key, $node>;
173    };
174}
175
176#[derive(Default, Debug, Copy, Clone)]
177struct NodeRelations<C: Cookie, I: NodeIndex, O: OptionalNodeIndex<I>> {
178    cookie: C,
179    parent: O,
180    first_child: O,
181    prev_sibling: I,
182    next_sibling: I,
183}
184
185/// A vector-backed tree where you can store nodes
186pub struct Tree<C: Cookie, I: NodeIndex, O: OptionalNodeIndex<I>, K: NodeKey<C, I>, N> {
187    nodes: Vec<N>,
188    relations: Vec<NodeRelations<C, I, O>>,
189    next_cookie: C,
190    free_slot: O,
191    _key: PhantomData<K>,
192}
193
194impl<C: Cookie, I: NodeIndex, O: OptionalNodeIndex<I>, K: NodeKey<C, I>, N> Tree<C, I, O, K, N> {
195    pub fn new() -> Self {
196        Self {
197            _key: PhantomData::default(),
198            nodes: Default::default(),
199            relations: Default::default(),
200            next_cookie: Default::default(),
201            free_slot: Default::default(),
202        }
203    }
204
205    fn rel(&self, index: I) -> &NodeRelations<C, I, O> {
206        &self.relations[index.into()]
207    }
208
209    fn rel_mut(&mut self, index: I) -> &mut NodeRelations<C, I, O> {
210        &mut self.relations[index.into()]
211    }
212
213    fn rel_key(&self, key: K) -> &NodeRelations<C, I, O> {
214        let node = self.rel(key.index());
215        assert_eq!(node.cookie, key.cookie());
216        node
217    }
218
219    fn rel_key_mut(&mut self, key: K) -> &mut NodeRelations<C, I, O> {
220        let node = self.rel_mut(key.index());
221        assert_eq!(node.cookie, key.cookie());
222        node
223    }
224
225    pub fn node_key(&self, index: I) -> K {
226        K::new(index, self.rel(index).cookie)
227    }
228
229    pub fn get(&self, key: K) -> Option<&N> {
230        if key.cookie() == self.rel(key.index()).cookie {
231            Some(&self.nodes[key.index().into()])
232        } else {
233            None
234        }
235    }
236
237    pub fn get_mut(&mut self, key: K) -> Option<&mut N> {
238        if key.cookie() == self.rel(key.index()).cookie {
239            Some(&mut self.nodes[key.index().into()])
240        } else {
241            None
242        }
243    }
244}
245
246impl<C: Cookie, I: NodeIndex, O: OptionalNodeIndex<I>, K: NodeKey<C, I>, N: Default> Tree<C, I, O, K, N> {
247    fn collect(&mut self, index: I) {
248        self.nodes[index.into()] = N::default();
249        *self.rel_mut(index) = NodeRelations::default();
250        self.rel_mut(index).next_sibling = index;
251        self.rel_mut(index).prev_sibling = index;
252        self.rel_mut(index).parent = self.free_slot;
253        self.rel_mut(index).cookie = C::invalid();
254        self.free_slot = Some(index).into();
255    }
256
257    /// Allocate a node in the internal vector
258    pub fn create(&mut self) -> K {
259        if let Some(free_slot) = self.free_slot.into() {
260            self.free_slot = self.rel(free_slot).parent;
261            let cookie = self.next_cookie.next();
262            self.rel_mut(free_slot).parent = None.into();
263            self.rel_mut(free_slot).cookie = cookie;
264            K::new(free_slot, cookie)
265        } else {
266            let index = self.nodes.len().try_into().unwrap();
267            self.nodes.push(N::default());
268            self.relations.push(NodeRelations::default());
269            self.collect(index);
270            self.create()
271        }
272    }
273
274    /// Remove the children of a node from the tree
275    ///
276    /// Internal references to the nodes are removed
277    /// (links with siblings & parent). This function
278    /// behaves recursively to remove nodes from leaf
279    /// to branch.
280    pub fn delete_children(&mut self, key: K) {
281        if let Some(child) = self.rel_key_mut(key).first_child.into() {
282            let mut current = self.rel(child).prev_sibling;
283            while current != key.index() {
284                if let Some(child) = self.rel(current).first_child.into() {
285                    // process children first, start with last child
286                    // (prev of first child is last child)
287                    current = self.rel(child).prev_sibling;
288                } else {
289                    // no child
290                    let parent = self.rel(current).parent.into().unwrap();
291                    let parent_fc = self.rel(parent).first_child.into().unwrap();
292
293                    // is `current` the first sibling?
294                    let next = if current == parent_fc {
295                        // no child, all siblings processed -> go up
296                        self.rel_mut(parent).first_child = None.into();
297                        parent
298                    } else {
299                        // no child, at least one sibling => go prev
300                        self.rel(current).prev_sibling
301                    };
302
303                    self.collect(current);
304                    current = next;
305                }
306            }
307        }
308
309        self.rel_key_mut(key).first_child = None.into();
310    }
311
312    /// Remove a node and its children from the tree
313    ///
314    /// Internal references to the node are removed
315    /// (links with siblings & parent). This function
316    /// behaves recursively to remove nodes from leaf
317    /// to branch.
318    pub fn delete(&mut self, key: K) {
319        self.delete_children(key);
320
321        let next = self.rel_key(key).next_sibling;
322        let has_a_sibling = next != key.index();
323        // if we had siblings:
324        if has_a_sibling {
325            // disconnect siblings
326            let prev = self.rel_key(key).prev_sibling;
327            self.rel_mut(next).prev_sibling = prev;
328            self.rel_mut(prev).next_sibling = next;
329        }
330
331        // if we had a parent:
332        if let Some(parent) = self.rel_key(key).parent.into() {
333            // if we were the first child:
334            let some_this_node = Some(key.index()).into();
335            let first_child = &mut self.rel_mut(parent).first_child;
336            if *first_child == some_this_node {
337                // disconnect parent
338                *first_child = match has_a_sibling {
339                    true => Some(next),
340                    false => None,
341                }.into();
342            }
343        }
344
345        self.collect(key.index());
346    }
347
348    /// Removes the children of a node and reset
349    /// its public properties to their defaults
350    pub fn reset(&mut self, key: K) {
351        self.delete_children(key);
352        self[key] = N::default();
353    }
354
355    /// Insert a node as the next sibling of `prev`
356    pub fn insert_after(&mut self, prev: K, key: K) {
357        assert_eq!(self.rel_key(key).parent.into(), None);
358        let parent = self.rel_key(prev).parent;
359
360        let prev = prev.index();
361        let this = key.index();
362        let next = self.rel(prev).next_sibling;
363        let last = self.rel(this).prev_sibling;
364
365        self.rel_mut(next).prev_sibling = last;
366        self.rel_mut(last).next_sibling = next;
367        self.rel_mut(this).prev_sibling = prev;
368        self.rel_mut(prev).next_sibling = this;
369
370        let mut current = this;
371        loop {
372            self.rel_mut(current).parent = parent;
373
374            current = self.rel(current).next_sibling;
375            if current == next {
376                break;
377            }
378        }
379    }
380
381    /// Adds one or multiple "children" nodes to a "parent" node
382    ///
383    /// `node` and all it's siblings are appended to the list
384    /// of children of `parent`.
385    pub fn append_children(&mut self, key: K, parent: K) {
386        assert_eq!(self.rel_key(key).parent.into(), None);
387
388        if let Some(last_child) = self.last_child(parent) {
389            self.insert_after(last_child, key);
390        } else {
391            self.rel_key_mut(parent).first_child = Some(key.index()).into();
392
393            let mut current = key.index();
394            loop {
395                self.rel_mut(current).parent = Some(parent.index()).into();
396
397                current = self.rel(current).next_sibling;
398                if current == key.index() {
399                    break;
400                }
401            }
402        }
403    }
404
405    /// Change references to `old` so that they
406    /// point to `new` instead and remove `old`
407    pub fn replace(&mut self, old: K, new: K) {
408        assert_eq!(self.rel_key(new).parent.into(), None);
409
410        if self.is_only_child(old) {
411            if let Some(parent) = self.parent(old) {
412                self.delete(old);
413                self.append_children(new, parent);
414            } else {
415                self.delete(old);
416            }
417        } else {
418            let prev = self.prev_sibling(old);
419            if let Some(parent) = self.parent(old) {
420                if self.first_child(parent) == Some(old) {
421                    // this is invalid until new is inserted
422                    self.rel_key_mut(parent).first_child = Some(new.index()).into();
423                }
424            }
425            self.delete(old);
426            self.insert_after(prev, new);
427        }
428    }
429
430    /// Detach a node from all its "children", leaving
431    /// them "orphans"
432    ///
433    /// If the node had no children, `None` is returned.
434    pub fn detach_children(&mut self, parent: K) -> Option<K> {
435        use core::mem::replace;
436
437        let parent_fc = replace(&mut self.rel_key_mut(parent).first_child, None.into()).into()?;
438
439        let mut current = parent_fc;
440        loop {
441            self.rel_mut(current).parent = None.into();
442
443            current = self.rel(current).next_sibling;
444            if current == parent_fc {
445                break;
446            }
447        }
448
449        Some(self.node_key(parent_fc))
450    }
451
452    /// Get the index of this node relative to its parent's children
453    pub fn child_index(&self, key: K) -> Option<usize> {
454        let parent = self.rel_key(key).parent.into()?;
455        let first_child = self.rel(parent).first_child.into().unwrap();
456        let mut i = 0;
457        let mut current = key.index();
458        while current != first_child {
459            current = self.rel(current).prev_sibling;
460            i += 1;
461        }
462        Some(i)
463    }
464
465    /// Locate the parent of a node
466    pub fn parent(&self, key: K) -> Option<K> {
467        Some(self.node_key(self.rel_key(key).parent.into()?))
468    }
469
470    /// Locate the first child of a node
471    pub fn first_child(&self, key: K) -> Option<K> {
472        Some(self.node_key(self.rel_key(key).first_child.into()?))
473    }
474
475    /// Locate the last child of a node
476    pub fn last_child(&self, key: K) -> Option<K> {
477        let first_child = self.first_child(key)?;
478        Some(self.prev_sibling(first_child))
479    }
480
481    /// Locate the previous sibling of a node
482    ///
483    /// Note: siblings are looping, meaning
484    /// that calling this for the first child
485    /// of a node will locate the last child of
486    /// the node.
487    pub fn prev_sibling(&self, key: K) -> K {
488        self.node_key(self.rel_key(key).prev_sibling)
489    }
490
491    /// Locate the next sibling of a node
492    ///
493    /// Note: siblings are looping, meaning
494    /// that calling this for the last child
495    /// of a node will locate the first child of
496    /// the node.
497    pub fn next_sibling(&self, key: K) -> K {
498        self.node_key(self.rel_key(key).next_sibling)
499    }
500
501    /// Returns true if node has no sibling except itself
502    pub fn is_only_child(&self, key: K) -> bool {
503        self.rel_key(key).prev_sibling == key.index()
504    }
505}
506
507/// Iterate on the children of a node
508///
509/// Note: if you add new children to `$parent` in
510/// the code block, they may not be iterated on, so this
511/// must not be done. Replacing or removing the child
512/// or one of its previous siblings is allowed. Replacing
513/// or removing its next sibling is not: that would result
514/// in an infinite loop. The relations of `$parent` must
515/// not be modified, too.
516#[macro_export]
517macro_rules! for_each_child {
518    ($tree:expr, $parent:expr, $child:ident, $code:block) => {
519        {
520            let parent = $parent;
521            let mut pending = $tree.first_child(parent);
522            while let Some($child) = pending {
523                let parent_fc = $tree.first_child(parent).unwrap();
524                let next_sibling = $tree.next_sibling($child);
525                pending = match next_sibling == parent_fc {
526                    true => None,
527                    false => Some(next_sibling),
528                };
529
530                $code
531            }
532        }
533    }
534}
535
536struct TreeNode<'a, C: Cookie, I: NodeIndex, O: OptionalNodeIndex<I>, K: NodeKey<C, I>, N> {
537    tree: &'a Tree<C, I, O, K, N>,
538    key: K,
539}
540
541impl<'a, C: Cookie, I: NodeIndex, O: OptionalNodeIndex<I>, K: NodeKey<C, I>, N: Default + ::core::fmt::Debug> ::core::fmt::Debug for TreeNode<'a, C, I, O, K, N> {
542    fn fmt(&self, fmt: &mut ::core::fmt::Formatter<'_>) -> ::core::fmt::Result {
543        let mut dbg = fmt.debug_struct("TreeNode");
544        dbg.field("node", &self.tree[self.key]);
545        for_each_child!(self.tree, self.key, child, {
546            dbg.field("child", &TreeNode {
547                tree: self.tree,
548                key: child,
549            });
550        });
551        dbg.finish()
552    }
553}
554
555impl<C: Cookie, I: NodeIndex, O: OptionalNodeIndex<I>, K: NodeKey<C, I>, N: Default + ::core::fmt::Debug> ::core::fmt::Debug for Tree<C, I, O, K, N> {
556    fn fmt(&self, fmt: &mut ::core::fmt::Formatter<'_>) -> ::core::fmt::Result {
557        let mut dbg = fmt.debug_list();
558
559        for i in 0..self.relations.len() {
560            let is_allocated = self.relations[i].cookie != C::invalid();
561            let has_no_parent = self.relations[i].parent.into().is_none();
562            if is_allocated && has_no_parent {
563                dbg.entry(&TreeNode {
564                    tree: self,
565                    key: self.node_key(i.into()),
566                });
567            }
568        }
569
570        dbg.finish()
571    }
572}
573
574impl<C: Cookie, I: NodeIndex, O: OptionalNodeIndex<I>, K: NodeKey<C, I>, N> Index<K> for Tree<C, I, O, K, N> {
575    type Output = N;
576
577    fn index(&self, key: K) -> &N {
578        let _ = self.rel_key(key);
579        &self.nodes[key.index().into()]
580    }
581}
582
583impl<C: Cookie, I: NodeIndex, O: OptionalNodeIndex<I>, K: NodeKey<C, I>, N> IndexMut<K> for Tree<C, I, O, K, N> {
584    fn index_mut(&mut self, key: K) -> &mut N {
585        let _ = self.rel_key(key);
586        &mut self.nodes[key.index().into()]
587    }
588}
oakwood/lib.rs

oakwood/
lib.rs
