Skip to main content

citadel_vector/vendored/prism/
graph.rs

1/// CSR (Compressed Sparse Row) graph for neighbor storage.
2///
3/// For n nodes, `offsets` has n+1 entries.
4/// Neighbors of node i are `neighbors[offsets[i]..offsets[i+1]]`.
5pub struct Graph {
6    pub offsets: Vec<u32>,
7    pub neighbors: Vec<u32>,
8    pub n: usize,
9}
10
11impl Graph {
12    /// Build from adjacency lists. Each entry in `adj` is the neighbor list for that node.
13    pub fn from_adj(adj: &[Vec<u32>]) -> Self {
14        let n = adj.len();
15        let mut offsets = Vec::with_capacity(n + 1);
16        let mut neighbors = Vec::new();
17        let mut offset = 0u32;
18        for list in adj {
19            offsets.push(offset);
20            neighbors.extend_from_slice(list);
21            offset += list.len() as u32;
22        }
23        offsets.push(offset);
24        Self {
25            offsets,
26            neighbors,
27            n,
28        }
29    }
30
31    /// Empty graph with n nodes and no edges.
32    pub fn empty(n: usize) -> Self {
33        Self {
34            offsets: vec![0; n + 1],
35            neighbors: Vec::new(),
36            n,
37        }
38    }
39
40    /// Degree of node i.
41    #[inline]
42    pub fn degree(&self, i: u32) -> usize {
43        let i = i as usize;
44        (self.offsets[i + 1] - self.offsets[i]) as usize
45    }
46
47    /// Neighbors of node i.
48    #[inline]
49    pub fn neighbors(&self, i: u32) -> &[u32] {
50        let i = i as usize;
51        let start = self.offsets[i] as usize;
52        let end = self.offsets[i + 1] as usize;
53        &self.neighbors[start..end]
54    }
55
56    /// Total number of edges (directed).
57    pub fn num_edges(&self) -> usize {
58        self.neighbors.len()
59    }
60}
61
62/// Mutable adjacency list builder that converts to CSR.
63pub struct AdjBuilder {
64    adj: Vec<Vec<u32>>,
65}
66
67impl AdjBuilder {
68    pub fn new(n: usize) -> Self {
69        Self {
70            adj: vec![Vec::new(); n],
71        }
72    }
73
74    /// Add a directed edge from `src` to `dst`.
75    #[inline]
76    pub fn add_edge(&mut self, src: u32, dst: u32) {
77        self.adj[src as usize].push(dst);
78    }
79
80    /// Add bidirectional edge.
81    #[inline]
82    pub fn add_undirected(&mut self, a: u32, b: u32) {
83        self.adj[a as usize].push(b);
84        self.adj[b as usize].push(a);
85    }
86
87    /// Get current neighbors (mutable).
88    pub fn neighbors_mut(&mut self, i: u32) -> &mut Vec<u32> {
89        &mut self.adj[i as usize]
90    }
91
92    /// Get current neighbors.
93    pub fn neighbors(&self, i: u32) -> &[u32] {
94        &self.adj[i as usize]
95    }
96
97    /// Total directed edges currently stored.
98    pub fn total_edges(&self) -> usize {
99        self.adj.iter().map(|v| v.len()).sum()
100    }
101
102    /// Create a read-only CSR graph snapshot without consuming the builder.
103    pub fn snapshot(&self) -> Graph {
104        Graph::from_adj(&self.adj)
105    }
106
107    /// Freeze into CSR graph, deduplicating edges.
108    pub fn build(mut self) -> Graph {
109        for list in &mut self.adj {
110            list.sort_unstable();
111            list.dedup();
112        }
113        Graph::from_adj(&self.adj)
114    }
115}
116
117#[cfg(test)]
118mod tests {
119    use super::*;
120
121    #[test]
122    fn test_graph_from_adj() {
123        let adj = vec![vec![1, 2], vec![0], vec![0, 1]];
124        let g = Graph::from_adj(&adj);
125        assert_eq!(g.n, 3);
126        assert_eq!(g.degree(0), 2);
127        assert_eq!(g.degree(1), 1);
128        assert_eq!(g.degree(2), 2);
129        assert_eq!(g.neighbors(0), &[1, 2]);
130        assert_eq!(g.neighbors(1), &[0]);
131        assert_eq!(g.neighbors(2), &[0, 1]);
132    }
133
134    #[test]
135    fn test_adj_builder() {
136        let mut builder = AdjBuilder::new(3);
137        builder.add_undirected(0, 1);
138        builder.add_undirected(0, 2);
139        let g = builder.build();
140        assert_eq!(g.neighbors(0), &[1, 2]);
141        assert_eq!(g.neighbors(1), &[0]);
142        assert_eq!(g.neighbors(2), &[0]);
143    }
144}