1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126

pub struct ArenaTree {
    nodes: Vec<Node>,
}
#[derive(Debug)]
struct Node {
    _id: NodeId,
    parent: Option<NodeId>,
}

#[derive(Debug, Clone, Copy)]
pub struct NodeId(pub usize);

impl ArenaTree {
    pub fn new() -> ArenaTree {
        ArenaTree { nodes: vec![] }
    }

    pub fn new_node(&mut self) -> NodeId {
        let next_idx = NodeId(self.nodes.len());

        self.nodes.push(Node {
            _id: next_idx,
            parent: None,
        });

        next_idx
    }

    pub fn connect(&mut self, id1: NodeId, id2: NodeId) {
        let root1 = self.root(id1);
        let root2 = self.root(id2);

        if root1.and(root2).is_none() {
            return;
        }

        if let Some(node) = self.nodes.get_mut(root2.unwrap().0) {
            node.parent = root1;
        }
    }

    pub fn connected(&self, id1: NodeId, id2: NodeId) -> bool {
        let root1 = self.root(id1);
        let root2 = self.root(id2);

        if root1.and(root2).is_none() {
            return false;
        }

        root1.unwrap().0 == root2.unwrap().0
    }

    fn root(&self, id: NodeId) -> Option<NodeId> {
        let node = self.nodes.get(id.0);
        node?;

        if let Some(parent) = node.unwrap().parent {
            self.root(parent)
        } else {
            Some(id)
        }
    }
}

impl Node {}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn unconnected_nodes() {
        let mut arena = ArenaTree::new();

        let node1 = arena.new_node();
        let node2 = arena.new_node();
        let node3 = arena.new_node();

        assert_eq!(arena.connected(node1, node2), false);
        assert_eq!(arena.connected(node1, node3), false);
    }

    #[test]
    fn connect_two_nodes() {
        let mut arena = ArenaTree::new();

        let node1 = arena.new_node();
        let node2 = arena.new_node();
        let node3 = arena.new_node();

        arena.connect(node1, node2);
        assert_eq!(arena.connected(node1, node2), true);
        assert_eq!(arena.connected(node1, node3), false);
    }

    #[test]
    fn connect_three_nodes() {
        let mut arena = ArenaTree::new();

        let node1 = arena.new_node();
        let node2 = arena.new_node();
        let node3 = arena.new_node();

        arena.connect(node1, node2);
        arena.connect(node3, node2);

        assert_eq!(arena.connected(node1, node2), true);
        assert_eq!(arena.connected(node1, node3), true);
        assert_eq!(arena.connected(node3, node2), true);
    }

    #[test]
    fn connect_one_none_node() {
        let mut arena = ArenaTree::new();
        let node1 = arena.new_node();

        arena.connect(node1, NodeId(2));
        arena.connect(NodeId(2), node1);
    }

    #[test]
    fn connect_two_none_node() {
        let mut arena = ArenaTree::new();
        arena.connect(NodeId(1), NodeId(2));
    }
}