1use std::cmp::{Ordering, Reverse};
2use std::collections::BinaryHeap;
3
4use rayon::prelude::*;
5use rustc_hash::{FxHashMap, FxHashSet};
6
7use crate::config::graph_filtering::GRAPH_FILTERING;
8use crate::config::scoring::EGO;
9use crate::types::{Fragment, FragmentId};
10
11#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
12pub enum EdgeCategory {
13 Semantic,
14 Structural,
15 Sibling,
16 Config,
17 ConfigGeneric,
18 Document,
19 Similarity,
20 History,
21 TestEdge,
22 Generic,
23}
24
25impl EdgeCategory {
26 pub fn from_str(s: &str) -> Self {
27 match s {
28 "semantic" => Self::Semantic,
29 "structural" => Self::Structural,
30 "sibling" => Self::Sibling,
31 "config" => Self::Config,
32 "config_generic" => Self::ConfigGeneric,
33 "document" => Self::Document,
34 "similarity" => Self::Similarity,
35 "history" => Self::History,
36 "test_edge" => Self::TestEdge,
37 _ => Self::Generic,
38 }
39 }
40
41 pub fn as_str(self) -> &'static str {
42 match self {
43 Self::Semantic => "semantic",
44 Self::Structural => "structural",
45 Self::Sibling => "sibling",
46 Self::Config => "config",
47 Self::ConfigGeneric => "config_generic",
48 Self::Document => "document",
49 Self::Similarity => "similarity",
50 Self::History => "history",
51 Self::TestEdge => "test_edge",
52 Self::Generic => "generic",
53 }
54 }
55
56 fn is_suppression_exempt(self) -> bool {
57 matches!(self, Self::Semantic | Self::Structural | Self::TestEdge)
58 }
59}
60
61pub struct CsrGraph {
62 pub n: usize,
63 pub indptr: Vec<u32>,
64 pub indices: Vec<u32>,
65 pub weights: Vec<f64>,
66 pub out_weight_sum: Vec<f64>,
67 pub node_to_idx: FxHashMap<FragmentId, u32>,
68 pub idx_to_node: Vec<FragmentId>,
69}
70
71#[derive(Default, Clone)]
76pub struct EdgeCapStats {
77 pub edges_before_cap: usize,
79 pub edges_after_cap: usize,
81 pub edges_dropped_by_cap: usize,
83 pub nodes_capped: usize,
86 pub max_out_edges_per_node: usize,
88 pub emissions_by_category: Vec<(EdgeCategory, u64, u64)>,
93}
94
95#[derive(Clone, Copy)]
100pub struct CompactEdge {
101 pub src: u32,
102 pub dst: u32,
103 pub weight: f64,
104 pub category: EdgeCategory,
105}
106
107pub struct CompactEdges {
112 pub node_to_idx: FxHashMap<FragmentId, u32>,
113 pub idx_to_node: Vec<FragmentId>,
114 pub edges: Vec<CompactEdge>,
115}
116
117pub fn intern_fragment_nodes(
118 fragments: &[Fragment],
119) -> (FxHashMap<FragmentId, u32>, Vec<FragmentId>) {
120 let mut idx_to_node: Vec<FragmentId> = fragments.iter().map(|f| f.id.clone()).collect();
121 idx_to_node.sort();
122 idx_to_node.dedup();
123 let node_to_idx = idx_to_node
124 .iter()
125 .enumerate()
126 .map(|(i, n)| (n.clone(), i as u32))
127 .collect();
128 (node_to_idx, idx_to_node)
129}
130
131pub fn dedup_compact_edges(edges: &mut Vec<CompactEdge>) {
136 edges.sort_by(|a, b| (a.src, a.dst).cmp(&(b.src, b.dst)));
137 let mut out = 0usize;
138 let mut i = 0usize;
139 while i < edges.len() {
140 let mut merged = edges[i];
141 let mut j = i + 1;
142 while j < edges.len() && edges[j].src == merged.src && edges[j].dst == merged.dst {
143 if edges[j].weight > merged.weight {
144 merged.weight = edges[j].weight;
145 }
146 j += 1;
147 }
148 edges[out] = merged;
149 out += 1;
150 i = j;
151 }
152 edges.truncate(out);
153}
154
155#[derive(Default)]
159pub struct EdgeCategoryTable {
160 node_to_idx: FxHashMap<FragmentId, u32>,
161 idx_to_node: Vec<FragmentId>,
162 entries: Vec<(u32, u32, EdgeCategory)>,
163 sorted: bool,
164}
165
166impl EdgeCategoryTable {
167 fn from_sorted_parts(
168 node_to_idx: FxHashMap<FragmentId, u32>,
169 idx_to_node: Vec<FragmentId>,
170 entries: Vec<(u32, u32, EdgeCategory)>,
171 ) -> Self {
172 debug_assert!(
173 entries
174 .windows(2)
175 .all(|w| (w[0].0, w[0].1) < (w[1].0, w[1].1))
176 );
177 Self {
178 node_to_idx,
179 idx_to_node,
180 entries,
181 sorted: true,
182 }
183 }
184
185 fn intern(&mut self, id: FragmentId) -> u32 {
186 if let Some(&i) = self.node_to_idx.get(&id) {
187 return i;
188 }
189 let i = self.idx_to_node.len() as u32;
190 self.idx_to_node.push(id.clone());
191 self.node_to_idx.insert(id, i);
192 i
193 }
194
195 pub fn insert(&mut self, src: FragmentId, dst: FragmentId, category: EdgeCategory) {
196 let s = self.intern(src);
197 let d = self.intern(dst);
198 self.entries.push((s, d, category));
199 self.sorted = false;
200 }
201
202 pub fn ensure_sorted(&mut self) {
205 if self.sorted {
206 return;
207 }
208 self.entries.sort_by_key(|e| (e.0, e.1));
209 let mut out = 0usize;
210 let mut i = 0usize;
211 while i < self.entries.len() {
212 let mut j = i;
213 while j + 1 < self.entries.len()
214 && self.entries[j + 1].0 == self.entries[i].0
215 && self.entries[j + 1].1 == self.entries[i].1
216 {
217 j += 1;
218 }
219 self.entries[out] = self.entries[j];
220 out += 1;
221 i = j + 1;
222 }
223 self.entries.truncate(out);
224 self.sorted = true;
225 }
226
227 pub fn get(&self, src: &FragmentId, dst: &FragmentId) -> Option<EdgeCategory> {
228 debug_assert!(self.sorted, "EdgeCategoryTable queried before freeze");
229 let s = *self.node_to_idx.get(src)?;
230 let d = *self.node_to_idx.get(dst)?;
231 self.entries
232 .binary_search_by_key(&(s, d), |e| (e.0, e.1))
233 .ok()
234 .map(|k| self.entries[k].2)
235 }
236
237 pub fn for_each<F: FnMut(&FragmentId, &FragmentId, EdgeCategory)>(&self, mut f: F) {
238 for &(s, d, c) in &self.entries {
239 f(
240 &self.idx_to_node[s as usize],
241 &self.idx_to_node[d as usize],
242 c,
243 );
244 }
245 }
246}
247
248pub struct Graph {
249 nodes: FxHashSet<FragmentId>,
250 fwd: FxHashMap<FragmentId, FxHashMap<FragmentId, f64>>,
251 rev: FxHashMap<FragmentId, FxHashMap<FragmentId, f64>>,
252 edge_categories: EdgeCategoryTable,
253 csr_cache: Option<(CsrGraph, CsrGraph)>,
254 pub cap_stats: EdgeCapStats,
255}
256
257impl Graph {
258 pub fn new() -> Self {
259 Self {
260 nodes: FxHashSet::default(),
261 fwd: FxHashMap::default(),
262 rev: FxHashMap::default(),
263 edge_categories: EdgeCategoryTable::default(),
264 csr_cache: None,
265 cap_stats: EdgeCapStats::default(),
266 }
267 }
268
269 pub fn edge_category(&self, src: &FragmentId, dst: &FragmentId) -> Option<EdgeCategory> {
270 self.edge_categories.get(src, dst)
271 }
272
273 pub fn for_each_categorized_edge<F: FnMut(&FragmentId, &FragmentId, EdgeCategory)>(
274 &self,
275 f: F,
276 ) {
277 self.edge_categories.for_each(f)
278 }
279
280 pub fn insert_edge_category(&mut self, src: FragmentId, dst: FragmentId, cat: EdgeCategory) {
281 self.edge_categories.insert(src, dst, cat);
282 }
283
284 pub fn categorized_edge_count(&self) -> usize {
285 self.edge_categories.entries.len()
286 }
287
288 pub fn add_node(&mut self, node: FragmentId) {
289 self.nodes.insert(node);
290 }
291
292 pub fn add_edge(&mut self, src: FragmentId, dst: FragmentId, weight: f64) {
293 if weight.is_nan() || weight.is_infinite() || weight <= 0.0 {
294 return;
295 }
296 if src == dst {
297 return;
298 }
299 debug_assert!(
300 self.csr_cache.is_none(),
301 "add_edge called after Graph was frozen"
302 );
303
304 let fwd_nbrs = self.fwd.entry(src.clone()).or_default();
305 let existing = fwd_nbrs.get(&dst).copied().unwrap_or(0.0);
306 let new_weight = existing.max(weight);
307 fwd_nbrs.insert(dst.clone(), new_weight);
308
309 let rev_nbrs = self.rev.entry(dst).or_default();
310 rev_nbrs.insert(src, new_weight);
311 }
312
313 pub fn node_count(&self) -> usize {
314 self.nodes.len()
315 }
316
317 pub fn nodes(&self) -> impl Iterator<Item = &FragmentId> {
318 self.nodes.iter()
319 }
320
321 pub fn edge_count(&self) -> usize {
322 if let Some((fwd, _)) = &self.csr_cache {
323 return fwd.indices.len();
324 }
325 self.fwd.values().map(|nbrs| nbrs.len()).sum()
326 }
327
328 pub fn freeze(&mut self) {
331 self.edge_categories.ensure_sorted();
332 if self.csr_cache.is_some() {
333 return;
334 }
335
336 let mut nodes: Vec<FragmentId> = self.nodes.iter().cloned().collect();
337 nodes.sort();
338
339 let node_to_idx: FxHashMap<FragmentId, u32> = nodes
340 .iter()
341 .enumerate()
342 .map(|(i, n)| (n.clone(), i as u32))
343 .collect();
344
345 let fwd = std::mem::take(&mut self.fwd);
346 let rev = std::mem::take(&mut self.rev);
347
348 let fwd_csr = build_csr_owned(fwd, &nodes, &node_to_idx);
349 let rev_csr = build_csr_owned(rev, &nodes, &node_to_idx);
350
351 self.csr_cache = Some((fwd_csr, rev_csr));
352 }
353
354 pub fn to_csr(&mut self) -> &(CsrGraph, CsrGraph) {
355 self.freeze();
356 self.csr_cache.as_ref().unwrap()
357 }
358
359 pub fn fwd_csr(&self) -> Option<&CsrGraph> {
360 self.csr_cache.as_ref().map(|(f, _)| f)
361 }
362
363 pub fn rev_csr(&self) -> Option<&CsrGraph> {
364 self.csr_cache.as_ref().map(|(_, r)| r)
365 }
366
367 pub fn forward_edge_weight(&self, src: &FragmentId, dst: &FragmentId) -> Option<f64> {
369 let fwd = self.fwd_csr()?;
370 let src_idx = *fwd.node_to_idx.get(src)? as usize;
371 let dst_idx = *fwd.node_to_idx.get(dst)?;
372 let s = fwd.indptr[src_idx] as usize;
373 let e = fwd.indptr[src_idx + 1] as usize;
374 for k in s..e {
375 if fwd.indices[k] == dst_idx {
376 return Some(fwd.weights[k]);
377 }
378 }
379 None
380 }
381
382 pub fn for_each_forward_neighbor<F: FnMut(&FragmentId, f64)>(
384 &self,
385 node: &FragmentId,
386 mut f: F,
387 ) {
388 let fwd = match self.fwd_csr() {
389 Some(c) => c,
390 None => return,
391 };
392 let idx = match fwd.node_to_idx.get(node) {
393 Some(&i) => i as usize,
394 None => return,
395 };
396 let s = fwd.indptr[idx] as usize;
397 let e = fwd.indptr[idx + 1] as usize;
398 for k in s..e {
399 let dst_idx = fwd.indices[k] as usize;
400 f(&fwd.idx_to_node[dst_idx], fwd.weights[k]);
401 }
402 }
403
404 pub fn ego_graph(
405 &self,
406 seeds: &FxHashSet<FragmentId>,
407 radius: usize,
408 ) -> FxHashMap<FragmentId, f64> {
409 let (fwd, rev) = match &self.csr_cache {
410 Some(c) => c,
411 None => return FxHashMap::default(),
412 };
413 if fwd.n == 0 {
414 return FxHashMap::default();
415 }
416
417 let mut valid_seed_idxs: Vec<u32> = seeds
418 .iter()
419 .filter_map(|s| fwd.node_to_idx.get(s).copied())
420 .collect();
421 valid_seed_idxs.sort_unstable();
422
423 let per_seed: Vec<Vec<(u32, u32, f64)>> = valid_seed_idxs
424 .par_iter()
425 .map(|&seed_idx| bfs_from_seed_with_path_weight(fwd, rev, seed_idx, radius))
426 .collect();
427
428 let gamma = EGO.per_hop_decay;
429 let mut scores: FxHashMap<u32, f64> = FxHashMap::default();
430 for visits in per_seed {
431 for (idx, dist, w_path) in visits {
432 let contribution = gamma.powi(dist as i32) * w_path;
433 *scores.entry(idx).or_insert(0.0) += contribution;
434 }
435 }
436
437 scores
438 .into_iter()
439 .map(|(idx, score)| (fwd.idx_to_node[idx as usize].clone(), score))
440 .collect()
441 }
442}
443
444fn bfs_from_seed_with_path_weight(
453 fwd: &CsrGraph,
454 rev: &CsrGraph,
455 seed_idx: u32,
456 radius: usize,
457) -> Vec<(u32, u32, f64)> {
458 let n = fwd.n;
459 let mut dist = vec![u32::MAX; n];
460 let mut max_w = vec![0.0_f64; n];
461 dist[seed_idx as usize] = 0;
462 max_w[seed_idx as usize] = 1.0;
463 let mut frontier: Vec<u32> = vec![seed_idx];
464
465 for step in 0..radius {
466 let new_dist = (step + 1) as u32;
467 let mut next: Vec<u32> = Vec::new();
468 for &u in &frontier {
469 let ui = u as usize;
470 let w_u = max_w[ui];
471 for csr in [fwd, rev] {
472 let s = csr.indptr[ui] as usize;
473 let e = csr.indptr[ui + 1] as usize;
474 for k in s..e {
475 let v = csr.indices[k];
476 let w_uv = csr.weights[k];
477 let candidate = w_u * w_uv;
478 let vi = v as usize;
479 if dist[vi] == u32::MAX {
480 dist[vi] = new_dist;
481 max_w[vi] = candidate;
482 next.push(v);
483 } else if dist[vi] == new_dist && candidate > max_w[vi] {
484 max_w[vi] = candidate;
485 }
486 }
487 }
488 }
489 frontier = next;
490 }
491
492 let mut result = Vec::new();
493 for i in 0..n {
494 if dist[i] != u32::MAX {
495 result.push((i as u32, dist[i], max_w[i]));
496 }
497 }
498 result
499}
500
501fn build_csr_owned(
502 adj: FxHashMap<FragmentId, FxHashMap<FragmentId, f64>>,
503 nodes: &[FragmentId],
504 node_to_idx: &FxHashMap<FragmentId, u32>,
505) -> CsrGraph {
506 let n = nodes.len();
507 let total_edges: usize = adj.values().map(|v| v.len()).sum();
508
509 let mut indptr = vec![0u32; n + 1];
510 let mut indices = Vec::with_capacity(total_edges);
511 let mut weights = Vec::with_capacity(total_edges);
512
513 for (i, node) in nodes.iter().enumerate() {
514 if let Some(nbrs) = adj.get(node) {
515 let mut edges: Vec<(u32, f64)> = nbrs
516 .iter()
517 .filter_map(|(dst, &w)| node_to_idx.get(dst).map(|&idx| (idx, w)))
518 .collect();
519 edges.sort_by_key(|&(idx, _)| idx);
520 for (idx, w) in edges {
521 indices.push(idx);
522 weights.push(w);
523 }
524 }
525 indptr[i + 1] = indices.len() as u32;
526 }
527
528 let mut out_weight_sum = vec![0.0f64; n];
529 for i in 0..n {
530 let s = indptr[i] as usize;
531 let e = indptr[i + 1] as usize;
532 if e > s {
533 out_weight_sum[i] = weights[s..e].iter().sum();
534 }
535 }
536
537 CsrGraph {
538 n,
539 indptr,
540 indices,
541 weights,
542 out_weight_sum,
543 node_to_idx: node_to_idx.clone(),
544 idx_to_node: nodes.to_vec(),
545 }
546}
547
548pub struct SuppressionFactors {
555 in_degree: Vec<u32>,
556 d_p95: f64,
557 sem_file_deg: Vec<u32>,
558}
559
560impl SuppressionFactors {
561 pub fn from_counters(in_degree: Vec<u32>, sem_file_deg: Vec<u32>) -> Self {
562 let mut degrees_sorted: Vec<u32> = in_degree.iter().copied().filter(|&d| d > 0).collect();
563 degrees_sorted.sort_unstable();
564 let d_p95 = if degrees_sorted.is_empty() {
565 0.0
566 } else {
567 let n = degrees_sorted.len();
568 let idx = ((n as f64 * 0.95).ceil() as usize)
569 .saturating_sub(1)
570 .min(n - 1);
571 degrees_sorted[idx] as f64
572 };
573 Self {
574 in_degree,
575 d_p95,
576 sem_file_deg,
577 }
578 }
579
580 fn from_edges(edges: &[CompactEdge], idx_to_node: &[FragmentId]) -> Self {
581 let n_nodes = idx_to_node.len();
582 let mut in_degree = vec![0u32; n_nodes];
583 for e in edges {
584 in_degree[e.dst as usize] += 1;
585 }
586
587 let mut sem_out_files: FxHashMap<u32, FxHashSet<&str>> = FxHashMap::default();
588 for e in edges {
589 if e.category == EdgeCategory::Semantic {
590 sem_out_files
591 .entry(e.src)
592 .or_default()
593 .insert(idx_to_node[e.dst as usize].path.as_ref());
594 }
595 }
596 let mut sem_file_deg = vec![0u32; n_nodes];
597 for (&src, files) in &sem_out_files {
598 sem_file_deg[src as usize] = files.len() as u32;
599 }
600
601 Self::from_counters(in_degree, sem_file_deg)
602 }
603
604 pub fn damp(&self, weight: f64, category: EdgeCategory, src: u32, dst: u32) -> f64 {
605 let mut w = weight;
606 let dst_deg = self.in_degree[dst as usize] as f64;
607 if dst_deg > self.d_p95 && !category.is_suppression_exempt() {
608 w /= dst_deg.ln_1p().max(1.0);
609 }
610 if category == EdgeCategory::Semantic {
611 let src_deg = self.sem_file_deg[src as usize];
612 if src_deg >= GRAPH_FILTERING.hub_out_degree_threshold as u32 {
613 w /= (src_deg as f64).sqrt();
614 }
615 }
616 w
617 }
618}
619
620fn apply_hub_suppression(edges: &mut [CompactEdge], idx_to_node: &[FragmentId]) {
621 if edges.is_empty() {
622 return;
623 }
624 let factors = SuppressionFactors::from_edges(edges, idx_to_node);
625 for e in edges.iter_mut() {
626 e.weight = factors.damp(e.weight, e.category, e.src, e.dst);
627 }
628}
629
630pub(crate) const DEFAULT_MAX_OUT_EDGES_PER_NODE: usize = 64;
637
638pub(crate) fn cap_out_edges_per_source(
647 edges: &mut Vec<CompactEdge>,
648 max_per_node: usize,
649) -> EdgeCapStats {
650 let edges_before = edges.len();
651
652 edges.sort_unstable_by(|a, b| {
653 a.src
654 .cmp(&b.src)
655 .then_with(|| {
656 b.weight
657 .partial_cmp(&a.weight)
658 .unwrap_or(std::cmp::Ordering::Equal)
659 })
660 .then_with(|| a.dst.cmp(&b.dst))
661 });
662
663 let mut nodes_capped = 0;
664 let mut out = 0usize;
665 let mut i = 0usize;
666 while i < edges.len() {
667 let src = edges[i].src;
668 let mut j = i;
669 while j < edges.len() && edges[j].src == src {
670 j += 1;
671 }
672 let group = j - i;
673 if group > max_per_node {
674 nodes_capped += 1;
675 }
676 let take = group.min(max_per_node);
677 for k in i..i + take {
678 edges[out] = edges[k];
679 out += 1;
680 }
681 i = j;
682 }
683 edges.truncate(out);
684
685 EdgeCapStats {
686 edges_before_cap: edges_before,
687 edges_after_cap: edges.len(),
688 edges_dropped_by_cap: edges_before - edges.len(),
689 nodes_capped,
690 max_out_edges_per_node: max_per_node,
691 emissions_by_category: Vec::new(),
692 }
693}
694
695pub(crate) fn read_max_out_edges_per_node() -> usize {
696 std::env::var("DIFFCTX_MAX_EDGES_PER_NODE")
697 .ok()
698 .and_then(|v| v.parse::<usize>().ok())
699 .filter(|&v| v > 0)
700 .unwrap_or(DEFAULT_MAX_OUT_EDGES_PER_NODE)
701}
702
703#[derive(Clone, Copy)]
708pub struct RankedCandidate {
709 pub weight: f64,
710 pub dst: u32,
711 pub category: EdgeCategory,
712}
713
714impl RankedCandidate {
715 fn rank(&self, other: &Self) -> Ordering {
716 self.weight
717 .partial_cmp(&other.weight)
718 .unwrap_or(Ordering::Equal)
719 .then_with(|| other.dst.cmp(&self.dst))
720 }
721}
722
723impl PartialEq for RankedCandidate {
724 fn eq(&self, other: &Self) -> bool {
725 self.rank(other) == Ordering::Equal
726 }
727}
728
729impl Eq for RankedCandidate {}
730
731impl PartialOrd for RankedCandidate {
732 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
733 Some(self.rank(other))
734 }
735}
736
737impl Ord for RankedCandidate {
738 fn cmp(&self, other: &Self) -> Ordering {
739 self.rank(other)
740 }
741}
742
743pub type SourceTopK = BinaryHeap<Reverse<RankedCandidate>>;
744
745pub fn push_bounded_top_k(heap: &mut SourceTopK, candidate: RankedCandidate, k: usize) {
746 if heap.len() < k {
747 heap.push(Reverse(candidate));
748 return;
749 }
750 if let Some(Reverse(worst)) = heap.peek() {
751 if candidate.rank(worst) == Ordering::Greater {
752 heap.pop();
753 heap.push(Reverse(candidate));
754 }
755 }
756}
757
758pub struct CappedEdges {
764 pub node_to_idx: FxHashMap<FragmentId, u32>,
765 pub idx_to_node: Vec<FragmentId>,
766 pub edges: Vec<CompactEdge>,
767 pub cap_stats: EdgeCapStats,
768}
769
770fn build_csr_from_pairs(
775 mut pairs: Vec<(u32, u32, f64)>,
776 idx_to_node: &[FragmentId],
777 node_to_idx: &FxHashMap<FragmentId, u32>,
778) -> CsrGraph {
779 pairs.sort_unstable_by_key(|p| (p.0, p.1));
780 let n = idx_to_node.len();
781
782 let mut indptr = vec![0u32; n + 1];
783 let mut indices = Vec::with_capacity(pairs.len());
784 let mut weights = Vec::with_capacity(pairs.len());
785
786 let mut row = 0usize;
787 for &(src, dst, w) in &pairs {
788 while row < src as usize {
789 row += 1;
790 indptr[row] = indices.len() as u32;
791 }
792 indices.push(dst);
793 weights.push(w);
794 }
795 while row < n {
796 row += 1;
797 indptr[row] = indices.len() as u32;
798 }
799
800 let mut out_weight_sum = vec![0.0f64; n];
801 for i in 0..n {
802 let s = indptr[i] as usize;
803 let e = indptr[i + 1] as usize;
804 if e > s {
805 out_weight_sum[i] = weights[s..e].iter().sum();
806 }
807 }
808
809 CsrGraph {
810 n,
811 indptr,
812 indices,
813 weights,
814 out_weight_sum,
815 node_to_idx: node_to_idx.clone(),
816 idx_to_node: idx_to_node.to_vec(),
817 }
818}
819
820pub fn build_graph_capped(fragments: &[Fragment], capped: CappedEdges) -> Graph {
824 let CappedEdges {
825 node_to_idx,
826 idx_to_node,
827 edges,
828 cap_stats,
829 } = capped;
830 tracing::debug!(
831 "edge cap K={}: {} -> {} (dropped {} from {} nodes)",
832 cap_stats.max_out_edges_per_node,
833 cap_stats.edges_before_cap,
834 cap_stats.edges_after_cap,
835 cap_stats.edges_dropped_by_cap,
836 cap_stats.nodes_capped,
837 );
838 assemble_graph(fragments, node_to_idx, idx_to_node, edges, cap_stats)
839}
840
841pub fn build_graph_compact(fragments: &[Fragment], compact: CompactEdges) -> Graph {
845 let CompactEdges {
846 node_to_idx,
847 idx_to_node,
848 mut edges,
849 } = compact;
850
851 apply_hub_suppression(&mut edges, &idx_to_node);
852
853 let max_per_node = read_max_out_edges_per_node();
854 let cap_stats = cap_out_edges_per_source(&mut edges, max_per_node);
855 tracing::debug!(
856 "edge cap K={}: {} -> {} (dropped {} from {} nodes)",
857 max_per_node,
858 cap_stats.edges_before_cap,
859 cap_stats.edges_after_cap,
860 cap_stats.edges_dropped_by_cap,
861 cap_stats.nodes_capped,
862 );
863
864 assemble_graph(fragments, node_to_idx, idx_to_node, edges, cap_stats)
865}
866
867fn is_valid_edge(e: &CompactEdge) -> bool {
874 e.src != e.dst && e.weight.is_finite() && e.weight > 0.0
875}
876
877fn assemble_graph(
886 fragments: &[Fragment],
887 node_to_idx: FxHashMap<FragmentId, u32>,
888 idx_to_node: Vec<FragmentId>,
889 edges: Vec<CompactEdge>,
890 cap_stats: EdgeCapStats,
891) -> Graph {
892 let valid_edges: Vec<CompactEdge> = edges.into_iter().filter(is_valid_edge).collect();
893
894 let mut category_entries: Vec<(u32, u32, EdgeCategory)> = valid_edges
895 .iter()
896 .map(|e| (e.src, e.dst, e.category))
897 .collect();
898 category_entries.sort_unstable_by_key(|e| (e.0, e.1));
899
900 let fwd_pairs: Vec<(u32, u32, f64)> = valid_edges
901 .iter()
902 .map(|e| (e.src, e.dst, e.weight))
903 .collect();
904 let rev_pairs: Vec<(u32, u32, f64)> = valid_edges
905 .iter()
906 .map(|e| (e.dst, e.src, e.weight))
907 .collect();
908 drop(valid_edges);
909
910 let fwd_csr = build_csr_from_pairs(fwd_pairs, &idx_to_node, &node_to_idx);
911 let rev_csr = build_csr_from_pairs(rev_pairs, &idx_to_node, &node_to_idx);
912
913 let mut graph = Graph::new();
914 for frag in fragments {
915 graph.nodes.insert(frag.id.clone());
916 }
917 graph.edge_categories =
918 EdgeCategoryTable::from_sorted_parts(node_to_idx, idx_to_node, category_entries);
919 graph.cap_stats = cap_stats;
920 graph.csr_cache = Some((fwd_csr, rev_csr));
921 graph
922}
923
924pub fn build_graph(
929 fragments: &[Fragment],
930 edges: FxHashMap<(FragmentId, FragmentId), f64>,
931 categories: FxHashMap<(FragmentId, FragmentId), EdgeCategory>,
932) -> Graph {
933 let (node_to_idx, idx_to_node) = intern_fragment_nodes(fragments);
934 let mut compact_edges = Vec::with_capacity(edges.len());
935 for ((src, dst), w) in &edges {
936 let s = match node_to_idx.get(src) {
937 Some(&i) => i,
938 None => continue,
939 };
940 let d = match node_to_idx.get(dst) {
941 Some(&i) => i,
942 None => continue,
943 };
944 let category = categories
945 .get(&(src.clone(), dst.clone()))
946 .copied()
947 .unwrap_or(EdgeCategory::Generic);
948 compact_edges.push(CompactEdge {
949 src: s,
950 dst: d,
951 weight: *w,
952 category,
953 });
954 }
955 dedup_compact_edges(&mut compact_edges);
956 build_graph_compact(
957 fragments,
958 CompactEdges {
959 node_to_idx,
960 idx_to_node,
961 edges: compact_edges,
962 },
963 )
964}
965
966#[cfg(test)]
967mod tests {
968 use super::*;
969 use std::sync::Arc;
970
971 fn fid(path: &str, start: u32, end: u32) -> FragmentId {
972 FragmentId::new(Arc::from(path), start, end)
973 }
974
975 fn collect_forward(g: &Graph, node: &FragmentId) -> Vec<(FragmentId, f64)> {
976 let mut out = Vec::new();
977 g.for_each_forward_neighbor(node, |nbr, w| out.push((nbr.clone(), w)));
978 out
979 }
980
981 #[test]
982 fn add_edge_takes_max_weight() {
983 let mut g = Graph::new();
984 let a = fid("a.rs", 1, 10);
985 let b = fid("b.rs", 1, 10);
986 g.add_node(a.clone());
987 g.add_node(b.clone());
988 g.add_edge(a.clone(), b.clone(), 0.5);
989 g.add_edge(a.clone(), b.clone(), 0.8);
990 g.add_edge(a.clone(), b.clone(), 0.3);
991 g.freeze();
992
993 let fwd = collect_forward(&g, &a);
994 assert_eq!(fwd.len(), 1);
995 assert!((fwd[0].1 - 0.8).abs() < 1e-9);
996 assert_eq!(fwd[0].0, b);
997 }
998
999 #[test]
1000 fn add_edge_drops_invalid_weights() {
1001 let mut g = Graph::new();
1002 let a = fid("a.rs", 1, 10);
1003 let b = fid("b.rs", 1, 10);
1004 g.add_node(a.clone());
1005 g.add_node(b.clone());
1006 g.add_edge(a.clone(), b.clone(), f64::NAN);
1007 g.add_edge(a.clone(), b.clone(), f64::INFINITY);
1008 g.add_edge(a.clone(), b.clone(), -1.0);
1009 g.add_edge(a.clone(), b.clone(), 0.0);
1010 g.freeze();
1011
1012 assert!(collect_forward(&g, &a).is_empty());
1013 assert_eq!(g.edge_count(), 0);
1014 }
1015
1016 #[test]
1017 fn csr_round_trip() {
1018 let mut g = Graph::new();
1019 let a = fid("a.rs", 1, 10);
1020 let b = fid("b.rs", 1, 10);
1021 let c = fid("c.rs", 1, 10);
1022 g.add_node(a.clone());
1023 g.add_node(b.clone());
1024 g.add_node(c.clone());
1025 g.add_edge(a.clone(), b.clone(), 1.0);
1026 g.add_edge(b.clone(), c.clone(), 2.0);
1027
1028 let (fwd, _rev) = g.to_csr();
1029 assert_eq!(fwd.n, 3);
1030 assert_eq!(fwd.indptr.len(), 4);
1031 assert!(fwd.out_weight_sum[fwd.node_to_idx[&a] as usize] > 0.0);
1032 }
1033
1034 #[test]
1035 fn ego_graph_scores() {
1036 let mut g = Graph::new();
1037 let a = fid("a.rs", 1, 10);
1038 let b = fid("b.rs", 1, 10);
1039 let c = fid("c.rs", 1, 10);
1040 g.add_node(a.clone());
1041 g.add_node(b.clone());
1042 g.add_node(c.clone());
1043 g.add_edge(a.clone(), b.clone(), 1.0);
1044 g.add_edge(b.clone(), c.clone(), 1.0);
1045 g.freeze();
1046
1047 let mut seeds = FxHashSet::default();
1048 seeds.insert(a.clone());
1049 let scores = g.ego_graph(&seeds, 2);
1050
1051 let gamma = crate::config::scoring::EGO.per_hop_decay;
1052 assert!((scores[&a] - 1.0).abs() < 1e-9);
1053 assert!((scores[&b] - gamma).abs() < 1e-9);
1054 assert!((scores[&c] - gamma * gamma).abs() < 1e-9);
1055 }
1056
1057 #[test]
1058 fn ego_graph_sums_over_seeds() {
1059 let mut g = Graph::new();
1060 let a = fid("a.rs", 1, 10);
1061 let b = fid("b.rs", 1, 10);
1062 let v = fid("v.rs", 1, 10);
1063 g.add_node(a.clone());
1064 g.add_node(b.clone());
1065 g.add_node(v.clone());
1066 g.add_edge(a.clone(), v.clone(), 1.0);
1067 g.add_edge(b.clone(), v.clone(), 1.0);
1068 g.freeze();
1069
1070 let mut seeds = FxHashSet::default();
1071 seeds.insert(a.clone());
1072 seeds.insert(b.clone());
1073 let scores = g.ego_graph(&seeds, 1);
1074
1075 let gamma = crate::config::scoring::EGO.per_hop_decay;
1076 assert!(
1077 (scores[&v] - 2.0 * gamma).abs() < 1e-9,
1078 "v reached by 2 seeds at d=1 must score 2·γ; got {}",
1079 scores[&v]
1080 );
1081 }
1082
1083 #[test]
1084 fn ego_graph_uses_path_weight() {
1085 let mut g = Graph::new();
1086 let a = fid("a.rs", 1, 10);
1087 let b = fid("b.rs", 1, 10);
1088 let c = fid("c.rs", 1, 10);
1089 g.add_node(a.clone());
1090 g.add_node(b.clone());
1091 g.add_node(c.clone());
1092 g.add_edge(a.clone(), b.clone(), 0.7);
1093 g.add_edge(b.clone(), c.clone(), 0.4);
1094 g.freeze();
1095
1096 let mut seeds = FxHashSet::default();
1097 seeds.insert(a.clone());
1098 let scores = g.ego_graph(&seeds, 2);
1099
1100 let gamma = crate::config::scoring::EGO.per_hop_decay;
1101 assert!(
1102 (scores[&b] - gamma * 0.7).abs() < 1e-9,
1103 "1-hop weighted score = γ·0.7; got {}",
1104 scores[&b]
1105 );
1106 assert!(
1107 (scores[&c] - gamma * gamma * 0.7 * 0.4).abs() < 1e-9,
1108 "2-hop product-of-weights score = γ²·0.7·0.4; got {}",
1109 scores[&c]
1110 );
1111 }
1112
1113 #[test]
1114 fn ego_graph_empty() {
1115 let mut g = Graph::new();
1116 g.freeze();
1117 let seeds = FxHashSet::default();
1118 let scores = g.ego_graph(&seeds, 2);
1119 assert!(scores.is_empty());
1120 }
1121
1122 #[test]
1123 fn dedup_compact_edges_max_weight_first_category() {
1124 let mut edges = vec![
1125 CompactEdge {
1126 src: 0,
1127 dst: 1,
1128 weight: 0.5,
1129 category: EdgeCategory::Semantic,
1130 },
1131 CompactEdge {
1132 src: 0,
1133 dst: 1,
1134 weight: 0.8,
1135 category: EdgeCategory::Similarity,
1136 },
1137 CompactEdge {
1138 src: 2,
1139 dst: 1,
1140 weight: 0.3,
1141 category: EdgeCategory::Sibling,
1142 },
1143 ];
1144 dedup_compact_edges(&mut edges);
1145 assert_eq!(edges.len(), 2);
1146 assert!((edges[0].weight - 0.8).abs() < 1e-9);
1147 assert_eq!(edges[0].category, EdgeCategory::Semantic);
1148 assert_eq!(edges[1].category, EdgeCategory::Sibling);
1149 }
1150
1151 #[test]
1152 fn cap_out_edges_per_source_keeps_top_k_ties_broken_by_ascending_dst() {
1153 let mut edges = vec![
1154 CompactEdge {
1155 src: 0,
1156 dst: 100,
1157 weight: 10.0,
1158 category: EdgeCategory::Semantic,
1159 },
1160 CompactEdge {
1161 src: 0,
1162 dst: 50,
1163 weight: 8.0,
1164 category: EdgeCategory::Semantic,
1165 },
1166 CompactEdge {
1167 src: 0,
1168 dst: 30,
1169 weight: 8.0,
1170 category: EdgeCategory::Semantic,
1171 },
1172 CompactEdge {
1173 src: 0,
1174 dst: 70,
1175 weight: 8.0,
1176 category: EdgeCategory::Semantic,
1177 },
1178 CompactEdge {
1179 src: 0,
1180 dst: 5,
1181 weight: 1.0,
1182 category: EdgeCategory::Semantic,
1183 },
1184 ];
1185
1186 let stats = cap_out_edges_per_source(&mut edges, 2);
1187
1188 assert_eq!(stats.edges_before_cap, 5);
1189 assert_eq!(stats.edges_after_cap, 2);
1190 assert_eq!(stats.edges_dropped_by_cap, 3);
1191 assert_eq!(stats.nodes_capped, 1);
1192 assert_eq!(stats.max_out_edges_per_node, 2);
1193
1194 let survivors: Vec<(u32, f64)> = edges.iter().map(|e| (e.dst, e.weight)).collect();
1195 assert_eq!(
1196 survivors,
1197 vec![(100, 10.0), (30, 8.0)],
1198 "of the three weight-8.0 ties (dst 30, 50, 70), only the lowest dst (30) may survive \
1199 the second slot"
1200 );
1201 }
1202
1203 #[test]
1204 fn ranked_candidate_tie_break_matches_cap_out_edges_tie_break() {
1205 let candidates = [
1211 RankedCandidate {
1212 weight: 5.0,
1213 dst: 30,
1214 category: EdgeCategory::Generic,
1215 },
1216 RankedCandidate {
1217 weight: 5.0,
1218 dst: 10,
1219 category: EdgeCategory::Generic,
1220 },
1221 RankedCandidate {
1222 weight: 5.0,
1223 dst: 20,
1224 category: EdgeCategory::Generic,
1225 },
1226 ];
1227
1228 let mut heap: SourceTopK = BinaryHeap::new();
1229 for &c in &candidates {
1230 push_bounded_top_k(&mut heap, c, 1);
1231 }
1232 let kept_by_heap: Vec<u32> = heap.into_iter().map(|Reverse(c)| c.dst).collect();
1233 assert_eq!(
1234 kept_by_heap,
1235 vec![10],
1236 "heap tie-break must keep the lowest dst"
1237 );
1238
1239 let mut edges: Vec<CompactEdge> = candidates
1240 .iter()
1241 .map(|c| CompactEdge {
1242 src: 0,
1243 dst: c.dst,
1244 weight: c.weight,
1245 category: c.category,
1246 })
1247 .collect();
1248 cap_out_edges_per_source(&mut edges, 1);
1249 let kept_by_sort: Vec<u32> = edges.iter().map(|e| e.dst).collect();
1250
1251 assert_eq!(
1252 kept_by_heap, kept_by_sort,
1253 "RankedCandidate::rank tie-break must agree with cap_out_edges_per_source's sort \
1254 tie-break"
1255 );
1256 }
1257
1258 #[test]
1259 fn two_pass_heap_cap_matches_materialized_cap() {
1260 let k = 3usize;
1271 type RawEdge = (u32, u32, f64, EdgeCategory);
1272
1273 let builder_logs: Vec<Vec<RawEdge>> = vec![
1274 vec![
1275 (0, 1, 9.0, EdgeCategory::Semantic),
1276 (0, 2, 9.0, EdgeCategory::Semantic),
1277 (0, 3, 7.0, EdgeCategory::Semantic),
1278 (0, 4, 1.0, EdgeCategory::Semantic),
1279 (1, 5, 4.0, EdgeCategory::Semantic),
1280 ],
1281 vec![
1282 (0, 4, 8.0, EdgeCategory::Structural),
1283 (0, 6, 2.0, EdgeCategory::Structural),
1284 (1, 5, 6.0, EdgeCategory::Structural),
1285 (1, 7, 3.0, EdgeCategory::Structural),
1286 ],
1287 vec![
1288 (0, 2, 9.0, EdgeCategory::Sibling),
1289 (0, 7, 9.0, EdgeCategory::Sibling),
1290 (1, 8, 0.5, EdgeCategory::Sibling),
1291 ],
1292 ];
1293
1294 let mut canonical_category: FxHashMap<(u32, u32), EdgeCategory> = FxHashMap::default();
1295 for log in &builder_logs {
1296 for &(src, dst, _w, cat) in log {
1297 canonical_category.entry((src, dst)).or_insert(cat);
1298 }
1299 }
1300 let canonicalize = |src: u32, dst: u32, fallback: EdgeCategory| {
1301 canonical_category
1302 .get(&(src, dst))
1303 .copied()
1304 .unwrap_or(fallback)
1305 };
1306
1307 let mut two_pass_edges: Vec<CompactEdge> = Vec::new();
1308 for log in &builder_logs {
1309 let mut per_source: FxHashMap<u32, SourceTopK> = FxHashMap::default();
1310 for &(src, dst, weight, cat) in log {
1311 let category = canonicalize(src, dst, cat);
1312 push_bounded_top_k(
1313 per_source.entry(src).or_default(),
1314 RankedCandidate {
1315 weight,
1316 dst,
1317 category,
1318 },
1319 k,
1320 );
1321 }
1322 for (src, heap) in per_source {
1323 for Reverse(c) in heap {
1324 two_pass_edges.push(CompactEdge {
1325 src,
1326 dst: c.dst,
1327 weight: c.weight,
1328 category: c.category,
1329 });
1330 }
1331 }
1332 }
1333 dedup_compact_edges(&mut two_pass_edges);
1334 cap_out_edges_per_source(&mut two_pass_edges, k);
1335
1336 let mut materialized_edges: Vec<CompactEdge> = builder_logs
1337 .iter()
1338 .flatten()
1339 .map(|&(src, dst, weight, cat)| CompactEdge {
1340 src,
1341 dst,
1342 weight,
1343 category: canonicalize(src, dst, cat),
1344 })
1345 .collect();
1346 dedup_compact_edges(&mut materialized_edges);
1347 cap_out_edges_per_source(&mut materialized_edges, k);
1348
1349 let normalize = |edges: &[CompactEdge]| -> Vec<(u32, u32, u64, EdgeCategory)> {
1350 let mut v: Vec<(u32, u32, u64, EdgeCategory)> = edges
1351 .iter()
1352 .map(|e| (e.src, e.dst, e.weight.to_bits(), e.category))
1353 .collect();
1354 v.sort_by_key(|t| (t.0, t.1));
1355 v
1356 };
1357
1358 let two_pass_normalized = normalize(&two_pass_edges);
1359 let materialized_normalized = normalize(&materialized_edges);
1360 assert_eq!(
1361 two_pass_normalized, materialized_normalized,
1362 "two-pass per-builder heap capping must be bit-identical to capping the fully \
1363 materialized edge set"
1364 );
1365 assert!(
1366 !materialized_normalized.is_empty() && materialized_normalized.len() <= 2 * k,
1367 "the cap must actually have fired in this fixture for the comparison to be meaningful"
1368 );
1369 }
1370
1371 fn plain_fragment(id: FragmentId) -> Fragment {
1372 Fragment {
1373 id,
1374 kind: crate::types::FragmentKind::Function,
1375 content: Arc::from(""),
1376 identifiers: FxHashSet::default(),
1377 token_count: 10,
1378 symbol_name: None,
1379 }
1380 }
1381
1382 #[test]
1383 fn assemble_graph_excludes_self_loops_and_non_finite_weights() {
1384 let a = fid("a.rs", 1, 10);
1385 let b = fid("b.rs", 1, 10);
1386 let frags = vec![plain_fragment(a.clone()), plain_fragment(b.clone())];
1387
1388 let mut edges: FxHashMap<(FragmentId, FragmentId), f64> = FxHashMap::default();
1389 let mut cats: FxHashMap<(FragmentId, FragmentId), EdgeCategory> = FxHashMap::default();
1390 edges.insert((a.clone(), b.clone()), 1.0);
1391 cats.insert((a.clone(), b.clone()), EdgeCategory::Semantic);
1392 edges.insert((a.clone(), a.clone()), 5.0);
1393 cats.insert((a.clone(), a.clone()), EdgeCategory::Semantic);
1394 edges.insert((b.clone(), a.clone()), f64::INFINITY);
1395 cats.insert((b.clone(), a.clone()), EdgeCategory::Semantic);
1396
1397 let graph = build_graph(&frags, edges, cats);
1398
1399 assert_eq!(
1400 graph.edge_count(),
1401 1,
1402 "only the valid a->b edge should survive"
1403 );
1404 assert_eq!(graph.categorized_edge_count(), 1);
1405 assert_eq!(graph.forward_edge_weight(&a, &b), Some(1.0));
1406 assert!(
1407 graph.forward_edge_weight(&a, &a).is_none(),
1408 "self-loop must be excluded from the live CSR"
1409 );
1410 assert!(
1411 graph.forward_edge_weight(&b, &a).is_none(),
1412 "infinite weight must be excluded from the live CSR"
1413 );
1414 }
1415
1416 #[test]
1417 fn category_table_stays_aligned_with_csr_when_cap_fires() {
1418 let hub = fid("hub.rs", 1, 10);
1419 let mut frags = vec![plain_fragment(hub.clone())];
1420 let mut edges: FxHashMap<(FragmentId, FragmentId), f64> = FxHashMap::default();
1421 let mut cats: FxHashMap<(FragmentId, FragmentId), EdgeCategory> = FxHashMap::default();
1422
1423 let n_leaves = DEFAULT_MAX_OUT_EDGES_PER_NODE + 10;
1424 for i in 0..n_leaves {
1425 let leaf = fid(&format!("leaf_{i}.rs"), 1, 10);
1426 frags.push(plain_fragment(leaf.clone()));
1427 edges.insert((hub.clone(), leaf.clone()), 1.0 + i as f64);
1428 cats.insert((hub.clone(), leaf.clone()), EdgeCategory::Semantic);
1429 }
1430
1431 let graph = build_graph(&frags, edges, cats);
1432
1433 assert!(
1434 graph.cap_stats.edges_dropped_by_cap > 0,
1435 "the cap must actually fire for this test to be meaningful"
1436 );
1437 assert_eq!(
1438 graph.categorized_edge_count(),
1439 graph.edge_count(),
1440 "the category table must be capped alongside the CSR"
1441 );
1442
1443 let mut checked = 0usize;
1444 graph.for_each_categorized_edge(|src, dst, _cat| {
1445 let w = graph.forward_edge_weight(src, dst);
1446 assert!(
1447 w.is_some_and(|w| w > 0.0),
1448 "every categorized edge must exist in the CSR with a positive weight; {src} -> {dst} did not"
1449 );
1450 checked += 1;
1451 });
1452 assert_eq!(checked, graph.edge_count());
1453 }
1454}