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