_diffctx/
provenance.rs

1use std::collections::VecDeque;
2use std::io::{BufWriter, Write};
3use std::path::Path;
4
5use rustc_hash::{FxHashMap, FxHashSet};
6
7use crate::pipeline::ScoredState;
8use crate::types::{Fragment, FragmentId};
9
10pub const PROVENANCE_ENV: &str = "DIFFCTX_PROVENANCE_DUMP";
11
12/// Env-gated per-candidate inclusion-provenance dump (#93). One JSONL line
13/// per scored candidate: relevance, seed distance in hops, per-edge-category
14/// incoming mass (`weight x rel(src)` summed over incoming edges), and the
15/// selection verdict. The default path costs one env probe and nothing else,
16/// which is what keeps the instrumentation E-class.
17pub fn maybe_dump(state: &ScoredState, selected: &[Fragment]) {
18    let Ok(path) = std::env::var(PROVENANCE_ENV) else {
19        return;
20    };
21    if path.is_empty() {
22        return;
23    }
24    if let Err(e) = dump(state, selected, Path::new(&path)) {
25        tracing::debug!("provenance dump to '{}' failed: {}", path, e);
26    }
27}
28
29pub fn seed_hops(state: &ScoredState) -> FxHashMap<FragmentId, u32> {
30    let graph = &state.scoring_result.graph;
31    let mut hops: FxHashMap<FragmentId, u32> = FxHashMap::default();
32    let mut queue: VecDeque<FragmentId> = VecDeque::new();
33    for core in &state.core_ids {
34        hops.insert(core.clone(), 0);
35        queue.push_back(core.clone());
36    }
37    // Relevance flows along edges in both directions (PPR blends forward and
38    // backward pushes), so distance is measured on the undirected graph.
39    let mut undirected: FxHashMap<FragmentId, Vec<FragmentId>> = FxHashMap::default();
40    graph.for_each_categorized_edge(|src, dst, _| {
41        undirected.entry(src.clone()).or_default().push(dst.clone());
42        undirected.entry(dst.clone()).or_default().push(src.clone());
43    });
44    while let Some(node) = queue.pop_front() {
45        let d = hops[&node];
46        if let Some(neighbors) = undirected.get(&node) {
47            for n in neighbors {
48                if !hops.contains_key(n) {
49                    hops.insert(n.clone(), d + 1);
50                    queue.push_back(n.clone());
51                }
52            }
53        }
54    }
55    hops
56}
57
58struct CatMass {
59    mass: f64,
60    top_source: FragmentId,
61    top_contribution: f64,
62}
63
64fn per_category_mass(
65    state: &ScoredState,
66) -> FxHashMap<FragmentId, FxHashMap<&'static str, CatMass>> {
67    let graph = &state.scoring_result.graph;
68    let rel = &state.scoring_result.rel_scores;
69    let mut mass: FxHashMap<FragmentId, FxHashMap<&'static str, CatMass>> = FxHashMap::default();
70    graph.for_each_categorized_edge(|src, dst, cat| {
71        let src_rel = rel.get(src).copied().unwrap_or(0.0);
72        if src_rel <= 0.0 {
73            return;
74        }
75        let w = graph.forward_edge_weight(src, dst).unwrap_or(0.0);
76        if w <= 0.0 {
77            return;
78        }
79        let contribution = w * src_rel;
80        let entry = mass
81            .entry(dst.clone())
82            .or_default()
83            .entry(cat.as_str())
84            .or_insert_with(|| CatMass {
85                mass: 0.0,
86                top_source: src.clone(),
87                top_contribution: 0.0,
88            });
89        entry.mass += contribution;
90        // Strict `>` with deterministic iteration would still tie-break by
91        // visit order; prefer the lexically-smaller source on equal
92        // contribution so the attribution is order-independent.
93        if contribution > entry.top_contribution
94            || (contribution == entry.top_contribution && src.path < entry.top_source.path)
95        {
96            entry.top_source = src.clone();
97            entry.top_contribution = contribution;
98        }
99    });
100    mass
101}
102
103/// Per-fragment incoming relevance mass grouped by edge category, sorted by
104/// mass descending: `(category, strongest_source_path, mass)`. Shared by the
105/// provenance dump and the locate renderer (#126) — one attribution pass,
106/// two consumers.
107pub fn incoming_attribution(
108    state: &ScoredState,
109) -> FxHashMap<FragmentId, Vec<(String, String, f64)>> {
110    per_category_mass(state)
111        .into_iter()
112        .map(|(id, cats)| {
113            let mut rows: Vec<(String, String, f64)> = cats
114                .into_iter()
115                .map(|(cat, m)| (cat.to_string(), m.top_source.path.to_string(), m.mass))
116                .collect();
117            rows.sort_by(|a, b| {
118                b.2.partial_cmp(&a.2)
119                    .unwrap_or(std::cmp::Ordering::Equal)
120                    .then(a.0.cmp(&b.0))
121            });
122            (id, rows)
123        })
124        .collect()
125}
126
127fn incoming_mass(state: &ScoredState) -> FxHashMap<FragmentId, FxHashMap<&'static str, f64>> {
128    per_category_mass(state)
129        .into_iter()
130        .map(|(id, cats)| (id, cats.into_iter().map(|(c, m)| (c, m.mass)).collect()))
131        .collect()
132}
133
134fn dump(state: &ScoredState, selected: &[Fragment], out_path: &Path) -> std::io::Result<()> {
135    let rel = &state.scoring_result.rel_scores;
136    let selected_ids: FxHashSet<&FragmentId> = selected.iter().map(|f| &f.id).collect();
137    let hops = seed_hops(state);
138    let mass = incoming_mass(state);
139
140    let mut fragments: Vec<&Fragment> = state.scoring_result.filtered_fragments.iter().collect();
141    fragments.sort_by(|a, b| {
142        a.id.path
143            .cmp(&b.id.path)
144            .then(a.id.start_line.cmp(&b.id.start_line))
145            .then(a.id.end_line.cmp(&b.id.end_line))
146    });
147
148    if let Some(parent) = out_path.parent() {
149        if !parent.as_os_str().is_empty() {
150            std::fs::create_dir_all(parent)?;
151        }
152    }
153    let mut w = BufWriter::new(std::fs::File::create(out_path)?);
154    for frag in fragments {
155        let contrib: serde_json::Map<String, serde_json::Value> = mass
156            .get(&frag.id)
157            .map(|per_cat| {
158                let mut sorted: Vec<_> = per_cat.iter().collect();
159                sorted.sort_by(|a, b| a.0.cmp(b.0));
160                sorted
161                    .into_iter()
162                    .map(|(cat, v)| ((*cat).to_string(), serde_json::json!(v)))
163                    .collect()
164            })
165            .unwrap_or_default();
166        let line = serde_json::json!({
167            "path": frag.id.path.as_ref(),
168            "start": frag.id.start_line,
169            "end": frag.id.end_line,
170            "kind": format!("{:?}", frag.kind).to_lowercase(),
171            "tokens": frag.token_count,
172            "relevance": rel.get(&frag.id).copied().unwrap_or(0.0),
173            "is_core": state.core_ids.contains(&frag.id),
174            "selected": selected_ids.contains(&frag.id),
175            "seed_hops": hops.get(&frag.id).map(|h| *h as i64).unwrap_or(-1),
176            // Which discovery strategy put this file in the universe at all.
177            // `null` for a changed file, which is never discovered — it is the
178            // seed. Splits "never surfaced" from "surfaced but not selected",
179            // which the selected set alone cannot distinguish (#130).
180            "discovery_source": state
181                .discovery_source
182                .get(&frag.id.path)
183                .map(|s| serde_json::json!(s))
184                .unwrap_or(serde_json::Value::Null),
185            "incoming_mass": serde_json::Value::Object(contrib),
186        });
187        writeln!(w, "{line}")?;
188    }
189    w.flush()
190}
_diffctx/provenance.rs

_diffctx/
provenance.rs
