nordisk/nornir

src/mimir_lens/room_feed/corpus.rs

//! The calibration corpus, from the code room: every function of one repo
//! with its OWN text — the room says where an item starts and ends (the
//! driver's spans), the checkout holds the lines.
//!
//! The extent is the ITEM's, never an embedding chunk's: an `impl` block is
//! one chunk, so a corpus cut at chunk borders would hand most functions a
//! block average. [`crate::mimir_calibrate`] measures similarity on these.

use std::collections::{BTreeMap, BTreeSet};
use std::path::Path;

use anyhow::Result;
use mimir_room::{EdgeSel, MaskSpec, Room};
use mimir_schema::{EdgeKind, Origin, kind as nk};

use crate::mimir_calibrate::CorpusFn;

/// The functions lifted out, and what was left behind (each with its reason).
#[derive(Debug, Clone, Default)]
pub struct Corpus {
    pub fns: Vec<CorpusFn>,
    /// Callable items of the repo the room holds.
    pub callables: usize,
    /// Produced by a macro expansion: their span is the invocation, not a body.
    pub macro_generated: usize,
    /// No span row (nothing says where they are).
    pub without_span: usize,
    /// Their file is not in the checkout, or is shorter than the span says.
    pub unreadable: usize,
    /// The files those lie in, each with how many functions it cost.
    pub unreadable_files: BTreeMap<String, usize>,
    /// Shorter than `min_lines` or longer than `max_lines`.
    pub outside_size: usize,
}

/// Lift every `fn` / `method` the repo defines out of `checkout` (the tree the
/// room's file paths are relative to).
///
/// `min_lines` drops functions too small to carry any signal (a 3-line getter
/// embeds close to every other 3-line getter). `max_lines` drops the giants a
/// 120-line chunk window never gives a vector of their own; `0` disables it.
///
/// `signature` is `fn name(param: Type, …)` from the room's parameter nodes,
/// in name order, WITHOUT a return type (the driver does not record one yet);
/// `signature_key` is the sorted parameter types — "same shape" for the hard
/// negatives. `callees` are the names of the items the function CALLS.
pub fn from_room(room: &Room, repo: &str, checkout: &Path, min_lines: usize, max_lines: usize) -> Result<Corpus> {
    let mut ords: Vec<u64> = Vec::new();
    for k in [nk::FN, nk::METHOD] {
        let spec = MaskSpec { repos: vec![repo.to_string()], node_kinds: vec![k.to_string()], origins: vec![Origin::Local], defined_only: true, ..Default::default() };
        ords.extend_from_slice(room.mask(&spec)?.value.ordinals());
    }
    let callables = room.mask_of(&ords);
    let out_edges = |kind: EdgeKind| -> Result<Vec<mimir_room::EdgeOut>> { Ok(room.edges(&EdgeSel { kinds: vec![kind], from: Some(&callables), to: None, cap: None, heaviest_first: false })?.value.edges) };

    let mut params: BTreeMap<u64, Vec<(String, String)>> = BTreeMap::new();
    for e in out_edges(EdgeKind::Contains)? {
        let p = room.node(e.to);
        if p.kind == Some(nk::PARAM) {
            params.entry(e.from).or_default().push((p.name.unwrap_or("_").to_string(), p.ty.unwrap_or("_").to_string()));
        }
    }
    let mut callees: BTreeMap<u64, BTreeSet<String>> = BTreeMap::new();
    for e in out_edges(EdgeKind::Calls)? {
        if let Some(name) = room.node(e.to).name {
            callees.entry(e.from).or_default().insert(name.to_string());
        }
    }
    let generated: BTreeSet<u64> = out_edges(EdgeKind::Expands)?.into_iter().map(|e| e.from).collect();

    let mut corpus = Corpus { callables: callables.len(), ..Default::default() };
    let mut files: BTreeMap<String, Option<Vec<String>>> = BTreeMap::new();
    for &o in callables.ordinals() {
        if generated.contains(&o) {
            corpus.macro_generated += 1;
            continue;
        }
        let n = room.node(o);
        let (Some(file), Some(start), Some(end)) = (n.file, n.start_line, n.end_line) else {
            corpus.without_span += 1;
            continue;
        };
        let (start, end) = (start as usize, end as usize);
        let lines = files.entry(file.to_string()).or_insert_with(|| std::fs::read_to_string(checkout.join(file)).ok().map(|t| t.lines().map(str::to_string).collect()));
        let Some(lines) = lines.as_ref().filter(|l| start >= 1 && end >= start && end <= l.len()) else {
            corpus.unreadable += 1;
            *corpus.unreadable_files.entry(file.to_string()).or_default() += 1;
            continue;
        };
        let span = end - start + 1;
        if span < min_lines || (max_lines != 0 && span > max_lines) {
            corpus.outside_size += 1;
            continue;
        }
        let container = n
            .parent
            .filter(|p| matches!(room.node(*p).kind, Some(nk::IMPL | nk::TRAIT)))
            .map(|p| room.display_path(p).rsplit("::").next().unwrap_or("").to_string())
            .unwrap_or_default();
        let mut ps = params.remove(&o).unwrap_or_default();
        ps.sort();
        let name = n.name.unwrap_or("").to_string();
        let mut types: Vec<&str> = ps.iter().map(|(_, t)| t.as_str()).collect();
        types.sort_unstable();
        corpus.fns.push(CorpusFn {
            file: file.to_string(),
            container,
            signature: format!("fn {name}({})", ps.iter().map(|(p, t)| format!("{p}: {t}")).collect::<Vec<_>>().join(", ")),
            signature_key: format!("fn({})", types.join(", ")),
            name,
            start_line: start,
            end_line: end,
            text: lines[start - 1..end].join("\n"),
            callees: callees.remove(&o).unwrap_or_default(),
        });
    }
    Ok(corpus)
}