examples/calibrate_group_space.rs
//! **Calibrate the GROUP-twin gates in the real embedding space** — the v2 run
//! the `mimir_groupclone` module docs promised ("the floor is model-relative —
//! a real embedder gets its own calibration before anyone acts on its number").
//!
//! The function-level precedent is `mimir_calibrate`'s RECORDED row (oden,
//! 2026-08-25, `examples/calibrate_embedding_space`); this is the same oracle
//! one level down: the unit is a SEMANTIC GROUP (def-use segment of a function
//! body), the embedded text is the group's NORMALISED body (`norm.text`) —
//! exactly what `clone_signals::feed::embed_groups` writes to
//! `embeddings_groups` — and the pairing gates under measurement are
//! `TwinOptions` (cosine floor + shared/rare role-class k-grams).
//!
//! Populations:
//!
//! * **positives, by construction** — a real group's SOURCE edited the way real
//! drifted copies drift, then re-normalised (a consistent rename or a retuned
//! literal normalises AWAY, so only edits that produce a DIFFERENT normalised
//! body are vector-path positives — the only pairs the cosine floor ever
//! judges):
//! - `reword-1stmt`: one statement's method / operator rewritten
//! (`push`→`insert`, `^`→`+`, …) — the v1 fixture's case (d), on real code;
//! - `drop-1stmt`: one middle statement removed;
//! - `swap-2stmt`: two adjacent statements reordered.
//! * **hard negatives** — pairs of real groups from different functions that
//! PASS the production structure gate (≥ `MIN_SHARED_KGRAMS` shared k-grams,
//! ≥ 1 rare) with different exact keys: exactly the candidate population the
//! cosine floor exists to judge. ⚠ Some of these are REAL twins wearing a
//! negative label (the function-level run measured that caveat at ≤ 16.3 %),
//! so the floor read off them is conservative, never optimistic.
//! * **random negatives** — groups from different files, for the report only.
//!
//! Usage (reads the trees READ-ONLY, writes nothing):
//!
//! ```text
//! calibrate_group_space <src-root>... [--sample N] [--negatives N] [--seed N]
//! ```
#[cfg(not(any(feature = "embed-tract", feature = "embed-ort")))]
fn main() {
eprintln!("built without an embedder (embed-tract / embed-ort): nothing to calibrate against");
std::process::exit(2);
}
#[cfg(any(feature = "embed-tract", feature = "embed-ort"))]
fn main() -> anyhow::Result<()> {
real::run()
}
#[cfg(any(feature = "embed-tract", feature = "embed-ort"))]
mod real {
use std::collections::{BTreeMap, BTreeSet, HashMap};
use std::path::{Path, PathBuf};
use nornir::clone_signals::norm::{self, NormBody};
use nornir::mimir_calibrate::{
Calibration, CosineDistribution, Lcg, PopulationSet, Provenance, HARD_NEGATIVE, RANDOM_NEGATIVE,
};
use nornir::mimir_groupclone::{self as gc, Group, GroupedItem, TwinOptions};
use nornir::vector::fnvec::{self, Embedding};
use nornir::vector::pair_index::dot;
/// One real group in the corpus.
struct CorpusGroup {
file: String,
/// Which function it came from (functions are the twin boundary).
item: usize,
/// The group's own SOURCE lines, joined (what the transforms edit).
source: String,
norm: NormBody,
kgrams: BTreeSet<u64>,
}
// ── corpus ───────────────────────────────────────────────────────────────
fn rust_files(root: &Path, out: &mut Vec<PathBuf>) {
let Ok(rd) = std::fs::read_dir(root) else { return };
for e in rd.flatten() {
let p = e.path();
let name = e.file_name().to_string_lossy().into_owned();
if p.is_dir() {
if name != "target" && !name.starts_with('.') {
rust_files(&p, out);
}
} else if name.ends_with(".rs") {
out.push(p);
}
}
}
/// Every `fn`/`method` item's text in `file_text`, via the one lexer this
/// repo allows (proc-macro2 with source lines). Nested items are their own
/// entries; text that does not tokenise yields nothing.
fn functions_in_file(file_text: &str) -> Vec<String> {
let Ok(ts) = file_text.parse::<proc_macro2::TokenStream>() else {
return Vec::new();
};
let lines: Vec<&str> = file_text.lines().collect();
let mut spans: Vec<(usize, usize)> = Vec::new();
collect_fn_spans(ts, &mut spans);
spans
.into_iter()
.filter(|(s, e)| *s >= 1 && *e <= lines.len() && e > s)
.map(|(s, e)| lines[s - 1..e].join("\n"))
.collect()
}
fn collect_fn_spans(ts: proc_macro2::TokenStream, out: &mut Vec<(usize, usize)>) {
use proc_macro2::{Delimiter, TokenTree};
let toks: Vec<TokenTree> = ts.into_iter().collect();
let mut i = 0usize;
while i < toks.len() {
if let TokenTree::Ident(id) = &toks[i] {
if *id == "fn" {
if let Some(TokenTree::Ident(_)) = toks.get(i + 1) {
let start = id.span().start().line;
// the body: the next Brace group at this level (a trait
// fn without one ends at `;` first)
let mut j = i + 2;
while j < toks.len() {
match &toks[j] {
TokenTree::Group(g) if g.delimiter() == Delimiter::Brace => {
out.push((start, g.span_close().end().line));
break;
}
TokenTree::Punct(p) if p.as_char() == ';' => break,
_ => {}
}
j += 1;
}
}
}
}
if let TokenTree::Group(g) = &toks[i] {
collect_fn_spans(g.stream(), out);
}
i += 1;
}
}
/// The groups of `fn_text` at or above the weight floor, with their SOURCE.
fn groups_of(fn_text: &str) -> Vec<(String, NormBody, BTreeSet<u64>)> {
let Some(groups) = gc::segment(fn_text) else { return Vec::new() };
let lines: Vec<&str> = fn_text.lines().collect();
groups
.into_iter()
.filter(|g: &Group| g.norm.tokens >= gc::MIN_GROUP_TOKENS)
.map(|g| {
let src = g
.lines
.iter()
.filter_map(|l| lines.get((*l - 1) as usize))
.copied()
.collect::<Vec<_>>()
.join("\n");
(src, g.norm, g.kgrams)
})
.collect()
}
// ── the constructed positives ────────────────────────────────────────────
/// Rewrite ONE statement: the first swap that applies. Method names on the
/// KEEP/METHODS lists and operators survive normalisation, so these edits
/// produce a DIFFERENT normalised body — the only kind of positive the
/// vector path ever judges.
fn reword_one(src: &str) -> String {
const SWAPS: &[(&str, &str)] = &[
(".push(", ".insert(0, "),
(".push_str(", ".insert_str(0, "),
(".max_by_key(", ".min_by_key("),
(".sort_by(", ".sort_unstable_by("),
(".iter(", ".into_iter("),
(".unwrap_or(", ".unwrap_or_else(|| "),
(" ^ ", " + "),
(" + ", " * "),
(" && ", " || "),
("+= 1", "+= 2"),
];
for (a, b) in SWAPS {
if src.contains(a) {
return src.replacen(a, b, 1);
}
}
// fallback: wrap the last line's expression — still one statement changed
let mut lines: Vec<String> = src.lines().map(str::to_string).collect();
if let Some(l) = lines.last_mut() {
*l = format!("{} // twin", l);
l.push_str(" ");
}
lines.join("\n")
}
/// Remove one middle statement line (groups with ≥ 4 lines).
fn drop_one(src: &str) -> Option<String> {
let lines: Vec<&str> = src.lines().collect();
if lines.len() < 4 {
return None;
}
let cut = lines.len() / 2;
Some(
lines
.iter()
.enumerate()
.filter(|(i, _)| *i != cut)
.map(|(_, l)| *l)
.collect::<Vec<_>>()
.join("\n"),
)
}
/// Swap two adjacent middle lines (groups with ≥ 4 lines).
fn swap_two(src: &str) -> Option<String> {
let mut lines: Vec<&str> = src.lines().collect();
if lines.len() < 4 {
return None;
}
let at = lines.len() / 2 - 1;
lines.swap(at, at + 1);
Some(lines.join("\n"))
}
// ── embedding: the SAME path as `embed_groups` (pieces → embed → pool) ──
fn embed_norm_texts(embedding: &mut dyn fnvec::Embedding, texts: &[&str]) -> anyhow::Result<Vec<Vec<f32>>> {
let mut out = Vec::with_capacity(texts.len());
let per = nornir::vector::store::embed_slice_chunks();
let t0 = std::time::Instant::now();
for slice in texts.chunks(per) {
let cut: Vec<Vec<&str>> = slice.iter().map(|t| fnvec::pieces(t)).collect();
let flat: Vec<&str> = cut.iter().flatten().copied().collect();
let embedded = embedding.embed(&flat)?;
anyhow::ensure!(embedded.len() == flat.len(), "embedder returned {} vectors for {} text(s)", embedded.len(), flat.len());
let mut next = 0usize;
for parts in &cut {
let of = &embedded[next..next + parts.len()];
next += parts.len();
out.push(fnvec::pool(of, &parts.iter().map(|p| p.len()).collect::<Vec<_>>()));
}
eprint!("\r embedded {}/{} ({:.0} s) ", out.len(), texts.len(), t0.elapsed().as_secs_f64());
}
eprintln!();
Ok(out)
}
// ── the run ──────────────────────────────────────────────────────────────
pub fn run() -> anyhow::Result<()> {
let mut roots: Vec<PathBuf> = Vec::new();
let (mut sample, mut negatives, mut seed) = (200usize, 500usize, 0x6e5fu64);
let mut args = std::env::args().skip(1);
while let Some(a) = args.next() {
match a.as_str() {
"--sample" => sample = args.next().unwrap_or_default().parse()?,
"--negatives" => negatives = args.next().unwrap_or_default().parse()?,
"--seed" => seed = args.next().unwrap_or_default().parse()?,
p => roots.push(PathBuf::from(p)),
}
}
anyhow::ensure!(!roots.is_empty(), "usage: calibrate_group_space <src-root>... [--sample N] [--negatives N]");
// 1. the corpus: every semantic group ≥ the weight floor, with source
let mut files: Vec<PathBuf> = Vec::new();
for r in &roots {
rust_files(r, &mut files);
}
files.sort();
let mut corpus: Vec<CorpusGroup> = Vec::new();
let mut functions = 0usize;
let mut item = 0usize;
for f in &files {
let Ok(text) = std::fs::read_to_string(f) else { continue };
for fn_text in functions_in_file(&text) {
functions += 1;
for (source, norm, kgrams) in groups_of(&fn_text) {
corpus.push(CorpusGroup { file: f.display().to_string(), item, source, norm, kgrams });
}
item += 1;
}
}
eprintln!(
"corpus: {} file(s), {functions} function(s), {} group(s) ≥ {} normalised tokens",
files.len(),
corpus.len(),
gc::MIN_GROUP_TOKENS
);
anyhow::ensure!(corpus.len() >= 50, "corpus too small to say anything ({} groups)", corpus.len());
// document frequency per k-gram over DISTINCT normalised bodies — the
// same rarity the production pairer uses
let mut df: HashMap<u64, BTreeSet<&str>> = HashMap::new();
for g in &corpus {
for k in &g.kgrams {
df.entry(*k).or_default().insert(g.norm.exact.as_str());
}
}
// 2. constructed positives from a spread sample
let mut rng = Lcg(seed);
let mut chosen: Vec<usize> = Vec::new();
let mut seen: BTreeSet<usize> = BTreeSet::new();
let want = sample.min(corpus.len());
let mut guard = 0usize;
while chosen.len() < want && guard < want * 100 {
guard += 1;
let i = rng.below(corpus.len());
if seen.insert(i) {
chosen.push(i);
}
}
type Xf = (&'static str, fn(&str) -> Option<String>);
let transforms: Vec<Xf> = vec![
("reword-1stmt", |s| Some(reword_one(s))),
("drop-1stmt", drop_one),
("swap-2stmt", swap_two),
];
// (label, original corpus idx, edited norm) — only edits that produced
// a DIFFERENT normalised body (the rest normalise away and are counted)
let mut pos_pairs: Vec<(&'static str, usize, NormBody)> = Vec::new();
let mut normalised_away = 0usize;
for &i in &chosen {
for (label, f) in &transforms {
let Some(edited) = f(&corpus[i].source) else { continue };
let Some(n) = norm::normalize(&edited) else { continue };
if n.exact == corpus[i].norm.exact {
normalised_away += 1;
continue;
}
if n.tokens < gc::MIN_GROUP_TOKENS {
continue;
}
pos_pairs.push((label, i, n));
}
}
eprintln!("positives: {} constructed pair(s) ({normalised_away} edit(s) normalised away — not vector-path pairs)", pos_pairs.len());
// 3. hard negatives: REAL pairs that pass the production STRUCTURE gate
// (different function, different body, ≥ MIN_SHARED_KGRAMS shared of
// which ≥ 1 rare) — the exact candidates the cosine floor judges
let opts = TwinOptions::default();
let is_rare = |k: &u64| df.get(k).is_some_and(|d| d.len() <= opts.rare_df);
let mut hard: Vec<(usize, usize, usize, usize)> = Vec::new(); // (a, b, shared, rare)
let mut index: HashMap<u64, Vec<usize>> = HashMap::new();
for (gi, g) in corpus.iter().enumerate() {
for k in &g.kgrams {
index.entry(*k).or_default().push(gi);
}
}
let mut cand: BTreeSet<(usize, usize)> = BTreeSet::new();
for (_, bucket) in &index {
if bucket.len() > opts.max_bucket {
continue;
}
for x in 0..bucket.len() {
for y in (x + 1)..bucket.len() {
let (a, b) = (bucket[x].min(bucket[y]), bucket[x].max(bucket[y]));
if corpus[a].item != corpus[b].item && corpus[a].norm.exact != corpus[b].norm.exact {
cand.insert((a, b));
}
}
}
}
let exact_twin_pairs = {
let mut by_key: BTreeMap<&str, BTreeSet<usize>> = BTreeMap::new();
for g in &corpus {
by_key.entry(g.norm.exact.as_str()).or_default().insert(g.item);
}
by_key.values().filter(|items| items.len() >= 2).count()
};
for (a, b) in &cand {
let shared: Vec<u64> = corpus[*a].kgrams.intersection(&corpus[*b].kgrams).copied().collect();
let rare = shared.iter().filter(|k| is_rare(k)).count();
if shared.len() >= opts.min_shared_kgrams && rare >= 1 {
hard.push((*a, *b, shared.len(), rare));
}
}
// spread-sample the hard negatives down to the budget
let mut rng2 = Lcg(seed ^ 0x51ed);
while hard.len() > negatives {
let i = rng2.below(hard.len());
hard.swap_remove(i);
}
eprintln!(
"hard negatives: {} structure-gate candidate pair(s) kept of {} in the corpus; {} normalised bodies are EXACT twins across functions (no vector needed for those)",
hard.len(),
cand.len(),
exact_twin_pairs
);
// random negatives: different files
let mut rand_pairs: Vec<(usize, usize)> = Vec::new();
let mut rng3 = Lcg(seed ^ 0xa17e);
let mut seen2: BTreeSet<(usize, usize)> = BTreeSet::new();
let mut guard = 0usize;
while rand_pairs.len() < negatives && guard < negatives * 50 {
guard += 1;
let a = rng3.below(corpus.len());
let b = rng3.below(corpus.len());
if a == b || corpus[a].file == corpus[b].file {
continue;
}
let key = (a.min(b), a.max(b));
if seen2.insert(key) {
rand_pairs.push(key);
}
}
// 4. embed every distinct normalised body ONCE, through the production
// seam (pieces → the real embedder → pool)
let mut slot: BTreeMap<&str, usize> = BTreeMap::new(); // exact key → texts idx
let mut texts: Vec<&str> = Vec::new();
let mut need: Vec<(&str, &str)> = Vec::new(); // (exact, text)
for &i in &chosen {
need.push((corpus[i].norm.exact.as_str(), corpus[i].norm.text.as_str()));
}
for (_, i, n) in &pos_pairs {
let _ = i;
need.push((n.exact.as_str(), n.text.as_str()));
}
for (a, b, _, _) in &hard {
need.push((corpus[*a].norm.exact.as_str(), corpus[*a].norm.text.as_str()));
need.push((corpus[*b].norm.exact.as_str(), corpus[*b].norm.text.as_str()));
}
for (a, b) in &rand_pairs {
need.push((corpus[*a].norm.exact.as_str(), corpus[*a].norm.text.as_str()));
need.push((corpus[*b].norm.exact.as_str(), corpus[*b].norm.text.as_str()));
}
for (exact, text) in need {
if !slot.contains_key(exact) {
slot.insert(exact, texts.len());
texts.push(text);
}
}
let mut embedding = fnvec::Compiled::default();
let profile = embedding.profile();
let profile_id = profile.id();
eprintln!(
"embedding {} distinct normalised group bodies — model `{}`, backend {}, profile {}",
texts.len(),
embedding.model_id(),
embedding.backend(),
profile_id
);
let vecs = embed_norm_texts(&mut embedding, &texts)?;
let v_of = |exact: &str| -> &[f32] { &vecs[slot[exact]] };
// 5. the populations
let mut by_label: BTreeMap<&'static str, Vec<f32>> = BTreeMap::new();
for (label, i, n) in &pos_pairs {
by_label.entry(*label).or_default().push(dot(v_of(&corpus[*i].norm.exact), v_of(&n.exact)));
}
let positives = PopulationSet::new(
"positives",
transforms
.iter()
.map(|(l, _)| CosineDistribution::measure(*l, by_label.get(*l).map(|v| v.as_slice()).unwrap_or(&[])))
.collect(),
);
let hard_c: Vec<f32> = hard.iter().map(|(a, b, _, _)| dot(v_of(&corpus[*a].norm.exact), v_of(&corpus[*b].norm.exact))).collect();
let rand_c: Vec<f32> = rand_pairs.iter().map(|(a, b)| dot(v_of(&corpus[*a].norm.exact), v_of(&corpus[*b].norm.exact))).collect();
let negatives = PopulationSet::new(
"negatives",
vec![CosineDistribution::measure(HARD_NEGATIVE, &hard_c), CosineDistribution::measure(RANDOM_NEGATIVE, &rand_c)],
);
let provenance = Provenance {
model_id: embedding.model_id(),
model_name: profile.model_name.clone(),
weights_sha: profile.weights_sha.clone(),
dim: profile.dim,
dtype: profile.dtype.clone(),
corpus: format!(
"{} — {} functions, {} semantic groups ≥ {} tokens, embedded per NORMALISED GROUP BODY",
roots.iter().map(|r| r.display().to_string()).collect::<Vec<_>>().join(","),
functions,
corpus.len(),
gc::MIN_GROUP_TOKENS
),
measured_at: chrono::Utc::now().format("%Y-%m-%d").to_string(),
};
let cal = Calibration::compute(provenance, positives, negatives, 1.0);
println!("\n{}\n", cal.report());
println!("model_profile (the id stored in embeddings_groups): {profile_id}");
println!("\nk sweep (k, floor, recall, hard-neg survival):");
for (k, floor, recall, hn) in cal.k_sweep(&[0.0, 0.5, 1.0, 1.5, 2.0, 3.0]) {
println!(" k={k:<4} floor={floor:.3} recall={:.1}% hard-neg={:.1}%", recall * 100.0, hn * 100.0);
}
println!("\nfloors read off the HARD-NEGATIVE quantile (rate, floor, recall):");
for rate in [0.05f32, 0.10, 0.20, 0.30] {
if let Some((floor, recall)) = cal.floor_for_hard_negative_rate(rate) {
println!(" FP≤{:>2.0}% floor={floor:.3} recall={:.1}%", rate * 100.0, recall * 100.0);
}
}
// 6. the structure gate, measured on the same pairs
let shared_of = |n: &NormBody, g: &CorpusGroup| {
let kg = gc::kgrams_of(&n.classes, opts.kgram);
let shared: Vec<u64> = kg.intersection(&g.kgrams).copied().collect();
let rare = shared.iter().filter(|k| is_rare(k)).count();
(shared.len(), rare)
};
let mut pos_shared: Vec<usize> = Vec::new();
let mut pos_rare_ok = 0usize;
for (_, i, n) in &pos_pairs {
let (s, r) = shared_of(n, &corpus[*i]);
if r >= 1 && s >= opts.min_shared_kgrams {
pos_rare_ok += 1;
}
pos_shared.push(s);
}
pos_shared.sort_unstable();
let q = |v: &[usize], p: f64| v.get(((v.len() as f64 - 1.0) * p) as usize).copied().unwrap_or(0);
println!(
"\nstructure gate on the constructed positives: shared {}-grams p05={} p50={} p95={}; {}/{} pass the production gate (≥{} shared, ≥1 rare)",
opts.kgram,
q(&pos_shared, 0.05),
q(&pos_shared, 0.50),
q(&pos_shared, 0.95),
pos_rare_ok,
pos_pairs.len(),
opts.min_shared_kgrams,
);
let hard_shared: Vec<usize> = {
let mut v: Vec<usize> = hard.iter().map(|(_, _, s, _)| *s).collect();
v.sort_unstable();
v
};
println!(
"structure gate on the hard negatives (all pass by construction): shared p05={} p50={} p95={}",
q(&hard_shared, 0.05),
q(&hard_shared, 0.50),
q(&hard_shared, 0.95)
);
// 7. survivors: the real pairer over the real corpus at candidate floors
let mut items: BTreeMap<usize, GroupedItem> = BTreeMap::new();
for g in &corpus {
let e = items.entry(g.item).or_insert_with(|| GroupedItem {
file: g.file.clone(),
ident: format!("fn#{}", g.item),
item_start: 1,
item_end: 1,
groups: Vec::new(),
});
e.groups.push(Group { lines: vec![1], norm: g.norm.clone(), kgrams: g.kgrams.clone() });
}
let items: Vec<GroupedItem> = items.into_values().collect();
let vectors: HashMap<String, Vec<f32>> = slot.iter().map(|(k, i)| (k.to_string(), vecs[*i].clone())).collect();
println!("\nsurvivors (the real pairer over the whole corpus, non-exact twins / exact twins):");
for floor in [0.80f32, 0.882, 0.90, 0.92, 0.95, 0.97] {
let o = TwinOptions { cosine_floor: floor, ..TwinOptions::default() };
let r = gc::pair(&items, &vectors, &o);
let exact = r.twins.iter().filter(|t| t.exact).count();
println!(
" floor {floor:.3}: {} twin(s) ({} exact, {} by vector), {} candidate(s) unjudged",
r.twins.len(),
exact,
r.twins.len() - exact,
r.candidates_unjudged
);
}
Ok(())
}
}