graph/
graph.rs

1use std::{
2    collections::{hash_map::Entry, HashMap},
3    fmt::Debug,
4    mem,
5    sync::{Arc, RwLock},
6};
7
8use token_ref_cell::{BoxToken, TokenRefCell};
9
10#[derive(Debug, Default)]
11pub struct Graph {
12    nodes: HashMap<NodeId, Arc<TokenRefCell<Node>>>,
13    token: BoxToken,
14}
15
16pub type NodeId = usize;
17
18#[derive(Debug, Default)]
19pub struct Node {
20    id: NodeId,
21    edges: HashMap<NodeId, Arc<TokenRefCell<Node>>>,
22}
23
24impl Graph {
25    pub fn new() -> Self {
26        Self::default()
27    }
28
29    pub fn get_node(&self, id: NodeId) -> Option<&Node> {
30        Some(self.nodes.get(&id)?.borrow(&self.token).into_ref())
31    }
32
33    pub fn insert_node(&mut self, id: NodeId) -> bool {
34        let entry = self.nodes.entry(id);
35        let inserted = matches!(entry, Entry::Vacant(_));
36        entry.or_insert_with(|| Node::new_cell(id, &self.token));
37        inserted
38    }
39
40    pub fn remove_node(&mut self, id: NodeId) -> bool {
41        let Some(node) = self.nodes.remove(&id) else {
42            return false;
43        };
44        let mut node_mut = node.borrow_mut(&mut self.token);
45        let mut reborrow_mut =
46            node_mut.reborrow_stateful_mut(|node| node.edges.values().map(AsRef::as_ref));
47        while let Some(mut other) = reborrow_mut.reborrow_opt_mut(|edges| edges.next()) {
48            other.edges.remove(&id);
49        }
50        true
51    }
52
53    pub fn insert_edge(&mut self, node1: NodeId, node2: NodeId) {
54        let mut node = |id| {
55            let entry = self.nodes.entry(id);
56            let node = entry.or_insert_with(|| Node::new_cell(node1, &self.token));
57            node.clone()
58        };
59        let node_cell1 = node(node1);
60        let node_cell2 = node(node2);
61        node_cell1
62            .borrow_mut(&mut self.token)
63            .edges
64            .insert(node2, node_cell2.clone());
65        node_cell2
66            .borrow_mut(&mut self.token)
67            .edges
68            .insert(node1, node_cell1);
69    }
70
71    pub fn remove_edge(&mut self, node1: NodeId, node2: NodeId) {
72        let mut remove = |a, b| {
73            let Some(node) = self.nodes.get(&a) else {
74                return false;
75            };
76            node.borrow_mut(&mut self.token).edges.remove(&b);
77            true
78        };
79        if remove(node1, node2) {
80            remove(node2, node1);
81        }
82    }
83
84    pub fn clear(&mut self) {
85        let mut nodes = mem::take(&mut self.nodes);
86        for (_, node) in nodes.drain() {
87            if let Ok(mut node) = node.try_borrow_mut(&mut self.token) {
88                node.edges.clear();
89            }
90        }
91        self.nodes = nodes;
92    }
93}
94
95impl Node {
96    fn new_cell(id: NodeId, token: &BoxToken) -> Arc<TokenRefCell<Node>> {
97        let edges = HashMap::new();
98        Arc::new(TokenRefCell::new(Node { id, edges }, token))
99    }
100
101    pub fn id(&self) -> NodeId {
102        self.id
103    }
104
105    pub fn edges(&self) -> impl Iterator<Item = NodeId> + '_ {
106        self.edges.keys().copied()
107    }
108}
109
110impl Drop for Graph {
111    fn drop(&mut self) {
112        self.clear();
113    }
114}
115
116fn main() {
117    let graph = Arc::new(RwLock::new(Graph::new()));
118    {
119        let mut graph_mut = graph.write().unwrap();
120        graph_mut.insert_edge(0, 1);
121        graph_mut.insert_edge(1, 2);
122        graph_mut.insert_edge(2, 0);
123        graph_mut.remove_node(1);
124    }
125    {
126        let graph_ref = graph.read().unwrap();
127        let node = graph_ref.get_node(0).unwrap();
128        assert_eq!(node.id(), 0);
129        let mut edges = node.edges();
130        assert_eq!(edges.next(), Some(2));
131        assert_eq!(edges.next(), None);
132    }
133}
graph/graph.rs

graph/
graph.rs
