src/vector/fnvec/room.rs
//! The function nodes of a code room, and the member half of the `embeddings`
//! feed: per consumed commit, the vectors of the functions whose text changed.
//!
//! ## Which nodes
//!
//! Node kinds `fn` and `method`, defined in the repo's own crates:
//!
//! | kind | what rustc calls it | why it is a chunk |
//! |---|---|---|
//! | `fn` | `DefKind::Fn` that is not a foreign item | a free function (a nested one is its own node) |
//! | `method` | `DefKind::AssocFn` | a method or associated fn of an `impl`, and a trait's own method |
//!
//! Not chunks: `ffi` (a declaration in an `extern` block has no body), closures
//! (no node of their own: their text is inside the function that holds them),
//! and a function a MACRO wrote — the item, or the `impl` / `trait` block it
//! sits in, has an EXPANDS edge. The driver positions such a node at the macro
//! invocation, so its "text" would be the invocation (`#[derive(Clone)]`,
//! `impl_int!(u8)`), shared by every function that invocation produced. Those
//! are counted ([`Census::macro_generated`]), never embedded. Neither is a
//! function in a file a build script generated into the target directory
//! ([`Census::outside_checkout`]): it is not source of the repo.
use std::collections::BTreeSet;
use std::path::{Component, Path};
use anyhow::Result;
use mimir_ingest::EventLog;
use mimir_room::{EdgeSel, MaskSpec, Room, Scope};
use mimir_schema::{EdgeKind, Origin, kind as nk};
use super::{At, FnNode, Run};
use crate::warehouse::iceberg::IcebergWarehouse;
/// The node kinds that are function chunks.
pub const KINDS: [&str; 2] = [nk::FN, nk::METHOD];
/// A function node of a room, with where the room holds it.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RoomFn {
pub ordinal: u64,
pub repo: String,
pub node: FnNode,
}
/// What the room holds of function-like nodes, and why some are no chunk.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct Census {
/// `fn` + `method` nodes defined in the repos' own crates.
pub callables: usize,
/// Written by a macro: positioned at the invocation, no text of their own.
pub macro_generated: usize,
/// No span row.
pub without_span: usize,
/// Positioned in a file that is not part of the checkout: code a build
/// script generated into the target directory (prost / bindgen output).
/// The driver records such a file by its absolute path.
pub outside_checkout: usize,
}
impl Census {
/// Function nodes that have no text of their own in the checkout — by
/// what they are, not because a feed did not run.
pub fn without_text(&self) -> usize {
self.macro_generated + self.without_span + self.outside_checkout
}
}
/// Is `file` (a path the driver wrote) inside the checkout it was dumped from?
pub fn in_checkout(file: &str) -> bool {
Path::new(file).components().all(|c| matches!(c, Component::Normal(_) | Component::CurDir))
}
/// The function nodes of `repos` that have a span of their own, and the census.
pub fn function_nodes(room: &Room, repos: &[String]) -> Result<(Vec<RoomFn>, Census)> {
let mut ords: Vec<u64> = Vec::new();
for k in KINDS {
let spec = MaskSpec { repos: repos.to_vec(), node_kinds: vec![k.to_string()], origins: vec![Origin::Local], defined_only: true, ..Default::default() };
ords.extend_from_slice(room.mask(&spec)?.value.ordinals());
}
ords.sort_unstable();
ords.dedup();
// the item, or the block it sits in, was written by a macro
let mut owners: BTreeSet<u64> = ords.iter().copied().collect();
owners.extend(ords.iter().filter_map(|o| room.node(*o).parent));
let owners = room.mask_of(&owners.into_iter().collect::<Vec<_>>());
let expands = room.edges(&EdgeSel { kinds: vec![EdgeKind::Expands], from: Some(&owners), to: None, cap: None, heaviest_first: false })?.value.edges;
let generated: BTreeSet<u64> = expands.into_iter().map(|e| e.from).collect();
let mut census = Census { callables: ords.len(), ..Default::default() };
let mut out = Vec::with_capacity(ords.len());
for o in ords {
let n = room.node(o);
if generated.contains(&o) || n.parent.is_some_and(|p| generated.contains(&p)) {
census.macro_generated += 1;
continue;
}
let (Some(file), Some(start), Some(end)) = (n.file, n.start_line, n.end_line) else {
census.without_span += 1;
continue;
};
if !in_checkout(file) {
census.outside_checkout += 1;
continue;
}
out.push(RoomFn {
ordinal: o,
repo: n.repo.unwrap_or("").to_string(),
node: FnNode { moniker: n.moniker.to_string(), file: file.to_string(), start_line: start, end_line: end },
});
}
Ok((out, census))
}
/// Read `file` (a path the driver wrote) under `root`. Only a path that stays
/// inside `root` is read.
pub fn read_under(root: &Path, file: &str) -> Option<String> {
in_checkout(file).then(|| std::fs::read_to_string(root.join(file)).ok()).flatten()
}
/// How the member half of the feed ended.
#[derive(Debug, Clone, PartialEq)]
pub enum Outcome {
Ran(Run),
/// This build cannot embed, and why. Recorded in the ledger.
Skipped(String),
}
/// **The member half of the `embeddings` feed.** `member` is caught up at
/// (`seq`, `sha`) and `checkout` is its tree at that commit: cut every function
/// node of the member out of the tree, embed the texts that have no vector
/// yet, write the link events of the functions that changed.
///
/// A build without an embedder answers [`Outcome::Skipped`] and says so in
/// the ledger; a run that fails (the budget gate, a model that does not load)
/// leaves its reason there too.
pub fn embed_member(wh: &IcebergWarehouse, log: &EventLog, member: &str, checkout: &Path, seq: i64, sha: &str) -> Result<Outcome> {
let at = At::member(member, seq, sha);
if let Some(why) = super::no_embedder() {
super::record_not_run(wh, &at, super::links::STATUS_SKIPPED, why)?;
return Ok(Outcome::Skipped(why.to_string()));
}
match embed_compiled(wh, log, member, checkout, seq, sha) {
Ok(run) => Ok(Outcome::Ran(run)),
Err(e) => {
// the reason stays readable next to the vectors
let _ = super::record_not_run(wh, &at, super::links::STATUS_ERROR, &format!("{e:#}"));
Err(e)
}
}
}
/// [`embed_member`] with the embedder given.
///
/// The room is built for this one question and not kept (it is one repo at one
/// ordinal; the workspace's own room stays in the cache).
pub fn embed_member_with(wh: &IcebergWarehouse, log: &EventLog, member: &str, checkout: &Path, seq: i64, sha: &str, embedding: &mut dyn super::Embedding) -> Result<Run> {
let scope = Scope::repo(member, Some(seq));
// on a thread of its own, on the lake's runtime: callable from any thread
let room = std::thread::scope(|s| s.spawn(|| log.block_on(Room::build_uncached(log, &scope))).join().unwrap_or_else(|p| std::panic::resume_unwind(p)))?;
let (nodes, census) = function_nodes(&room, &[member.to_string()])?;
drop(room);
let (chunks, unreadable) = super::chunks(nodes.into_iter().map(|n| n.node).collect(), |f| read_under(checkout, f));
let without = census.without_text() + unreadable.len();
let run = super::embed(wh, &At::member(member, seq, sha), &chunks, without, embedding)?;
eprintln!(
"nornir-embed: {member}@{seq}: {} (macro-generated {}, generated outside the checkout {}, no span {}, unreadable {})",
run.line(),
census.macro_generated,
census.outside_checkout,
census.without_span,
unreadable.len()
);
Ok(run)
}
#[cfg(any(feature = "embed-tract", feature = "embed-ort"))]
fn embed_compiled(wh: &IcebergWarehouse, log: &EventLog, member: &str, checkout: &Path, seq: i64, sha: &str) -> Result<Run> {
embed_member_with(wh, log, member, checkout, seq, sha, &mut super::Compiled::default())
}
#[cfg(not(any(feature = "embed-tract", feature = "embed-ort")))]
fn embed_compiled(_: &IcebergWarehouse, _: &EventLog, _: &str, _: &Path, _: i64, _: &str) -> Result<Run> {
anyhow::bail!("{}", super::no_embedder().unwrap_or("no embedder"))
}
/// What [`read_member`] cut out of the tree: the function chunks of a member
/// at one commit. Reading needs the checkout AT that commit; embedding what
/// was read ([`embed_read`]) needs the tree no more — the two halves the
/// consumer's deferred feed runs on either side of its hold on the tree.
pub struct MemberRead {
at: At,
chunks: Vec<super::FnChunk>,
without: usize,
}
/// The READING half of [`embed_member`]: the room of `member` at `seq` and the
/// text of every function node, out of `checkout`. No model is loaded.
pub fn read_member(log: &EventLog, member: &str, checkout: &Path, seq: i64, sha: &str) -> Result<MemberRead> {
let scope = Scope::repo(member, Some(seq));
// on a thread of its own, on the lake's runtime: callable from any thread
let room = std::thread::scope(|s| s.spawn(|| log.block_on(Room::build_uncached(log, &scope))).join().unwrap_or_else(|p| std::panic::resume_unwind(p)))?;
let (nodes, census) = function_nodes(&room, &[member.to_string()])?;
drop(room);
let (chunks, unreadable) = super::chunks(nodes.into_iter().map(|n| n.node).collect(), |f| read_under(checkout, f));
Ok(MemberRead { at: At::member(member, seq, sha), without: census.without_text() + unreadable.len(), chunks })
}
/// The EMBEDDING half of [`embed_member`], over what [`read_member`] read:
/// the same vectors, link events and ledger row, with the tree not looked at.
pub fn embed_read(wh: &IcebergWarehouse, read: MemberRead) -> Result<Outcome> {
if let Some(why) = super::no_embedder() {
super::record_not_run(wh, &read.at, super::links::STATUS_SKIPPED, why)?;
return Ok(Outcome::Skipped(why.to_string()));
}
match embed_read_compiled(wh, &read) {
Ok(run) => Ok(Outcome::Ran(run)),
Err(e) => {
let _ = super::record_not_run(wh, &read.at, super::links::STATUS_ERROR, &format!("{e:#}"));
Err(e)
}
}
}
#[cfg(any(feature = "embed-tract", feature = "embed-ort"))]
fn embed_read_compiled(wh: &IcebergWarehouse, read: &MemberRead) -> Result<Run> {
super::embed(wh, &read.at, &read.chunks, read.without, &mut super::Compiled::default())
}
#[cfg(not(any(feature = "embed-tract", feature = "embed-ort")))]
fn embed_read_compiled(_: &IcebergWarehouse, _: &MemberRead) -> Result<Run> {
anyhow::bail!("{}", super::no_embedder().unwrap_or("no embedder"))
}
/// One member's line of an on-demand run.
#[derive(Debug, Clone, PartialEq)]
pub enum MemberRun {
Ran(Run),
/// Not run, and why (this build cannot embed; the checkout is not at the
/// commit the room is of).
NotRun(String),
Failed(String),
}
/// **The on-demand door** (`nornir workspace embed`): the function half of the
/// feed for every consumed member, now — no fetch, no consume. A member is
/// embedded only when its checkout under `git_dir` is AT the commit the room
/// holds: the spans are that commit's, and cutting another tree by them would
/// hand a function the wrong lines.
pub fn embed_consumed(code_room: &crate::code_room::CodeRoom, wh: &IcebergWarehouse, git_dir: &Path) -> std::result::Result<Vec<(String, MemberRun)>, crate::code_room::Refusal> {
use crate::code_room::{Refusal, State};
let log = code_room.log()?;
let offsets = std::thread::scope(|s| s.spawn(|| log.block_on(log.offsets())).join().unwrap_or_else(|p| std::panic::resume_unwind(p)))
.map_err(|e| Refusal::new(State::Internal, format!("code room of `{}`: {e:#}", code_room.workspace())))?;
if offsets.is_empty() {
return Err(Refusal::new(State::NotConsumed, format!("code room of `{}`: no repo has been consumed yet — there is no function node to embed", code_room.workspace())));
}
let mut out = Vec::new();
for (member, at) in offsets {
let checkout = git_dir.join(&member);
let head = crate::gitio::head_sha(&checkout).unwrap_or_default();
let run = if head != at.sha {
MemberRun::NotRun(format!(
"its checkout {} is at {}, the code room at seq {} ({}): consume first (`nornir workspace populate {}`)",
checkout.display(),
if head.is_empty() { "no readable HEAD".to_string() } else { head[..head.len().min(10)].to_string() },
at.seq,
&at.sha[..at.sha.len().min(10)],
code_room.workspace()
))
} else {
match embed_member(wh, &log, &member, &checkout, at.seq, &at.sha) {
Ok(Outcome::Ran(run)) => MemberRun::Ran(run),
Ok(Outcome::Skipped(why)) => MemberRun::NotRun(why),
Err(e) => MemberRun::Failed(format!("{e:#}")),
}
};
out.push((member, run));
}
Ok(out)
}
/// The dependency ring, now and in this process (`nornir workspace embed <ws>
/// --ring`): the same plan and run the background job does, with the crates a
/// former run skipped tried again. Returns the plan's line and the report.
#[cfg(any(feature = "embed-tract", feature = "embed-ort"))]
pub fn embed_ring(code_room: &crate::code_room::CodeRoom, wh: &IcebergWarehouse) -> Result<(String, super::ring::Report)> {
use super::Embedding as _;
let log = code_room.log().map_err(|r| anyhow::anyhow!("[{}] {r}", r.state))?;
let mut embedding = super::Compiled::default();
let (dirs, unasked) = super::ring::crate_dirs(&log, &code_room.layout().base().join("clones"));
for why in &unasked {
eprintln!("nornir-embed: {}: ring sources: {why}", code_room.workspace());
}
let plan = super::ring::plan(wh, &log, &dirs, &embedding.profile().id(), true)?;
let line = plan.line();
eprintln!("nornir-embed: {}: {line}", code_room.workspace());
let ws = code_room.workspace().to_string();
let rep = super::ring::run(wh, &plan, &mut embedding, &mut |krate, run| eprintln!("nornir-embed: {ws}: ring {krate}: {}", run.line()))?;
Ok((line, rep))
}
/// Without an embedder there is no ring run.
#[cfg(not(any(feature = "embed-tract", feature = "embed-ort")))]
pub fn embed_ring(_: &crate::code_room::CodeRoom, _: &IcebergWarehouse) -> Result<(String, super::ring::Report)> {
anyhow::bail!("{}", super::no_embedder().unwrap_or("no embedder"))
}