Skip to main content

_diffctx/
graph.rs

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/// Statistics from the per-source out-edge cap that runs in
72/// `build_graph` after `apply_hub_suppression`. Surfaced to Python
73/// via `LatencyBreakdown` so calibration runs can quantify how often
74/// the cap fires and how many edges it discards.
75#[derive(Default, Clone)]
76pub struct EdgeCapStats {
77    /// Edge count after merge + hub suppression, before cap.
78    pub edges_before_cap: usize,
79    /// Edge count after cap.
80    pub edges_after_cap: usize,
81    /// Edges discarded by the cap (lowest-weight neighbors of overfull nodes).
82    pub edges_dropped_by_cap: usize,
83    /// Number of source nodes that had > `max_per_node` outgoing edges
84    /// and therefore had their neighbor list truncated.
85    pub nodes_capped: usize,
86    /// The `max_per_node` value actually applied (after env-var override).
87    pub max_out_edges_per_node: usize,
88    /// Per-category (raw emissions, first-seen deduped) counts from
89    /// pass 1 of the two-pass edge build, sorted by category name.
90    /// Names the builder category responsible for near-dense emission
91    /// blowups (#116). Empty on the materialized-edge path.
92    pub emissions_by_category: Vec<(EdgeCategory, u64, u64)>,
93}
94
95/// A single node-interned edge. 24 bytes instead of a string-keyed
96/// hashmap entry — the difference between ~400 MB and several GB on
97/// multi-million-edge repositories (vendored Go monorepos, vscode),
98/// which is what used to OOM-kill memory-limited benchmark runners.
99#[derive(Clone, Copy)]
100pub struct CompactEdge {
101    pub src: u32,
102    pub dst: u32,
103    pub weight: f64,
104    pub category: EdgeCategory,
105}
106
107/// Node-interned edge list: the memory-bounded intermediate between
108/// edge collection and graph construction. `idx_to_node` is the sorted,
109/// deduplicated fragment-id universe, so CSR node indexing built from it
110/// is identical to the ordering `Graph::freeze` produces.
111pub 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
131/// Merge duplicate (src, dst) entries: weight = max across duplicates,
132/// category = first occurrence in input order (builder order), matching
133/// the historical `EdgeDict` max-merge + `or_insert` category semantics.
134/// Requires a stable sort so first-in-input stays first-in-group.
135pub 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/// Compact (src, dst) -> category lookup over the pre-cap edge set.
156/// Replaces the FragmentId-pair-keyed hashmap that used to retain the
157/// full uncapped edge universe for the lifetime of the pipeline.
158#[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    /// Stable-sort + last-wins dedup, matching hashmap overwrite
203    /// semantics for repeated `insert` of the same key.
204    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    /// Convert the build-time hashmap representation into CSR and drop the hashmaps.
329    /// After freeze, fwd/rev are empty and all reads go through CSR.
330    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    /// Look up the weight of edge `src -> dst` in the forward CSR.
368    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    /// Invoke `f(neighbor_id, weight)` for each forward neighbor of `node`.
383    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
444/// BFS over `fwd ∪ rev` from a single seed, tracking both the shortest
445/// hop distance and the max-product edge-weight path of that length.
446///
447/// Implements the paper's `R_ego` kernel (§4.4.2):
448/// `R_ego(v) = Σ_{u∈E_0} 1[d_hop(u,v) ≤ L] · γ^{d_hop} · W_path(u,v)`
449/// where `W_path(u,v) = max_π ∏_{(a,b)∈π} w_{ab}` over paths of length
450/// equal to `d_hop`. The `Σ` over seeds is performed in `ego_graph`;
451/// per-seed shortest-distance + max-product is computed here.
452fn 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
548/// Hub-suppression damping factors, precomputed from dedup-aware degree
549/// counters so the damped weight of any edge can be evaluated without
550/// materializing the edge universe. The per-edge arithmetic is identical
551/// to the historical in-place suppression loops: at most one division,
552/// by `ln(1+in_degree)` for non-exempt edges into hub destinations or by
553/// `sqrt(fan)` for Semantic edges out of wide-fan sources.
554pub 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
630/// Default top-K out-edges per source node. Calibrated against the
631/// observed edge density distribution: typical Python file emits
632/// 5-20 semantic + 2-5 structural + ≤20 sibling edges (~30-50 normal),
633/// so K=64 preserves all legitimate edges while clamping pathological
634/// dense nodes (e.g. utility hubs in django/material-ui that radiate
635/// into thousands of dependents).
636pub(crate) const DEFAULT_MAX_OUT_EDGES_PER_NODE: usize = 64;
637
638/// Truncate each node's outgoing edge list to the top-K by weight.
639/// Run AFTER `apply_hub_suppression` so the suppression pass sees
640/// the true in-degree distribution; otherwise its IDF damping is
641/// computed against a graph that has already been thinned.
642///
643/// Returns the cap stats for diagnostic surfacing into `LatencyBreakdown`.
644/// Keeps each source's top-K neighbors by weight (ties broken by dst
645/// index for determinism), truncating the rest in place.
646pub(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/// Candidate out-edge ranked by post-suppression weight. The ordering
704/// replicates the sort in `cap_out_edges_per_source` (descending weight,
705/// ties by ascending dst index), so a bounded min-heap of these keeps
706/// exactly the edges the full sort-and-truncate would keep.
707#[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
758/// Output of the two-pass edge construction: the final damped, deduped,
759/// per-source-capped edge set plus cap stats derived from pass-1 counters.
760/// The category table is *not* carried here — `assemble_graph` derives it
761/// from this same post-cap `edges` vector, which is the only way to
762/// guarantee the exported category table and the CSR never disagree.
763pub 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
770/// Build a CSR from (src, dst, weight) triples over the interned node
771/// universe. Sorting by (src, dst) reproduces the neighbor ordering the
772/// hashmap-based `freeze` path produced (neighbors sorted by dst index),
773/// so weights, out-weight sums, and traversal results are bit-identical.
774fn 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
820/// Graph-construction path for pre-capped edges from the two-pass
821/// pipeline: suppression and the per-source cap already ran, so only
822/// the category table and CSR assembly remain.
823pub 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
841/// Materialized-edge construction path kept for the map-based adapter
842/// and small callers: hub suppression, per-source cap, and CSR assembly
843/// all run on the interned edge array.
844pub 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
867/// Self-loops and non-finite/non-positive weights must never reach a live
868/// graph: they used to be rejected only by `Graph::add_edge`'s guard, which
869/// has no production caller (both real construction paths route through
870/// here). An infinite weight makes `out_weight_sum` infinite, every
871/// normalized PPR transition probability NaN, and the score map come back
872/// empty with exit 0 -- silently shipping core-only context.
873fn is_valid_edge(e: &CompactEdge) -> bool {
874    e.src != e.dst && e.weight.is_finite() && e.weight > 0.0
875}
876
877/// Builds both the CSR and the category table from the *same* filtered,
878/// already-capped `edges` vector. Deriving `edge_categories` here instead
879/// of accepting it as a separately pre-cap parameter is what guarantees
880/// `Graph::categorized_edge_count() == Graph::edge_count()` always --
881/// previously the category table was built from the pre-cap edge set in
882/// both callers, so it disagreed with the CSR whenever the cap actually
883/// fired (capped-away edges exported with `weight: 0.0`, `edge_count`
884/// mismatching `len(edges)` in the same JSON document).
885fn 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
924/// Map-based adapter kept for tests and small callers; converts to the
925/// compact representation and delegates. Edges whose endpoints are not
926/// in `fragments` are dropped (the hashmap path silently dropped them
927/// at CSR construction).
928pub 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        // Pins that `RankedCandidate::rank`'s dst tie-break agrees with
1206        // `cap_out_edges_per_source`'s sort tie-break: both must favor the
1207        // lower dst index on equal weight. A regression flipping either
1208        // side would silently change which edges the two-pass path keeps
1209        // relative to the materialized path.
1210        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        // Faithful to `edges::collect_capped_edges`'s two passes: pass 1
1261        // fixes a single canonical category per (src, dst) pair (first
1262        // builder, in registration order, to emit it) *before* pass 2's
1263        // per-builder bounded-heap capping runs. This test verifies the
1264        // correctness claim documented there -- per-builder top-K heaps,
1265        // merged + deduped + capped again, must select exactly the same
1266        // surviving edges as capping the fully materialized union
1267        // directly, including the case where the pair's canonical
1268        // category came from a builder whose own (locally weaker) copy
1269        // gets evicted from its own per-builder heap before the merge.
1270        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}