src/vector/fnvec.rs
//! **One vector per function node.**
//!
//! A chunk here is the source text of ONE function-like node of the code room
//! — a free `fn`, a method, an associated fn — cut by the span the rustc driver
//! recorded for that node at that commit. Never a line window, never a parser
//! of ours: where a function starts and ends is the compiler's answer
//! (`tests/no_second_parser.rs`).
//!
//! ```text
//! node id (SCIP moniker) ──▶ content_hash ──▶ vector
//! embedding_node_events embeddings (model_profile)
//! ```
//!
//! * The vector lives where every vector lives: `embeddings`, keyed
//! `(content_hash, model_profile)`. Two functions with the same text share
//! one row — a body is embedded ONCE, in any number of crates and commits.
//! * The link node → `content_hash` is an append-only change log in the same
//! catalog ([`links`]): per (repo, moniker) the last event up to a consumer
//! ordinal is the state. A commit writes events for the functions whose text
//! changed and embeds only content that has no vector yet.
//! * A run writes one ledger row (`embedding_node_runs`): which state the
//! vectors are of, in which model space, on which backend — or why it did
//! not embed. A reader needs no embedder to know what it is reading.
//!
//! The line-window chunks ([`super::chunk`]) are untouched: they stay the unit
//! of free-text search (`vector search`, `Funnel.Retrieve`) and of every file
//! that is not Rust.
use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet};
use anyhow::Result;
use crate::warehouse::iceberg::IcebergWarehouse;
pub mod links;
#[cfg(feature = "code-room")]
pub mod ring;
#[cfg(feature = "code-room")]
pub mod room;
/// The cut rule, part of every link row: a change of the rule is a new link
/// set (the old one is not read), never a silent mix of two cuts.
///
/// `fn-v1`: the lines `start_line..=end_line` of the node's driver span, the
/// indentation all of them share removed, trailing whitespace trimmed, joined
/// with `\n`. The indentation goes so that the same body nested one level
/// deeper (a method, a fn inside a `mod`) is the same content.
///
/// The vector of a text longer than [`PIECE_LINES`] lines or [`PIECE_BYTES`]
/// bytes is the length-weighted mean of the vectors of its [`pieces`],
/// normalized — still ONE vector per function, keyed by the hash of the whole
/// text.
pub const CHUNKER: &str = "fn-v1";
/// The most lines one forward pass is handed (the line-window chunker's own
/// size: comfortably inside the model's context).
pub const PIECE_LINES: usize = 120;
/// The most bytes one forward pass is handed. Lines are not always short (a
/// generated table, a minified literal), and a transformer's cost and memory
/// grow with the SQUARE of the token count: MEASURED on znippy-zoomies
/// (2026-10-13, tract CPU), functions embedded whole made one forward hold
/// 9.5 GB and run for more than twenty minutes on one core.
pub const PIECE_BYTES: usize = 8_000;
/// Cut `text` into the pieces it is embedded as: runs of whole lines, each at
/// most [`PIECE_LINES`] lines and [`PIECE_BYTES`] bytes; a single line longer
/// than that is cut at character boundaries. A text within both limits is one
/// piece, itself.
pub fn pieces(text: &str) -> Vec<&str> {
let mut out: Vec<&str> = Vec::new();
let (mut start, mut lines) = (0usize, 0usize);
let mut at = 0usize;
for line in text.split_inclusive('\n') {
let end = at + line.len();
if lines > 0 && (lines == PIECE_LINES || end - start > PIECE_BYTES) {
out.push(text[start..at].trim_end_matches('\n'));
(start, lines) = (at, 0);
}
if line.len() > PIECE_BYTES {
// one line over the limit: cut it, on character boundaries
let mut from = at;
while end - from > PIECE_BYTES {
let mut to = from + PIECE_BYTES;
while !text.is_char_boundary(to) {
to -= 1;
}
out.push(&text[from..to]);
from = to;
}
(start, lines) = (from, 0);
}
lines += 1;
at = end;
}
if start < text.len() {
out.push(text[start..].trim_end_matches('\n'));
}
out.retain(|p| !p.trim().is_empty());
out
}
/// The vector of a text from the vectors of its pieces: the mean weighted by
/// piece length, L2-normalized. One piece is its own vector, untouched.
pub fn pool(vectors: &[Vec<f32>], weights: &[usize]) -> Vec<f32> {
if vectors.len() == 1 {
return vectors[0].clone();
}
let dim = vectors.first().map_or(0, Vec::len);
let mut out = vec![0f32; dim];
for (v, w) in vectors.iter().zip(weights) {
for (o, x) in out.iter_mut().zip(v) {
*o += *x * *w as f32;
}
}
let norm = out.iter().map(|x| x * x).sum::<f32>().sqrt();
if norm > 0.0 {
out.iter_mut().for_each(|x| *x /= norm);
}
out
}
/// `repo` value of the dependency ring's rows (the code room's own name for it).
pub const RING: &str = "@universe";
/// The feed that writes the member rows, as answers name it.
pub const FEED: &str = "embeddings";
/// One function-like node, positioned.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct FnNode {
/// The node id (SCIP moniker).
pub moniker: String,
/// The file its span lies in, as the driver wrote it (relative to the checkout).
pub file: String,
pub start_line: u32,
pub end_line: u32,
}
/// A function node with its text.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct FnChunk {
pub node: FnNode,
pub text: String,
/// SHA-256 of [`Self::text`] — the key of the vector.
pub content_hash: String,
}
/// Cut `start..=end` (1-based, inclusive) out of a file's lines. `None` when
/// the span does not lie in the file or covers only whitespace.
pub fn cut(lines: &[&str], start: u32, end: u32) -> Option<String> {
let (s, e) = (start as usize, end as usize);
if s == 0 || e < s || e > lines.len() {
return None;
}
let body = &lines[s - 1..e];
let indent = body
.iter()
.filter(|l| !l.trim().is_empty())
.map(|l| l.len() - l.trim_start_matches([' ', '\t']).len())
.min()?;
let text = body.iter().map(|l| if l.trim().is_empty() { "" } else { l[indent..].trim_end() }).collect::<Vec<_>>().join("\n");
let text = text.trim_matches('\n');
(!text.is_empty()).then(|| text.to_string())
}
/// Cut every node of `nodes`. `read` hands back a file's text by the path the
/// driver wrote (`None` = not there). Nodes whose file is missing or shorter
/// than the span says come back in the second list.
pub fn chunks(nodes: Vec<FnNode>, mut read: impl FnMut(&str) -> Option<String>) -> (Vec<FnChunk>, Vec<FnNode>) {
let mut files: HashMap<String, Option<String>> = HashMap::new();
let (mut out, mut unreadable) = (Vec::with_capacity(nodes.len()), Vec::new());
for node in nodes {
let text = files.entry(node.file.clone()).or_insert_with(|| read(&node.file));
let body = text.as_deref().and_then(|t| cut(&t.lines().collect::<Vec<_>>(), node.start_line, node.end_line));
match body {
Some(text) => out.push(FnChunk { content_hash: super::chunk::hash_text(&text), text, node }),
None => unreadable.push(node),
}
}
(out, unreadable)
}
/// One embedding run at a time in a process. A consume's `embeddings` feed
/// (both halves) and the background ring job ([`ring`]) would otherwise load
/// two models and fan two forward passes across the same cores or the same
/// GPU. The feed holds the turn for its run; the ring job takes it per crate,
/// so a feed waits for at most one crate.
pub fn embed_turn() -> std::sync::MutexGuard<'static, ()> {
static TURN: std::sync::Mutex<()> = std::sync::Mutex::new(());
TURN.lock().unwrap_or_else(|e| e.into_inner())
}
/// Where a run's rows belong: a member at a consumer ordinal, or one crate of
/// the dependency ring.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct At {
/// The member, or [`RING`].
pub repo: String,
/// `""` for a member; `crate@version` for a ring crate.
pub unit: String,
pub seq: i64,
pub sha: String,
}
impl At {
pub fn member(repo: &str, seq: i64, sha: &str) -> Self {
Self { repo: repo.to_string(), unit: String::new(), seq, sha: sha.to_string() }
}
pub fn ring(krate: &str) -> Self {
Self { repo: RING.to_string(), unit: krate.to_string(), seq: 0, sha: String::new() }
}
}
/// What one run did.
#[derive(Debug, Clone, Default, PartialEq)]
pub struct Run {
/// Function nodes that have a chunk at this state.
pub functions: usize,
/// Distinct texts among them.
pub distinct_bodies: usize,
/// Vectors this run embedded and wrote.
pub new_vectors: usize,
/// Link events written: nodes new or changed (`+`), nodes gone (`-`).
pub plus: usize,
pub minus: usize,
/// Function nodes without a chunk (macro-generated, no span, unreadable).
pub without_chunk: usize,
/// Forward-pass inputs behind those vectors (a long function is several pieces).
pub pieces: usize,
/// Wall clock of the embed calls alone.
pub embed_ms: u64,
pub ms: u64,
pub backend: String,
}
impl Run {
/// Functions per second of the embed calls (`None` when nothing was embedded).
pub fn rate(&self) -> Option<f64> {
(self.new_vectors > 0 && self.embed_ms > 0).then(|| self.new_vectors as f64 * 1000.0 / self.embed_ms as f64)
}
pub fn line(&self) -> String {
let rate = self
.rate()
.map(|r| format!(" ({} piece(s)) at {r:.2} function/s in {:.1}s on {}", self.pieces, self.embed_ms as f64 / 1000.0, self.backend))
.unwrap_or_default();
format!(
"{} function(s), {} distinct bodies, {} new vector(s){rate}; +{} −{} link(s); {} without a chunk",
self.functions, self.distinct_bodies, self.new_vectors, self.plus, self.minus, self.without_chunk
)
}
}
/// The delta of `chunks` against the open link state of their scope.
#[derive(Debug, Default)]
pub struct Delta<'a> {
/// Nodes whose text hashes to something else than the state says (or that are new).
pub plus: Vec<&'a FnChunk>,
/// Nodes of the state that are no longer function nodes with a chunk.
pub minus: Vec<String>,
}
/// `state` = the open links of the scope (moniker → content hash).
pub fn delta<'a>(state: &BTreeMap<String, String>, chunks: &'a [FnChunk]) -> Delta<'a> {
let now: BTreeSet<&str> = chunks.iter().map(|c| c.node.moniker.as_str()).collect();
Delta {
plus: chunks.iter().filter(|c| state.get(&c.node.moniker) != Some(&c.content_hash)).collect(),
minus: state.keys().filter(|m| !now.contains(m.as_str())).cloned().collect(),
}
}
/// The distinct texts of `chunks` that have no vector in `stored` yet, first
/// occurrence first.
pub fn missing<'a>(chunks: &'a [FnChunk], stored: &HashSet<String>) -> Vec<&'a FnChunk> {
let mut seen: HashSet<&str> = HashSet::new();
chunks.iter().filter(|c| !stored.contains(&c.content_hash) && seen.insert(c.content_hash.as_str())).collect()
}
/// What a run needs to turn text into vectors, loaded only when a text has no
/// vector yet (a commit that changes no function text loads no model).
pub trait Embedding {
/// The model space the vectors are stored under — known without loading.
fn profile(&self) -> super::store::ModelProfile;
fn model_id(&self) -> String;
fn backend(&self) -> String;
/// Embed `texts`, in order. Called once per slice.
fn embed(&mut self, texts: &[&str]) -> Result<Vec<Vec<f32>>>;
/// Price the whole job before the first slice (the budget gate).
fn price(&mut self, _texts: &[&str]) -> Result<()> {
Ok(())
}
/// [`Self::price`] for a job too large to hold as texts: `sample` is a
/// spread of the texts, `total` how many there are.
fn price_sampled(&mut self, _sample: &[&str], _total: usize) -> Result<()> {
Ok(())
}
}
/// The compiled embedder: the profile comes from the build, the model is
/// loaded on the first text that needs a vector.
#[cfg(any(feature = "embed-tract", feature = "embed-ort"))]
#[derive(Default)]
pub struct Compiled {
loaded: Option<Box<dyn super::store::Embedder>>,
}
#[cfg(any(feature = "embed-tract", feature = "embed-ort"))]
impl Compiled {
fn get(&mut self) -> Result<&dyn super::store::Embedder> {
if self.loaded.is_none() {
self.loaded = Some(super::load_embedder()?);
}
Ok(&**self.loaded.as_ref().expect("just loaded"))
}
}
#[cfg(any(feature = "embed-tract", feature = "embed-ort"))]
impl Embedding for Compiled {
fn profile(&self) -> super::store::ModelProfile {
match &self.loaded {
Some(e) => e.profile(),
None => super::embed_support::profile(),
}
}
fn model_id(&self) -> String {
super::selected_model_id().to_string()
}
fn backend(&self) -> String {
super::embedder_backend().to_string()
}
fn embed(&mut self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
self.get()?.embed(texts)
}
fn price(&mut self, texts: &[&str]) -> Result<()> {
if let Some(eta) = super::store::check_embed_budget(self.get()?, texts)? {
eprintln!("nornir-embed: {}", eta.summary());
}
Ok(())
}
fn price_sampled(&mut self, sample: &[&str], total: usize) -> Result<()> {
let Some(budget) = super::store::embed_budget()? else { return Ok(()) };
if sample.len() < 2 {
return Ok(());
}
let backend = self.backend();
// a job that cannot be priced is let through (it fails on its own next line if the backend is broken)
let Ok(eta) = self.get()?.estimate(sample, total) else { return Ok(()) };
eprintln!("nornir-embed: {}", eta.summary());
anyhow::ensure!(
eta.fits(budget),
"embedding {total} function bodies would take ~{:.1} h at {:.2} chunk/s on {backend}, over the {:.0} min budget. \
Nothing was embedded. Run it where the GPU provider is, or accept the cost with {}=off (or e.g. {}=48h)",
eta.eta.as_secs_f64() / 3600.0,
eta.rate,
budget.as_secs_f64() / 60.0,
super::store::NORNIR_EMBED_BUDGET_ENV,
super::store::NORNIR_EMBED_BUDGET_ENV,
);
Ok(())
}
}
/// Any [`super::store::Embedder`] as an [`Embedding`] (tests, benches).
pub struct Loaded<'a, E: super::store::Embedder + ?Sized> {
pub embedder: &'a E,
pub backend: &'a str,
}
impl<E: super::store::Embedder + ?Sized> Embedding for Loaded<'_, E> {
fn profile(&self) -> super::store::ModelProfile {
self.embedder.profile()
}
fn model_id(&self) -> String {
self.embedder.profile().model_name
}
fn backend(&self) -> String {
self.backend.to_string()
}
fn embed(&mut self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
self.embedder.embed(texts)
}
}
/// Why this BINARY cannot embed, or `None` when it can.
pub fn no_embedder() -> Option<&'static str> {
if cfg!(any(feature = "embed-tract", feature = "embed-ort")) { None } else { Some("built without embedder (embed-tract/embed-ort)") }
}
/// Bring the scope `at` to `chunks`: embed the texts that have no vector yet
/// (deduped by content hash, sliced, each slice durable before the next),
/// then write the link events of the nodes that changed, then the ledger row.
///
/// The order is what makes a crash harmless: vectors without links are
/// content-addressed and reused by the next run; a link is written only when
/// its vector is. `without_chunk` = function nodes of the scope that have no
/// chunk (counted in the ledger, never linked).
pub fn embed(wh: &IcebergWarehouse, at: &At, chunks: &[FnChunk], without_chunk: usize, embedding: &mut dyn Embedding) -> Result<Run> {
// the delta is against the newest state written, whatever its ordinal
let state = links::open(&links::read_events_of(wh, &[at.repo.as_str()])?, &at.repo, i64::MAX);
let mut stored = wh.block_on(super::store::existing_content_hashes_in(wh, vector_table(at), &embedding.profile().id()))?;
embed_against(wh, at, chunks, without_chunk, embedding, &state, &mut stored)
}
/// The table a scope's vectors live in: `embeddings` for a member's own
/// functions, a table of its own for the dependency ring. A ring is 10^5–10^6
/// bodies, and every reader of `embeddings` (the lenses, free-text search)
/// scans the model's rows: the ring must not sit among a workspace's vectors.
pub fn vector_table(at: &At) -> &'static str {
if at.repo == RING { crate::warehouse::iceberg::TABLE_EMBEDDINGS_RING } else { crate::warehouse::iceberg::TABLE_EMBEDDINGS }
}
/// [`embed`] for a caller that already holds the scope's open links (`state`)
/// and the content hashes stored for the model in the scope's vector table
/// ([`vector_table`]; `stored` is extended by what this run writes) — the ring
/// job, which runs once per crate over one read.
pub fn embed_against(
wh: &IcebergWarehouse,
at: &At,
chunks: &[FnChunk],
without_chunk: usize,
embedding: &mut dyn Embedding,
state: &BTreeMap<String, String>,
stored: &mut HashSet<String>,
) -> Result<Run> {
let t0 = std::time::Instant::now();
let model = embedding.profile();
let profile = model.id();
let d = delta(state, chunks);
let todo = missing(chunks, stored);
let mut run = Run {
functions: chunks.len(),
distinct_bodies: chunks.iter().map(|c| c.content_hash.as_str()).collect::<HashSet<_>>().len(),
without_chunk,
backend: embedding.backend(),
..Default::default()
};
if !todo.is_empty() {
// priced by what a forward pass is handed: the pieces
let all: Vec<&str> = todo.iter().flat_map(|c| pieces(&c.text)).collect();
embedding.price(&all)?;
drop(all);
let total = todo.len();
for slice in todo.chunks(super::store::embed_slice_chunks()) {
let cut: Vec<Vec<&str>> = slice.iter().map(|c| pieces(&c.text)).collect();
let texts: Vec<&str> = cut.iter().flatten().copied().collect();
let t = std::time::Instant::now();
let embedded = embedding.embed(&texts)?;
run.embed_ms += t.elapsed().as_millis() as u64;
anyhow::ensure!(embedded.len() == texts.len(), "embedder returned {} vectors for {} text(s)", embedded.len(), texts.len());
let mut vectors: Vec<Vec<f32>> = Vec::with_capacity(slice.len());
let mut next = 0usize;
for (c, parts) in slice.iter().zip(&cut) {
anyhow::ensure!(!parts.is_empty(), "{} has no text to embed", c.node.moniker);
let of = &embedded[next..next + parts.len()];
next += parts.len();
for v in of {
anyhow::ensure!(v.len() == model.dim, "embedding of {} has dim {} != model dim {}", c.node.moniker, v.len(), model.dim);
}
vectors.push(pool(of, &parts.iter().map(|p| p.len()).collect::<Vec<_>>()));
}
let rows: Vec<(&str, &[f32])> = slice.iter().zip(&vectors).map(|(c, v)| (c.content_hash.as_str(), v.as_slice())).collect();
wh.block_on(super::store::append_embeddings_to(wh, vector_table(at), &profile, model.dim, model.normalize, &rows))?;
run.new_vectors += slice.len();
run.pieces += texts.len();
stored.extend(slice.iter().map(|c| c.content_hash.clone()));
if total > slice.len() {
let secs = run.embed_ms as f64 / 1000.0;
let rate = if secs > 0.0 { run.new_vectors as f64 / secs } else { 0.0 };
eprintln!(
"nornir-embed: {}{}: {}/{total} function(s) · {rate:.2} function/s · elapsed {secs:.0}s · eta {:.0}s",
at.repo,
if at.unit.is_empty() { String::new() } else { format!(" {}", at.unit) },
run.new_vectors,
if rate > 0.0 { (total - run.new_vectors) as f64 / rate } else { f64::NAN },
);
}
}
}
run.plus = d.plus.len();
run.minus = d.minus.len();
let mut events: Vec<links::Event> = d.plus.iter().map(|c| links::Event::plus(at, &c.node.moniker, &c.content_hash)).collect();
events.extend(d.minus.iter().map(|m| links::Event::minus(at, m)));
links::append_events(wh, &events)?;
run.ms = t0.elapsed().as_millis() as u64;
links::append_run(wh, &links::RunRow::of(at, &profile, &embedding.model_id(), &run, links::STATUS_OK, None))?;
Ok(run)
}
/// Record that a run did not embed, and why (no embedder in this build, the
/// budget gate, a failure) — so a reader can say it instead of "no vectors".
pub fn record_not_run(wh: &IcebergWarehouse, at: &At, status: &str, why: &str) -> Result<()> {
let run = Run::default();
links::append_run(wh, &links::RunRow::of(at, "", "", &run, status, Some(why)))
}
/// What a reader gets: per node its vector, and what the vectors are of.
#[derive(Debug, Default)]
pub struct Vectors {
/// moniker → content hash, per repo (the open links at the ordinal asked for).
pub links: BTreeMap<String, BTreeMap<String, String>>,
/// content hash → vector, in [`Self::model_profile`]'s space.
pub by_hash: HashMap<String, Vec<f32>>,
pub dim: usize,
/// The model space read (`""` = no run has ever embedded).
pub model_profile: String,
pub model_id: String,
/// The newest run row of each repo (member rows only).
pub last_run: BTreeMap<String, links::RunRow>,
/// Link events / ledger rows read.
pub events: usize,
pub runs: usize,
}
impl Vectors {
/// The vector of a node, when its link has one.
pub fn of(&self, repo: &str, moniker: &str) -> Option<(&str, &[f32])> {
let hash = self.links.get(repo)?.get(moniker)?;
self.by_hash.get(hash).map(|v| (hash.as_str(), v.as_slice()))
}
}
/// Read the links of `repos` (each up to its ordinal) and the vectors they
/// point at. The model space is the one the newest green run wrote: a reader
/// has no embedder to ask, and must never mix two spaces.
/// [`read`], kept by the process: the last answer beside the snapshots of
/// the three tables it reads (link events, run ledger, vectors) and the
/// repos asked for. An unchanged stamp IS the read (the live workspace:
/// 1.8 GB of link events and 2.6 GB of runs, 12-21 s per ask, read by
/// similar_code and semantic_twins on every ask); any commit to one of the
/// tables reads again. One entry: the newest ask's.
pub fn read_kept(wh: &IcebergWarehouse, repos: &BTreeMap<String, i64>) -> Result<std::sync::Arc<Vectors>> {
type Key = (std::path::PathBuf, BTreeMap<String, i64>, [Option<i64>; 3]);
static KEPT: std::sync::Mutex<Option<(Key, std::sync::Arc<Vectors>)>> = std::sync::Mutex::new(None);
use iceberg::Catalog as _;
let stamp = wh.block_on(async {
let mut out = [None; 3];
for (i, name) in [crate::warehouse::iceberg::TABLE_EMBEDDING_NODE_EVENTS, crate::warehouse::iceberg::TABLE_EMBEDDING_NODE_RUNS, crate::warehouse::iceberg::TABLE_EMBEDDINGS].into_iter().enumerate() {
let ident = wh.table_ident(name);
if wh.catalog().table_exists(&ident).await? {
out[i] = wh.catalog().load_table(&ident).await?.metadata().current_snapshot().map(|s| s.snapshot_id());
}
}
anyhow::Ok(out)
})?;
let key: Key = (wh.root().to_path_buf(), repos.clone(), stamp);
if let Ok(k) = KEPT.lock()
&& let Some((kk, v)) = k.as_ref()
&& *kk == key
{
return Ok(v.clone());
}
let v = std::sync::Arc::new(read(wh, repos)?);
if let Ok(mut k) = KEPT.lock() {
*k = Some((key, v.clone()));
}
Ok(v)
}
pub fn read(wh: &IcebergWarehouse, repos: &BTreeMap<String, i64>) -> Result<Vectors> {
let events = links::read_events_of(wh, &repos.keys().map(String::as_str).collect::<Vec<_>>())?;
let runs = links::read_runs(wh)?;
let mut out = Vectors { events: events.len(), runs: runs.len(), ..Default::default() };
for r in &runs {
if r.unit.is_empty() && repos.contains_key(&r.repo) {
let newer = out.last_run.get(&r.repo).is_none_or(|p| (p.seq, p.ts_micros) <= (r.seq, r.ts_micros));
if newer {
out.last_run.insert(r.repo.clone(), r.clone());
}
}
}
if let Some(r) = runs.iter().filter(|r| r.status == links::STATUS_OK && !r.model_profile.is_empty() && (repos.contains_key(&r.repo) || r.repo == RING)).max_by_key(|r| r.ts_micros) {
out.model_profile = r.model_profile.clone();
out.model_id = r.model_id.clone();
}
for (repo, seq) in repos {
out.links.insert(repo.clone(), links::open(&events, repo, *seq));
}
if out.model_profile.is_empty() {
return Ok(out);
}
let needed: HashSet<&str> = out.links.values().flat_map(|m| m.values().map(String::as_str)).collect();
if !needed.is_empty() {
let (by_hash, dim) = wh.block_on(super::store::read_vectors_present(wh, &out.model_profile, &needed))?;
out.by_hash = by_hash;
out.dim = dim;
}
Ok(out)
}
#[cfg(test)]
mod tests {
use super::*;
/// The kept read is the read: an empty warehouse answers the same twice,
/// the second from the kept entry (the same Arc).
#[test]
fn the_kept_vector_read_is_the_read_until_a_table_moves() {
let dir = tempfile::tempdir().unwrap();
let wh = IcebergWarehouse::open(dir.path()).unwrap();
let repos: BTreeMap<String, i64> = [("r".to_string(), 1)].into();
let a = read_kept(&wh, &repos).unwrap();
let b = read_kept(&wh, &repos).unwrap();
assert!(std::sync::Arc::ptr_eq(&a, &b), "an unchanged stamp: the kept answer");
assert_eq!((a.events, a.runs, a.dim), (read(&wh, &repos).unwrap().events, 0, 0));
let other: BTreeMap<String, i64> = [("r".to_string(), 2)].into();
assert!(!std::sync::Arc::ptr_eq(&a, &read_kept(&wh, &other).unwrap()), "another ask reads again");
}
fn node(m: &str, file: &str, s: u32, e: u32) -> FnNode {
FnNode { moniker: m.into(), file: file.into(), start_line: s, end_line: e }
}
#[test]
fn a_cut_is_the_spans_lines_without_the_shared_indent() {
let src = "impl T {\n fn a(&self) -> u8 {\n 1\n }\n\n fn b() {} \n}\n";
let lines: Vec<&str> = src.lines().collect();
assert_eq!(cut(&lines, 2, 4).as_deref(), Some("fn a(&self) -> u8 {\n 1\n}"));
assert_eq!(cut(&lines, 6, 6).as_deref(), Some("fn b() {}"));
// a span that is not in the file, or holds no text, is no chunk
assert_eq!(cut(&lines, 6, 9), None);
assert_eq!(cut(&lines, 0, 2), None);
assert_eq!(cut(&lines, 5, 5), None);
}
/// The same body at another nesting depth is the same content: ONE vector.
#[test]
fn the_same_body_nested_deeper_has_the_same_content_hash() {
let free = "fn total(xs: &[u32]) -> u32 {\n xs.iter().sum()\n}\n";
let method = "impl Cart {\n fn total(xs: &[u32]) -> u32 {\n xs.iter().sum()\n }\n}\n";
let (cs, bad) = chunks(vec![node("a", "free.rs", 1, 3), node("b", "method.rs", 2, 4)], |f| Some(if f == "free.rs" { free } else { method }.to_string()));
assert!(bad.is_empty());
assert_eq!(cs[0].content_hash, cs[1].content_hash);
assert_eq!(missing(&cs, &HashSet::new()).len(), 1, "identical bodies are embedded once");
}
#[test]
fn a_node_whose_file_is_gone_or_too_short_is_reported_not_cut() {
let (cs, bad) = chunks(vec![node("a", "gone.rs", 1, 2), node("b", "short.rs", 3, 9), node("c", "short.rs", 1, 1)], |f| (f == "short.rs").then(|| "fn c() {}\n".to_string()));
assert_eq!(cs.iter().map(|c| c.node.moniker.as_str()).collect::<Vec<_>>(), ["c"]);
assert_eq!(bad.iter().map(|n| n.moniker.as_str()).collect::<Vec<_>>(), ["a", "b"]);
}
/// A text within the limits is one piece (itself); a long one is cut into
/// runs of whole lines; one over-long line is cut on character boundaries.
/// Nothing is lost and no piece is over a limit.
#[test]
fn a_long_text_is_cut_into_pieces_within_the_limits_and_nothing_is_lost() {
let short = "fn a() {\n 1\n}";
assert_eq!(pieces(short), [short]);
let long: String = (0..300).map(|i| format!(" let x{i} = {i};\n")).collect();
let cut = pieces(long.trim_end());
assert_eq!(cut.len(), 3);
assert!(cut.iter().all(|p| p.lines().count() <= PIECE_LINES && p.len() <= PIECE_BYTES));
assert_eq!(cut.join("\n"), long.trim_end(), "the pieces are the text");
// many bytes on few lines: the byte limit cuts first
let wide: String = (0..40).map(|i| format!("{}{i}\n", "é".repeat(400))).collect();
let cut = pieces(&wide);
assert!(cut.len() > 1 && cut.iter().all(|p| p.len() <= PIECE_BYTES), "{:?}", cut.iter().map(|p| p.len()).collect::<Vec<_>>());
assert_eq!(cut.join("\n") + "\n", wide);
// ONE line over the limit
let line = "é".repeat(9_001);
let cut = pieces(&line);
assert!(cut.len() >= 3 && cut.iter().all(|p| p.len() <= PIECE_BYTES));
assert_eq!(cut.concat(), line);
}
#[test]
fn a_pooled_vector_is_the_length_weighted_mean_normalized() {
assert_eq!(pool(&[vec![0.6, 0.8]], &[5]), [0.6, 0.8], "one piece is its own vector");
let v = pool(&[vec![1.0, 0.0], vec![0.0, 1.0]], &[3, 1]);
let norm = (v[0] * v[0] + v[1] * v[1]).sqrt();
assert!((norm - 1.0).abs() < 1e-6 && (v[0] / v[1] - 3.0).abs() < 1e-5, "{v:?}");
}
/// Vector = `[bytes, lines, 1]`, normalized; records what it was asked for.
#[derive(Default)]
struct Shape(std::sync::Mutex<Vec<usize>>);
impl crate::vector::store::Embedder for Shape {
fn profile(&self) -> crate::vector::store::ModelProfile {
crate::vector::store::ModelProfile { model_name: "test-shape".into(), weights_sha: "w".into(), tokenizer_sha: "t".into(), pooling: "none".into(), normalize: true, dim: 3, dtype: "f32".into() }
}
fn embed(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
self.0.lock().unwrap().extend(texts.iter().map(|t| t.lines().count()));
Ok(texts.iter().map(|t| pool(&[vec![t.len() as f32, t.lines().count() as f32, 1.0], vec![0.0; 3]], &[1, 0])).collect())
}
}
/// A function longer than a forward pass may be is embedded piece by piece
/// — no forward is handed more than [`PIECE_LINES`] lines — and still gets
/// ONE vector, under the hash of its whole text.
#[test]
fn a_long_function_is_embedded_in_pieces_and_stored_as_one_vector() {
let dir = tempfile::tempdir().unwrap();
let wh = IcebergWarehouse::open(dir.path()).unwrap();
let long: String = std::iter::once("fn table() {".to_string()).chain((0..298).map(|i| format!(" row({i});"))).chain(std::iter::once("}".to_string())).collect::<Vec<_>>().join("\n");
let (cs, _) = chunks(vec![node("short", "f.rs", 1, 1), node("long", "g.rs", 1, 300)], |f| Some(if f == "f.rs" { "fn a() {}\n".to_string() } else { long.clone() }));
let shape = Shape::default();
let run = embed(&wh, &At::member("r", 1, "sha"), &cs, 0, &mut Loaded { embedder: &shape, backend: "test" }).unwrap();
assert_eq!((run.functions, run.new_vectors, run.pieces), (2, 2, 4), "{run:?}");
let asked = shape.0.lock().unwrap().clone();
// the budget gate has no say here (a `Loaded` embedder is not priced): 1 + 3 forwards
assert_eq!(asked.len(), 4);
assert!(asked.iter().all(|lines| *lines <= PIECE_LINES), "a forward was handed {asked:?} lines");
let read = read(&wh, &[("r".to_string(), 1)].into()).unwrap();
let (hash, v) = read.of("r", "long").unwrap();
assert_eq!(hash, crate::vector::chunk::hash_text(&long), "keyed by the WHOLE text");
assert!((v.iter().map(|x| x * x).sum::<f32>().sqrt() - 1.0).abs() < 1e-5, "one normalized vector: {v:?}");
assert_eq!(read.by_hash.len(), 2);
}
/// The delta is exactly the nodes whose text changed, and the nodes that left.
#[test]
fn the_delta_names_only_the_nodes_whose_text_changed() {
let read = |v: u8| move |_: &str| Some(format!("fn a() {{ 1 }}\nfn b() {{ {v} }}\nfn c() {{ 3 }}\n"));
let nodes = || vec![node("a", "f.rs", 1, 1), node("b", "f.rs", 2, 2), node("c", "f.rs", 3, 3)];
let (first, _) = chunks(nodes(), read(2));
let state: BTreeMap<String, String> = first.iter().map(|c| (c.node.moniker.clone(), c.content_hash.clone())).collect();
assert!(delta(&state, &first).plus.is_empty() && delta(&state, &first).minus.is_empty(), "nothing changed: no event");
// `b` changes, `c` is deleted
let (second, _) = chunks(vec![node("a", "f.rs", 1, 1), node("b", "f.rs", 2, 2)], read(9));
let d = delta(&state, &second);
assert_eq!(d.plus.iter().map(|c| c.node.moniker.as_str()).collect::<Vec<_>>(), ["b"]);
assert_eq!(d.minus, ["c"]);
}
}