src/clone_signals/feed.rs
//! The `clone_shapes` feed's work: per consumed commit of a member, the shape
//! of every function whose normalised body changed — no model, no GPU, a few
//! seconds per member — and the normalised bodies' vectors, embedded by the
//! `embeddings` feed's function half through the same embedder (the GPU
//! worker of `$NORNIR_EMBED_REMOTE`).
use std::collections::{BTreeMap, HashSet};
use std::path::Path;
use anyhow::Result;
use mimir_ingest::EventLog;
use mimir_room::{Room, Scope};
use super::norm::{NormBody, normalize};
use super::store::{self, Event, Shape};
use crate::vector::fnvec::{self, FnChunk, room::function_nodes, room::read_under};
use crate::warehouse::iceberg::{IcebergWarehouse, TABLE_EMBEDDINGS_GROUPS, TABLE_EMBEDDINGS_NORM};
/// What one member run did.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct ShapeRun {
/// Function chunks read at this state.
pub functions: usize,
/// … of which did not tokenise (no shape).
pub untokenised: usize,
/// Distinct exact keys among them.
pub distinct: usize,
/// Events written: new or changed shapes, functions gone.
pub plus: usize,
pub minus: usize,
pub ms: u64,
}
impl ShapeRun {
pub fn line(&self) -> String {
format!(
"{} function(s), {} distinct normalised bodies, +{} −{} shape event(s), {} untokenised, {} ms",
self.functions, self.distinct, self.plus, self.minus, self.untokenised, self.ms
)
}
}
/// A body written for one CPU feature or vector width: `#[target_feature]`,
/// `std::arch` / `core::arch`, or x86 (`_mm*`) / NEON (`vld1*`, `vst1*`)
/// intrinsics. Its look-alikes (`…_avx2` / `…_avx512`) are deliberate.
pub fn is_simd(text: &str) -> bool {
text.contains("target_feature")
|| text.contains("::arch::")
|| text.split(|c: char| !(c.is_alphanumeric() || c == '_')).any(|w| w.starts_with("_mm") || w.starts_with("vld1") || w.starts_with("vst1"))
}
/// The shape of one chunk, or `None` when its text does not tokenise.
pub fn shape_of(text: &str) -> Option<(NormBody, Shape)> {
let n = normalize(text)?;
let minhash = super::minhash::signature(&n.classes).unwrap_or_else(|| vec![u32::MAX; super::minhash::PERMS]);
let s = Shape { exact: n.exact.clone(), tokens: n.tokens as i64, minhash, simd: is_simd(text) };
Some((n, s))
}
/// Bring the shapes of `member` to `chunks` (its functions at `seq`): events
/// only for the functions whose shape changed and the ones that left.
pub fn write_shapes(wh: &IcebergWarehouse, member: &str, seq: i64, sha: &str, chunks: &[FnChunk]) -> Result<ShapeRun> {
let t0 = std::time::Instant::now();
let events = store::read(wh, &[member])?;
let state = store::open(&events, member, i64::MAX);
let mut run = ShapeRun { functions: chunks.len(), ..Default::default() };
let mut out: Vec<Event> = Vec::new();
let mut now: HashSet<String> = HashSet::new();
let mut distinct: HashSet<String> = HashSet::new();
for c in chunks {
let Some((_, shape)) = shape_of(&c.text) else {
run.untokenised += 1;
continue;
};
// keyed by the STABLE id: a version bump of the crate is not a new function
let id = mimir_schema::moniker::stable(&c.node.moniker).into_owned();
distinct.insert(shape.exact.clone());
if state.get(id.as_str()).is_some_and(|s| **s == shape) {
now.insert(id);
continue;
}
out.push(Event { repo: member.to_string(), seq, sha: sha.to_string(), moniker: id.clone(), shape: Some(shape), ts_micros: 0 });
now.insert(id);
}
run.plus = out.len();
for m in state.keys().filter(|m| !now.contains(**m)) {
out.push(Event { repo: member.to_string(), seq, sha: sha.to_string(), moniker: m.to_string(), shape: None, ts_micros: 0 });
}
run.minus = out.len() - run.plus;
run.distinct = distinct.len();
store::append(wh, &out)?;
run.ms = t0.elapsed().as_millis() as u64;
Ok(run)
}
/// The function chunks of `member` at `seq`, cut out of `checkout` (the
/// embeddings feed's own cut).
pub fn read_chunks(log: &EventLog, member: &str, checkout: &Path, seq: i64) -> Result<Vec<FnChunk>> {
let scope = Scope::repo(member, Some(seq));
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, _) = function_nodes(&room, &[member.to_string()])?;
drop(room);
Ok(fnvec::chunks(nodes.into_iter().map(|n| n.node).collect(), |f| read_under(checkout, f)).0)
}
/// **The member half of the `clone_shapes` feed.**
pub fn shapes_member(wh: &IcebergWarehouse, log: &EventLog, member: &str, checkout: &Path, seq: i64, sha: &str) -> Result<ShapeRun> {
let chunks = read_chunks(log, member, checkout, seq)?;
write_shapes(wh, member, seq, sha, &chunks)
}
/// What the normalised half of an embed run did.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct NormRun {
pub distinct: usize,
pub new_vectors: usize,
pub embed_ms: u64,
}
/// Embed the normalised bodies of `chunks` that have no vector yet in
/// `embeddings_norm` (keyed by the exact key, in the embedder's model space),
/// through `embedding` — the same embedder, and the same worker, as the raw
/// bodies. Content-addressed: a body any function had before costs nothing.
pub fn embed_normalised(wh: &IcebergWarehouse, chunks: &[FnChunk], embedding: &mut dyn fnvec::Embedding) -> Result<NormRun> {
let profile = embedding.profile();
let id = profile.id();
let mut stored = wh.block_on(crate::vector::store::existing_content_hashes_in(wh, TABLE_EMBEDDINGS_NORM, &id))?;
let mut todo: BTreeMap<String, String> = BTreeMap::new();
for c in chunks {
if let Some(n) = normalize(&c.text)
&& !stored.contains(&n.exact)
{
todo.entry(n.exact).or_insert(n.text);
}
}
let mut run = NormRun { distinct: todo.len(), ..Default::default() };
let todo: Vec<(String, String)> = todo.into_iter().collect();
for slice in todo.chunks(crate::vector::store::embed_slice_chunks()) {
let cut: Vec<Vec<&str>> = slice.iter().map(|(_, t)| fnvec::pieces(t)).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 next = 0usize;
let mut vectors = Vec::with_capacity(slice.len());
for parts in &cut {
let of = &embedded[next..next + parts.len()];
next += parts.len();
vectors.push(fnvec::pool(of, &parts.iter().map(|p| p.len()).collect::<Vec<_>>()));
}
let rows: Vec<(&str, &[f32])> = slice.iter().zip(&vectors).filter(|(_, v)| !v.is_empty()).map(|((h, _), v)| (h.as_str(), v.as_slice())).collect();
wh.block_on(crate::vector::store::append_embeddings_to(wh, TABLE_EMBEDDINGS_NORM, &id, profile.dim, profile.normalize, &rows))?;
run.new_vectors += rows.len();
stored.extend(slice.iter().map(|(h, _)| h.clone()));
}
Ok(run)
}
/// Embed the normalised SEMANTIC GROUPS of `chunks` that have no vector yet in
/// `embeddings_groups` ([`crate::mimir_groupclone`]: def-use segments of a
/// function body, keyed by the group's exact key) — the same seam, the same
/// embedder and the same content-addressing as [`embed_normalised`], one table
/// down. A group below [`crate::mimir_groupclone::MIN_GROUP_TOKENS`] is never
/// embedded: the pairer never reads it, so a vector of it would be dead rows.
pub fn embed_groups(wh: &IcebergWarehouse, chunks: &[FnChunk], embedding: &mut dyn fnvec::Embedding) -> Result<NormRun> {
let profile = embedding.profile();
let id = profile.id();
let mut stored = wh.block_on(crate::vector::store::existing_content_hashes_in(wh, TABLE_EMBEDDINGS_GROUPS, &id))?;
let mut todo: BTreeMap<String, String> = BTreeMap::new();
for c in chunks {
for g in crate::mimir_groupclone::segment(&c.text).unwrap_or_default() {
if g.norm.tokens >= crate::mimir_groupclone::MIN_GROUP_TOKENS && !stored.contains(&g.norm.exact) {
todo.entry(g.norm.exact).or_insert(g.norm.text);
}
}
}
let mut run = NormRun { distinct: todo.len(), ..Default::default() };
let todo: Vec<(String, String)> = todo.into_iter().collect();
for slice in todo.chunks(crate::vector::store::embed_slice_chunks()) {
let cut: Vec<Vec<&str>> = slice.iter().map(|(_, t)| fnvec::pieces(t)).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 next = 0usize;
let mut vectors = Vec::with_capacity(slice.len());
for parts in &cut {
let of = &embedded[next..next + parts.len()];
next += parts.len();
vectors.push(fnvec::pool(of, &parts.iter().map(|p| p.len()).collect::<Vec<_>>()));
}
let rows: Vec<(&str, &[f32])> = slice.iter().zip(&vectors).filter(|(_, v)| !v.is_empty()).map(|((h, _), v)| (h.as_str(), v.as_slice())).collect();
wh.block_on(crate::vector::store::append_embeddings_to(wh, TABLE_EMBEDDINGS_GROUPS, &id, profile.dim, profile.normalize, &rows))?;
run.new_vectors += rows.len();
stored.extend(slice.iter().map(|(h, _)| h.clone()));
}
Ok(run)
}