1use std::collections::{HashMap, VecDeque};
2
3use petgraph::algo;
4use petgraph::graph::{DiGraph, NodeIndex};
5use petgraph::visit::EdgeRef;
6use serde::{Deserialize, Serialize};
7
8use crate::edge::{EdgeKind, GraphEdge};
9use crate::node::{GraphNode, NodeKind};
10
11#[derive(Debug, Clone, Serialize, Deserialize)]
12pub struct CodeGraph {
13 pub nodes: HashMap<u64, GraphNode>,
14 pub edges: Vec<GraphEdge>,
15 pub next_id: u64,
16}
17
18pub struct PetGraphWrapper {
19 pub graph: DiGraph<GraphNode, EdgeKind>,
20 pub node_id_to_index: HashMap<u64, NodeIndex>,
21 pub index_to_node_id: HashMap<NodeIndex, u64>,
22}
23
24impl CodeGraph {
25 pub fn new() -> Self {
26 Self {
27 nodes: HashMap::new(),
28 edges: Vec::new(),
29 next_id: 1,
30 }
31 }
32
33 pub fn add_node(&mut self, kind: NodeKind, name: &str) -> u64 {
34 let id = self.next_id;
35 self.next_id += 1;
36 let node = GraphNode::new(id, name.to_string(), kind);
37 self.nodes.insert(id, node);
38 id
39 }
40
41 pub fn add_edge(&mut self, source_id: u64, target_id: u64, kind: EdgeKind) {
42 let edge = GraphEdge::new(source_id, target_id, kind);
43 self.edges.push(edge);
44 }
45
46 pub fn get_node(&self, id: u64) -> Option<&GraphNode> {
47 self.nodes.get(&id)
48 }
49
50 pub fn find_nodes_by_name(&self, name: &str) -> Vec<&GraphNode> {
51 self.nodes
52 .values()
53 .filter(|n| n.name.contains(name))
54 .collect()
55 }
56
57 pub fn find_nodes_by_path(&self, path: &str) -> Vec<&GraphNode> {
58 self.nodes
59 .values()
60 .filter(|n| n.file_path.as_deref().map_or(false, |p| p == path))
61 .collect()
62 }
63
64 pub fn reverse_dependencies(&self, node_id: u64) -> Vec<&GraphEdge> {
65 self.edges
66 .iter()
67 .filter(|e| e.target_id == node_id)
68 .collect()
69 }
70
71 pub fn dependencies(&self, node_id: u64) -> Vec<&GraphEdge> {
72 self.edges
73 .iter()
74 .filter(|e| e.source_id == node_id)
75 .collect()
76 }
77
78 pub fn find_path(&self, from_id: u64, to_id: u64) -> Option<Vec<u64>> {
79 let pet = self.to_petgraph();
80 let from = *pet.node_id_to_index.get(&from_id)?;
81 let to = *pet.node_id_to_index.get(&to_id)?;
82 let path = algo::dijkstra(&pet.graph, from, Some(to), |_| 1);
83 if path.contains_key(&to) {
84 let result = vec![to_id];
85 Some(result)
86 } else {
87 None
88 }
89 }
90
91 pub fn detect_cycles(&self) -> Vec<Vec<u64>> {
92 let pet = self.to_petgraph();
93 let mut cycles = Vec::new();
94 if algo::is_cyclic_directed(&pet.graph) {
95 let mut visited = HashMap::new();
96 let mut stack = Vec::new();
97 for node in pet.graph.node_indices() {
98 if visited.contains_key(&node) {
99 continue;
100 }
101 if let Some(cycle) = self.dfs_find_cycle(node, &pet, &mut visited, &mut stack) {
102 cycles.push(cycle);
103 }
104 }
105 }
106 cycles
107 }
108
109 fn dfs_find_cycle(
110 &self,
111 start: NodeIndex,
112 pet: &PetGraphWrapper,
113 visited: &mut HashMap<NodeIndex, bool>,
114 stack: &mut Vec<u64>,
115 ) -> Option<Vec<u64>> {
116 visited.insert(start, true);
117 stack.push(pet.index_to_node_id.get(&start).copied().unwrap_or(0));
118 for edge in pet.graph.edges(start) {
119 let target = edge.target();
120 if !visited.contains_key(&target) {
121 if let Some(cycle) = self.dfs_find_cycle(target, pet, visited, stack) {
122 return Some(cycle);
123 }
124 } else if stack.contains(&pet.index_to_node_id.get(&target).copied().unwrap_or(0)) {
125 let cycle_start = stack
126 .iter()
127 .position(|&x| x == pet.index_to_node_id.get(&target).copied().unwrap_or(0))
128 .unwrap_or(0);
129 return Some(stack[cycle_start..].to_vec());
130 }
131 }
132 stack.pop();
133 None
134 }
135
136 pub fn to_petgraph(&self) -> PetGraphWrapper {
137 let mut graph = DiGraph::new();
138 let mut node_id_to_index = HashMap::new();
139 let mut index_to_node_id = HashMap::new();
140 for (id, node) in &self.nodes {
141 let idx = graph.add_node(node.clone());
142 node_id_to_index.insert(*id, idx);
143 index_to_node_id.insert(idx, *id);
144 }
145 for edge in &self.edges {
146 if let (Some(&source_idx), Some(&target_idx)) = (
147 node_id_to_index.get(&edge.source_id),
148 node_id_to_index.get(&edge.target_id),
149 ) {
150 graph.add_edge(source_idx, target_idx, edge.kind.clone());
151 }
152 }
153 PetGraphWrapper {
154 graph,
155 node_id_to_index,
156 index_to_node_id,
157 }
158 }
159
160 pub fn importance_scores(&self) -> HashMap<u64, f64> {
161 let pet = self.to_petgraph();
162 let mut scores: HashMap<u64, f64> = HashMap::new();
163 let num_nodes = pet.graph.node_count();
164 if num_nodes == 0 {
165 return scores;
166 }
167 let damping = 0.85;
168 let initial = 1.0 / num_nodes as f64;
169 for &id in self.nodes.keys() {
170 scores.insert(id, initial);
171 }
172 for _ in 0..20 {
173 let mut new_scores: HashMap<u64, f64> = HashMap::new();
174 let dangling_sum: f64 = scores.values().sum();
175 for (&id, _) in &self.nodes {
176 let mut score = (1.0 - damping) / num_nodes as f64;
177 let mut incoming_weight = 0.0;
178 for edge in &self.edges {
179 if edge.target_id == id {
180 let source_score = scores.get(&edge.source_id).copied().unwrap_or(initial);
181 let out_degree = self
182 .edges
183 .iter()
184 .filter(|e| e.source_id == edge.source_id)
185 .count();
186 if out_degree > 0 {
187 incoming_weight += source_score * edge.weight / out_degree as f64;
188 }
189 }
190 }
191 score += damping * incoming_weight;
192 score += damping * dangling_sum / num_nodes as f64;
193 new_scores.insert(id, score);
194 }
195 scores = new_scores;
196 }
197 scores
198 }
199
200 pub fn bfs_reachable(&self, start_id: u64, max_depth: usize) -> Vec<u64> {
201 let mut visited = HashMap::new();
202 let mut queue = VecDeque::new();
203 let mut result = Vec::new();
204 queue.push_back((start_id, 0));
205 visited.insert(start_id, true);
206 while let Some((node_id, depth)) = queue.pop_front() {
207 if depth > max_depth {
208 continue;
209 }
210 result.push(node_id);
211 for edge in &self.edges {
212 if edge.source_id == node_id && !visited.contains_key(&edge.target_id) {
213 visited.insert(edge.target_id, true);
214 queue.push_back((edge.target_id, depth + 1));
215 }
216 }
217 }
218 result
219 }
220}
221
222impl Default for CodeGraph {
223 fn default() -> Self {
224 Self::new()
225 }
226}