1use crate::node::Node;
2use crate::slab;
3use crate::NodeId;
4use snowflake::ProcessUniqueId;
5
6#[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; let result = tree.get(id);
133
134 assert!(result.is_none());
135 }
136}