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}