1use crate::core::id::{Eid, Vid};
13use fxhash::FxBuildHasher;
14use std::collections::HashMap;
15
16#[derive(Clone, Copy, Debug)]
20pub struct EdgeEntry {
21 pub eid: Eid,
22 pub src_vid: Vid,
23 pub dst_vid: Vid,
24 pub edge_type: u32,
25}
26
27type FxHashMap<K, V> = HashMap<K, V, FxBuildHasher>;
29
30#[derive(Clone, Debug)]
35pub struct SimpleGraph {
36 vertices: FxHashMap<Vid, ()>,
38 outgoing: FxHashMap<Vid, Vec<EdgeEntry>>,
40 incoming: FxHashMap<Vid, Vec<EdgeEntry>>,
42 edge_map: FxHashMap<Eid, EdgeEntry>,
44}
45
46#[derive(Clone, Copy, Debug, PartialEq, Eq)]
48pub enum Direction {
49 Outgoing,
50 Incoming,
51}
52
53impl Default for SimpleGraph {
54 fn default() -> Self {
55 Self {
56 vertices: HashMap::with_hasher(FxBuildHasher::default()),
57 outgoing: HashMap::with_hasher(FxBuildHasher::default()),
58 incoming: HashMap::with_hasher(FxBuildHasher::default()),
59 edge_map: HashMap::with_hasher(FxBuildHasher::default()),
60 }
61 }
62}
63
64impl SimpleGraph {
65 pub fn new() -> Self {
67 Self::default()
68 }
69
70 pub fn with_capacity(vertices: usize, edges: usize) -> Self {
72 Self {
73 vertices: HashMap::with_capacity_and_hasher(vertices, FxBuildHasher::default()),
74 outgoing: HashMap::with_capacity_and_hasher(vertices, FxBuildHasher::default()),
75 incoming: HashMap::with_capacity_and_hasher(vertices, FxBuildHasher::default()),
76 edge_map: HashMap::with_capacity_and_hasher(edges, FxBuildHasher::default()),
77 }
78 }
79
80 pub fn add_vertex(&mut self, vid: Vid) -> bool {
82 self.vertices.insert(vid, ()).is_none()
85 }
86
87 pub fn remove_vertex(&mut self, vid: Vid) {
89 self.vertices.remove(&vid);
90
91 if let Some(edges) = self.outgoing.remove(&vid) {
93 for edge in edges {
95 self.edge_map.remove(&edge.eid);
96 if let Some(incoming_list) = self.incoming.get_mut(&edge.dst_vid) {
97 incoming_list.retain(|e| e.eid != edge.eid);
98 }
99 }
100 }
101
102 if let Some(edges) = self.incoming.remove(&vid) {
104 for edge in edges {
106 self.edge_map.remove(&edge.eid);
107 if let Some(outgoing_list) = self.outgoing.get_mut(&edge.src_vid) {
108 outgoing_list.retain(|e| e.eid != edge.eid);
109 }
110 }
111 }
112 }
113
114 pub fn contains_vertex(&self, vid: Vid) -> bool {
116 self.vertices.contains_key(&vid)
117 }
118
119 pub fn vertex_count(&self) -> usize {
121 self.vertices.len()
122 }
123
124 pub fn add_edge(&mut self, src_vid: Vid, dst_vid: Vid, eid: Eid, edge_type: u32) {
126 self.add_vertex(src_vid);
128 self.add_vertex(dst_vid);
129
130 self.add_edge_unchecked(src_vid, dst_vid, eid, edge_type);
131 }
132
133 #[inline]
135 pub fn add_edge_unchecked(&mut self, src_vid: Vid, dst_vid: Vid, eid: Eid, edge_type: u32) {
136 let entry = EdgeEntry {
137 eid,
138 src_vid,
139 dst_vid,
140 edge_type,
141 };
142
143 self.outgoing.entry(src_vid).or_default().push(entry);
145
146 self.incoming.entry(dst_vid).or_default().push(entry);
148
149 self.edge_map.insert(eid, entry);
151 }
152
153 pub fn remove_edge(&mut self, eid: Eid) -> Option<EdgeEntry> {
157 let entry = self.edge_map.remove(&eid)?;
159
160 if let Some(outgoing_list) = self.outgoing.get_mut(&entry.src_vid) {
162 outgoing_list.retain(|e| e.eid != eid);
163 }
164
165 if let Some(incoming_list) = self.incoming.get_mut(&entry.dst_vid) {
167 incoming_list.retain(|e| e.eid != eid);
168 }
169
170 Some(entry)
171 }
172
173 pub fn neighbors(&self, vid: Vid, direction: Direction) -> &[EdgeEntry] {
176 let map = match direction {
177 Direction::Outgoing => &self.outgoing,
178 Direction::Incoming => &self.incoming,
179 };
180 map.get(&vid).map(|v| v.as_slice()).unwrap_or(&[])
181 }
182
183 pub fn edges(&self) -> impl Iterator<Item = &EdgeEntry> {
185 self.outgoing.values().flat_map(|v| v.iter())
186 }
187
188 pub fn edge(&self, eid: Eid) -> Option<&EdgeEntry> {
190 self.edge_map.get(&eid)
191 }
192
193 pub fn vertex(&self, vid: Vid) -> Option<Vid> {
195 self.vertices.contains_key(&vid).then_some(vid)
196 }
197
198 pub fn edge_count(&self) -> usize {
200 self.outgoing.values().map(|v| v.len()).sum()
201 }
202
203 pub fn vertices(&self) -> impl Iterator<Item = Vid> + '_ {
205 self.vertices.keys().copied()
206 }
207
208 pub fn clear(&mut self) {
210 self.vertices.clear();
211 self.outgoing.clear();
212 self.incoming.clear();
213 self.edge_map.clear();
214 }
215}
216
217#[cfg(test)]
218mod tests {
219 use super::*;
220
221 #[test]
222 fn test_add_remove_vertex() {
223 let mut g = SimpleGraph::new();
224 let v1 = Vid::new(1);
225 let v2 = Vid::new(2);
226
227 assert!(g.add_vertex(v1));
228 assert!(!g.add_vertex(v1)); assert!(g.add_vertex(v2));
230
231 assert_eq!(g.vertex_count(), 2);
232 assert!(g.contains_vertex(v1));
233
234 g.remove_vertex(v1);
235 assert_eq!(g.vertex_count(), 1);
236 assert!(!g.contains_vertex(v1));
237 }
238
239 #[test]
240 fn test_add_remove_edge() {
241 let mut g = SimpleGraph::new();
242 let v1 = Vid::new(1);
243 let v2 = Vid::new(2);
244 let v3 = Vid::new(3);
245 let e1 = Eid::new(101);
246 let e2 = Eid::new(102);
247
248 g.add_edge(v1, v2, e1, 1);
249 g.add_edge(v1, v3, e2, 1);
250
251 assert_eq!(g.edge_count(), 2);
252
253 let out = g.neighbors(v1, Direction::Outgoing);
255 assert_eq!(out.len(), 2);
256
257 let inc = g.neighbors(v2, Direction::Incoming);
259 assert_eq!(inc.len(), 1);
260 assert_eq!(inc[0].src_vid, v1);
261
262 let removed = g.remove_edge(e1);
264 assert!(removed.is_some());
265 assert_eq!(removed.unwrap().dst_vid, v2);
266
267 assert_eq!(g.edge_count(), 1);
268 assert_eq!(g.neighbors(v1, Direction::Outgoing).len(), 1);
269 assert_eq!(g.neighbors(v2, Direction::Incoming).len(), 0);
270 }
271
272 #[test]
273 fn test_remove_vertex_with_edges() {
274 let mut g = SimpleGraph::new();
275 let v1 = Vid::new(1);
276 let v2 = Vid::new(2);
277 let v3 = Vid::new(3);
278 let e1 = Eid::new(101);
279 let e2 = Eid::new(102);
280
281 g.add_edge(v1, v2, e1, 1);
282 g.add_edge(v2, v3, e2, 1);
283
284 g.remove_vertex(v2);
286
287 assert_eq!(g.vertex_count(), 2);
288 assert_eq!(g.edge_count(), 0);
289 assert_eq!(g.neighbors(v1, Direction::Outgoing).len(), 0);
290 assert_eq!(g.neighbors(v3, Direction::Incoming).len(), 0);
291 }
292
293 #[test]
294 fn test_edge_iteration() {
295 let mut g = SimpleGraph::new();
296 let v1 = Vid::new(1);
297 let v2 = Vid::new(2);
298 let v3 = Vid::new(3);
299 let e1 = Eid::new(101);
300 let e2 = Eid::new(102);
301
302 g.add_edge(v1, v2, e1, 1);
303 g.add_edge(v2, v3, e2, 2);
304
305 let edges: Vec<_> = g.edges().collect();
306 assert_eq!(edges.len(), 2);
307 }
308
309 #[test]
310 fn test_neighbor_iteration_after_deletion() {
311 let mut g = SimpleGraph::new();
313 let v1 = Vid::new(1);
314 let v2 = Vid::new(2);
315 let v3 = Vid::new(3);
316 let e1 = Eid::new(101);
317 let e2 = Eid::new(102);
318
319 g.add_edge(v1, v2, e1, 1);
320 g.add_edge(v1, v3, e2, 1);
321
322 g.remove_edge(e1);
324
325 let neighbors = g.neighbors(v1, Direction::Outgoing);
327 assert_eq!(neighbors.len(), 1);
328 assert_eq!(neighbors[0].dst_vid, v3);
329 }
330
331 #[test]
332 fn test_edge_lookup_o1() {
333 let mut g = SimpleGraph::new();
335 let v1 = Vid::new(1);
336 let v2 = Vid::new(2);
337 let v3 = Vid::new(3);
338 let e1 = Eid::new(101);
339 let e2 = Eid::new(102);
340
341 g.add_edge(v1, v2, e1, 1);
342 g.add_edge(v2, v3, e2, 2);
343
344 let edge = g.edge(e1);
346 assert!(edge.is_some());
347 let edge = edge.unwrap();
348 assert_eq!(edge.src_vid, v1);
349 assert_eq!(edge.dst_vid, v2);
350 assert_eq!(edge.edge_type, 1);
351
352 let non_existent = Eid::new(999);
354 assert!(g.edge(non_existent).is_none());
355
356 g.remove_edge(e1);
358 assert!(g.edge(e1).is_none());
359 assert!(g.edge(e2).is_some()); }
361}