knossos/utils/
arena.rs

1pub struct ArenaTree {
2    nodes: Vec<Node>,
3}
4#[derive(Debug)]
5struct Node {
6    _id: NodeId,
7    parent: Option<NodeId>,
8}
9
10#[derive(Debug, Clone, Copy)]
11pub struct NodeId(pub usize);
12
13impl ArenaTree {
14    pub const fn new() -> ArenaTree {
15        ArenaTree { nodes: vec![] }
16    }
17
18    pub fn new_node(&mut self) -> NodeId {
19        let next_idx = NodeId(self.nodes.len());
20
21        self.nodes.push(Node {
22            _id: next_idx,
23            parent: None,
24        });
25
26        next_idx
27    }
28
29    pub fn connect(&mut self, id1: NodeId, id2: NodeId) {
30        let root1 = self.root(id1);
31        let root2 = self.root(id2);
32
33        if root1.and(root2).is_none() {
34            return;
35        }
36
37        if let Some(node) = self.nodes.get_mut(root2.unwrap().0) {
38            node.parent = root1;
39        }
40    }
41
42    pub fn connected(&self, id1: NodeId, id2: NodeId) -> bool {
43        let root1 = self.root(id1);
44        let root2 = self.root(id2);
45
46        if root1.and(root2).is_none() {
47            return false;
48        }
49
50        root1.unwrap().0 == root2.unwrap().0
51    }
52
53    fn root(&self, id: NodeId) -> Option<NodeId> {
54        let node = self.nodes.get(id.0);
55        node?;
56
57        if let Some(parent) = node.unwrap().parent {
58            self.root(parent)
59        } else {
60            Some(id)
61        }
62    }
63}
64
65impl Node {}
66
67#[cfg(test)]
68mod tests {
69    use super::*;
70
71    #[test]
72    fn unconnected_nodes() {
73        let mut arena = ArenaTree::new();
74
75        let node1 = arena.new_node();
76        let node2 = arena.new_node();
77        let node3 = arena.new_node();
78
79        assert!(!arena.connected(node1, node2));
80        assert!(!arena.connected(node1, node3));
81    }
82
83    #[test]
84    fn connect_two_nodes() {
85        let mut arena = ArenaTree::new();
86
87        let node1 = arena.new_node();
88        let node2 = arena.new_node();
89        let node3 = arena.new_node();
90
91        arena.connect(node1, node2);
92        assert!(arena.connected(node1, node2));
93        assert!(!arena.connected(node1, node3));
94    }
95
96    #[test]
97    fn connect_three_nodes() {
98        let mut arena = ArenaTree::new();
99
100        let node1 = arena.new_node();
101        let node2 = arena.new_node();
102        let node3 = arena.new_node();
103
104        arena.connect(node1, node2);
105        arena.connect(node3, node2);
106
107        assert!(arena.connected(node1, node2));
108        assert!(arena.connected(node1, node3));
109        assert!(arena.connected(node3, node2));
110    }
111
112    #[test]
113    fn connect_one_none_node() {
114        let mut arena = ArenaTree::new();
115        let node1 = arena.new_node();
116
117        arena.connect(node1, NodeId(2));
118        assert!(!arena.connected(node1, NodeId(2)));
119
120        arena.connect(NodeId(2), node1);
121        assert!(!arena.connected(node1, NodeId(2)));
122    }
123
124    #[test]
125    fn connect_two_none_node() {
126        let arena = ArenaTree::new();
127        arena.connected(NodeId(1), NodeId(2));
128        assert!(!arena.connected(NodeId(1), NodeId(2)));
129    }
130}