1use std::collections::{BTreeMap, HashMap, HashSet, VecDeque};
18
19use crate::model::{Edge, EdgeKind, Symbol};
20use petgraph::graph::{DiGraph, NodeIndex};
21use petgraph::visit::EdgeRef;
22use petgraph::Direction::{Incoming, Outgoing};
23
24pub const DEPENDENCY: &[EdgeKind] = &[
26 EdgeKind::Calls,
27 EdgeKind::Imports,
28 EdgeKind::Instantiates,
29 EdgeKind::Inherits,
30 EdgeKind::Implements,
31 EdgeKind::Consumes,
32 EdgeKind::Subscribes,
33 EdgeKind::Reads,
34];
35
36pub const PROVIDES: &[EdgeKind] = &[
39 EdgeKind::Serves,
40 EdgeKind::Publishes,
41 EdgeKind::Writes,
42 EdgeKind::Defines,
43];
44
45#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
47pub struct Node {
48 pub id: String,
49 pub name: String,
50 pub kind: String,
51 pub repo: String,
52 pub location: String,
54 pub signature: Option<String>,
55}
56
57#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
59pub struct FlowStep {
60 pub node: Node,
61 pub depth: usize,
62 pub via: String,
63 pub from: String,
64}
65
66#[derive(Debug, Clone, serde::Serialize)]
68pub struct Hop {
69 pub node: Node,
70 pub depth: usize,
71 pub via: EdgeKind,
72}
73
74#[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize)]
76pub struct Zoom {
77 pub focus: Option<Node>,
78 pub container: Option<Node>,
79 pub members: Vec<Node>,
80 pub callers: Vec<Node>,
81 pub callees: Vec<Node>,
82 pub importers: Vec<Node>,
83 pub contracts: Vec<ContractRole>,
85 pub providers: Vec<Node>,
87 pub consumers: Vec<Node>,
88}
89
90#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
92pub struct ContractRole {
93 pub role: EdgeKind,
94 pub contract: Node,
95 pub counterparts: Vec<Node>,
96}
97
98#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
100pub struct ContractView {
101 pub contract: Node,
102 pub providers: Vec<Node>,
103 pub consumers: Vec<Node>,
104}
105
106pub fn is_contract_role(k: &EdgeKind) -> bool {
108 PROVIDES.contains(k)
109 || matches!(
110 k,
111 EdgeKind::Consumes | EdgeKind::Subscribes | EdgeKind::Reads
112 )
113}
114
115struct NodeData {
116 sym: Symbol,
117}
118
119pub struct CodeGraph {
120 g: DiGraph<NodeData, EdgeKind>,
121 by_id: HashMap<String, NodeIndex>,
122 by_name: HashMap<String, Vec<NodeIndex>>,
123 pub dangling_edges: usize,
125}
126
127impl CodeGraph {
128 pub fn build(symbols: &[Symbol], edges: &[Edge]) -> Self {
130 let mut g = DiGraph::new();
131 let mut by_id = HashMap::new();
132 let mut by_name: HashMap<String, Vec<NodeIndex>> = HashMap::new();
133
134 for s in symbols {
135 let key = s.id.render();
136 if by_id.contains_key(&key) {
137 continue; }
139 let idx = g.add_node(NodeData { sym: s.clone() });
140 by_id.insert(key, idx);
141 by_name.entry(s.name.clone()).or_default().push(idx);
142 }
143
144 let mut dangling_edges = 0;
145 for e in edges {
146 match (by_id.get(&e.src.render()), by_id.get(&e.dst.render())) {
147 (Some(&a), Some(&b)) => {
148 g.add_edge(a, b, e.kind.clone());
149 }
150 _ => dangling_edges += 1,
151 }
152 }
153
154 Self {
155 g,
156 by_id,
157 by_name,
158 dangling_edges,
159 }
160 }
161
162 pub fn node_count(&self) -> usize {
163 self.g.node_count()
164 }
165 pub fn edge_count(&self) -> usize {
166 self.g.edge_count()
167 }
168
169 pub fn scope_of(&self, focus: &str) -> Option<(String, Vec<String>)> {
171 let idx = self.lookup(focus)?;
172 let mut ids = vec![self.node_of(idx).id];
173 let mut cur = idx;
174 for _ in 0..6 {
175 let parent = self
176 .g
177 .edges_directed(cur, Incoming)
178 .find(|e| *e.weight() == EdgeKind::Contains)
179 .map(|e| e.source());
180 match parent {
181 Some(p) => {
182 ids.push(self.node_of(p).id);
183 cur = p;
184 }
185 None => break,
186 }
187 }
188 Some((self.node_of(idx).name, ids))
189 }
190
191 pub fn node(&self, idx: NodeIndex) -> Node {
192 self.node_of(idx)
193 }
194
195 fn node_of(&self, idx: NodeIndex) -> Node {
197 let mut n = self.plain(idx);
198 if self.g[idx].sym.kind == crate::model::SymbolKind::Method {
199 if let Some(owner) = self
200 .g
201 .edges_directed(idx, Incoming)
202 .find(|e| *e.weight() == EdgeKind::Contains)
203 .map(|e| &self.g[e.source()].sym)
204 .filter(|o| o.kind != crate::model::SymbolKind::File)
205 {
206 n.name = format!("{}.{}", owner.name, n.name);
207 }
208 }
209 n
210 }
211
212 fn plain(&self, idx: NodeIndex) -> Node {
213 let s = &self.g[idx].sym;
214 Node {
215 id: s.id.render(),
216 name: s.name.clone(),
217 kind: format!("{:?}", s.kind),
218 repo: s.id.package.clone(),
219 location: format!("{}:{}", s.file_path, s.range.start.line),
220 signature: s.signature.clone(),
221 }
222 }
223
224 pub fn lookup(&self, focus: &str) -> Option<NodeIndex> {
228 if let Some(&i) = self.by_id.get(focus) {
229 return Some(i);
230 }
231 let mut candidates: Vec<NodeIndex> = self
232 .g
233 .node_indices()
234 .filter(|&i| {
235 let d = &self.g[i].sym.id.descriptor;
236 d == focus || d.contains(focus)
237 })
238 .collect();
239 if candidates.is_empty() {
240 if let Some(v) = self.by_name.get(focus) {
241 candidates = v.clone();
242 }
243 }
244 if let (true, Some((_, method))) = (candidates.is_empty(), focus.rsplit_once('.')) {
246 candidates = self
247 .by_name
248 .get(method)
249 .into_iter()
250 .flatten()
251 .copied()
252 .filter(|&i| self.node_of(i).name == focus)
253 .collect();
254 }
255 candidates.sort_by_key(|&i| {
256 let s = &self.g[i].sym;
257 let inexact = (s.name != focus && self.node_of(i).name != focus) as usize;
258 let is_file = matches!(s.kind, crate::model::SymbolKind::File) as usize;
259 let path = s.file_path.to_ascii_lowercase();
260 let is_test = (path.contains("test") || path.contains("spec")) as usize;
261 (inexact, is_file, is_test, s.id.descriptor.len())
262 });
263 candidates.first().copied()
264 }
265
266 pub fn repos(&self) -> Vec<(String, usize)> {
268 let mut m: BTreeMap<String, usize> = BTreeMap::new();
269 for i in self.g.node_indices() {
270 if self.g[i].sym.kind != crate::model::SymbolKind::Contract {
271 *m.entry(self.g[i].sym.id.package.clone()).or_default() += 1;
272 }
273 }
274 m.into_iter().collect()
275 }
276
277 pub fn dependents_count(&self, rendered_id: &str) -> usize {
281 let Some(&idx) = self.by_id.get(rendered_id) else {
282 return 0;
283 };
284 self.g
285 .edges_directed(idx, Incoming)
286 .filter(|e| DEPENDENCY.contains(e.weight()))
287 .count()
288 }
289
290 pub fn find(&self, query: &str, limit: usize) -> Vec<Node> {
293 let mut v: Vec<NodeIndex> = self
294 .g
295 .node_indices()
296 .filter(|&i| {
297 let s = &self.g[i].sym;
298 s.name == query || s.name.contains(query) || s.id.descriptor.contains(query)
299 })
300 .collect();
301 v.sort_by_key(|&i| {
302 let s = &self.g[i].sym;
303 let is_file = matches!(s.kind, crate::model::SymbolKind::File) as usize;
304 (is_file, s.id.descriptor.len())
305 });
306 v.into_iter().take(limit).map(|i| self.node_of(i)).collect()
307 }
308
309 pub fn ripple(&self, focus: NodeIndex, max_depth: usize) -> Vec<Hop> {
313 self.bfs(focus, max_depth, Incoming, DEPENDENCY, PROVIDES)
314 }
315
316 pub fn flow(&self, focus: NodeIndex, max_depth: usize) -> Vec<FlowStep> {
319 use crate::model::SymbolKind::Contract;
320 let mut seen: HashSet<NodeIndex> = HashSet::from([focus]);
321 let mut q: VecDeque<(NodeIndex, usize)> = VecDeque::from([(focus, 0)]);
322 let mut out = Vec::new();
323 while let Some((n, depth)) = q.pop_front() {
324 if depth >= max_depth {
325 continue;
326 }
327 let next: Vec<(NodeIndex, EdgeKind)> = if self.g[n].sym.kind == Contract {
328 self.g
329 .edges_directed(n, Incoming)
330 .filter(|e| matches!(e.weight(), EdgeKind::Serves | EdgeKind::Subscribes))
331 .map(|e| (e.source(), e.weight().clone()))
332 .collect()
333 } else {
334 self.g
335 .edges_directed(n, Outgoing)
336 .filter(|e| {
337 matches!(
338 e.weight(),
339 EdgeKind::Calls
340 | EdgeKind::Instantiates
341 | EdgeKind::Consumes
342 | EdgeKind::Publishes
343 | EdgeKind::Writes
344 )
345 })
346 .map(|e| (e.target(), e.weight().clone()))
347 .collect()
348 };
349 for (m, via) in next {
350 if seen.insert(m) {
351 out.push(FlowStep {
352 node: self.node_of(m),
353 depth: depth + 1,
354 via: via.as_str().to_string(),
355 from: self.g[n].sym.id.render(),
356 });
357 q.push_back((m, depth + 1));
358 }
359 }
360 }
361 out
362 }
363
364 pub fn entry_points(&self, top: usize) -> Vec<(NodeIndex, usize)> {
367 use crate::model::SymbolKind::{Function, Method};
368 let mut out: Vec<(NodeIndex, usize, usize)> = self
369 .g
370 .node_indices()
371 .filter(|&i| matches!(self.g[i].sym.kind, Function | Method))
372 .filter(|&i| crate::embed::rank::path_penalty(&self.g[i].sym.file_path) > 0.3)
373 .filter_map(|i| {
374 let steps = self.flow(i, 6);
375 let repos: HashSet<&str> = steps.iter().map(|s| s.node.repo.as_str()).collect();
376 (!steps.is_empty()).then_some((i, repos.len(), steps.len()))
377 })
378 .collect();
379 out.sort_by_key(|x| std::cmp::Reverse((x.1, x.2)));
380 out.truncate(top);
381 out.into_iter().map(|(i, _, n)| (i, n)).collect()
382 }
383
384 pub fn thread(&self, focus: NodeIndex, max_depth: usize) -> Vec<Hop> {
386 self.bfs(
387 focus,
388 max_depth,
389 Outgoing,
390 &[EdgeKind::Calls, EdgeKind::Instantiates],
391 &[],
392 )
393 }
394
395 fn bfs(
396 &self,
397 start: NodeIndex,
398 max_depth: usize,
399 dir: petgraph::Direction,
400 kinds: &[EdgeKind],
401 reversed: &[EdgeKind],
402 ) -> Vec<Hop> {
403 let mut seen: HashSet<NodeIndex> = HashSet::from([start]);
404 let mut q: VecDeque<(NodeIndex, usize)> = VecDeque::from([(start, 0)]);
405 let mut out = Vec::new();
406 while let Some((n, depth)) = q.pop_front() {
407 if depth >= max_depth {
408 continue;
409 }
410 let forward = self
411 .g
412 .edges_directed(n, dir)
413 .filter(|e| kinds.contains(e.weight()))
414 .map(|e| {
415 (
416 e.weight(),
417 if dir == Incoming {
418 e.source()
419 } else {
420 e.target()
421 },
422 )
423 });
424 let backward = self
425 .g
426 .edges_directed(n, dir.opposite())
427 .filter(|e| reversed.contains(e.weight()))
428 .map(|e| {
429 (
430 e.weight(),
431 if dir == Incoming {
432 e.target()
433 } else {
434 e.source()
435 },
436 )
437 });
438 for (kind, next) in forward.chain(backward).collect::<Vec<_>>() {
439 if seen.insert(next) {
440 out.push(Hop {
441 node: self.node_of(next),
442 depth: depth + 1,
443 via: kind.clone(),
444 });
445 q.push_back((next, depth + 1));
446 }
447 }
448 }
449 out
450 }
451
452 fn sides(&self, contract: NodeIndex) -> (Vec<Node>, Vec<Node>) {
454 let (mut p, mut c) = (Vec::new(), Vec::new());
455 let (mut seen_p, mut seen_c) = (HashSet::new(), HashSet::new());
456 for e in self.g.edges_directed(contract, Incoming) {
457 if PROVIDES.contains(e.weight()) {
458 if seen_p.insert(e.source()) {
459 p.push(self.node_of(e.source()));
460 }
461 } else if is_contract_role(e.weight()) && seen_c.insert(e.source()) {
462 c.push(self.node_of(e.source()));
463 }
464 }
465 (p, c)
466 }
467
468 pub fn contracts(&self, kind: Option<&str>, dangling: bool) -> Vec<ContractView> {
470 let mut out: Vec<ContractView> = self
471 .g
472 .node_indices()
473 .filter(|&i| self.g[i].sym.kind == crate::model::SymbolKind::Contract)
474 .filter(|&i| kind.is_none_or(|k| self.g[i].sym.name.starts_with(&format!("{k}:"))))
475 .map(|i| {
476 let (providers, consumers) = self.sides(i);
477 ContractView {
478 contract: self.node_of(i),
479 providers,
480 consumers,
481 }
482 })
483 .filter(|v| !dangling || v.providers.is_empty() || v.consumers.is_empty())
484 .collect();
485 out.sort_by(|a, b| a.contract.name.cmp(&b.contract.name));
486 out
487 }
488
489 pub fn zoom(&self, focus: NodeIndex) -> Zoom {
491 let mut z = Zoom {
492 focus: Some(self.node_of(focus)),
493 ..Default::default()
494 };
495 for e in self.g.edges_directed(focus, Incoming) {
496 match e.weight() {
497 EdgeKind::Contains => z.container = Some(self.node_of(e.source())),
498 EdgeKind::Calls | EdgeKind::Instantiates => {
499 z.callers.push(self.node_of(e.source()))
500 }
501 EdgeKind::Imports => z.importers.push(self.node_of(e.source())),
502 _ => {}
503 }
504 }
505 for e in self.g.edges_directed(focus, Incoming) {
506 if is_contract_role(e.weight()) {
507 if PROVIDES.contains(e.weight()) {
508 z.providers.push(self.node_of(e.source()));
509 } else {
510 z.consumers.push(self.node_of(e.source()));
511 }
512 }
513 }
514 let mut roles = HashSet::new();
515 let me = self.g[focus].sym.id.render();
516 for e in self.g.edges_directed(focus, Outgoing) {
517 if is_contract_role(e.weight()) && roles.insert((e.weight().clone(), e.target())) {
518 let provider = PROVIDES.contains(e.weight());
519 let (providers, consumers) = self.sides(e.target());
520 let mut counterparts = if provider { consumers } else { providers };
521 counterparts.retain(|n| n.id != me);
522 z.contracts.push(ContractRole {
523 role: e.weight().clone(),
524 contract: self.node_of(e.target()),
525 counterparts,
526 });
527 }
528 }
529 for e in self.g.edges_directed(focus, Outgoing) {
530 match e.weight() {
531 EdgeKind::Contains => z.members.push(self.node_of(e.target())),
532 EdgeKind::Calls | EdgeKind::Instantiates => {
533 z.callees.push(self.node_of(e.target()))
534 }
535 _ => {}
536 }
537 }
538 z
539 }
540
541 pub fn context_pack(&self, focus: NodeIndex, token_budget: usize) -> Vec<(Node, usize)> {
547 let ranks = self.personalized_pagerank(focus, 0.85, 40);
548 let mut ranked: Vec<(NodeIndex, f32)> = ranks
549 .into_iter()
550 .enumerate()
551 .filter(|&(i, _)| NodeIndex::new(i) != focus)
552 .map(|(i, r)| (NodeIndex::new(i), r))
553 .collect();
554 ranked.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
555
556 let mut pack = Vec::new();
559 let mut spent = est_tokens(&self.g[focus].sym);
560 pack.push((self.node_of(focus), spent));
561 for (idx, rank) in ranked {
562 if rank <= 0.0 {
563 break; }
565 if matches!(
566 self.g[idx].sym.kind,
567 crate::model::SymbolKind::File | crate::model::SymbolKind::Contract
568 ) {
569 continue;
570 }
571 let t = est_tokens(&self.g[idx].sym);
572 if spent + t > token_budget {
573 continue;
574 }
575 spent += t;
576 pack.push((self.node_of(idx), t));
577 }
578 pack
579 }
580
581 fn personalized_pagerank(&self, seed: NodeIndex, damping: f32, iters: usize) -> Vec<f32> {
585 let n = self.g.node_count();
586 if n == 0 {
587 return vec![];
588 }
589 let mut neighbors: Vec<Vec<usize>> = vec![Vec::new(); n];
591 for e in self.g.edge_references() {
592 if !DEPENDENCY.contains(e.weight()) && !PROVIDES.contains(e.weight()) {
593 continue;
594 }
595 let (a, b) = (e.source().index(), e.target().index());
596 neighbors[a].push(b);
597 neighbors[b].push(a);
598 }
599 let s = seed.index();
600 let mut rank = vec![0.0f32; n];
601 rank[s] = 1.0;
602 for _ in 0..iters {
603 let mut next = vec![0.0f32; n];
604 next[s] += 1.0 - damping;
606 for u in 0..n {
607 if rank[u] == 0.0 || neighbors[u].is_empty() {
608 continue;
609 }
610 let share = damping * rank[u] / neighbors[u].len() as f32;
611 for &v in &neighbors[u] {
612 next[v] += share;
613 }
614 }
615 rank = next;
616 }
617 rank
618 }
619}
620
621fn est_tokens(s: &Symbol) -> usize {
623 let sig = s.signature.as_deref().map_or(0, str::len);
624 (sig + s.file_path.len() + 24) / 4
625}
626
627#[cfg(test)]
628mod tests {
629 use super::*;
630 use crate::model::{GlobalSymbolId, Language, Position, Range, RepoId, SymbolKind};
631
632 fn id(pkg: &str, desc: &str) -> GlobalSymbolId {
633 GlobalSymbolId {
634 scheme: "t".into(),
635 package_manager: ".".into(),
636 package: pkg.into(),
637 version: ".".into(),
638 descriptor: desc.into(),
639 }
640 }
641 fn sym(pkg: &str, desc: &str, name: &str, kind: SymbolKind) -> Symbol {
642 Symbol {
643 id: id(pkg, desc),
644 name: name.into(),
645 kind,
646 language: Language::Python,
647 repo: RepoId(pkg.into()),
648 file_path: format!("{pkg}/f.py"),
649 range: Range {
650 start: Position { line: 3, column: 0 },
651 end: Position { line: 9, column: 0 },
652 },
653 signature: Some(format!("def {name}(): ...")),
654 docstring: None,
655 fingerprint: None,
656 }
657 }
658 fn edge(s: &Symbol, d: &Symbol, k: EdgeKind) -> Edge {
659 Edge {
660 src: s.id.clone(),
661 dst: d.id.clone(),
662 kind: k,
663 confidence: 1.0,
664 cross_repo: s.id.package != d.id.package,
665 }
666 }
667
668 #[test]
669 fn ripple_finds_transitive_cross_repo_callers() {
670 let util = sym("lib", "lib/util().", "util", SymbolKind::Function);
672 let mid = sym("app", "app/mid().", "mid", SymbolKind::Function);
673 let top = sym("app", "app/top().", "top", SymbolKind::Function);
674 let syms = vec![util.clone(), mid.clone(), top.clone()];
675 let edges = vec![
676 edge(&mid, &util, EdgeKind::Calls),
677 edge(&top, &mid, EdgeKind::Calls),
678 ];
679 let g = CodeGraph::build(&syms, &edges);
680 let focus = g.lookup("lib/util().").unwrap();
681 let hits = g.ripple(focus, 5);
682 let names: Vec<_> = hits.iter().map(|h| h.node.name.as_str()).collect();
683 assert!(names.contains(&"mid"));
684 assert!(names.contains(&"top")); assert!(hits.iter().any(|h| h.node.repo == "app")); }
687
688 #[test]
689 fn context_pack_respects_budget_and_leads_with_focus() {
690 let a = sym("lib", "lib/a().", "a", SymbolKind::Function);
691 let b = sym("lib", "lib/b().", "b", SymbolKind::Function);
692 let c = sym("lib", "lib/c().", "c", SymbolKind::Function);
693 let syms = vec![a.clone(), b.clone(), c.clone()];
694 let edges = vec![edge(&a, &b, EdgeKind::Calls), edge(&b, &c, EdgeKind::Calls)];
695 let g = CodeGraph::build(&syms, &edges);
696 let focus = g.lookup("lib/a().").unwrap();
697 let pack = g.context_pack(focus, 10_000);
698 assert_eq!(pack[0].0.name, "a"); assert!(pack.len() >= 2); let tiny = g.context_pack(focus, 1); assert_eq!(tiny.len(), 1);
703 }
704}