1pub mod builder;
34pub mod edge;
35pub mod naming;
36pub mod node;
37pub mod resolver;
38pub mod scc;
39
40use std::collections::HashMap;
41
42pub use builder::{AnalysisParts, GraphBuilder};
43pub use edge::{ConfidenceTier, EdgeData, EdgeKind, confidence_tier};
44pub use node::{
45 DataGraphNode, ExternalClassification, ExternalNode, FileNode, NodeData, SymbolNode,
46};
47pub use scc::{DeployabilityHint, Scc, SccAnalysis};
48
49use crate::language::LangId;
50use crate::model::{FileId, SnapshotId, SymbolId};
51use petgraph::graph::{DiGraph, EdgeIndex, NodeIndex};
52
53#[derive(Debug, Clone)]
55pub struct CodeGraph {
56 graph: DiGraph<NodeData, EdgeData>,
58
59 edge_index: HashMap<(NodeIndex, NodeIndex, EdgeKind), EdgeIndex>,
63
64 pub(crate) file_to_index: HashMap<FileId, NodeIndex>,
66
67 pub(crate) symbol_to_index: HashMap<SymbolId, NodeIndex>,
69
70 pub(crate) external_index: HashMap<String, NodeIndex>,
72
73 pub snapshot_id: SnapshotId,
75}
76
77impl CodeGraph {
78 pub fn new(snapshot_id: SnapshotId) -> Self {
80 Self {
81 graph: DiGraph::new(),
82 edge_index: HashMap::new(),
83 file_to_index: HashMap::new(),
84 symbol_to_index: HashMap::new(),
85 external_index: HashMap::new(),
86 snapshot_id,
87 }
88 }
89
90 pub(crate) fn from_parts(
95 graph: DiGraph<NodeData, EdgeData>,
96 edge_index: HashMap<(NodeIndex, NodeIndex, EdgeKind), EdgeIndex>,
97 file_to_index: HashMap<FileId, NodeIndex>,
98 symbol_to_index: HashMap<SymbolId, NodeIndex>,
99 external_index: HashMap<String, NodeIndex>,
100 snapshot_id: SnapshotId,
101 ) -> Self {
102 Self {
103 graph,
104 edge_index,
105 file_to_index,
106 symbol_to_index,
107 external_index,
108 snapshot_id,
109 }
110 }
111 pub fn graph(&self) -> &DiGraph<NodeData, EdgeData> {
114 &self.graph
115 }
116
117 pub fn add_node(&mut self, node: NodeData) -> NodeIndex {
126 self.graph.add_node(node)
127 }
128
129 pub fn add_file_node(&mut self, node: FileNode) -> NodeIndex {
131 let file_id = node.id;
132 let idx = self.graph.add_node(NodeData::File(node));
133 self.file_to_index.insert(file_id, idx);
134 idx
135 }
136
137 pub fn add_symbol_node(&mut self, node: SymbolNode) -> NodeIndex {
139 let symbol_id = node.id;
140 let idx = self.graph.add_node(NodeData::Symbol(node));
141 self.symbol_to_index.insert(symbol_id, idx);
142 idx
143 }
144 pub fn file_node_index(&self, file_id: FileId) -> Option<NodeIndex> {
145 self.file_to_index.get(&file_id).copied()
146 }
147 pub fn symbol_node_index(&self, symbol_id: SymbolId) -> Option<NodeIndex> {
148 self.symbol_to_index.get(&symbol_id).copied()
149 }
150 pub fn file_node(&self, file_id: FileId) -> Option<&FileNode> {
151 let idx = self.file_node_index(file_id)?;
152 self.graph.node_weight(idx).and_then(|data| data.as_file())
153 }
154 pub fn symbol_node(&self, symbol_id: SymbolId) -> Option<&SymbolNode> {
155 let idx = self.symbol_node_index(symbol_id)?;
156 self.graph
157 .node_weight(idx)
158 .and_then(|data| data.as_symbol())
159 }
160 pub fn file_count(&self) -> usize {
161 self.file_to_index.len()
162 }
163 pub fn symbol_count(&self) -> usize {
164 self.symbol_to_index.len()
165 }
166 pub fn external_count(&self) -> usize {
167 self.external_index.len()
168 }
169 pub fn node_count(&self) -> usize {
170 self.graph.node_count()
171 }
172 pub fn edge_count(&self) -> usize {
173 self.graph.edge_count()
174 }
175 pub fn get_or_create_external_node(&mut self, raw_path: String, language: LangId) -> NodeIndex {
178 if let Some(&idx) = self.external_index.get(&raw_path) {
179 return idx;
180 }
181 let node = NodeData::External(ExternalNode {
182 raw_path: raw_path.clone(),
183 language,
184 classification: None,
185 });
186 let idx = self.graph.add_node(node);
187 self.external_index.insert(raw_path, idx);
188 idx
189 }
190 pub fn add_edge_normalized(
194 &mut self,
195 source: NodeIndex,
196 target: NodeIndex,
197 kind: EdgeKind,
198 confidence: f32,
199 ) {
200 self.add_edge_normalized_with_flow(source, target, kind, confidence, None);
201 }
202
203 pub fn add_edge_normalized_with_flow(
208 &mut self,
209 source: NodeIndex,
210 target: NodeIndex,
211 kind: EdgeKind,
212 confidence: f32,
213 flow_kind: Option<crate::model::FlowKind>,
214 ) {
215 let confidence = confidence.clamp(0.0, 1.0);
216 let key = (source, target, kind);
217 if let Some(&edge_idx) = self.edge_index.get(&key) {
218 self.graph[edge_idx].merge_repeated(confidence, flow_kind);
219 return;
220 }
221 let edge_idx = self.graph.add_edge(
222 source,
223 target,
224 EdgeData {
225 kind,
226 confidence,
227 flow_kind,
228 },
229 );
230 self.edge_index.insert(key, edge_idx);
231 }
232
233 pub fn files(&self) -> impl Iterator<Item = (FileId, &FileNode)> + '_ {
234 let mut files: Vec<(FileId, &FileNode)> = self
235 .file_to_index
236 .iter()
237 .filter_map(|(file_id, &idx)| {
238 self.graph
239 .node_weight(idx)
240 .and_then(|data| data.as_file().map(|f| (*file_id, f)))
241 })
242 .collect();
243 files.sort_by(|a, b| a.1.path.cmp(&b.1.path));
244 files.into_iter()
245 }
246
247 pub fn symbols(&self) -> impl Iterator<Item = (SymbolId, &SymbolNode)> + '_ {
252 let mut symbols: Vec<(SymbolId, &SymbolNode)> = self
253 .symbol_to_index
254 .iter()
255 .filter_map(|(symbol_id, &idx)| {
256 self.graph
257 .node_weight(idx)
258 .and_then(|data| data.as_symbol().map(|s| (*symbol_id, s)))
259 })
260 .collect();
261 symbols.sort_by(|a, b| {
262 let left = self.file_node(a.1.file_id).map(|file| &file.path);
263 let right = self.file_node(b.1.file_id).map(|file| &file.path);
264 left.cmp(&right)
265 .then(a.1.name.cmp(&b.1.name))
266 .then(a.0.to_raw().cmp(&b.0.to_raw()))
267 });
268 symbols.into_iter()
269 }
270 pub fn edges_of_kind(
271 &self,
272 kind: EdgeKind,
273 ) -> impl Iterator<Item = (NodeIndex, NodeIndex)> + '_ {
274 self.graph.edge_indices().filter_map(move |edge_idx| {
275 let (source, target) = self.graph.edge_endpoints(edge_idx)?;
276 let weight = self.graph.edge_weight(edge_idx)?;
277 if weight.kind == kind {
278 Some((source, target))
279 } else {
280 None
281 }
282 })
283 }
284
285 pub fn reference_edges(&self) -> impl Iterator<Item = (SymbolId, SymbolId, f32)> + '_ {
290 self.graph.edge_indices().filter_map(move |edge_idx| {
291 let weight = self.graph.edge_weight(edge_idx)?;
292 if weight.kind != EdgeKind::Reference {
293 return None;
294 }
295 let (source, target) = self.graph.edge_endpoints(edge_idx)?;
296 let source_id = self.graph.node_weight(source)?.as_symbol()?.id;
297 let target_id = self.graph.node_weight(target)?.as_symbol()?.id;
298 Some((source_id, target_id, weight.confidence))
299 })
300 }
301}
302
303#[cfg(test)]
304mod tests {
305 use super::*;
306 use crate::language::LangId;
307 use crate::model::{LineColumn, SourceRange, Symbol, SymbolKind, Visibility, ids::SnapshotId};
308 use std::path::PathBuf;
309
310 fn test_range() -> SourceRange {
311 SourceRange {
312 byte_start: 0,
313 byte_end: 10,
314 start: LineColumn { line: 1, column: 0 },
315 end: LineColumn {
316 line: 1,
317 column: 10,
318 },
319 }
320 }
321
322 fn test_file_node(id: u32, path: &str) -> NodeData {
323 NodeData::File(FileNode {
324 id: FileId::new(id).unwrap(),
325 path: PathBuf::from(path),
326 language: LangId::Rust,
327 snapshot_id: SnapshotId::new(1).unwrap(),
328 })
329 }
330
331 fn test_symbol(id: u32, name: &str) -> Symbol {
332 Symbol {
333 id: SymbolId::new(id).unwrap(),
334 name: name.to_string(),
335 kind: SymbolKind::Function,
336 language: LangId::Rust,
337 file_path: PathBuf::from("test.rs"),
338 source_range: test_range(),
339 name_range: None,
340 visibility: Some(Visibility::Public),
341 signature: None,
342 docstring: None,
343 is_async: false,
344 }
345 }
346
347 #[test]
348 fn code_graph_new_empty() {
349 let graph = CodeGraph::new(SnapshotId::new(1).unwrap());
350 assert_eq!(graph.node_count(), 0);
351 assert_eq!(graph.edge_count(), 0);
352 assert_eq!(graph.snapshot_id.to_raw(), 1);
353 }
354
355 #[test]
356 fn code_graph_file_lookup_returns_none_for_missing() {
357 let graph = CodeGraph::new(SnapshotId::new(1).unwrap());
358 assert!(graph.file_node(FileId::new(1).unwrap()).is_none());
359 assert!(graph.file_node_index(FileId::new(1).unwrap()).is_none());
360 }
361
362 #[test]
363 fn code_graph_symbol_lookup_returns_none_for_missing() {
364 let graph = CodeGraph::new(SnapshotId::new(1).unwrap());
365 assert!(graph.symbol_node(SymbolId::new(1).unwrap()).is_none());
366 assert!(graph.symbol_node_index(SymbolId::new(1).unwrap()).is_none());
367 }
368
369 #[test]
370 fn builder_produces_valid_code_graph() {
371 let mut builder = GraphBuilder::new(SnapshotId::new(1).unwrap());
372 let file_id = builder.add_file(PathBuf::from("test.rs"), LangId::Rust);
373 let symbol = test_symbol(1, "main");
374 let _sym_idx = builder.add_symbol(&symbol).unwrap();
375
376 let graph = builder.build();
377
378 assert_eq!(graph.file_count(), 1);
379 assert_eq!(graph.symbol_count(), 1);
380 assert_eq!(graph.node_count(), 2);
381 assert_eq!(graph.edge_count(), 1); let file_lookup = graph.file_node(file_id);
385 assert!(file_lookup.is_some());
386 assert_eq!(file_lookup.unwrap().language, LangId::Rust);
387
388 let sym_lookup = graph.symbol_node(SymbolId::new(1).unwrap());
389 assert!(sym_lookup.is_some());
390 assert_eq!(sym_lookup.unwrap().name, "main");
391 }
392
393 #[test]
394 fn code_graph_iteration_over_files() {
395 let mut builder = GraphBuilder::new(SnapshotId::new(1).unwrap());
396 builder.add_file(PathBuf::from("a.rs"), LangId::Rust);
397 builder.add_file(PathBuf::from("b.py"), LangId::Python);
398
399 let graph = builder.build();
400 let files: Vec<_> = graph.files().collect();
401
402 assert_eq!(files.len(), 2);
403 }
404
405 #[test]
406 fn file_iteration_is_path_sorted() {
407 let mut builder = GraphBuilder::new(SnapshotId::new(1).unwrap());
408 for name in ["e.rs", "a.rs", "d.rs", "b.rs", "c.rs"] {
409 builder.add_file(PathBuf::from(name), LangId::Rust);
410 }
411
412 let graph = builder.build();
413 let paths: Vec<PathBuf> = graph.files().map(|(_, file)| file.path.clone()).collect();
414 let mut sorted = paths.clone();
415 sorted.sort();
416
417 assert_eq!(paths, sorted, "file iteration must be path sorted");
418 let again: Vec<PathBuf> = graph.files().map(|(_, file)| file.path.clone()).collect();
419 assert_eq!(paths, again, "file iteration must be stable");
420 }
421
422 #[test]
423 fn symbol_iteration_is_path_then_name_sorted() {
424 let mut builder = GraphBuilder::new(SnapshotId::new(1).unwrap());
425 let entries = [
426 ("b.rs", 3, "zeta"),
427 ("a.rs", 1, "beta"),
428 ("a.rs", 2, "alpha"),
429 ("c.rs", 4, "gamma"),
430 ];
431 for (path, id, name) in entries.iter() {
432 builder.add_file(PathBuf::from(path), LangId::Rust);
433 let mut symbol = test_symbol(*id, name);
434 symbol.file_path = PathBuf::from(path);
435 builder.add_symbol(&symbol).unwrap();
436 }
437 let mut sorted_entries = entries;
438 sorted_entries.sort_by_key(|(path, _, name)| (*path, *name));
439
440 let graph = builder.build();
441 let ordered: Vec<(String, String)> = graph
442 .symbols()
443 .map(|(_, symbol)| {
444 let path = graph
445 .file_node(symbol.file_id)
446 .map(|file| file.path.display().to_string())
447 .unwrap_or_default();
448 (path, symbol.name.clone())
449 })
450 .collect();
451 let entries = sorted_entries;
452 let expected: Vec<(String, String)> = entries
453 .iter()
454 .map(|(path, _, name)| (path.to_string(), name.to_string()))
455 .collect();
456
457 assert_eq!(ordered, expected, "symbols must be path then name sorted");
458 }
459
460 #[test]
461 fn code_graph_iteration_over_symbols() {
462 let mut builder = GraphBuilder::new(SnapshotId::new(1).unwrap());
463 let _file_id = builder.add_file(PathBuf::from("test.rs"), LangId::Rust);
464 let sym1 = test_symbol(1, "func_a");
465 let sym2 = test_symbol(2, "func_b");
466 builder.add_symbol(&sym1).unwrap();
467 builder.add_symbol(&sym2).unwrap();
468
469 let graph = builder.build();
470 let symbols: Vec<_> = graph.symbols().collect();
471
472 assert_eq!(symbols.len(), 2);
473 }
474
475 #[test]
476 fn code_graph_edges_of_kind_filtering() {
477 let mut builder = GraphBuilder::new(SnapshotId::new(1).unwrap());
478 let file1 = builder.add_file(PathBuf::from("a.rs"), LangId::Rust);
479 let _file2 = builder.add_file(PathBuf::from("b.rs"), LangId::Rust);
480
481 builder.add_import(file1, PathBuf::from("b.rs"));
483
484 let graph = builder.build();
485
486 let ownership_edges: Vec<_> = graph.edges_of_kind(EdgeKind::Ownership).collect();
487 let import_edges: Vec<_> = graph.edges_of_kind(EdgeKind::Import).collect();
488
489 assert_eq!(ownership_edges.len(), 0); assert_eq!(import_edges.len(), 1);
491 }
492
493 #[test]
494 fn build_populated_edge_index_normalizes_post_build_edges() {
495 let mut builder = GraphBuilder::new(SnapshotId::new(1).unwrap());
496 let file_a = builder.add_file(PathBuf::from("a.rs"), LangId::Rust);
497 let file_b = builder.add_file(PathBuf::from("b.rs"), LangId::Rust);
498 builder.add_import(file_a, PathBuf::from("b.rs"));
499 let mut graph = builder.build();
500 assert_eq!(graph.edge_count(), 1);
501
502 let a_idx = graph.file_node_index(file_a).unwrap();
503 let b_idx = graph.file_node_index(file_b).unwrap();
504
505 graph.add_edge_normalized(a_idx, b_idx, EdgeKind::Import, 0.5);
507 graph.add_edge_normalized(a_idx, b_idx, EdgeKind::Import, 0.8);
508 assert_eq!(graph.edge_count(), 1);
509
510 graph.add_edge_normalized(a_idx, b_idx, EdgeKind::Reference, 0.9);
512 assert_eq!(graph.edge_count(), 2);
513 }
514
515 #[test]
516 fn add_edge_normalized_handles_multiple_edge_kinds_between_same_nodes() {
517 let mut graph = CodeGraph::new(SnapshotId::new(1).unwrap());
518 let n1 = graph.add_node(test_file_node(1, "a.rs"));
519 let n2 = graph.add_node(test_file_node(2, "b.rs"));
520
521 graph.add_edge_normalized(n1, n2, EdgeKind::Reference, 0.7);
523 graph.add_edge_normalized(n1, n2, EdgeKind::Import, 0.5);
525 graph.add_edge_normalized(n1, n2, EdgeKind::Reference, 0.9);
527
528 let ref_count = graph
529 .graph()
530 .edges_connecting(n1, n2)
531 .filter(|e| e.weight().kind == EdgeKind::Reference)
532 .count();
533 assert_eq!(
534 ref_count, 1,
535 "Expected 1 Reference edge, but found {ref_count}"
536 );
537 assert_eq!(graph.edge_count(), 2);
538 }
539
540 #[test]
541 fn add_edge_normalized_with_flow_preserves_first_flow_kind() {
542 use crate::model::{DataNodeId, DataScope, FlowKind};
543 let mut graph = CodeGraph::new(SnapshotId::new(1).unwrap());
544 let n1 = graph.add_node(NodeData::Data(DataGraphNode {
545 id: DataNodeId::new(1).unwrap(),
546 symbol_id: None,
547 name: Some("x".into()),
548 scope: DataScope::Local,
549 type_hint: None,
550 source_range: test_range(),
551 }));
552 let n2 = graph.add_node(NodeData::Data(DataGraphNode {
553 id: DataNodeId::new(2).unwrap(),
554 symbol_id: None,
555 name: Some("y".into()),
556 scope: DataScope::Local,
557 type_hint: None,
558 source_range: test_range(),
559 }));
560
561 graph.add_edge_normalized_with_flow(n1, n2, EdgeKind::Flow, 0.9, Some(FlowKind::DefUse));
562 graph.add_edge_normalized_with_flow(n1, n2, EdgeKind::Flow, 0.8, Some(FlowKind::Argument));
563
564 assert_eq!(graph.edge_count(), 1);
565 let edge = graph.graph().edges_connecting(n1, n2).next().unwrap();
566 assert_eq!(edge.weight().flow_kind, Some(FlowKind::DefUse));
567 assert_eq!(edge.weight().confidence, 0.9);
568 }
569
570 #[test]
571 fn add_edge_normalized_duplicate_never_grows_edge_count() {
572 let mut graph = CodeGraph::new(SnapshotId::new(1).unwrap());
573 let n1 = graph.add_node(test_file_node(1, "a.rs"));
574 let n2 = graph.add_node(test_file_node(2, "b.rs"));
575
576 graph.add_edge_normalized(n1, n2, EdgeKind::Import, 0.5);
577 for confidence in [0.6, 0.3, 0.9, 0.4, 0.7] {
578 graph.add_edge_normalized(n1, n2, EdgeKind::Import, confidence);
579 }
580
581 assert_eq!(graph.edge_count(), 1);
582 let edge = graph.graph().edges_connecting(n1, n2).next().unwrap();
583 assert_eq!(edge.weight().confidence, 0.9);
584 }
585
586 #[test]
587 fn add_edge_normalized_counts_only_distinct_triples() {
588 let mut graph = CodeGraph::new(SnapshotId::new(1).unwrap());
589 let n1 = graph.add_node(test_file_node(1, "a.rs"));
590 let n2 = graph.add_node(test_file_node(2, "b.rs"));
591 let n3 = graph.add_node(test_file_node(3, "c.rs"));
592
593 graph.add_edge_normalized(n1, n2, EdgeKind::Import, 0.5);
594 graph.add_edge_normalized(n1, n2, EdgeKind::Reference, 0.6);
595 graph.add_edge_normalized(n1, n3, EdgeKind::Reference, 0.7);
596 graph.add_edge_normalized(n1, n2, EdgeKind::Import, 0.9);
597
598 assert_eq!(graph.edge_count(), 3);
599 }
600}