eidos_kernel/retrieval/

scoring.rs

use std::collections::{BTreeMap, HashMap};

use crate::schema::{
    Edge, EdgeBasis, Graph, Node, RelationProfile, core_relation_profile, relation,
};
use crate::trigram::TrigramIndex;

use super::{terms_of, within_edit1};

/// Precomputed, graph-invariant inputs to grounding — the per-node BM25F token fields and
/// lowercased titles. Built ONCE per graph via [`GroundIndex::build`]; grounding many queries against
/// one graph reuses it instead of rebuilding it every call.
#[derive(Clone)]
pub struct GroundIndex {
    pub(super) lc_titles: Vec<String>,
    /// BM25F prototype: per-node tokenized fields
    /// [id, title, summary, aliases, query_examples, relations].
    pub(super) bm25_fields: Vec<[Vec<String>; 6]>,
    /// Average token count per field across the graph — BM25 length normalization.
    pub(super) bm25_avglen: [f64; 6],
    /// The BM25F ranker parameters for this index. Injected at build time; defaults to the frozen
    /// calibration. `ground_scoped` reads it — the kernel never consults the environment.
    pub(super) config: RankerConfig,
    /// Trigram inverted index for fast candidate narrowing.
    trigram_index: TrigramIndex,
}

impl GroundIndex {
    /// Precompute the ground inputs for `graph` with the FROZEN ranker calibration (one O(N) pass).
    pub fn build(graph: &Graph) -> Self {
        Self::build_with_config(graph, RankerConfig::default())
    }

    /// Precompute the ground inputs with an explicit [`RankerConfig`] — the injection seam for the
    /// dev-only parameter sweep. Production/eval callers use [`GroundIndex::build`] (frozen).
    pub fn build_with_config(graph: &Graph, config: RankerConfig) -> Self {
        let relation_surfaces = relation_surfaces(graph);
        let bm25_fields: Vec<[Vec<String>; 6]> = graph
            .nodes
            .iter()
            .map(|node| bm25f_fields(node, &relation_surfaces))
            .collect();
        let n = bm25_fields.len().max(1) as f64;
        let mut bm25_avglen = [0.0f64; 6];
        for f in &bm25_fields {
            for j in 0..6 {
                bm25_avglen[j] += f[j].len() as f64;
            }
        }
        for a in &mut bm25_avglen {
            *a = (*a / n).max(1.0);
        }
        GroundIndex {
            lc_titles: graph.nodes.iter().map(|n| n.title.to_lowercase()).collect(),
            // D0.6: trigram narrowing is dormant (scoring.rs:835-841 in retrieval.rs does not
            // read it; D3 will re-enable narrowing against the SAME tokenized forms `terms_of`
            // produces). Skip the (large, redundant) index construction here — costs nothing in
            // scores (candidates() is never called today) and saves an O(nodes × avg_haystack)
            // memory+time hit per build. The module + tests are preserved for D3.
            trigram_index: TrigramIndex::default(),
            bm25_fields,
            bm25_avglen,
            config,
        }
    }

    /// Narrow the candidate set using the trigram index.
    /// Returns `None` if narrowing isn't possible (short terms) — caller falls back to scanning all.
    /// Returns `Some(indices)` — the node indices that MIGHT match (superset; final scoring decides).
    pub fn candidates(&self, terms: &[String]) -> Option<Vec<usize>> {
        self.trigram_index.candidates(terms)
    }
}

/// Frozen BM25F field weights `[id, title, summary, aliases, query_examples, relations]`. These
/// are calibration constants — moving one moves the frozen eval baselines. See
/// [`crate::calibration`] for the full calibration surface and which test pins each value.
pub const DEFAULT_BM25F_WEIGHTS: [f64; 6] = [5.0, 8.0, 2.0, 6.0, 4.0, 3.0];
/// Frozen BM25 term-frequency saturation. Near-optimal for short docs (measured).
pub const DEFAULT_BM25F_K1: f64 = 1.2;
/// Frozen BM25 length-normalization. Near-optimal for short docs (measured).
pub const DEFAULT_BM25F_B: f64 = 0.75;

/// BM25F ranker parameters. INJECTED, never read from the environment: the kernel is pure, so
/// `Default` is the frozen calibration and any override is passed in by the caller (the dev-only
/// parameter sweep via `eidos-eval-runner --k1/--b/--weights`). This is what makes "same inputs →
/// same scores" hold without depending on ambient process state.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct RankerConfig {
    pub k1: f64,
    pub b: f64,
    pub weights: [f64; 6],
}

impl Default for RankerConfig {
    fn default() -> Self {
        Self {
            k1: DEFAULT_BM25F_K1,
            b: DEFAULT_BM25F_B,
            weights: DEFAULT_BM25F_WEIGHTS,
        }
    }
}

/// Back-compat alias for the internal scorer name.
pub(super) type Bm25fParams = RankerConfig;

pub(super) struct Bm25fScorer<'a> {
    pub(super) terms: &'a [String],
    pub(super) df: &'a HashMap<&'a str, usize>,
    pub(super) n: usize,
    pub(super) avglen: &'a [f64; 6],
    pub(super) params: Bm25fParams,
}

#[derive(Clone, Copy)]
pub(super) struct Bm25fNodeShape {
    pub(super) is_code: bool,
    pub(super) is_module: bool,
}

fn bm25f_fields(
    node: &Node,
    relation_surfaces: &BTreeMap<String, Vec<String>>,
) -> [Vec<String>; 6] {
    [
        terms_of(&node.id.replace(['.', '-', '_', '/'], " ").to_lowercase()),
        terms_of(&node.title.to_lowercase()),
        terms_of(&node.summary.to_lowercase()),
        terms_of(&node.aliases.join(" ").to_lowercase()),
        terms_of(&node.query_examples.join(" ").to_lowercase()),
        terms_of(
            &relation_surfaces
                .get(&node.id)
                .map(|surfaces| surfaces.join(" "))
                .unwrap_or_default()
                .to_lowercase(),
        ),
    ]
}

pub(super) fn bm25f_score(
    fields: &[Vec<String>; 6],
    scorer: &Bm25fScorer<'_>,
    shape: Bm25fNodeShape,
) -> (f64, f64, usize, usize, usize, f64) {
    let mut score = 0.0;
    let mut identity = 0.0;
    let mut matched = 0usize;
    let mut matched_identity = 0usize;
    let mut matched_name = 0usize;
    // Highest IDF among query terms that EXACTLY name this node (exact hit in an id/title/alias
    // field). A high value means a RARE, distinctive term names the node — "carbonara" names one
    // doc; "egg" names five. Lets confidence anchor on a distinctive name even when common co-terms
    // dilute id_coverage below the usual bar.
    let mut max_exact_name_idf = 0.0_f64;
    for t in scorer.terms {
        let dft = *scorer.df.get(t.as_str()).unwrap_or(&0);
        if dft == 0 {
            continue;
        }
        let idf = ((scorer.n as f64 - dft as f64 + 0.5) / (dft as f64 + 0.5) + 1.0).ln();
        let mut wtf = 0.0;
        let mut id_wtf = 0.0;
        let mut name_hit = false;
        let mut exact_name_hit = false;
        for (j, field) in fields.iter().enumerate() {
            let exact_tf = field.iter().filter(|x| x.as_str() == t.as_str()).count() as f64;
            let mut tf = exact_tf;
            if tf == 0.0 && t.len() >= 5 {
                tf = 0.5
                    * field
                        .iter()
                        .filter(|x| x.len() >= 5 && within_edit1(t.as_bytes(), x.as_bytes()))
                        .count() as f64;
            }
            if tf > 0.0 {
                let norm = 1.0 - scorer.params.b
                    + scorer.params.b * (field.len() as f64) / scorer.avglen[j];
                let c = scorer.params.weights[j] * tf / norm;
                wtf += c;
                if j == 0 || j == 1 || (!shape.is_module && j == 2) || (!shape.is_code && j == 3) {
                    id_wtf += c;
                    name_hit = true;
                }
                // NAME-GRADE evidence: an exact (non-fuzzy) token hit in a NAME field — id,
                // title, or (for docs) aliases. Summary hits describe; they do not name. A
                // fuzzy edit-1 hit ("scheduler"~"schedule") is a typo hypothesis, not a name
                // claim. This is what lets id_coverage anchor confidence without letting a
                // description scatter or near-miss impersonate identity.
                if exact_tf > 0.0 && (j == 0 || j == 1 || (!shape.is_code && j == 3)) {
                    exact_name_hit = true;
                }
            }
        }
        if wtf > 0.0 {
            matched += 1;
            score += idf * wtf / (scorer.params.k1 + wtf);
            if id_wtf > 0.0 {
                identity += idf * id_wtf / (scorer.params.k1 + id_wtf);
            }
            if name_hit {
                matched_identity += 1;
            }
            if exact_name_hit {
                matched_name += 1;
                max_exact_name_idf = max_exact_name_idf.max(idf);
            }
        }
    }
    (
        score,
        identity,
        matched,
        matched_identity,
        matched_name,
        max_exact_name_idf,
    )
}

pub(super) fn bm25f_df<'a>(
    terms: &'a [String],
    fields: &[[Vec<String>; 6]],
) -> HashMap<&'a str, usize> {
    terms
        .iter()
        .map(|t| {
            let exact = fields
                .iter()
                .filter(|f| f.iter().any(|fl| fl.iter().any(|tok| tok == t)))
                .count();
            let c = if exact > 0 || t.len() < 5 {
                exact
            } else {
                fields
                    .iter()
                    .filter(|f| {
                        f.iter().any(|fl| {
                            fl.iter().any(|tok| {
                                tok.len() >= 5 && within_edit1(t.as_bytes(), tok.as_bytes())
                            })
                        })
                    })
                    .count()
            };
            (t.as_str(), c)
        })
        .collect()
}

fn relation_surfaces(graph: &Graph) -> BTreeMap<String, Vec<String>> {
    let nodes = graph
        .nodes
        .iter()
        .map(|node| (node.id.as_str(), node))
        .collect::<HashMap<_, _>>();
    let mut surfaces: BTreeMap<String, Vec<String>> = BTreeMap::new();
    for edge in &graph.edges {
        if !relation_is_searchable(edge, &graph.relation_profiles) {
            continue;
        }
        let Some(from) = nodes.get(edge.from.as_str()) else {
            continue;
        };
        let Some(to) = nodes.get(edge.to.as_str()) else {
            continue;
        };
        push_relation_surface(
            &mut surfaces,
            &edge.from,
            relation_phrase(&edge.relation, &graph.relation_profiles),
            to,
        );
        push_relation_surface(
            &mut surfaces,
            &edge.to,
            reverse_relation_phrase(&edge.relation, &graph.relation_profiles),
            from,
        );
    }
    surfaces
}

fn relation_is_searchable(
    edge: &Edge,
    profiles: &std::collections::BTreeMap<String, RelationProfile>,
) -> bool {
    if edge.basis != EdgeBasis::Resolved {
        return false;
    }
    // Traversal-only call-graph edges power callers/impact but must not enrich grounding surfaces.
    if edge.evidence == crate::schema::TRAVERSAL_ONLY_EVIDENCE {
        return false;
    }
    profiles
        .get(&edge.relation)
        .cloned()
        .or_else(|| core_relation_profile(&edge.relation))
        .is_none_or(|profile| profile.searchable)
}

fn push_relation_surface(
    surfaces: &mut BTreeMap<String, Vec<String>>,
    node_id: &str,
    relation_phrase: String,
    other: &Node,
) {
    let mut surface = format!("{relation_phrase} {} {}", other.id, other.title);
    for alias in &other.aliases {
        surface.push(' ');
        surface.push_str(alias);
    }
    let entry = surfaces.entry(node_id.to_string()).or_default();
    if !entry.iter().any(|existing| existing == &surface) {
        entry.push(surface);
    }
}

pub(super) fn relation_phrase(
    relation: &str,
    profiles: &std::collections::BTreeMap<String, RelationProfile>,
) -> String {
    profiles
        .get(relation)
        .cloned()
        .or_else(|| core_relation_profile(relation))
        .map_or_else(
            || relation.replace('_', " "),
            |profile| profile.forward_phrase.clone(),
        )
}

pub(super) fn reverse_relation_phrase(
    relation: &str,
    profiles: &std::collections::BTreeMap<String, RelationProfile>,
) -> String {
    if let Some(profile) = profiles
        .get(relation)
        .cloned()
        .or_else(|| core_relation_profile(relation))
    {
        return profile.reverse_phrase.clone();
    }
    match relation {
        relation::LINKS_TO => "linked from".to_string(),
        relation::REFERENCES => "referenced by".to_string(),
        relation::IMPLEMENTS => "implemented by".to_string(),
        "depends_on" => "required by".to_string(),
        "produces" => "produced by".to_string(),
        "consumes" => "consumed by".to_string(),
        "uses_tool" => "used by".to_string(),
        "requires_contract" => "required contract for".to_string(),
        "dispatches" => "dispatched by".to_string(),
        other => format!("{} by", other.replace('_', " ")),
    }
}

// ─── D0.6 — trigram narrowing dormant, scores unaffected ───────────────

#[cfg(test)]
mod dormant_trigram_tests {
    use super::*;
    use crate::retrieval::ground_with;
    use crate::schema::Node;

    fn mk_node(id: &str, title: &str, summary: &str) -> Node {
        Node {
            id: id.into(),
            kind: crate::schema::Kind::Doc,
            partition: None,
            subkind: None,
            title: title.into(),
            summary: summary.into(),
            aliases: Vec::new(),
            tags: Vec::new(),
            query_examples: Vec::new(),
            source_files: Vec::new(),
            span: None,
        }
    }

    fn two_node_graph() -> Graph {
        let mut g = Graph::default();
        g.nodes.push(mk_node("doc.a", "A", "alpha"));
        g.nodes.push(mk_node("doc.b", "B", "beta"));
        g
    }

    #[test]
    fn ground_index_does_not_construct_populated_trigram() {
        // D0.6: dormant state is observable. `trigram_index` must be the empty default,
        // not a populated one — guarantees scoring.rs:835-841's full scan isn't shadowed
        // by data we never read. The trigram index is populated by walking per-node
        // haystacks; on an empty index every 3+-char term returns `Some(vec![])`
        // (trigram doesn't exist anywhere → no candidates). A populated index would return
        // `Some(non_empty)` for matching terms. Pin the contract: the index we build
        // gives back the empty-candidate result.
        let g = two_node_graph();
        let idx = GroundIndex::build(&g);
        let result = idx.candidates(&[String::from("alpha")]);
        assert!(
            matches!(result, Some(ref v) if v.is_empty()),
            "GroundIndex::build must leave trigram_index empty (D0.6 dormant state); \
             candidates() returned {result:?} — non-empty would mean trigram narrowing is \
             silently providing data retrieval.rs:835-841 never reads"
        );
        // The same is true for any 3+-char term: no postings means no narrowing,
        // so the scoring loop never depends on this index.
        let result2 = idx.candidates(&[String::from("beta")]);
        assert!(
            matches!(result2, Some(ref v) if v.is_empty()),
            "dormant trigram index must return Some(vec![]) for any term (no postings); got {result2:?}"
        );
    }

    #[test]
    fn ground_score_is_identical_to_dormant_trigram_world() {
        // The actual byte-identical-pinning property: today's dormant state must not have
        // changed any score. Build a graph + index, ground a query, capture the result.
        // The frozen eval baselines will catch any deviation; this test is the local pin.
        let g = two_node_graph();
        let idx = GroundIndex::build(&g);
        let hits = ground_with(&g, &idx, "alpha", 10);
        assert!(
            !hits.is_empty(),
            "ground must still return at least one hit when the index is dormant"
        );
        // The first hit should be the alpha-node; band is unambiguous-or-better because
        // the query is a single short term.
        assert_eq!(hits[0].id, "doc.a");
    }
}