Skip to main content

tsift_algorithms/
graph_builder.rs

1use std::collections::{HashMap, HashSet};
2
3pub struct Graph {
4    pub node_vec: Vec<String>,
5    pub node_idx: HashMap<String, usize>,
6    pub out_adj: Vec<HashSet<usize>>,
7    pub in_adj: Vec<HashSet<usize>>,
8}
9
10impl Graph {
11    pub fn node_count(&self) -> usize {
12        self.node_vec.len()
13    }
14
15    pub fn edge_count(&self) -> usize {
16        self.out_adj.iter().map(|s| s.len()).sum()
17    }
18
19    pub fn ensure_nodes(&mut self, names: &[String]) {
20        for name in names {
21            if !self.node_idx.contains_key(name) {
22                let idx = self.node_vec.len();
23                self.node_idx.insert(name.clone(), idx);
24                self.node_vec.push(name.clone());
25                self.out_adj.push(HashSet::new());
26                self.in_adj.push(HashSet::new());
27            }
28        }
29    }
30}
31
32pub fn build_graph(edges: &[(String, String)]) -> Graph {
33    let (node_vec, node_idx) = build_node_index(edges);
34    let n = node_vec.len();
35    let mut out_adj: Vec<HashSet<usize>> = vec![HashSet::new(); n];
36    let mut in_adj: Vec<HashSet<usize>> = vec![HashSet::new(); n];
37    for (a, b) in edges {
38        let ai = node_idx[a];
39        let bi = node_idx[b];
40        if ai != bi {
41            out_adj[ai].insert(bi);
42            in_adj[bi].insert(ai);
43        }
44    }
45    Graph {
46        node_vec,
47        node_idx,
48        out_adj,
49        in_adj,
50    }
51}
52
53pub fn build_node_index(edges: &[(String, String)]) -> (Vec<String>, HashMap<String, usize>) {
54    let mut node_vec: Vec<String> = Vec::new();
55    let mut node_idx: HashMap<String, usize> = HashMap::new();
56    for (a, b) in edges {
57        for name in [a, b] {
58            if !node_idx.contains_key(name) {
59                node_idx.insert(name.clone(), node_vec.len());
60                node_vec.push(name.clone());
61            }
62        }
63    }
64    (node_vec, node_idx)
65}
66
67#[cfg(test)]
68mod tests {
69    use super::*;
70
71    fn e(a: &str, b: &str) -> (String, String) {
72        (a.to_string(), b.to_string())
73    }
74
75    #[test]
76    fn empty_graph() {
77        let g = build_graph(&[]);
78        assert_eq!(g.node_count(), 0);
79        assert_eq!(g.edge_count(), 0);
80    }
81
82    #[test]
83    fn single_edge() {
84        let g = build_graph(&[e("a", "b")]);
85        assert_eq!(g.node_count(), 2);
86        assert_eq!(g.edge_count(), 1);
87        assert!(g.out_adj[g.node_idx["a"]].contains(&g.node_idx["b"]));
88        assert!(g.in_adj[g.node_idx["b"]].contains(&g.node_idx["a"]));
89    }
90
91    #[test]
92    fn self_loop_excluded() {
93        let g = build_graph(&[e("a", "a")]);
94        assert_eq!(g.node_count(), 1);
95        assert_eq!(g.edge_count(), 0);
96    }
97
98    #[test]
99    fn build_node_index_basic() {
100        let (nodes, idx) = build_node_index(&[e("a", "b"), e("b", "c")]);
101        assert_eq!(nodes.len(), 3);
102        assert!(idx.contains_key("a"));
103        assert!(idx.contains_key("b"));
104        assert!(idx.contains_key("c"));
105    }
106
107    #[test]
108    fn ensure_nodes_adds_missing() {
109        let mut g = build_graph(&[e("a", "b")]);
110        g.ensure_nodes(&["c".to_string(), "d".to_string()]);
111        assert_eq!(g.node_count(), 4);
112        assert!(g.node_idx.contains_key("c"));
113        assert!(g.node_idx.contains_key("d"));
114    }
115
116    #[test]
117    fn ensure_nodes_idempotent() {
118        let mut g = build_graph(&[e("a", "b")]);
119        g.ensure_nodes(&["a".to_string()]);
120        assert_eq!(g.node_count(), 2);
121    }
122
123    #[test]
124    fn multi_edge_graph() {
125        let g = build_graph(&[e("a", "b"), e("a", "c"), e("b", "c")]);
126        assert_eq!(g.node_count(), 3);
127        assert_eq!(g.edge_count(), 3);
128    }
129}