howudoin/
flat_tree.rs

1//! A common abstraction of a 'flat' tree which tracks the nodes with an ordered map and the set of
2//! root nodes.
3//!
4//! This pattern is common with implementors of [`crate::Consume`] so is presented here, but is by
5//! no means a feature complete implementation.
6use std::{
7    borrow::Borrow,
8    collections::{BTreeMap, BTreeSet},
9};
10
11/// A common abstraction of a 'flat' tree which tracks the nodes with an ordered map and the set of
12/// root nodes.
13///
14/// Note that the data structures use ordered maps and sets which work well with the incrementing
15/// [`crate::Id`] key.
16pub struct FlatTree<K, V> {
17    /// The nodes.
18    pub nodes: BTreeMap<K, V>,
19    /// The set of root nodes.
20    pub roots: BTreeSet<K>,
21}
22
23impl<K, V> Default for FlatTree<K, V> {
24    fn default() -> Self {
25        Self {
26            nodes: Default::default(),
27            roots: Default::default(),
28        }
29    }
30}
31
32impl<K, V> FlatTree<K, V>
33where
34    K: Clone + Eq + Ord,
35{
36    /// Get a node's value using the key.
37    pub fn get<Q>(&self, k: &Q) -> Option<&V>
38    where
39        K: Borrow<Q>,
40        Q: Eq + Ord,
41    {
42        self.nodes.get(k)
43    }
44
45    /// Get a node's value (mutable reference) using the key.
46    pub fn get_mut<Q>(&mut self, k: &Q) -> Option<&mut V>
47    where
48        K: Borrow<Q>,
49        Q: Eq + Ord,
50    {
51        self.nodes.get_mut(k)
52    }
53
54    /// Insert a root node.
55    pub fn insert_root(&mut self, k: K, v: V) -> Option<V> {
56        let x = self.insert(k.clone(), v);
57        self.roots.insert(k);
58
59        x
60    }
61
62    /// Insert a node.
63    ///
64    /// If the node is a root node, [`insert_root`] should be used.
65    ///
66    /// [`insert_root`]: Self::insert_root
67    pub fn insert(&mut self, k: K, v: V) -> Option<V> {
68        self.nodes.insert(k, v)
69    }
70
71    /// Remove a node with the key.
72    pub fn remove<Q>(&mut self, k: &Q) -> Option<V>
73    where
74        K: Borrow<Q>,
75        Q: Eq + Ord,
76    {
77        let x = self.nodes.remove(k);
78        self.roots.remove(k);
79
80        x
81    }
82
83    /// Does the structure contain a node with the key.
84    pub fn contains_node<Q>(&self, k: &Q) -> bool
85    where
86        K: Borrow<Q>,
87        Q: Eq + Ord,
88    {
89        self.nodes.contains_key(k)
90    }
91
92    /// Check if a key is a root node.
93    pub fn is_root<Q>(&self, k: &Q) -> bool
94    where
95        K: Borrow<Q>,
96        Q: Eq + Ord,
97    {
98        self.roots.contains(k)
99    }
100
101    /// An iterator over the roots, fetching the nodes' values.
102    pub fn roots(&self) -> impl Iterator<Item = (&K, &V)> {
103        self.roots
104            .iter()
105            .filter_map(|k| self.get(k).map(|v| (k, v)))
106    }
107}
howudoin/flat_tree.rs

howudoin/
flat_tree.rs
