Skip to main content

slab_tree/
core_tree.rs

1use crate::node::Node;
2use crate::slab;
3use crate::NodeId;
4use snowflake::ProcessUniqueId;
5
6///
7/// A wrapper around a Slab containing Node<T> values.
8///
9/// Groups a collection of Node<T>s with a process unique id.
10///
11#[derive(Debug, PartialEq)]
12pub(crate) struct CoreTree<T> {
13    id: ProcessUniqueId,
14    slab: slab::Slab<Node<T>>,
15}
16
17impl<T> CoreTree<T> {
18    pub(crate) fn new(capacity: usize) -> CoreTree<T> {
19        CoreTree {
20            id: ProcessUniqueId::new(),
21            slab: slab::Slab::new(capacity),
22        }
23    }
24
25    pub(crate) fn capacity(&self) -> usize {
26        self.slab.capacity()
27    }
28
29    pub(crate) fn insert(&mut self, data: T) -> NodeId {
30        let key = self.slab.insert(Node::new(data));
31        self.new_node_id(key)
32    }
33
34    pub(crate) fn remove(&mut self, node_id: NodeId) -> Option<T> {
35        self.filter_by_tree_id(node_id)
36            .and_then(|id| self.slab.remove(id.index))
37            .map(|node| node.data)
38    }
39
40    pub(crate) fn get(&self, node_id: NodeId) -> Option<&Node<T>> {
41        self.filter_by_tree_id(node_id)
42            .and_then(|id| self.slab.get(id.index))
43    }
44
45    pub(crate) fn get_mut(&mut self, node_id: NodeId) -> Option<&mut Node<T>> {
46        self.filter_by_tree_id(node_id)
47            .and_then(move |id| self.slab.get_mut(id.index))
48    }
49
50    fn new_node_id(&self, index: slab::Index) -> NodeId {
51        NodeId {
52            tree_id: self.id,
53            index,
54        }
55    }
56
57    fn filter_by_tree_id(&self, node_id: NodeId) -> Option<NodeId> {
58        if node_id.tree_id != self.id {
59            return None;
60        }
61        Some(node_id)
62    }
63}
64
65#[cfg_attr(tarpaulin, skip)]
66#[cfg(test)]
67mod tests {
68    use super::*;
69
70    #[test]
71    fn capacity() {
72        let capacity = 5;
73        let tree = CoreTree::<i32>::new(capacity);
74        assert_eq!(tree.capacity(), capacity);
75    }
76
77    #[test]
78    fn insert() {
79        let mut tree = CoreTree::new(0);
80
81        let id = tree.insert(1);
82        let id2 = tree.insert(3);
83
84        assert_eq!(tree.get(id).unwrap().data, 1);
85        assert_eq!(tree.get(id2).unwrap().data, 3);
86    }
87
88    #[test]
89    fn remove() {
90        let mut tree = CoreTree::new(0);
91
92        let id = tree.insert(1);
93        assert_eq!(tree.get(id).unwrap().data, 1);
94
95        let one = tree.remove(id);
96        assert!(one.is_some());
97
98        let one = one.unwrap();
99        assert_eq!(one, 1);
100    }
101
102    #[test]
103    fn get() {
104        let mut tree = CoreTree::new(0);
105
106        let id = tree.insert(1);
107        let id2 = tree.insert(3);
108
109        assert_eq!(tree.get(id).unwrap().data, 1);
110        assert_eq!(tree.get(id2).unwrap().data, 3);
111    }
112
113    #[test]
114    fn get_mut() {
115        let mut tree = CoreTree::new(0);
116
117        let id = tree.insert(1);
118        let id2 = tree.insert(3);
119
120        assert_eq!(tree.get_mut(id).unwrap().data, 1);
121        assert_eq!(tree.get_mut(id2).unwrap().data, 3);
122    }
123
124    #[test]
125    fn get_with_bad_id() {
126        let mut tree = CoreTree::new(0);
127        let tree2: CoreTree<i32> = CoreTree::new(0);
128
129        let mut id = tree.insert(1);
130        id.tree_id = tree2.id; // oops, wrong tree id.
131
132        let result = tree.get(id);
133
134        assert!(result.is_none());
135    }
136}