pub mod accel;
pub mod add;
pub mod agentfiles;
pub mod brief;
pub mod cache;
pub mod chunk;
pub mod clientfile;
pub mod clientid;
pub mod clients;
pub mod clock;
pub mod compact;
pub mod coreml;
pub mod cuda;
pub mod daemon;
pub mod doctor;
pub mod embed;
pub mod filter;
pub mod fleet;
mod formats;
pub mod gpu;
pub mod graph;
pub mod home;
pub mod hook;
pub mod http;
pub mod image;
pub mod index;
pub mod keyscan;
pub mod ledger;
pub mod leftovers;
pub mod llama;
pub mod lock;
pub mod mcp;
pub mod openvino;
pub mod packs;
pub mod pattern;
pub mod pipeline;
pub mod portal;
pub mod prices;
pub mod priority;
pub mod proxy;
pub mod replay;
pub mod report;
pub mod rerank;
pub mod routes;
pub mod schedule;
pub mod service;
pub mod session;
pub mod setup;
pub mod store;
pub mod system;
pub mod tree;
pub mod trt;
pub mod upgrade;
pub mod usage;
pub mod watch;
use anyhow::{Result, bail};
use embed::Model;
use fastembed::TextEmbedding;
use filter::Filter;
use index::{Allowlist, VectorIndex};
use rusqlite::Connection;
use serde::Serialize;
use std::path::{Path, PathBuf};
use std::time::{SystemTime, UNIX_EPOCH};
const BIT_WIDTH: usize = 4;
const EMBED_BATCH: usize = 8;
pub(crate) const UNSORTED_ENV: &str = "SEMLITH_UNSORTED";
const VARIANTS_KEY: &str = "variants";
static SESSIONS: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
pub fn writer_sessions() -> usize {
SESSIONS.load(std::sync::atomic::Ordering::Relaxed)
}
const SORT_WINDOW: usize = 256;
const ROW_COMMIT: std::time::Duration = std::time::Duration::from_millis(250);
const PROGRESS_TICK: std::time::Duration = std::time::Duration::from_millis(250);
const PROGRESS_HEARTBEAT: std::time::Duration = std::time::Duration::from_secs(2);
struct Progress {
seen: usize,
at: std::time::Instant,
}
impl Default for Progress {
fn default() -> Self {
Self {
seen: 0,
at: std::time::Instant::now(),
}
}
}
impl Progress {
fn due(&mut self, batched: &std::sync::atomic::AtomicUsize) -> Option<usize> {
let n = batched.load(std::sync::atomic::Ordering::Relaxed);
let since = self.at.elapsed();
if since < PROGRESS_TICK || (n == self.seen && since < PROGRESS_HEARTBEAT) {
return None;
}
self.seen = n;
self.at = std::time::Instant::now();
Some(n)
}
}
const PENDING_EMBED: usize = 16;
pub fn default_model() -> Model {
Model::Granite
}
const PENDING: &str = "";
const GENERATION: &str = "index_generation";
const RRF_K: f32 = 60.0;
const DEFINITION_LIFT: usize = 3;
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize)]
#[serde(rename_all = "lowercase")]
pub enum Shape {
Identifier,
Question,
}
impl Shape {
pub fn as_str(self) -> &'static str {
match self {
Self::Identifier => "identifier-shaped",
Self::Question => "question-shaped",
}
}
pub fn weighting(self) -> &'static str {
match self {
Self::Identifier => "keyword weighted 2×",
Self::Question => "vector and keyword equal",
}
}
fn keyword_weight(self) -> f32 {
match self {
Self::Identifier => 2.0,
Self::Question => 1.0,
}
}
fn list_constant(self, list: &str) -> f32 {
const STEEP: f32 = 12.0;
match (self, list) {
(Self::Identifier, "keyword") => STEEP,
(Self::Question, "vector") => STEEP,
(_, "image") => STEEP,
_ => RRF_K,
}
}
}
type CandidateList = (&'static str, bool, Vec<(u64, f32)>);
type Scored = ((bool, u64), f32);
fn fuse(
lists: &[CandidateList],
constant_of: impl Fn(&str) -> f32,
) -> (Vec<Scored>, Vec<Vec<&'static str>>) {
let mut fused: Vec<((bool, u64), f32)> = Vec::new();
let mut seen: std::collections::HashMap<(bool, u64), usize> = std::collections::HashMap::new();
let mut badges: Vec<Vec<&'static str>> = Vec::new();
for (name, is_image, ranking) in lists {
let constant = constant_of(name);
let derived = *name == "graph";
for (rank, (id, weight)) in ranking.iter().enumerate() {
let key = (*is_image, *id);
let contribution = weight / (constant + rank as f32 + 1.0);
match seen.get(&key) {
Some(&slot) => {
if !derived {
fused[slot].1 += contribution;
}
badges[slot].push(name);
}
None => {
seen.insert(key, fused.len());
fused.push((key, contribution));
badges.push(vec![name]);
}
}
}
}
(fused, badges)
}
pub fn shape_of(query: &str) -> Shape {
let trimmed = query.trim();
if trimmed.is_empty() || trimmed.split_whitespace().count() > 1 {
return Shape::Question;
}
let identifier = trimmed
.chars()
.all(|c| c.is_alphanumeric() || c == '_' || c == ':' || c == '.' || c == '!');
if identifier && trimmed.chars().any(|c| c.is_alphabetic()) {
Shape::Identifier
} else {
Shape::Question
}
}
pub fn code_shaped(query: &str) -> bool {
const CODE_WORDS: [&str; 22] = [
"function",
"functions",
"method",
"methods",
"call",
"calls",
"called",
"caller",
"callers",
"callee",
"struct",
"class",
"impl",
"trait",
"enum",
"type",
"field",
"variable",
"signature",
"parameter",
"return",
"returns",
];
if shape_of(query) == Shape::Identifier {
return true;
}
query.split_whitespace().any(|word| {
let word = word.trim_matches(|c: char| !c.is_alphanumeric() && c != '_' && c != ':');
let inner_capital =
word.chars().skip(1).any(char::is_uppercase) && word.chars().any(char::is_lowercase);
word.contains('_')
|| word.contains("::")
|| inner_capital
|| CODE_WORDS.contains(&word.to_lowercase().as_str())
}) || query.contains("()")
}
fn names_tests(query: &str) -> bool {
query.split(|c: char| !c.is_alphanumeric()).any(|w| {
let w = w.to_lowercase();
w.starts_with("test") || w.starts_with("spec") || w.starts_with("fixture")
})
}
pub fn is_test_path(path: &str) -> bool {
Path::new(path).components().any(|c| {
matches!(
c.as_os_str().to_str(),
Some("tests" | "test" | "fixtures" | "spec" | "__tests__")
)
})
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, serde::Serialize)]
#[serde(rename_all = "lowercase")]
pub enum Prefer {
#[default]
Any,
Code,
Docs,
}
impl Prefer {
pub fn as_str(self) -> &'static str {
match self {
Self::Any => "any",
Self::Code => "code",
Self::Docs => "docs",
}
}
pub fn parse(raw: &str) -> Result<Self> {
match raw.trim().to_ascii_lowercase().as_str() {
"any" | "" => Ok(Self::Any),
"code" => Ok(Self::Code),
"docs" | "doc" => Ok(Self::Docs),
other => anyhow::bail!("prefer is code, docs or any, not {other:?}"),
}
}
fn multiplier(self, is_code: bool) -> f32 {
const LIFT: f32 = 1.5;
match (self, is_code) {
(Self::Any, _) => 1.0,
(Self::Code, true) | (Self::Docs, false) => LIFT,
_ => 1.0,
}
}
}
const GRAPH_PROXIMITY: f32 = 0.15;
const STALE_PENALTY: f32 = 0.10;
const TEST_PENALTY: f32 = 0.10;
type Candidate = (((bool, u64), f32), Vec<&'static str>);
const TEXT_LISTS: f32 = 3.0;
const WEAK_IMAGE: f32 = 0.25;
const IMAGE_FLOOR: f32 = 0.28;
pub const IMAGE_FLOOR_ENV: &str = "SEMLITH_IMAGE_FLOOR";
fn image_floor() -> f32 {
std::env::var(IMAGE_FLOOR_ENV)
.ok()
.and_then(|v| v.parse::<f32>().ok())
.unwrap_or(IMAGE_FLOOR)
}
const RANK_DEPTH: usize = 4;
const CHECKPOINT_FILES: usize = 200;
const CHECKPOINT_FILES_ENV: &str = "SEMLITH_CHECKPOINT_FILES";
fn checkpoint_files() -> usize {
match std::env::var(CHECKPOINT_FILES_ENV)
.ok()
.and_then(|v| v.parse().ok())
{
Some(files) if files > 0 => files,
_ => CHECKPOINT_FILES,
}
}
const INFERRED_EXPANSION: f32 = 0.5;
struct Reached {
id: u64,
weight: f32,
tier: String,
proximity: f32,
}
fn provenance_of(graph: &[Reached], id: u64) -> Option<String> {
graph.iter().find(|r| r.id == id).map(|r| r.tier.clone())
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Target {
Span { path: String, start: u32, end: u32 },
Symbol(String),
}
impl Target {
pub fn parse(raw: &str) -> Self {
let raw = raw.trim();
if let Some((path, lines)) = raw.rsplit_once(':')
&& !path.is_empty()
{
let (start, end) = match lines.split_once('-') {
Some((a, b)) => (a.parse::<u32>().ok(), b.parse::<u32>().ok()),
None => (lines.parse::<u32>().ok(), lines.parse::<u32>().ok()),
};
if let (Some(start), Some(end)) = (start, end) {
return Self::Span {
path: path.to_string(),
start: start.min(end),
end: start.max(end),
};
}
}
Self::Symbol(raw.to_string())
}
}
#[derive(Debug, Clone, Serialize)]
pub struct Span {
#[serde(serialize_with = "serialize_plain")]
pub path: String,
pub start_line: u32,
pub end_line: u32,
pub text: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub symbol: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub symbol_kind: Option<String>,
pub fresh: bool,
#[serde(skip_serializing_if = "std::ops::Not::not")]
pub from_disk: bool,
#[serde(skip_serializing_if = "Option::is_none")]
pub store: Option<String>,
}
#[derive(Debug, Clone, Serialize)]
#[serde(untagged)]
pub enum Read {
One(Span),
Choose(Vec<store::SymbolRow>),
}
#[derive(Debug, Clone, Serialize)]
pub struct Hit {
pub score: f32,
#[serde(serialize_with = "serialize_plain")]
pub path: String,
pub start_line: u32,
pub end_line: u32,
pub text: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub store: Option<String>,
#[serde(skip_serializing_if = "Vec::is_empty")]
pub lists: Vec<&'static str>,
#[serde(skip_serializing_if = "Option::is_none")]
pub image: Option<Pixels>,
pub fresh: bool,
#[serde(skip_serializing_if = "Option::is_none")]
pub symbol: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub symbol_kind: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub symbol_line: Option<u32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub provenance: Option<String>,
#[serde(skip_serializing_if = "Vec::is_empty")]
#[serde(serialize_with = "serialize_plain_list")]
pub copies: Vec<String>,
}
#[derive(Debug, Clone, Copy, serde::Serialize)]
pub struct Pixels {
pub width: u32,
pub height: u32,
}
#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
pub enum FileOutcome {
#[default]
Indexing,
Unchanged,
Skipped,
Removed,
Refused,
Writing,
Failed,
Progress,
}
impl FileOutcome {
pub fn as_str(self) -> &'static str {
match self {
Self::Indexing => "indexing",
Self::Writing => "writing",
Self::Unchanged => "unchanged",
Self::Skipped => "skipped",
Self::Removed => "removed",
Self::Refused => "refused",
Self::Failed => "failed",
Self::Progress => "progress",
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum SkipReason {
Empty,
TooLarge,
NotRegular,
Unreadable(String),
Binary,
NoText,
NotDecodableImage,
}
impl SkipReason {
pub fn as_str(&self) -> String {
match self {
Self::Empty => "empty".to_string(),
Self::TooLarge => format!("over {} MiB", chunk::MAX_FILE_BYTES / (1024 * 1024)),
Self::NotRegular => "not a regular file".to_string(),
Self::Unreadable(e) => format!("unreadable: {e}"),
Self::Binary => "binary".to_string(),
Self::NoText => "no text in this document".to_string(),
Self::NotDecodableImage => "not a decodable image".to_string(),
}
}
pub fn class(&self) -> &'static str {
match self {
Self::TooLarge => store::class::POLICY,
_ => store::class::UNINDEXABLE,
}
}
pub fn kind(&self) -> &'static str {
match self {
Self::Empty => "empty",
Self::TooLarge => "too large",
Self::NotRegular => "not a regular file",
Self::Unreadable(_) => "unreadable",
Self::Binary => "binary",
Self::NoText => "no text in this document",
Self::NotDecodableImage => "not a decodable image",
}
}
}
#[derive(Debug, Default, Clone)]
pub struct Boundary {
pub roots: Option<Vec<PathBuf>>,
pub allow_secrets: bool,
}
impl Boundary {
pub fn within(roots: Vec<PathBuf>) -> Self {
Self {
roots: Some(roots),
allow_secrets: false,
}
}
pub fn refuses(&self, path: &Path, walked: bool, home: Option<&Path>) -> Option<Refusal> {
self.resolved(home).refuses(path, walked)
}
pub fn resolved(&self, home: Option<&Path>) -> ResolvedBoundary {
ResolvedBoundary {
within: self.roots.as_deref().map(filter::resolve_boundary),
boundary: self.clone(),
home: home.map(Path::to_path_buf),
}
}
}
#[derive(Debug, Clone)]
pub struct ResolvedBoundary {
boundary: Boundary,
within: Option<Vec<String>>,
home: Option<PathBuf>,
}
impl ResolvedBoundary {
pub fn refuses(&self, path: &Path, walked: bool) -> Option<Refusal> {
if let (Some(roots), Some(within)) = (&self.boundary.roots, &self.within)
&& !filter::within_resolved(path, within)
{
let named = if roots.is_empty() {
"this store has no registered roots, so only the home directory is inside \
its boundary"
.to_string()
} else {
format!(
"this store's roots are {}",
roots
.iter()
.map(|r| r.display().to_string())
.collect::<Vec<_>>()
.join(", ")
)
};
return Some(Refusal {
why: format!(
"is outside this store's roots, so semlith will not index it into this \
store — {named}. Index it into the store that covers it, add it as a \
root first, or index it from the command line."
),
credential: false,
});
}
if !self.boundary.allow_secrets
&& let Some(why) = filter::denied_against(path, self.home.as_deref())
{
if walked && why == filter::Denied::Hidden {
return None;
}
return Some(Refusal {
why: why.reason(),
credential: matches!(why, filter::Denied::Name(_) | filter::Denied::Directory(_)),
});
}
None
}
}
#[derive(Debug, Clone)]
pub struct Refusal {
pub why: String,
pub credential: bool,
}
#[derive(Debug, Default, Clone)]
pub struct IndexProgress {
pub outcome: FileOutcome,
pub scanned: usize,
pub indexed: usize,
pub chunks: usize,
pub total: usize,
pub symbols: usize,
pub why: Option<String>,
pub threads: usize,
pub lanes: std::collections::BTreeMap<String, usize>,
pub bytes: u64,
pub bytes_total: u64,
pub rows: usize,
}
#[derive(Debug, Default, Clone, Serialize)]
pub struct IndexReport {
pub scanned: usize,
pub indexed: usize,
pub unchanged: usize,
pub skipped: usize,
pub removed: usize,
pub chunks: usize,
pub remaining: usize,
#[serde(skip)]
pub pending: Vec<PathBuf>,
#[serde(skip_serializing_if = "Vec::is_empty")]
pub generated: Vec<String>,
pub rechunked: bool,
#[serde(skip_serializing_if = "Vec::is_empty")]
pub refused: Vec<(String, String)>,
#[serde(skip_serializing_if = "Vec::is_empty")]
pub failed: Vec<(String, String)>,
#[serde(skip_serializing_if = "std::collections::BTreeMap::is_empty")]
pub skipped_reasons: std::collections::BTreeMap<String, usize>,
pub secrets_indexed: usize,
pub dummies_indexed: usize,
pub accepted_indexed: usize,
pub symbols: usize,
pub edges: usize,
pub images: usize,
pub stopped: bool,
pub threads: usize,
pub embedded: usize,
#[serde(skip)]
pub batched: usize,
#[serde(skip)]
pub lanes: std::collections::BTreeMap<String, usize>,
#[serde(skip)]
pub bytes: u64,
#[serde(skip)]
pub bytes_total: u64,
pub written: Vec<String>,
pub stages: pipeline::Stages,
pub cache_lookups: usize,
pub cache_hits: usize,
pub rows: usize,
}
#[derive(Default)]
struct CacheWrites {
fresh: Vec<([u8; 32], &'static str, Vec<f32>)>,
hits: Vec<([u8; 32], &'static str)>,
}
const SCAN_RULES: u32 = 2;
const GRAPH_RULES: u32 = 3;
#[derive(Debug, Clone, Serialize)]
pub struct Review {
#[serde(serialize_with = "serialize_plain")]
pub path: String,
pub class: String,
pub rule: String,
#[serde(skip_serializing_if = "Vec::is_empty")]
pub matches: Vec<keyscan::Match>,
#[serde(skip_serializing_if = "Option::is_none")]
pub confidence: Option<u8>,
}
#[derive(Debug, Clone, Default, Serialize)]
pub struct Plan {
pub embed: usize,
pub embed_bytes: u64,
pub languages: std::collections::BTreeMap<String, usize>,
pub unchanged: usize,
pub not_indexed: std::collections::BTreeMap<String, usize>,
pub review: Vec<Review>,
#[serde(serialize_with = "serialize_plain_list")]
pub credential: Vec<String>,
pub dummies: usize,
#[serde(skip_serializing_if = "Option::is_none")]
pub eta_ms: Option<u64>,
pub seconds: f64,
}
fn embeddable_bytes(meta: Option<std::fs::Metadata>) -> u64 {
meta.filter(|m| m.is_file() && m.len() <= chunk::MAX_FILE_BYTES)
.map_or(0, |m| m.len())
}
fn say_file(
on_file: &mut dyn FnMut(&Path, IndexProgress),
report: &IndexReport,
total: usize,
path: &Path,
outcome: FileOutcome,
why: Option<String>,
) {
debug_assert!(
why.is_some()
|| !matches!(
outcome,
FileOutcome::Skipped | FileOutcome::Refused | FileOutcome::Failed
),
"{} must say why",
outcome.as_str()
);
on_file(
path,
IndexProgress {
outcome,
scanned: report.scanned,
indexed: report.indexed,
chunks: report.embedded.max(report.batched),
total,
symbols: report.symbols,
why,
threads: report.threads,
lanes: report.lanes.clone(),
bytes: report.bytes,
bytes_total: report.bytes_total,
rows: report.rows,
},
);
}
fn skip(report: &mut IndexReport, why: &SkipReason) {
report.skipped += 1;
*report
.skipped_reasons
.entry(why.kind().to_string())
.or_insert(0) += 1;
}
pub const FAULT_PANIC_ENV: &str = "SEMLITH_FAULT_PANIC";
fn fault_panic(path: &Path) {
if let Ok(name) = std::env::var(FAULT_PANIC_ENV)
&& path
.file_name()
.is_some_and(|n| n.to_string_lossy() == name)
{
panic!("{FAULT_PANIC_ENV} names {name}");
}
}
impl Drop for Semlith {
fn drop(&mut self) {
self.release_embedder();
}
}
fn contained<T>(work: impl FnOnce() -> T) -> Result<T> {
std::panic::catch_unwind(std::panic::AssertUnwindSafe(work)).map_err(|panic| {
let message = panic
.downcast_ref::<&str>()
.map(|s| s.to_string())
.or_else(|| panic.downcast_ref::<String>().cloned())
.unwrap_or_else(|| "no message".to_string());
anyhow::anyhow!("the reader panicked on this file: {message}")
})
}
fn failed(report: &mut IndexReport, path: &Path, e: &anyhow::Error) {
report
.failed
.push((path.display().to_string(), format!("{e:#}")));
}
#[derive(Debug, Clone, Serialize)]
pub struct Finding {
#[serde(serialize_with = "serialize_plain")]
pub path: String,
pub why: String,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Flow {
Run,
Pause,
Yield,
Stop,
}
const PAUSE_TICK: std::time::Duration = std::time::Duration::from_millis(120);
pub struct Semlith {
dir: PathBuf,
db: Connection,
index: VectorIndex,
images: VectorIndex,
exact: index::Exact,
model: Model,
dim: usize,
embedder: Option<TextEmbedding>,
embedder_threads: usize,
index_session: Option<pipeline::CpuWork>,
index_threads: usize,
embedder_alt: Option<TextEmbedding>,
query_embedder: Option<TextEmbedding>,
batches: u64,
last_variant: &'static str,
last_embed: Option<std::time::Instant>,
budget_generation: u64,
lane_chunks: std::collections::BTreeMap<String, usize>,
tokenizer: Option<tokenizers::Tokenizer>,
clip: image::Clip,
generation: u64,
pub quiet: bool,
pub boundary: Boundary,
prehashed: std::collections::HashMap<PathBuf, (u64, i64, String)>,
scrub: bool,
}
impl Semlith {
pub fn open(dir: impl AsRef<Path>, model: Option<Model>) -> Result<Self> {
let dir = dir.as_ref().to_path_buf();
home::secure_dir(&dir)?;
let db_path = dir.join("store.db");
let db = store::open(&db_path)?;
home::tighten_file(&db_path);
for beside in ["store.db-wal", "store.db-shm"] {
home::tighten_file(&dir.join(beside));
}
let model = match store::get_meta(&db, "model")? {
Some(existing) => {
let existing: Model = existing.parse().map_err(anyhow::Error::msg)?;
if let Some(want) = model
&& want != existing
{
bail!(
"store was built with {existing}, not {want}; \
delete {} to rebuild with a different model",
dir.display()
);
}
existing
}
None => {
let m = model.unwrap_or_else(default_model);
let _writing = store::Writing::begin(&db)?;
store::set_meta(&db, "model", &m.to_string())?;
store::set_meta(&db, store::FORMAT_KEY, &store::FORMAT_VERSION.to_string())?;
m
}
};
let dim = model.dim()?;
let sharded = store::format(&db)? >= store::SHARDED_FORMAT;
let index = VectorIndex::open(&dir, dim, BIT_WIDTH, sharded)?;
let images = VectorIndex::open(&dir.join(image::INDEX_DIR), image::DIM, BIT_WIDTH, true)?;
let generation = generation(&db)?;
Ok(Self {
exact: index::Exact::open(&dir, dim),
dir,
db,
index,
images,
model,
dim,
embedder: None,
embedder_threads: 0,
index_session: None,
index_threads: 0,
embedder_alt: None,
query_embedder: None,
batches: 0,
last_variant: embed::Variant::Int8.name(),
last_embed: None,
budget_generation: index::budget_generation(),
lane_chunks: std::collections::BTreeMap::new(),
tokenizer: None,
clip: image::Clip::default(),
generation,
quiet: false,
boundary: Boundary::default(),
prehashed: Default::default(),
scrub: false,
})
}
pub fn open_existing(dir: impl AsRef<Path>) -> Result<Self> {
let dir = dir.as_ref();
if !dir.join("store.db").exists() {
bail!(
"{} is not a semlith store — no store.db in it; index it first",
dir.display()
);
}
Self::open(dir, None)
}
pub fn refresh(&mut self) -> Result<()> {
let current = generation(&self.db)?;
if current == self.generation {
return Ok(());
}
let was_resident = self.index.is_resident();
self.reopen_indexes()?;
if was_resident {
self.index.prepare()?;
}
self.generation = current;
Ok(())
}
fn reopen_indexes(&mut self) -> Result<()> {
let sharded = store::format(&self.db)? >= store::SHARDED_FORMAT;
self.index = VectorIndex::open(&self.dir, self.dim, BIT_WIDTH, sharded)?;
self.images = VectorIndex::open(
&self.dir.join(image::INDEX_DIR),
image::DIM,
BIT_WIDTH,
true,
)?;
Ok(())
}
fn search_vectors(
&mut self,
vector: &[f32],
depth: usize,
allowlist: &Allowlist,
) -> Result<(Vec<f32>, Vec<u64>)> {
for _ in 0..40 {
if !compact::swapping(&self.db)? {
self.refresh()?;
let seen = self.generation;
let found = self.index.search(vector, depth, allowlist)?;
if generation(&self.db)? == seen && !compact::swapping(&self.db)? {
return Ok(found);
}
}
std::thread::sleep(std::time::Duration::from_millis(50));
}
self.refresh()?;
self.index.search(vector, depth, allowlist)
}
pub fn model(&self) -> &Model {
&self.model
}
pub fn dim(&self) -> usize {
self.dim
}
pub fn len(&self) -> usize {
if let Some(resident) = self.index.resident_len() {
return resident;
}
store::durable_chunks(&self.db).unwrap_or(0) as usize
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
pub fn db(&self) -> &Connection {
&self.db
}
pub fn dir(&self) -> &Path {
&self.dir
}
pub fn shards(&self) -> Option<(usize, usize)> {
self.index
.max_resident()
.map(|max| (self.index.shards(), max))
}
pub fn evictions(&self) -> u64 {
self.index.evictions()
}
pub fn tokenizer(&self) -> Option<&tokenizers::Tokenizer> {
self.tokenizer.as_ref()
}
fn embedder(&mut self) -> Result<&mut TextEmbedding> {
let threads = embed::threads_in_force();
if self.embedder.is_some() && self.embedder_threads != threads {
self.release_embedder();
}
if self.embedder.is_none() {
let cache = model_cache_dir()?;
self.tokenizer = self.model.tokenizer(&cache);
let (variant, _) = embed::index_variant();
self.embedder = Some(self.model.load_variant(
cache,
chunk::MAX_CHARS / 2,
self.quiet,
threads,
variant,
)?);
self.embedder_threads = threads;
SESSIONS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
}
Ok(self.embedder.as_mut().unwrap())
}
fn release_embedder(&mut self) {
if self.embedder.take().is_some() {
SESSIONS.fetch_sub(1, std::sync::atomic::Ordering::Relaxed);
}
if self.index_session.take().is_some() {
SESSIONS.fetch_sub(1, std::sync::atomic::Ordering::Relaxed);
}
self.embedder_alt = None;
self.query_embedder = None;
self.clip = image::Clip::default();
}
pub fn release_if_idle(&mut self, idle: std::time::Duration) -> bool {
match self.last_embed {
Some(at)
if (self.embedder.is_some() || self.index_session.is_some())
&& at.elapsed() >= idle =>
{
self.release_embedder();
true
}
_ => false,
}
}
pub fn embedder_threads(&self) -> Option<usize> {
self.embedder.as_ref().map(|_| self.embedder_threads)
}
pub fn follow_budget(&mut self) {
let now = index::budget_generation();
if now != self.budget_generation {
self.budget_generation = now;
self.index.apply_budget();
self.images.apply_budget();
}
}
pub fn resident_shards(&self) -> usize {
self.index.resident_shards()
}
pub fn warm(&mut self) -> Result<()> {
self.embedder()?;
self.index.prepare()?;
Ok(())
}
pub fn warm_index(&mut self) -> Result<()> {
self.index.prepare()
}
fn embed(&mut self, texts: Vec<String>) -> Result<Vec<Vec<f32>>> {
let _lifted = priority::embedding();
self.last_embed = Some(std::time::Instant::now());
self.embedder()?;
self.batches += 1;
let (main, alt) = embed::index_variant();
let use_alt = alt.is_some() && self.batches.is_multiple_of(2);
let session = if let (true, Some(variant)) = (use_alt, alt) {
if self.embedder_alt.is_none() {
self.embedder_alt = Some(self.model.load_variant(
model_cache_dir()?,
chunk::MAX_CHARS / 2,
self.quiet,
self.embedder_threads,
variant,
)?);
}
self.last_variant = variant.name();
self.embedder_alt.as_mut().expect("loaded above")
} else {
self.last_variant = main.name();
self.embedder.as_mut().expect("loaded above")
};
let mut out = session
.embed(texts, Some(EMBED_BATCH))
.map_err(|e| anyhow::anyhow!("{e}"))?;
for v in &mut out {
normalize(v);
}
Ok(out)
}
pub fn pin_snapshot(&self) -> Result<()> {
self.db.execute_batch("BEGIN")?;
let _: i64 = self
.db
.query_row("SELECT COUNT(*) FROM meta", [], |r| r.get(0))?;
Ok(())
}
pub fn release_snapshot(&self) {
let _ = self.db.execute_batch("COMMIT");
}
fn accepted_folders(&self) -> Vec<PathBuf> {
store::acceptances(&self.db)
.unwrap_or_default()
.into_iter()
.filter(|a| a.class == store::class::POLICY)
.map(|a| PathBuf::from(a.path))
.collect()
}
fn walk(&self, roots: &[PathBuf]) -> Walked {
let allowed = self.accepted_folders();
let started = std::time::Instant::now();
let mut walked = walk_allowing(roots, &allowed);
walked.walk_ms = started.elapsed().as_millis() as u64;
walked
}
pub fn plan(&mut self, roots: &[PathBuf]) -> Result<Plan> {
let started = std::time::Instant::now();
let walked = self.walk(roots);
let mut plan = Plan::default();
let home = crate::home::user_home().ok().map(|h| canonical(&h));
let rechunk = store::format(&self.db)? < store::CODE_CONTEXT;
let rescan = self.rules_outdated()?;
let all = walked
.named
.into_iter()
.map(|p| (p, false))
.chain(walked.files.into_iter().map(|p| (p, true)));
let not = |plan: &mut Plan, class: &str| {
*plan.not_indexed.entry(class.to_string()).or_insert(0) += 1;
};
for dir in &walked.generated {
plan.review.push(Review {
path: dir.to_string_lossy().into_owned(),
class: store::class::POLICY.to_string(),
rule: "a generated or vendored folder the walk steps over".to_string(),
matches: Vec::new(),
confidence: None,
});
not(&mut plan, store::class::POLICY);
}
for _ in &walked.unreadable {
not(&mut plan, store::class::UNINDEXABLE);
}
for path in &walked.credentials {
plan.credential.push(path.to_string_lossy().into_owned());
not(&mut plan, store::class::CREDENTIAL);
}
for _ in &walked.excluded {
not(&mut plan, store::class::EXCLUDED);
}
let boundary = self.boundary.resolved(home.as_deref());
for (path, walked) in all {
let key = path.to_string_lossy().into_owned();
if let Some(refusal) = boundary.refuses(&path, walked) {
if refusal.credential {
plan.credential.push(key);
not(&mut plan, store::class::CREDENTIAL);
} else {
not(&mut plan, store::class::EXCLUDED);
}
continue;
}
let Ok(meta) = std::fs::metadata(&path) else {
not(&mut plan, store::class::UNINDEXABLE);
continue;
};
if !meta.is_file() || meta.len() == 0 {
not(&mut plan, store::class::UNINDEXABLE);
continue;
}
let accepted = store::acceptance(&self.db, &key)?;
if meta.len() > chunk::MAX_FILE_BYTES
&& !accepted
.as_ref()
.is_some_and(|a| a.class == store::class::POLICY)
{
plan.review.push(Review {
path: key,
class: store::class::POLICY.to_string(),
rule: format!(
"over {} MiB ({} MiB)",
chunk::MAX_FILE_BYTES / (1024 * 1024),
meta.len() / (1024 * 1024)
),
matches: Vec::new(),
confidence: None,
});
not(&mut plan, store::class::POLICY);
continue;
}
let Ok(bytes) = std::fs::read(&path) else {
not(&mut plan, store::class::UNINDEXABLE);
continue;
};
let hash = blake3::hash(&bytes).to_hex().to_string();
let mtime = meta
.modified()
.ok()
.and_then(|t| t.duration_since(std::time::UNIX_EPOCH).ok())
.map_or(0, |d| d.as_secs() as i64);
self.prehashed
.insert(path.clone(), (meta.len(), mtime, hash.clone()));
let unchanged =
!rechunk && store::file_hash(&self.db, &key)?.as_deref() == Some(hash.as_str());
if unchanged && !rescan {
plan.unchanged += 1;
continue;
}
if image::is_image(&path) {
if unchanged {
plan.unchanged += 1;
} else {
plan.embed += 1;
plan.embed_bytes += meta.len();
}
continue;
}
let text = match chunk::extract(&path, &bytes) {
Ok(text) => text,
Err(_) => {
not(&mut plan, store::class::UNINDEXABLE);
continue;
}
};
if self.boundary.allow_secrets {
plan.embed += usize::from(!unchanged);
plan.unchanged += usize::from(unchanged);
continue;
}
let found = keyscan::scan(&key, &text);
match keyscan::decide(&self.db, &key, &text, &found)? {
keyscan::Decision::Index { .. } => {
if !found.is_empty() {
plan.dummies += 1;
}
if unchanged {
plan.unchanged += 1;
} else {
plan.embed += 1;
plan.embed_bytes += meta.len();
if let Some(lang) = filter::language_of_path(&key) {
*plan.languages.entry(lang.name.to_string()).or_insert(0) += 1;
}
}
}
keyscan::Decision::Refuse(rule) => {
let live: Vec<keyscan::Match> =
found.into_iter().filter(|m| m.dummy.is_none()).collect();
plan.review.push(Review {
path: key,
class: store::class::CONTENT.to_string(),
confidence: live.iter().map(|m| m.confidence).max(),
rule,
matches: live,
});
not(&mut plan, store::class::CONTENT);
}
}
}
plan.eta_ms = store::get_meta(&self.db, "embed_bytes_per_sec")?
.and_then(|v| v.parse::<f64>().ok())
.filter(|r| *r > 0.0)
.map(|rate| (plan.embed_bytes as f64 / rate * 1000.0) as u64);
plan.seconds = started.elapsed().as_secs_f64();
Ok(plan)
}
fn rules_outdated(&self) -> Result<bool> {
Ok(store::get_meta(&self.db, "scan_rules")?
.and_then(|v| v.parse::<u32>().ok())
.unwrap_or(0)
< SCAN_RULES)
}
pub fn accept(&mut self, path: &str, mode: &str, source: &str) -> Result<store::Acceptance> {
let key = self.stored_key(path);
if let Some(why) = filter::denied(Path::new(&key))
&& why != filter::Denied::Hidden
{
anyhow::bail!(
"{key} is a credential file ({}) and is never accepted; `semlith index \
--include-secrets` is the only way to index one",
why.reason()
);
}
if !matches!(mode, "redacted" | "as-is" | "refused") {
anyhow::bail!("mode is redacted, as-is or refused, not {mode:?}");
}
let row = store::refusal(&self.db, &key)?;
let class = row
.as_ref()
.map_or(store::class::CONTENT, |r| r.class.as_str())
.to_string();
if !(store::class::reviewable(&class) || class == store::class::DUMMY) {
anyhow::bail!(
"{key} is not indexable ({}); accepting cannot change that",
class
);
}
if mode == "refused" && class != store::class::DUMMY {
anyhow::bail!("only a file let through as test dummies can be refused instead");
}
let accepted = self.writing(|me| {
let salt = store::salt(&me.db)?;
let mut prints = Vec::new();
let mut confidence = None;
if class == store::class::CONTENT || class == store::class::DUMMY {
let bytes = anyhow::Context::with_context(std::fs::read(&key), || {
format!("reading {key}")
})?;
let text = chunk::extract(Path::new(&key), &bytes)
.map_err(|why| anyhow::anyhow!("{key}: {}", why.as_str()))?;
let found = keyscan::scan(&key, &text);
for m in found.iter().filter(|m| m.dummy.is_none()) {
prints.push(keyscan::fingerprint(&salt, &text, m));
}
confidence = found.iter().map(|m| m.confidence).max();
}
let a = store::Acceptance {
path: key.clone(),
class: if class == store::class::DUMMY {
store::class::DUMMY.to_string()
} else {
class.clone()
},
mode: mode.to_string(),
fingerprints: prints,
confidence,
at: now(),
source: source.to_string(),
};
store::accept(&me.db, &a)?;
if mode == "refused" {
let (gone, images) = me.evict(&key)?;
if gone + images > 0 {
me.save()?;
}
}
crate::ledger::acceptance(&me.db, &a, "accept")?;
Ok(a)
})?;
Ok(accepted)
}
pub fn revoke(&mut self, path: &str) -> Result<bool> {
let key = self.stored_key(path);
self.writing(|me| {
let Some(a) = store::acceptance(&me.db, &key)? else {
return Ok(false);
};
store::revoke(&me.db, &key)?;
let (gone, images) = me.evict(&key)?;
if gone + images > 0 {
me.save()?;
}
crate::ledger::acceptance(&me.db, &a, "revoke")?;
Ok(true)
})
}
fn stored_key(&self, path: &str) -> String {
let given = Path::new(path);
if given.exists() {
canonical(given).to_string_lossy().into_owned()
} else {
path.to_string()
}
}
pub fn index_paths(
&mut self,
roots: &[PathBuf],
on_file: impl FnMut(&Path, IndexProgress),
) -> Result<IndexReport> {
let _lock = lock::StoreLock::acquire(&self.dir)?;
self.index_walk(roots, on_file)
}
pub fn index_paths_within(
&mut self,
roots: &[PathBuf],
budget: std::time::Duration,
on_file: impl FnMut(&Path, IndexProgress),
) -> Result<IndexReport> {
let _lock = lock::StoreLock::acquire(&self.dir)?;
self.index_within_held(roots, budget, on_file)
}
pub(crate) fn index_within_held(
&mut self,
roots: &[PathBuf],
budget: std::time::Duration,
on_file: impl FnMut(&Path, IndexProgress),
) -> Result<IndexReport> {
let deadline = std::time::Instant::now() + budget;
self.index_set(
Walk::Roots(roots.to_vec()),
true,
Handed::Walked(Some(deadline)),
None,
on_file,
)
}
pub(crate) fn index_walk_under(
&mut self,
roots: &[PathBuf],
control: &dyn Fn() -> Flow,
on_file: impl FnMut(&Path, IndexProgress),
) -> Result<IndexReport> {
self.index_set(
Walk::Roots(roots.to_vec()),
true,
Handed::Walked(None),
Some(control),
on_file,
)
}
pub(crate) fn index_within_held_under(
&mut self,
roots: &[PathBuf],
budget: std::time::Duration,
control: &dyn Fn() -> Flow,
on_file: impl FnMut(&Path, IndexProgress),
) -> Result<IndexReport> {
let deadline = std::time::Instant::now() + budget;
self.index_set(
Walk::Roots(roots.to_vec()),
true,
Handed::Walked(Some(deadline)),
Some(control),
on_file,
)
}
pub(crate) fn index_rest_held_under(
&mut self,
files: Vec<PathBuf>,
budget: std::time::Duration,
bytes_total: Option<u64>,
control: &dyn Fn() -> Flow,
on_file: impl FnMut(&Path, IndexProgress),
) -> Result<IndexReport> {
self.index_set(
Walk::Done(Walked {
walk_ms: 0,
files,
named: Vec::new(),
unreadable: Vec::new(),
generated: Vec::new(),
credentials: Vec::new(),
excluded: Vec::new(),
}),
true,
Handed::Rest {
budget,
bytes_total,
},
Some(control),
on_file,
)
}
pub(crate) fn index_walk(
&mut self,
roots: &[PathBuf],
on_file: impl FnMut(&Path, IndexProgress),
) -> Result<IndexReport> {
self.index_set(
Walk::Roots(roots.to_vec()),
true,
Handed::Walked(None),
None,
on_file,
)
}
pub(crate) fn index_changed(
&mut self,
paths: Vec<PathBuf>,
on_file: impl FnMut(&Path, IndexProgress),
) -> Result<IndexReport> {
self.index_set(
Walk::Done(Walked {
walk_ms: 0,
files: paths,
named: Vec::new(),
unreadable: Vec::new(),
generated: Vec::new(),
credentials: Vec::new(),
excluded: Vec::new(),
}),
false,
Handed::Walked(None),
None,
on_file,
)
}
fn index_set(
&mut self,
walk: Walk,
sweep: bool,
handed: Handed,
control: Option<&dyn Fn() -> Flow>,
on_file: impl FnMut(&Path, IndexProgress),
) -> Result<IndexReport> {
let _lifted = priority::embedding();
let _class = priority::indexing_thread();
let _writer = embed::writer();
self.writing(move |me| me.index_set_writing(walk, sweep, handed, control, on_file))
}
fn writing<T>(&mut self, body: impl FnOnce(&mut Self) -> Result<T>) -> Result<T> {
store::read_only(&self.db, false)?;
let out = body(self);
let _ = self.tx_commit();
let _ = store::read_only(&self.db, true);
out
}
fn refuse_file(
&mut self,
report: &mut IndexReport,
total: usize,
path: &Path,
refusal: &Refusal,
on_file: &mut dyn FnMut(&Path, IndexProgress),
) -> Result<()> {
let evicted = if refusal.credential {
let key = path.to_string_lossy().into_owned();
let (chunks, images) = self.evict(&key)?;
chunks + images
} else {
0
};
self.scrub |= evicted > 0;
report.removed += usize::from(evicted > 0);
let why = if evicted > 0 {
format!(
"{} Its earlier contents have been removed from this store.",
refusal.why
)
} else {
refusal.why.clone()
};
report
.refused
.push((path.display().to_string(), why.clone()));
report.scanned += 1;
let class = if refusal.credential {
store::class::CREDENTIAL
} else {
store::class::EXCLUDED
};
store::refuse(
&self.db,
&path.to_string_lossy(),
class,
&refusal.why,
&[],
1,
now(),
)?;
say_file(
on_file,
report,
total,
path,
FileOutcome::Refused,
Some(why),
);
Ok(())
}
fn index_set_writing(
&mut self,
walk: Walk,
sweep: bool,
handed: Handed,
control: Option<&dyn Fn() -> Flow>,
mut on_file: impl FnMut(&Path, IndexProgress),
) -> Result<IndexReport> {
self.index.clean();
let mut report = IndexReport::default();
let mut written: Vec<String> = Vec::new();
let mut completed: Vec<(i64, String)> = Vec::new();
let checkpointing = matches!(self.index, index::VectorIndex::Sharded(_));
let every = checkpoint_files();
let mut since_checkpoint = 0usize;
let rechunk = store::format(&self.db)? < store::CODE_CONTEXT;
report.rechunked = rechunk;
let rescan = self.rules_outdated()?;
let regraph = store::get_meta(&self.db, "graph_rules")?
.and_then(|v| v.parse::<u32>().ok())
.unwrap_or(0)
< GRAPH_RULES;
let prehashed = std::mem::take(&mut self.prehashed);
let run_started = std::time::Instant::now();
let home = crate::home::user_home().ok().map(|h| canonical(&h));
let boundary = self.boundary.resolved(home.as_deref());
let (deadline, budget, bytes_known) = match handed {
Handed::Walked(deadline) => (deadline, None, None),
Handed::Rest {
budget,
bytes_total,
} => (None, Some(budget), bytes_total),
};
let each = budget.is_some().then_some(&boundary);
let (head, pending_walk, done_walk) = match walk {
Walk::Done(walked) => (Vec::new(), None, Some(walked)),
Walk::Roots(roots) => {
let accepted = self.accepted_folders();
let head = if accel::managed() {
recent_head(&roots, &boundary, &accepted)
} else {
Vec::new()
};
(head, Some((roots, accepted)), None)
}
};
let mut total;
let mut walk_ms = 0u64;
let first = match done_walk {
Some(walked) => {
walk_ms = walked.walk_ms;
let (paths, counted) =
self.walk_setup(walked, &boundary, each.is_some(), &mut report, &mut on_file)?;
total = counted;
report.bytes_total = bytes_known.unwrap_or_else(|| bytes_of(&paths));
paths
}
None => {
total = head.len();
report.bytes_total = bytes_of(&head);
head.clone()
}
};
let head_set: std::collections::HashSet<PathBuf> = head.iter().cloned().collect();
let deadline = deadline.or_else(|| budget.map(|b| std::time::Instant::now() + b));
let granite = self.model == Model::Granite;
let tokenizer = if granite && (!first.is_empty() || pending_walk.is_some()) {
let cache = model_cache_dir()?;
embed::granite_graph(&cache, self.quiet, embed::index_variant().0)?;
if self.tokenizer.is_none() {
self.tokenizer = self.model.tokenizer(&cache);
}
Some(session::tokenizer(&cache)?)
} else {
None
};
let ctx = pipeline::Context {
rechunk,
rescan,
regraph,
allow_secrets: self.boundary.allow_secrets,
hashes: store::all_hashes(&self.db)?,
uncounted: store::uncounted(&self.db)?,
over_cap: store::acceptances(&self.db)?
.into_iter()
.filter(|a| a.class == store::class::POLICY)
.map(|a| a.path)
.collect(),
prehashed,
each: each.cloned(),
tokenizer,
clocks: pipeline::Clocks::default(),
cache: accel::cache_in_use().then(|| cache::Scope::of(&self.model)),
variants: accel::cache_variants(),
lookups: Default::default(),
hits: Default::default(),
};
let mut cached = CacheWrites::default();
let mut clock = pipeline::WriterClock::start(walk_ms);
let sealed = pending_walk.is_none();
let mut run_lanes: std::collections::BTreeMap<String, usize> = Default::default();
let (cpu_back, cpu_returned) = std::sync::mpsc::channel();
let prepare_threads = pipeline::prepare_threads();
std::thread::scope(|scope| -> Result<()> {
let prefetch = pipeline::Prefetch::start(scope, first, sealed, &ctx, prepare_threads);
let mut walker = pending_walk.map(|(roots, accepted)| {
scope.spawn(move || {
let started = std::time::Instant::now();
let mut walked = walk_allowing(&roots, &accepted);
walked.walk_ms = started.elapsed().as_millis() as u64;
walked
})
});
let mut stage: Option<pipeline::Embedder> = None;
let mut window = pipeline::Window {
ids: Vec::new(),
pieces: Vec::new(),
hashes: Vec::new(),
};
let paused = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
let ask = |paused: &std::sync::atomic::AtomicBool| -> bool {
let Some(ask) = control else { return true };
match ask() {
Flow::Pause => {
paused.store(true, std::sync::atomic::Ordering::Relaxed);
std::thread::sleep(PAUSE_TICK);
true
}
Flow::Stop => false,
Flow::Run | Flow::Yield => {
paused.store(false, std::sync::atomic::Ordering::Relaxed);
true
}
}
};
let mut committed = std::time::Instant::now();
let finish_walk = |me: &mut Self,
walker: &mut Option<std::thread::ScopedJoinHandle<'_, Walked>>,
report: &mut IndexReport,
total: &mut usize,
clock: &mut pipeline::WriterClock,
on_file: &mut dyn FnMut(&Path, IndexProgress)|
-> Result<()> {
let Some(handle) = walker.take() else {
return Ok(());
};
let waiting = std::time::Instant::now();
let walked = handle
.join()
.map_err(|_| anyhow::anyhow!("the walk of the tree panicked"))?;
clock.waited_on_walk(waiting);
let (rest, counted) = me.walk_setup(walked, &boundary, false, report, on_file)?;
let rest: Vec<PathBuf> =
rest.into_iter().filter(|p| !head_set.contains(p)).collect();
*total = head_set.len() + counted.saturating_sub(head_set.len());
report.bytes_total += bytes_of(&rest);
prefetch.extend_and_seal(rest);
Ok(())
};
let mut next = 0usize;
loop {
let seen = next;
let (listed, complete) = prefetch.len();
if seen >= listed {
if complete {
break;
}
finish_walk(
self,
&mut walker,
&mut report,
&mut total,
&mut clock,
&mut on_file,
)?;
continue;
}
next += 1;
let path = prefetch.path(seen).expect("inside the list");
if let Some(deadline) = deadline
&& report.indexed > 0
&& std::time::Instant::now() >= deadline
{
finish_walk(
self,
&mut walker,
&mut report,
&mut total,
&mut clock,
&mut on_file,
)?;
report.pending = prefetch.from(seen);
report.remaining = report.pending.len();
break;
}
if let Some(ask) = control {
let mut stop = false;
let mut yielded = false;
loop {
match ask() {
Flow::Run => break,
Flow::Pause => {
paused.store(true, std::sync::atomic::Ordering::Relaxed);
std::thread::sleep(PAUSE_TICK)
}
Flow::Yield => {
yielded = true;
break;
}
Flow::Stop => {
stop = true;
break;
}
}
}
paused.store(false, std::sync::atomic::Ordering::Relaxed);
if yielded {
finish_walk(
self,
&mut walker,
&mut report,
&mut total,
&mut clock,
&mut on_file,
)?;
report.pending = prefetch.from(seen);
report.remaining = report.pending.len();
break;
}
if stop {
report.remaining = prefetch.len().0 - seen;
report.stopped = true;
break;
}
}
if let Some(stage) = stage.as_mut() {
while let Some(done) = stage.next(std::time::Duration::ZERO) {
let done = done.map_err(anyhow::Error::msg)?;
self.land(done, &mut report, &mut run_lanes, &mut cached)?;
say_file(
&mut on_file,
&report,
total,
&path,
FileOutcome::Progress,
None,
);
}
}
let waiting = std::time::Instant::now();
let prepared = prefetch.take(seen);
clock.waited_on_prepare(waiting);
let key = path.to_string_lossy().into_owned();
if let pipeline::Prepared::Refused(refusal) = &prepared {
self.refuse_file(&mut report, total, &path, refusal, &mut on_file)?;
continue;
}
report.scanned += 1;
let bytes_before = report.bytes;
let (hash, len, text, units, found, ready) = match prepared {
pipeline::Prepared::Refused(_) => unreachable!("handled above"),
pipeline::Prepared::Unchanged {
size,
mtime,
file_bytes,
} => {
report.bytes += file_bytes;
self.tx_begin()?;
store::heal_stamps(&self.db, &key, size as i64, mtime, now())?;
report.unchanged += 1;
say_file(
&mut on_file,
&report,
total,
&path,
FileOutcome::Unchanged,
None,
);
continue;
}
pipeline::Prepared::Unusable {
why,
gone,
file_bytes,
} => {
report.bytes += file_bytes;
self.tx_begin()?;
if gone {
let ids = store::delete_file(&self.db, &key, now())?;
if !ids.is_empty() {
for id in ids {
self.index.remove(id)?;
}
report.removed += 1;
say_file(
&mut on_file,
&report,
total,
&path,
FileOutcome::Removed,
None,
);
continue;
}
}
skip(&mut report, &why);
store::refuse(&self.db, &key, why.class(), &why.as_str(), &[], 1, now())?;
say_file(
&mut on_file,
&report,
total,
&path,
FileOutcome::Skipped,
Some(why.as_str()),
);
continue;
}
pipeline::Prepared::Same {
len,
mtime,
file_bytes,
rescan: rescanned,
regraph: regraphed,
units,
} => {
report.bytes += file_bytes;
self.tx_begin()?;
if let Some(units) = units {
store::set_units_by_path(&self.db, &key, units)?;
}
if let Some((text, found)) = rescanned
&& let keyscan::Decision::Refuse(why) =
keyscan::decide(&self.db, &key, &text, &found)?
{
let live: Vec<keyscan::Match> =
found.into_iter().filter(|m| m.dummy.is_none()).collect();
store::refuse(
&self.db,
&key,
store::class::CONTENT,
&why,
&live,
1,
now(),
)?;
let (gone, images) = self.evict(&key)?;
report.removed += usize::from(gone + images > 0);
report
.refused
.push((path.display().to_string(), why.clone()));
say_file(
&mut on_file,
&report,
total,
&path,
FileOutcome::Refused,
Some(why),
);
continue;
}
if let Some(extraction) = regraphed
&& let Some((file_id, spans)) = store::graph_input(&self.db, &key)?
{
store::delete_graph(&self.db, file_id)?;
let (symbols, edges) =
self.write_graph(Some(extraction), file_id, &spans)?;
report.symbols += symbols;
report.edges += edges;
}
if let Some(mtime) = mtime {
store::heal_stamps(&self.db, &key, len as i64, mtime, now())?;
}
report.unchanged += 1;
say_file(
&mut on_file,
&report,
total,
&path,
FileOutcome::Unchanged,
None,
);
continue;
}
pipeline::Prepared::Image {
bytes,
hash,
file_bytes,
} => {
report.bytes += file_bytes;
self.tx_begin()?;
if let Some(why) = image::too_large(&bytes) {
report
.refused
.push((path.display().to_string(), why.clone()));
say_file(
&mut on_file,
&report,
total,
&path,
FileOutcome::Refused,
Some(why),
);
continue;
}
let Some((width, height)) = image::dimensions(&bytes) else {
let why = SkipReason::NotDecodableImage;
skip(&mut report, &why);
store::refuse(
&self.db,
&key,
why.class(),
&why.as_str(),
&[],
1,
now(),
)?;
say_file(
&mut on_file,
&report,
total,
&path,
FileOutcome::Skipped,
Some(why.as_str()),
);
continue;
};
say_file(
&mut on_file,
&report,
total,
&path,
FileOutcome::Indexing,
None,
);
let vector = match self.clip.embed_image(&path, self.quiet) {
Ok(vector) => vector,
Err(e) => {
failed(&mut report, &path, &e);
say_file(
&mut on_file,
&report,
total,
&path,
FileOutcome::Failed,
Some(format!("{e:#}")),
);
continue;
}
};
for id in store::image_ids_of(&self.db, &key)? {
self.images.remove(id as u64)?;
}
for id in store::delete_file(&self.db, &key, now())? {
self.index.remove(id)?;
}
let file_id =
store::insert_file(&self.db, &key, PENDING, bytes.len() as u64, now())?;
written.push(key.clone());
let image_id = store::insert_image(&self.db, file_id, width, height)?;
self.images.add(&vector, &[image_id as u64])?;
completed.push((file_id, hash));
report.indexed += 1;
report.images += 1;
continue;
}
pipeline::Prepared::Failed { error, file_bytes } => {
report.bytes += file_bytes;
failed(&mut report, &path, &error);
say_file(
&mut on_file,
&report,
total,
&path,
FileOutcome::Failed,
Some(format!("{error:#}")),
);
continue;
}
pipeline::Prepared::Skipped { why, file_bytes } => {
report.bytes += file_bytes;
self.tx_begin()?;
skip(&mut report, &why);
store::refuse(&self.db, &key, why.class(), &why.as_str(), &[], 1, now())?;
say_file(
&mut on_file,
&report,
total,
&path,
FileOutcome::Skipped,
Some(why.as_str()),
);
continue;
}
pipeline::Prepared::Text {
len,
hash,
file_bytes,
text,
units,
found,
ready,
} => {
report.bytes += file_bytes;
(hash, len, text, units, found, ready)
}
};
let file_bytes = report.bytes - bytes_before;
self.tx_begin()?;
let decided = if self.boundary.allow_secrets {
if keyscan::refuses(&found) {
report.secrets_indexed += 1;
}
store::unrefuse(&self.db, &key)?;
None
} else {
match keyscan::decide(&self.db, &key, &text, &found)? {
keyscan::Decision::Index {
text: decided,
accepted,
} => {
if accepted {
report.accepted_indexed += 1;
store::unrefuse(&self.db, &key)?;
} else if !found.is_empty() {
report.dummies_indexed += 1;
store::refuse(
&self.db,
&key,
store::class::DUMMY,
"every match is a declared test dummy",
&found,
1,
now(),
)?;
} else {
store::unrefuse(&self.db, &key)?;
}
(decided != text).then_some(decided)
}
keyscan::Decision::Refuse(why) => {
let live: Vec<keyscan::Match> = found
.iter()
.filter(|m| m.dummy.is_none())
.cloned()
.collect();
store::refuse(
&self.db,
&key,
store::class::CONTENT,
&why,
if live.is_empty() { &found } else { &live },
1,
now(),
)?;
let (gone, images) = self.evict(&key)?;
let evicted = gone + images;
report.removed += usize::from(evicted > 0);
let why = if evicted > 0 {
format!(
"{why}. Its earlier contents have been removed from this store."
)
} else {
why
};
report
.refused
.push((path.display().to_string(), why.clone()));
say_file(
&mut on_file,
&report,
total,
&path,
FileOutcome::Refused,
Some(why),
);
continue;
}
}
};
let ready = match (decided, ready) {
(Some(redacted), _) => pipeline::take_ready(
&path,
&redacted,
&ctx,
None,
std::time::Instant::now(),
),
(None, Some(ready)) => ready,
(None, None) => {
pipeline::take_ready(&path, &text, &ctx, None, std::time::Instant::now())
}
};
drop(text);
let pipeline::Ready {
extraction,
chunks,
pieces,
hashes,
} = match ready {
Ok(ready) => ready,
Err(e) => {
failed(&mut report, &path, &e);
say_file(
&mut on_file,
&report,
total,
&path,
FileOutcome::Failed,
Some(format!("{e:#}")),
);
continue;
}
};
if chunks.is_empty() {
let why = SkipReason::NoText;
skip(&mut report, &why);
store::refuse(&self.db, &key, why.class(), &why.as_str(), &[], 1, now())?;
say_file(
&mut on_file,
&report,
total,
&path,
FileOutcome::Skipped,
Some(why.as_str()),
);
continue;
}
let file_end = report.bytes;
report.bytes = bytes_before;
say_file(
&mut on_file,
&report,
total,
&path,
FileOutcome::Indexing,
None,
);
for id in store::delete_file(&self.db, &key, now())? {
self.index.remove(id)?;
}
let file_id = store::insert_file(&self.db, &key, PENDING, len, now())?;
written.push(key.clone());
let parsed = match (&extraction, graph::language_of(&path)) {
(Some(_), _) => "parsed",
(None, Some(lang)) if graph::has_graph(lang) => "timeout",
(None, _) => "none",
};
store::set_file_graph(&self.db, file_id, parsed)?;
if let Some(units) = units {
store::set_file_units(&self.db, file_id, units)?;
}
let mut spans: Vec<(u32, u32, i64)> = Vec::with_capacity(chunks.len());
let mut halted = false;
let count = chunks.len();
for (((ord, c), piece), hash) in chunks.iter().enumerate().zip(pieces).zip(hashes) {
let id = store::insert_chunk(
&self.db,
file_id,
ord,
c.start_line,
c.end_line,
&c.text,
)?;
spans.push((c.start_line, c.end_line, id));
report.rows += 1;
window.ids.push(id as u64);
window.pieces.push(piece);
window.hashes.push(hash);
if window.ids.len() >= SORT_WINDOW {
report.bytes = bytes_before + file_bytes * (ord as u64 + 1) / count as u64;
if !self.hand_over(
scope,
&mut stage,
&mut window,
&cpu_back,
&paused,
&ask,
&mut clock,
&mut |n| {
report.batched = n;
say_file(
&mut on_file,
&report,
total,
&path,
FileOutcome::Progress,
None,
)
},
)? {
halted = true;
break;
}
}
}
report.threads = self.index_threads();
report.lanes = self.lane_chunks.clone();
report.bytes = file_end;
if halted {
report.remaining = prefetch.len().0 - seen;
report.stopped = true;
break;
}
let (symbols, edges) = self.write_graph(extraction, file_id, &spans)?;
report.symbols += symbols;
report.edges += edges;
completed.push((file_id, hash));
report.indexed += 1;
report.chunks += count;
if committed.elapsed() >= ROW_COMMIT {
self.tx_commit()?;
committed = std::time::Instant::now();
}
since_checkpoint += 1;
if checkpointing && since_checkpoint >= every {
say_file(
&mut on_file,
&report,
total,
&path,
FileOutcome::Writing,
Some("writing the index to disk".to_string()),
);
let drained = self.hand_over(
scope,
&mut stage,
&mut window,
&cpu_back,
&paused,
&ask,
&mut clock,
&mut |n| {
report.batched = n;
say_file(
&mut on_file,
&report,
total,
&path,
FileOutcome::Progress,
None,
)
},
)? && self.drain(
&mut stage,
&mut report,
&mut run_lanes,
&mut cached,
&paused,
&ask,
&mut clock,
&mut |report: &IndexReport| {
say_file(
&mut on_file,
report,
total,
&path,
FileOutcome::Progress,
None,
)
},
)?;
if !drained {
report.remaining = prefetch.len().0 - seen - 1;
report.stopped = true;
break;
}
self.checkpoint(&mut completed)?;
since_checkpoint = 0;
say_file(
&mut on_file,
&report,
total,
&path,
FileOutcome::Indexing,
None,
);
}
}
let last = prefetch
.path(prefetch.len().0.saturating_sub(1))
.unwrap_or_default();
drop(prefetch);
if !report.stopped {
let drained = self.hand_over(
scope,
&mut stage,
&mut window,
&cpu_back,
&paused,
&ask,
&mut clock,
&mut |n| {
report.batched = n;
say_file(
&mut on_file,
&report,
total,
&last,
FileOutcome::Progress,
None,
)
},
)? && self.drain(
&mut stage,
&mut report,
&mut run_lanes,
&mut cached,
&paused,
&ask,
&mut clock,
&mut |report: &IndexReport| {
say_file(
&mut on_file,
report,
total,
&last,
FileOutcome::Progress,
None,
)
},
)?;
if !drained {
report.stopped = true;
}
}
if let Some(stage) = stage.as_mut() {
if report.stopped {
stage.cancel();
} else {
stage.close();
}
}
drop(stage);
drop(cpu_back);
Ok(())
})?;
if let Ok(work) = cpu_returned.try_recv() {
if let pipeline::CpuWork::Ids { main, .. } = &work {
self.index_threads = main.threads();
}
self.index_session = Some(work);
}
self.tx_commit()?;
report.stages = clock.stages(&ctx.clocks, &run_lanes);
report.threads = self.index_threads();
report.cache_lookups = ctx.lookups.load(std::sync::atomic::Ordering::Relaxed) as usize;
report.cache_hits = ctx.hits.load(std::sync::atomic::Ordering::Relaxed) as usize;
if let Some(scope) = &ctx.cache
&& report.cache_lookups > 0
&& let Some(mut cache) = cache::Cache::open()
{
let fresh: Vec<([u8; 32], &'static str, &[f32])> = cached
.fresh
.iter()
.map(|(hash, variant, vector)| {
(scope.key(hash, variant), *variant, vector.as_slice())
})
.collect();
let hits: Vec<[u8; 32]> = cached
.hits
.iter()
.map(|(hash, variant)| scope.key(hash, variant))
.collect();
let _ = cache.record(&fresh, &hits, report.cache_lookups as u64);
}
report.written = written;
if sweep && report.pending.is_empty() && !report.stopped {
for key in store::all_paths(&self.db)? {
if !Path::new(&key).exists() {
for id in store::delete_file(&self.db, &key, now())? {
self.index.remove(id)?;
}
report.removed += 1;
}
}
for row in store::refusals(&self.db)? {
if !Path::new(&row.path).exists() {
store::unrefuse(&self.db, &row.path)?;
}
}
if rescan {
store::set_meta(&self.db, "scan_rules", &SCAN_RULES.to_string())?;
}
if regraph {
store::set_meta(&self.db, "graph_rules", &GRAPH_RULES.to_string())?;
}
}
if report.indexed > 0 || report.removed > 0 || !self.index.exists() {
self.save()?;
}
self.commit_hashes(&mut completed)?;
let secs = run_started.elapsed().as_secs_f64();
if report.indexed > 0 && report.bytes >= 64 * 1024 && secs >= 0.25 {
let sample = report.bytes as f64 / secs;
let rate = match store::get_meta(&self.db, "embed_bytes_per_sec")?
.and_then(|v| v.parse::<f64>().ok())
.filter(|r| *r > 0.0)
{
Some(known) => (known + sample) / 2.0,
None => sample,
};
store::set_meta(&self.db, "embed_bytes_per_sec", &format!("{rate:.0}"))?;
}
if rechunk && sweep && report.pending.is_empty() && !report.stopped {
store::set_meta(
&self.db,
store::FORMAT_KEY,
&store::FORMAT_VERSION.to_string(),
)?;
}
if std::mem::take(&mut self.scrub) {
store::reclaim(&self.db)?;
}
Ok(report)
}
fn walk_setup(
&mut self,
walked: Walked,
boundary: &ResolvedBoundary,
each: bool,
report: &mut IndexReport,
on_file: &mut dyn FnMut(&Path, IndexProgress),
) -> Result<(Vec<PathBuf>, usize)> {
let Walked {
walk_ms: _,
files: walked_paths,
named,
unreadable: unwalkable,
generated,
credentials: hidden_credentials,
excluded,
} = walked;
self.tx_begin()?;
let (paths, refused): (Vec<PathBuf>, Vec<(PathBuf, Refusal)>) = if each {
(walked_paths, Vec::new())
} else {
let mut allowed = Vec::with_capacity(walked_paths.len() + named.len());
let mut refused = Vec::new();
let all = named
.into_iter()
.map(|p| (p, false))
.chain(walked_paths.into_iter().map(|p| (p, true)));
for (path, walked) in all {
match boundary.refuses(&path, walked) {
Some(why) => refused.push((path, why)),
None => allowed.push(path),
}
}
allowed.sort();
if accel::managed() {
recent_first(&mut allowed);
}
(allowed, refused)
};
let total = paths.len() + refused.len() + unwalkable.len();
for (path, refusal) in &refused {
self.refuse_file(report, total, path, refusal, &mut *on_file)?;
}
report.generated = generated.iter().map(|p| p.display().to_string()).collect();
for (path, rule) in &excluded {
let folder = path.is_dir();
if rule == IGNORE_FILE {
report.skipped += 1;
*report
.skipped_reasons
.entry(IGNORE_FILE.to_string())
.or_insert(0) += 1;
}
store::refuse(
&self.db,
&path.to_string_lossy(),
store::class::EXCLUDED,
&format!("left out by {rule}; change the rule rather than accept the file"),
&[],
if folder { 0 } else { 1 },
now(),
)?;
}
if !excluded.is_empty() {
let out: Vec<&Path> = excluded.iter().map(|(p, _)| p.as_path()).collect();
for key in store::all_paths(&self.db)? {
if out.iter().any(|o| Path::new(&key).starts_with(o)) {
let (chunks, images) = self.evict(&key)?;
report.removed += usize::from(chunks + images > 0);
}
}
}
for path in &hidden_credentials {
let why = filter::denied(path).map(|d| d.reason()).unwrap_or_default();
store::refuse(
&self.db,
&path.to_string_lossy(),
store::class::CREDENTIAL,
&why,
&[],
1,
now(),
)?;
}
for dir in &generated {
store::refuse(
&self.db,
&dir.to_string_lossy(),
store::class::POLICY,
"a generated or vendored folder the walk steps over",
&[],
0,
now(),
)?;
}
for (path, why) in &unwalkable {
store::refuse(
&self.db,
&path.to_string_lossy(),
store::class::UNINDEXABLE,
why,
&[],
1,
now(),
)?;
report
.failed
.push((path.display().to_string(), why.clone()));
report.scanned += 1;
say_file(
&mut *on_file,
report,
total,
path,
FileOutcome::Failed,
Some(why.clone()),
);
}
self.tx_commit()?;
Ok((paths, total))
}
fn checkpoint(&mut self, completed: &mut Vec<(i64, String)>) -> Result<()> {
self.tx_commit()?;
if completed.is_empty() {
return Ok(());
}
self.save()?;
self.commit_hashes(completed)
}
fn expansion_weight(confidence: &str) -> f32 {
match confidence {
graph::EXTRACTED | graph::RESOLVED => 1.0,
_ => INFERRED_EXPANSION,
}
}
fn with_pending(
&mut self,
query: &[f32],
keyword: &[u64],
scores: Vec<f32>,
ids: Vec<u64>,
) -> Result<(Vec<f32>, Vec<u64>)> {
if store::pending_share(&self.db)?.is_none() {
return Ok((scores, ids));
}
let candidates: Vec<u64> = keyword
.iter()
.filter(|id| !ids.contains(id))
.copied()
.collect();
let pending: Vec<(u64, String)> = store::pending_among(&self.db, &candidates)?
.into_iter()
.take(PENDING_EMBED)
.collect();
if pending.is_empty() {
return Ok((scores, ids));
}
let texts: Vec<String> = pending.iter().map(|(_, text)| text.clone()).collect();
let vectors = self.embed(texts)?;
let mut merged: Vec<(u64, f32)> = ids.into_iter().zip(scores).collect();
for ((id, _), vector) in pending.iter().zip(&vectors) {
merged.push((*id, index::cosine(query, vector)));
}
merged.sort_by(|a, b| b.1.total_cmp(&a.1));
let (ids, scores) = merged.into_iter().unzip();
Ok((scores, ids))
}
pub fn pending_share(&self) -> Result<Option<f64>> {
store::pending_share(&self.db)
}
fn rescored(&self, query: &[f32], scores: Vec<f32>, ids: Vec<u64>) -> (Vec<f32>, Vec<u64>) {
if ids.len() < 2 || !self.exact.exists() {
return (scores, ids);
}
let mut exact: Vec<(u64, f32)> = Vec::with_capacity(ids.len());
for id in &ids {
match self.exact.get(*id) {
Some(vector) => exact.push((*id, index::cosine(query, &vector))),
None => return (scores, ids),
}
}
exact.sort_by(|a, b| b.1.total_cmp(&a.1));
let (rescored_ids, rescored_scores): (Vec<u64>, Vec<f32>) = exact.into_iter().unzip();
(rescored_scores, rescored_ids)
}
fn graph_expansion(
&self,
dense: &[u64],
keyword: &[u64],
depth: usize,
filter: &Filter,
) -> Result<Vec<Reached>> {
const SEEDS: usize = 8;
const VISITS: usize = 128;
const HUB: usize = 32;
const RESOLVE_BATCH: usize = 256;
let mut mass: std::collections::BTreeMap<u64, f32> = std::collections::BTreeMap::new();
for list in [dense, keyword] {
for (rank, id) in list.iter().take(SEEDS).enumerate() {
*mass.entry(*id).or_default() += 1.0 / (RRF_K + rank as f32 + 1.0);
}
}
if mass.is_empty() {
return Ok(Vec::new());
}
let mut personal: std::collections::BTreeMap<String, f32> =
std::collections::BTreeMap::new();
for (id, mass) in &mass {
let names = store::symbols_in_chunks(&self.db, &[*id])?;
if names.is_empty() {
continue;
}
let share = mass / names.len() as f32;
for name in names {
*personal.entry(name).or_default() += share;
}
}
if personal.is_empty() {
return Ok(Vec::new());
}
let kinds = graph::dependency_kinds();
let ranked = graph::expand(&personal, VISITS, |name| {
if store::defined_more_than(&self.db, name, HUB)? {
return Ok(Vec::new());
}
let mut out = Vec::new();
for end in store::edges_out_short_of_hubs(&self.db, name, &kinds, HUB)?
.into_iter()
.chain(store::edges_in(&self.db, name, &kinds)?)
{
if end.confidence == graph::AMBIGUOUS {
continue;
}
if !crate::filter::is_code(&end.symbol.path) {
continue;
}
let weight = Self::expansion_weight(&end.confidence);
out.push((end.symbol.name, weight, end.confidence));
}
Ok(out)
})?;
let best = ranked.first().map(|(_, mass, _)| *mass);
let mut ids: Vec<Reached> = Vec::new();
'names: for batch in ranked.chunks(RESOLVE_BATCH) {
let names: Vec<String> = batch.iter().map(|(name, _, _)| name.clone()).collect();
let place: std::collections::HashMap<&str, usize> = names
.iter()
.enumerate()
.map(|(i, name)| (name.as_str(), i))
.collect();
let mut symbols = store::symbols_by_names(&self.db, &names, filter.groups())?;
symbols.sort_by_key(|s| place.get(s.name.as_str()).copied().unwrap_or(usize::MAX));
for symbol in symbols {
let Some(chunk_id) = symbol.chunk_id else {
continue;
};
let id = chunk_id as u64;
let found = batch.iter().find(|(name, _, _)| *name == symbol.name);
let tier = found
.map(|(_, _, tier)| tier.clone())
.unwrap_or_else(|| graph::INFERRED.to_string());
let weight = Self::expansion_weight(&tier);
let proximity = match (found, best) {
(Some((_, mass, _)), Some(best)) if best > 0.0 => mass / best,
_ => 0.0,
};
if !ids.iter().any(|r| r.id == id) {
ids.push(Reached {
id,
weight,
tier,
proximity,
});
}
if ids.len() >= depth {
break 'names;
}
}
}
Ok(ids)
}
fn write_graph(
&self,
extraction: Option<graph::Extraction>,
file_id: i64,
spans: &[(u32, u32, i64)],
) -> Result<(usize, usize)> {
let Some(extraction) = extraction else {
return Ok((0, 0));
};
let mut ids: std::collections::HashMap<&str, (i64, bool)> =
std::collections::HashMap::new();
for symbol in &extraction.symbols {
let chunk_id = spans
.iter()
.find(|(start, end, _)| symbol.start_line >= *start && symbol.start_line <= *end)
.map(|(_, _, id)| *id);
let id = store::insert_symbol(&self.db, file_id, chunk_id, symbol)?;
let module = symbol.kind == "module";
match ids.get(symbol.name.as_str()) {
None => {
ids.insert(symbol.name.as_str(), (id, module));
}
Some((_, true)) if !module => {
ids.insert(symbol.name.as_str(), (id, false));
}
_ => {}
}
}
let ids: std::collections::HashMap<&str, i64> =
ids.into_iter().map(|(k, (id, _))| (k, id)).collect();
let mut written = 0;
for edge in &extraction.edges {
let Some(src) = ids.get(edge.from.as_str()) else {
continue;
};
store::insert_edge(
&self.db,
*src,
&edge.to,
&edge.kind,
&edge.confidence,
edge.hint.as_deref(),
edge.line,
)?;
written += 1;
}
Ok((extraction.symbols.len(), written))
}
fn commit_hashes(&mut self, completed: &mut Vec<(i64, String)>) -> Result<()> {
if completed.is_empty() {
return Ok(());
}
let tx = self.db.unchecked_transaction()?;
for (file_id, hash) in completed.iter() {
tx.execute(
"UPDATE files SET hash = ?1 WHERE id = ?2",
rusqlite::params![hash, file_id],
)?;
}
tx.commit()?;
completed.clear();
Ok(())
}
#[allow(clippy::too_many_arguments)]
fn hand_over<'scope>(
&mut self,
scope: &'scope std::thread::Scope<'scope, '_>,
stage: &mut Option<pipeline::Embedder>,
window: &mut pipeline::Window,
cpu_back: &std::sync::mpsc::Sender<pipeline::CpuWork>,
paused: &std::sync::Arc<std::sync::atomic::AtomicBool>,
ask: &dyn Fn(&std::sync::atomic::AtomicBool) -> bool,
clock: &mut pipeline::WriterClock,
tick: &mut dyn FnMut(usize),
) -> Result<bool> {
if window.ids.is_empty() {
return Ok(true);
}
self.tx_commit()?;
if stage.is_none() {
let work = self.index_work()?;
*stage = Some(pipeline::Embedder::start(
scope,
work,
self.model == Model::Granite,
cpu_back.clone(),
std::sync::Arc::clone(paused),
));
}
let running = stage.as_mut().expect("started above");
let full = std::mem::replace(
window,
pipeline::Window {
ids: Vec::new(),
pieces: Vec::new(),
hashes: Vec::new(),
},
);
let waiting = std::time::Instant::now();
let mut stopped = false;
let batched = std::sync::Arc::clone(&running.batched);
let mut progress = Progress::default();
let sent = running.send(full, || {
if !ask(paused) {
stopped = true;
return false;
}
if !paused.load(std::sync::atomic::Ordering::Relaxed)
&& let Some(n) = progress.due(&batched)
{
tick(n);
}
std::thread::sleep(std::time::Duration::from_millis(1));
true
});
clock.waited_on_embed(waiting);
if !sent && !stopped {
bail!("the embed stage ended before the run did");
}
Ok(sent)
}
#[allow(clippy::too_many_arguments)]
fn drain(
&mut self,
stage: &mut Option<pipeline::Embedder>,
report: &mut IndexReport,
run_lanes: &mut std::collections::BTreeMap<String, usize>,
cached: &mut CacheWrites,
paused: &std::sync::atomic::AtomicBool,
ask: &dyn Fn(&std::sync::atomic::AtomicBool) -> bool,
clock: &mut pipeline::WriterClock,
tick: &mut dyn FnMut(&IndexReport),
) -> Result<bool> {
let Some(running) = stage.as_mut() else {
return Ok(true);
};
let mut progress = Progress::default();
while running.outstanding > 0 {
let waiting = std::time::Instant::now();
let got = running.next(std::time::Duration::from_millis(20));
clock.waited_on_embed(waiting);
match got {
Some(done) => {
self.land(done.map_err(anyhow::Error::msg)?, report, run_lanes, cached)?;
if !paused.load(std::sync::atomic::Ordering::Relaxed) {
tick(report);
}
}
None => {
if !ask(paused) {
return Ok(false);
}
if !paused.load(std::sync::atomic::Ordering::Relaxed)
&& let Some(n) = progress.due(&running.batched)
{
report.batched = n;
tick(report);
}
}
}
}
Ok(true)
}
fn land(
&mut self,
done: pipeline::Embedded,
report: &mut IndexReport,
run_lanes: &mut std::collections::BTreeMap<String, usize>,
cached: &mut CacheWrites,
) -> Result<()> {
let flat: Vec<f32> = done.vectors.iter().flatten().copied().collect();
anyhow::ensure!(
flat.len() == done.ids.len() * self.dim,
"the embed stage handed back {} values for {} chunks",
flat.len(),
done.ids.len()
);
self.follow_budget();
self.index.add(&flat, &done.ids)?;
self.exact.append(&flat, &done.ids)?;
self.record_variants(&done.variants)?;
for (lane, n) in &done.lanes {
self.note_lane(lane, *n);
*run_lanes.entry(lane.to_string()).or_default() += n;
}
report.embedded += done.ids.len();
report.threads = self.index_threads();
report.lanes = self.lane_chunks.clone();
self.last_embed = Some(std::time::Instant::now());
for (at, ((variant, from_cache), hash)) in done.rows.iter().zip(&done.hashes).enumerate() {
let Some(hash) = hash else { continue };
if *from_cache {
cached.hits.push((*hash, variant));
} else if !variant.is_empty() {
cached
.fresh
.push((*hash, variant, done.vectors[at].clone()));
}
}
Ok(())
}
fn index_work(&mut self) -> Result<pipeline::CpuWork> {
let threads = embed::threads_in_force();
if let Some(work) = self.index_session.take() {
if self.index_threads == threads {
return Ok(work);
}
SESSIONS.fetch_sub(1, std::sync::atomic::Ordering::Relaxed);
}
let cache = model_cache_dir()?;
let work = match self.model {
Model::Granite => {
let (main, alt) = embed::index_variant();
pipeline::CpuWork::Ids {
main: session::CpuSession::open(&cache, main, threads, self.quiet)?,
alt: alt
.map(|variant| {
session::CpuSession::open(&cache, variant, threads, self.quiet)
})
.transpose()?,
}
}
Model::Builtin(_) => pipeline::CpuWork::Text(Box::new(self.model.load_variant(
cache,
chunk::MAX_CHARS / 2,
self.quiet,
threads,
embed::Variant::Int8,
)?)),
};
self.index_threads = threads;
SESSIONS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
Ok(work)
}
fn index_threads(&self) -> usize {
match pipeline::CPU_THREADS.load(std::sync::atomic::Ordering::Relaxed) {
0 => self.index_threads,
live => live,
}
}
fn tx_begin(&self) -> Result<()> {
if self.db.is_autocommit() {
self.db.execute_batch("BEGIN")?;
}
Ok(())
}
fn tx_commit(&self) -> Result<()> {
if !self.db.is_autocommit() {
self.db.execute_batch("COMMIT")?;
}
Ok(())
}
fn note_lane(&mut self, lane: &'static str, chunks: usize) {
*self.lane_chunks.entry(lane.to_string()).or_default() += chunks;
}
fn record_variants(&self, counts: &[(&'static str, usize)]) -> Result<()> {
if counts.is_empty() {
return Ok(());
}
let mut variants: std::collections::BTreeMap<String, u64> =
store::get_meta(&self.db, VARIANTS_KEY)?
.and_then(|text| serde_json::from_str(&text).ok())
.unwrap_or_default();
for (variant, n) in counts {
*variants.entry((*variant).to_string()).or_default() += *n as u64;
}
store::set_meta(&self.db, VARIANTS_KEY, &serde_json::to_string(&variants)?)
}
pub fn variants(&self) -> std::collections::BTreeMap<String, u64> {
store::get_meta(&self.db, VARIANTS_KEY)
.ok()
.flatten()
.and_then(|text| serde_json::from_str(&text).ok())
.unwrap_or_default()
}
pub(crate) fn undo_held(&mut self, keys: &[String]) -> Result<usize> {
self.writing(|me| {
for key in keys {
me.evict(key)?;
}
me.save()?;
Ok(keys.len())
})
}
pub fn forget(&mut self, path: &Path) -> Result<(usize, usize)> {
let _lock = lock::StoreLock::acquire(&self.dir)?;
self.forget_held(path)
}
pub(crate) fn forget_held(&mut self, path: &Path) -> Result<(usize, usize)> {
self.writing(|me| me.forget_writing(path))
}
fn forget_writing(&mut self, path: &Path) -> Result<(usize, usize)> {
let key = canonical(path).to_string_lossy().into_owned();
let (chunks, images) = self.evict(&key)?;
if chunks > 0 || images > 0 {
self.save()?;
}
Ok((chunks, images))
}
fn evict(&mut self, key: &str) -> Result<(usize, usize)> {
let images = store::image_ids_of(&self.db, key)?;
let ids = store::delete_file(&self.db, key, now())?;
for id in &ids {
self.index.remove(*id)?;
}
for id in &images {
self.images.remove(*id as u64)?;
}
Ok((ids.len(), images.len()))
}
pub fn out_of_root(&self, roots: &[PathBuf]) -> Result<Vec<String>> {
if roots.is_empty() {
return Ok(Vec::new());
}
let resolved = filter::resolve_boundary(roots);
Ok(store::all_paths(&self.db)?
.into_iter()
.filter(|key| !filter::within_resolved(Path::new(key), &resolved))
.collect())
}
pub fn prune_out_of_root(&mut self, roots: &[PathBuf]) -> Result<usize> {
let strays = self.out_of_root(roots)?;
if strays.is_empty() {
return Ok(0);
}
self.writing(|me| {
for key in &strays {
me.evict(key)?;
}
me.index.save()?;
Ok(())
})?;
Ok(strays.len())
}
pub fn scan(&self) -> Result<Vec<Finding>> {
let home = match crate::home::user_home() {
Ok(home) => canonical(&home),
Err(e) => bail!(
"semlith cannot tell where your home directory is, so it cannot say which \
of this store's files sit under a credential directory: {e}. Set \
SEMLITH_HOME, or HOME, and scan again."
),
};
let home = Some(home);
let mut out = Vec::new();
for key in store::all_paths(&self.db)? {
let path = Path::new(&key);
if let Some(why) = filter::denied_against(path, home.as_deref())
&& why != filter::Denied::Hidden
{
out.push(Finding {
path: key,
why: why.reason(),
});
continue;
}
if let Some(found) = filter::scan_text(&store::text_of(&self.db, &key)?) {
out.push(Finding {
path: key,
why: found.reason(),
});
}
}
Ok(out)
}
pub fn forget_findings(&mut self, findings: &[Finding]) -> Result<usize> {
self.writing(|me| {
let mut gone = 0;
for finding in findings {
let (chunks, images) = me.evict(&finding.path)?;
if chunks > 0 || images > 0 {
gone += 1;
}
}
if gone > 0 {
me.save()?;
}
Ok(gone)
})
}
pub fn search(&mut self, query: &str, k: usize) -> Result<Vec<Hit>> {
self.search_filtered(query, k, &Filter::default())
}
fn allowlist(&mut self, filter: &Filter) -> Result<Allowlist> {
if filter.is_empty() {
return Ok(Allowlist::All);
}
let candidates = store::filtered_chunk_ids(&self.db, filter.groups())?;
let mut ids = Vec::with_capacity(candidates.len());
for id in candidates {
if self.index.contains(id)? {
ids.push(id);
}
}
Ok(if ids.is_empty() {
Allowlist::Empty
} else if ids.len() == self.len() {
Allowlist::All
} else {
Allowlist::Subset(ids)
})
}
pub fn matching_files(&self, filter: &Filter) -> Result<i64> {
store::matching_files(&self.db, filter.groups())
}
pub fn matching_paths(&self, filter: &Filter) -> Result<Vec<String>> {
store::filtered_paths(&self.db, filter.groups())
}
pub fn search_filtered(&mut self, query: &str, k: usize, filter: &Filter) -> Result<Vec<Hit>> {
let vector = self.embed_query(query)?;
self.search_with_vector(query, &vector, k, filter)
}
pub fn embed_query(&mut self, query: &str) -> Result<Vec<f32>> {
let text = self.model.query_text(query);
let wanted = embed::query_variant();
if embed::index_variant() == (wanted, None) {
return Ok(self.embed(vec![text])?.remove(0));
}
let _lifted = priority::embedding();
if self.query_embedder.is_none() {
self.query_embedder = Some(self.model.load_variant(
model_cache_dir()?,
chunk::MAX_CHARS / 2,
self.quiet,
embed::embed_threads(),
wanted,
)?);
}
let mut out = self
.query_embedder
.as_mut()
.expect("loaded above")
.embed(vec![text], Some(1))
.map_err(|e| anyhow::anyhow!("{e}"))?;
let mut vector = out.remove(0);
normalize(&mut vector);
Ok(vector)
}
pub fn search_with_vector(
&mut self,
query: &str,
vector: &[f32],
k: usize,
filter: &Filter,
) -> Result<Vec<Hit>> {
Ok(self
.search_ranked(query, vector, k, filter)?
.into_iter()
.map(|(hit, _)| hit)
.collect())
}
pub fn search_ranked(
&mut self,
query: &str,
vector: &[f32],
k: usize,
filter: &Filter,
) -> Result<Vec<(Hit, f32)>> {
self.search_preferring(query, vector, k, filter, Prefer::default())
}
pub fn search_preferring(
&mut self,
query: &str,
vector: &[f32],
k: usize,
filter: &Filter,
prefer: Prefer,
) -> Result<Vec<(Hit, f32)>> {
self.refresh()?;
if k == 0
|| (self.is_empty()
&& store::chunk_count(&self.db)? == 0
&& store::image_count(&self.db)? == 0)
{
return Ok(Vec::new());
}
let depth = (k * RANK_DEPTH).max(k);
let shape = shape_of(query);
let allowlist = self.allowlist(filter)?;
if matches!(allowlist, Allowlist::Empty) {
return Ok(Vec::new());
}
let definitions: Vec<u64> = if shape == Shape::Identifier {
let name = query.trim().to_string();
let mut ids = Vec::new();
for row in
store::symbols_by_names(&self.db, std::slice::from_ref(&name), filter.groups())?
{
if let Some(chunk) = row.chunk_id {
let id = chunk as u64;
if !ids.contains(&id) {
ids.push(id);
}
}
if ids.len() >= DEFINITION_LIFT {
break;
}
}
ids
} else {
Vec::new()
};
let (dense_scores, dense_ids) = self.search_vectors(vector, depth, &allowlist)?;
let seeds = dense_ids.clone();
let (dense_scores, dense_ids) = self.rescored(vector, dense_scores, dense_ids);
let keyword_ids = store::keyword_search(&self.db, query, depth, filter.groups())?;
let (dense_scores, dense_ids) = if shape == Shape::Identifier {
(dense_scores, dense_ids)
} else {
self.with_pending(vector, &keyword_ids, dense_scores, dense_ids)?
};
let images = self.image_search(query, depth, filter)?;
let graph_ids = self
.graph_expansion(&seeds, &keyword_ids, depth, filter)
.unwrap_or_else(|e| {
eprintln!("the graph list failed, searching without it: {e:#}");
Vec::new()
});
let image_list: Vec<(u64, f32)> = images
.iter()
.map(|(id, similarity)| {
let weight = if *similarity >= image_floor() {
TEXT_LISTS
} else {
WEAK_IMAGE
};
(*id, weight)
})
.collect();
let candidate_lists = [
(
"definition",
false,
definitions.iter().map(|id| (*id, 0.0)).collect(),
),
("image", true, image_list),
(
"vector",
false,
dense_ids.iter().map(|id| (*id, 1.0)).collect(),
),
(
"keyword",
false,
keyword_ids
.iter()
.map(|id| (*id, shape.keyword_weight()))
.collect(),
),
(
"graph",
false,
graph_ids.iter().map(|r| (r.id, r.weight)).collect(),
),
];
let (mut fused, lists) = fuse(&candidate_lists, |name| shape.list_constant(name));
if !definitions.is_empty() {
let top = fused.iter().map(|(_, score)| *score).fold(0.0f32, f32::max);
for (at, id) in definitions.iter().enumerate() {
if let Some(slot) = fused.iter().position(|(key, _)| *key == (false, *id)) {
fused[slot].1 = top + (definitions.len() - at) as f32;
}
}
}
let mut ranked: Vec<Candidate> = fused.into_iter().zip(lists).collect();
ranked.sort_by(|a, b| b.0.1.total_cmp(&a.0.1).then_with(|| a.0.0.cmp(&b.0.0)));
ranked.truncate(depth.max(k));
let mut hits = Vec::with_capacity(ranked.len());
let mut proximity: Vec<f32> = Vec::with_capacity(ranked.len());
for (((is_image, id), score), found_by) in ranked {
if is_image {
if let Some(row) = store::image(&self.db, id as i64)? {
hits.push((
Hit {
score,
path: row.path,
start_line: 0,
end_line: 0,
text: String::new(),
store: None,
lists: found_by,
image: Some(Pixels {
width: row.width,
height: row.height,
}),
fresh: true,
symbol: None,
symbol_kind: None,
symbol_line: None,
provenance: None,
copies: Vec::new(),
},
0.0,
));
proximity.push(0.0);
}
continue;
}
if let Some(row) = store::chunk(&self.db, id)? {
let similarity = dense_ids
.iter()
.position(|d| *d == id)
.and_then(|i| dense_scores.get(i).copied())
.unwrap_or(0.0);
hits.push((
Hit {
score,
path: row.path,
start_line: row.start_line,
end_line: row.end_line,
text: row.text,
store: None,
lists: found_by,
image: None,
fresh: true,
symbol: None,
symbol_kind: None,
symbol_line: None,
provenance: provenance_of(&graph_ids, id),
copies: Vec::new(),
},
similarity,
));
proximity.push(
graph_ids
.iter()
.find(|r| r.id == id)
.map(|r| r.proximity)
.unwrap_or(0.0),
);
}
}
self.mark_freshness(&mut hits)?;
self.name_enclosing_symbols(&mut hits)?;
if shape != Shape::Identifier && rerank::enabled() && hits.len() > 1 {
let head = hits.len().min(rerank::RERANK_DEPTH);
let texts: Vec<String> = hits[..head]
.iter()
.map(|(hit, _)| {
let text: String = hit.text.chars().take(rerank::RERANK_CHARS).collect();
format!("{}\n{text}", hit.path)
})
.collect();
let quiet = self.quiet;
let cache = model_cache_dir().unwrap_or_default();
let order = rerank::with(&cache, quiet, |model| {
model.map(|model| rerank::order(model, query, &texts))
});
if let Some(order) = order {
match order {
Ok(order) if order.len() == head => {
let mut scores: Vec<f32> =
hits[..head].iter().map(|(hit, _)| hit.score).collect();
scores.sort_by(|a, b| b.total_cmp(a));
let mut by_fusion: Vec<usize> = (0..head).collect();
by_fusion.sort_by(|a, b| hits[*b].0.score.total_cmp(&hits[*a].0.score));
let mut fused_rank = vec![0usize; head];
for (rank, at) in by_fusion.into_iter().enumerate() {
fused_rank[at] = rank;
}
let mut blended: Vec<(usize, f32)> = order
.into_iter()
.enumerate()
.map(|(rescored, at)| {
let vote = 1.0 / (RRF_K + rescored as f32 + 1.0)
+ 1.0 / (RRF_K + fused_rank[at] as f32 + 1.0);
(at, vote)
})
.collect();
blended.sort_by(|a, b| b.1.total_cmp(&a.1).then_with(|| a.0.cmp(&b.0)));
for (rank, (at, _)) in blended.into_iter().enumerate() {
hits[at].0.score = scores[rank];
hits[at].0.lists.push("rerank");
}
}
Ok(_) => {}
Err(e) => {
if !quiet {
eprintln!("the rescoring pass failed, keeping the fused order: {e}");
}
}
}
}
}
let about_tests = names_tests(query);
for ((hit, _), proximity) in hits.iter_mut().zip(&proximity) {
hit.score *= 1.0 + GRAPH_PROXIMITY * proximity;
if !hit.fresh {
hit.score *= 1.0 - STALE_PENALTY;
}
if !about_tests && is_test_path(&hit.path) {
hit.score *= 1.0 - TEST_PENALTY;
}
hit.score *= prefer.multiplier(filter::is_code(&hit.path));
}
hits.sort_by(|a, b| {
b.0.score
.total_cmp(&a.0.score)
.then_with(|| a.0.path.cmp(&b.0.path))
.then_with(|| a.0.start_line.cmp(&b.0.start_line))
});
collapse_copies(&mut hits);
if !about_tests {
product_first(&mut hits);
}
hits.truncate(k);
Ok(hits)
}
pub fn read(&self, target: &Target, filter: &Filter) -> Result<Option<Read>> {
self.read_within(target, filter, &[])
}
pub fn read_within(
&self,
target: &Target,
filter: &Filter,
roots: &[PathBuf],
) -> Result<Option<Read>> {
let (path, start, end) = match target {
Target::Span { path, start, end } => {
let exact: Vec<String> = if Path::new(path).is_relative() {
let mut found = Vec::new();
for root in roots {
let joined = root.join(path);
let joined = joined.to_string_lossy();
for hit in store::files_ending_with(&self.db, &joined, 2)? {
if plain(&hit) == plain(&joined) && !found.contains(&hit) {
found.push(hit);
}
}
}
found
} else {
Vec::new()
};
let candidates = if exact.is_empty() {
store::files_ending_with(&self.db, path, 8)?
} else {
exact
};
match candidates.len() {
0 => return Ok(None),
1 => (candidates[0].clone(), *start, *end),
_ => anyhow::bail!(
"{path:?} matches {} indexed files: {}. Name more of the path.",
candidates.len(),
candidates.join(", ")
),
}
}
Target::Symbol(name) => {
let (qualifier, bare) = crate::graph::split_qualified(name);
let mut found = store::symbols_named(&self.db, bare, 200)?;
if found
.iter()
.any(|r| !crate::graph::NAVIGATIONAL_KINDS.contains(&r.kind.as_str()))
{
found.retain(|r| !crate::graph::NAVIGATIONAL_KINDS.contains(&r.kind.as_str()));
}
if let Some(q) = qualifier
&& found.iter().any(|r| crate::graph::owned_by(r, q))
{
found.retain(|r| crate::graph::owned_by(r, q));
}
match found.len() {
0 => {
let whole = Target::Span {
path: name.clone(),
start: 1,
end: u32::MAX,
};
return self.read_within(&whole, filter, roots);
}
1 => {
let only = &found[0];
(only.path.clone(), only.start_line, only.end_line)
}
_ => return Ok(Some(Read::Choose(found))),
}
}
};
let allowed = if filter.is_empty() {
None
} else {
Some(store::filtered_chunk_ids(&self.db, filter.groups())?)
};
let chunks: Vec<store::ChunkRow> = store::chunks_overlapping(&self.db, &path, start, end)?
.into_iter()
.filter(|c| match &allowed {
Some(ids) => ids.contains(&(c.id as u64)),
None => true,
})
.collect();
if chunks.is_empty() {
return Ok(None);
}
let mut lines: Vec<(u32, String)> = Vec::new();
for chunk in &chunks {
for (offset, line) in chunk.text.lines().enumerate() {
let number = chunk.start_line + offset as u32;
if lines.iter().any(|(n, _)| *n == number) {
continue;
}
lines.push((number, line.to_string()));
}
}
lines.sort_by_key(|(number, _)| *number);
lines.retain(|(number, _)| *number >= start && *number <= end);
if lines.is_empty() {
return Ok(None);
}
let first = lines.first().map(|(n, _)| *n).unwrap_or(start);
let last = lines.last().map(|(n, _)| *n).unwrap_or(end);
let text = lines
.into_iter()
.map(|(_, line)| line)
.collect::<Vec<_>>()
.join("\n");
let mut span = Span {
path,
start_line: first,
end_line: last,
text,
symbol: None,
symbol_kind: None,
fresh: true,
from_disk: false,
store: None,
};
self.name_enclosing_symbol(&mut span)?;
self.mark_span_freshness(&mut span)?;
if !span.fresh {
self.read_from_disk(&mut span, start, end)?;
}
Ok(Some(Read::One(span)))
}
fn read_from_disk(&self, span: &mut Span, start: u32, end: u32) -> Result<()> {
let Ok(bytes) = std::fs::read(&span.path) else {
return Ok(());
};
if bytes.len() as u64 > chunk::MAX_FILE_BYTES {
return Ok(());
}
let hash = blake3::hash(&bytes).to_hex().to_string();
if store::file_hash(&self.db, &span.path)?.as_deref() == Some(hash.as_str()) {
span.fresh = true;
return Ok(());
}
let Ok(text) = String::from_utf8(bytes) else {
return Ok(());
};
let text = match keyscan::readable(&self.db, &span.path, &text)? {
Some(text) => text,
None => return Ok(()),
};
let lines: Vec<&str> = text.lines().collect();
let first = start.max(1);
let last = end.min(lines.len() as u32);
if first > last {
return Ok(());
}
span.text = lines[(first - 1) as usize..last as usize].join("\n");
span.start_line = first;
span.end_line = last;
span.from_disk = true;
Ok(())
}
fn name_enclosing_symbol(&self, span: &mut Span) -> Result<()> {
let by_file = store::symbols_in_files(&self.db, std::slice::from_ref(&span.path))?;
let Some(symbols) = by_file.get(&span.path) else {
return Ok(());
};
if let Some((_, name, kind)) = enclosing_definition(symbols, span.start_line, span.end_line)
{
span.symbol = Some(name);
span.symbol_kind = Some(kind);
}
Ok(())
}
fn mark_span_freshness(&self, span: &mut Span) -> Result<()> {
let stamps = store::file_stamps(&self.db, std::slice::from_ref(&span.path))?;
let Some((bytes, indexed_at)) = stamps.get(&span.path) else {
return Ok(());
};
span.fresh = match std::fs::metadata(&span.path) {
Ok(meta) => {
meta.len() as i64 == *bytes
&& meta
.modified()
.ok()
.and_then(|t| t.duration_since(std::time::UNIX_EPOCH).ok())
.is_none_or(|d| d.as_secs() as i64 <= *indexed_at)
}
Err(_) => false,
};
Ok(())
}
fn name_enclosing_symbols(&self, hits: &mut [(Hit, f32)]) -> Result<()> {
let mut paths: Vec<String> = Vec::new();
for (hit, _) in hits.iter() {
if hit.image.is_none() && !paths.contains(&hit.path) {
paths.push(hit.path.clone());
}
}
let by_file = store::symbols_in_files(&self.db, &paths)?;
for (hit, _) in hits.iter_mut() {
let Some(symbols) = by_file.get(&hit.path) else {
continue;
};
if let Some((start, name, kind)) =
enclosing_definition(symbols, hit.start_line, hit.end_line)
{
hit.symbol = Some(name);
hit.symbol_kind = Some(kind);
hit.symbol_line = Some(start);
}
}
Ok(())
}
fn mark_freshness(&self, hits: &mut [(Hit, f32)]) -> Result<()> {
if hits.is_empty() {
return Ok(());
}
let paths: Vec<String> = {
let mut seen: Vec<String> = Vec::new();
for (hit, _) in hits.iter() {
if !seen.contains(&hit.path) {
seen.push(hit.path.clone());
}
}
seen
};
let stamps = store::file_stamps(&self.db, &paths)?;
let mut fresh: std::collections::HashMap<&str, bool> = std::collections::HashMap::new();
for path in &paths {
let Some((bytes, indexed_at)) = stamps.get(path) else {
fresh.insert(path.as_str(), true);
continue;
};
let state = match std::fs::metadata(path) {
Ok(meta) => {
let size_matches = meta.len() as i64 == *bytes;
let unmodified = meta
.modified()
.ok()
.and_then(|t| t.duration_since(std::time::UNIX_EPOCH).ok())
.is_none_or(|d| d.as_secs() as i64 <= *indexed_at);
size_matches && unmodified
}
Err(_) => false,
};
fresh.insert(path.as_str(), state);
}
for (hit, _) in hits.iter_mut() {
hit.fresh = fresh.get(hit.path.as_str()).copied().unwrap_or(true);
}
Ok(())
}
fn image_search(
&mut self,
query: &str,
depth: usize,
filter: &Filter,
) -> Result<Vec<(u64, f32)>> {
if store::image_count(&self.db)? == 0 {
return Ok(Vec::new());
}
let allowlist = if filter.is_empty() {
index::Allowlist::All
} else {
let mut ids = Vec::new();
for id in store::filtered_image_ids(&self.db, filter.groups())? {
if self.images.contains(id as u64)? {
ids.push(id as u64);
}
}
if ids.is_empty() {
return Ok(Vec::new());
}
index::Allowlist::Subset(ids)
};
let mut vector = self.clip.embed_query(query, self.quiet)?;
normalize(&mut vector);
let (scores, ids) = self.images.search(&vector, depth, &allowlist)?;
Ok(ids
.into_iter()
.zip(scores.into_iter().chain(std::iter::repeat(0.0)))
.collect())
}
pub fn image_count(&self) -> Result<i64> {
store::image_count(&self.db)
}
pub fn stats(&self) -> Result<(i64, i64, i64)> {
store::stats(&self.db)
}
pub fn readable(&self) -> Result<()> {
store::readable(&self.db)
}
pub fn exact_bytes(&self) -> Option<u64> {
self.exact.exists().then(|| self.exact.bytes())
}
pub fn retired_symbols(&self) -> Result<i64> {
store::symbols_past_count(&self.db)
}
pub fn save(&mut self) -> Result<()> {
self.index.save()?;
self.images.save()?;
self.generation = generation(&self.db)? + 1;
store::set_meta(&self.db, GENERATION, &self.generation.to_string())?;
Ok(())
}
}
pub(crate) fn normalize(v: &mut [f32]) {
let norm = v.iter().map(|x| x * x).sum::<f32>().sqrt();
if norm > 0.0 {
for x in v {
*x /= norm;
}
}
}
fn generation(db: &Connection) -> Result<u64> {
Ok(store::get_meta(db, GENERATION)?
.and_then(|v| v.parse().ok())
.unwrap_or(0))
}
fn now() -> i64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.as_secs() as i64)
.unwrap_or(0)
}
pub fn human_bytes(bytes: i64) -> String {
const KB: f64 = 1024.0;
let b = bytes as f64;
if b >= KB * KB {
format!("{:.1} MB", b / (KB * KB))
} else if b >= KB {
format!("{:.0} KB", b / KB)
} else {
format!("{bytes} B")
}
}
pub fn serialize_plain_list<S: serde::Serializer>(
paths: &[String],
s: S,
) -> Result<S::Ok, S::Error> {
s.collect_seq(paths.iter().map(|p| plain(p)))
}
pub fn serialize_plain<S: serde::Serializer>(path: &str, s: S) -> Result<S::Ok, S::Error> {
s.serialize_str(&plain(path))
}
thread_local! {
static CANONICAL_CALLS: std::cell::Cell<u64> = const { std::cell::Cell::new(0) };
}
pub fn canonical_calls() -> u64 {
CANONICAL_CALLS.with(|c| c.get())
}
pub fn canonical(path: &Path) -> PathBuf {
CANONICAL_CALLS.with(|c| c.set(c.get() + 1));
std::fs::canonicalize(path).unwrap_or_else(|_| path.to_path_buf())
}
fn product_first(hits: &mut Vec<(Hit, f32)>) {
let product = |h: &Hit| filter::is_code(&h.path) && !is_test_path(&h.path);
let Some(first) = hits.iter().position(|(h, _)| product(h)) else {
return;
};
let held: Vec<usize> = (0..first)
.filter(|&i| is_test_path(&hits[i].0.path))
.collect();
if held.is_empty() {
return;
}
let mut moved: Vec<(Hit, f32)> = Vec::with_capacity(held.len());
for &i in held.iter().rev() {
moved.push(hits.remove(i));
}
moved.reverse();
let at = hits
.iter()
.position(|(h, _)| product(h))
.map_or(hits.len(), |p| p + 1);
for (offset, hit) in moved.into_iter().enumerate() {
hits.insert(at + offset, hit);
}
}
fn collapse_copies(hits: &mut Vec<(Hit, f32)>) {
let mut first: std::collections::HashMap<blake3::Hash, usize> =
std::collections::HashMap::new();
let mut keep: Vec<bool> = Vec::with_capacity(hits.len());
let mut copies: Vec<(usize, String)> = Vec::new();
for (i, (hit, _)) in hits.iter().enumerate() {
if hit.image.is_some() || hit.text.trim().is_empty() {
keep.push(true);
continue;
}
let key = blake3::hash(hit.text.as_bytes());
match first.get(&key) {
Some(&j) if hits[j].0.path != hit.path => {
copies.push((j, hit.path.clone()));
keep.push(false);
}
Some(_) => keep.push(true),
None => {
first.insert(key, i);
keep.push(true);
}
}
}
for (j, path) in copies {
if !hits[j].0.copies.contains(&path) {
hits[j].0.copies.push(path);
}
}
let mut flags = keep.into_iter();
hits.retain(|_| flags.next().unwrap_or(true));
}
pub(crate) fn enclosing_definition(
symbols: &[(u32, u32, String, String)],
start: u32,
end: u32,
) -> Option<(u32, String, String)> {
let starts_inside = symbols
.iter()
.filter(|(s, _, _, _)| *s >= start && *s <= end)
.min_by_key(|(s, e, _, _)| (*s, std::cmp::Reverse(e.saturating_sub(*s))));
let contains = || {
symbols
.iter()
.filter(|(s, e, _, _)| *s <= start && *e >= start)
.min_by_key(|(s, e, _, _)| e.saturating_sub(*s))
};
let overlaps = || {
symbols
.iter()
.filter(|(s, e, _, _)| *s <= end && *e >= start)
.min_by_key(|(s, e, _, _)| e.saturating_sub(*s))
};
starts_inside
.or_else(contains)
.or_else(overlaps)
.map(|(s, _, name, kind)| (*s, name.clone(), kind.clone()))
}
pub struct Shortener {
roots: Vec<std::path::PathBuf>,
used: std::cell::RefCell<std::collections::BTreeSet<std::path::PathBuf>>,
}
impl Shortener {
pub fn new(mut roots: Vec<std::path::PathBuf>) -> Self {
roots.sort_by_key(|r| std::cmp::Reverse(r.as_os_str().len()));
Self {
roots,
used: Default::default(),
}
}
pub fn short(&self, path: &str) -> String {
let plain_path = plain(path);
for root in &self.roots {
let root_text = plain(&root.to_string_lossy());
if let Some(rest) = plain_path.strip_prefix(root_text.as_str())
&& let Some(rest) = rest.strip_prefix(['/', '\\'])
{
self.used.borrow_mut().insert(root.clone());
return rest.to_string();
}
}
plain_path
}
pub fn header(&self) -> Option<String> {
let used = self.used.borrow();
if used.is_empty() {
return None;
}
let roots: Vec<String> = used.iter().map(|r| plain(&r.to_string_lossy())).collect();
Some(if roots.len() == 1 {
format!("root {}", roots[0])
} else {
format!("roots {}", roots.join(", "))
})
}
pub fn with_header(&self, text: String) -> String {
match self.header() {
Some(h) if !text.is_empty() => format!("{h}\n{text}"),
_ => text,
}
}
}
pub fn plain(text: &str) -> String {
if let Some(rest) = text.strip_prefix(r"\\?\UNC\") {
return format!(r"\\{rest}");
}
match text.strip_prefix(r"\\?\") {
Some(rest)
if rest.len() >= 2
&& rest.as_bytes()[0].is_ascii_alphabetic()
&& rest.as_bytes()[1] == b':' =>
{
rest.to_string()
}
_ => text.to_string(),
}
}
pub const MODEL_CACHE_ENV: &str = "SEMLITH_MODEL_CACHE";
pub fn model_cache_dir() -> Result<PathBuf> {
if let Ok(dir) = std::env::var(MODEL_CACHE_ENV) {
return Ok(PathBuf::from(dir));
}
Ok(home::user_home()?
.join(".cache")
.join("semlith")
.join("models"))
}
pub fn check_roots(roots: &[PathBuf]) -> (Vec<PathBuf>, Vec<(PathBuf, String)>) {
let mut readable = Vec::new();
let mut refused = Vec::new();
for root in roots {
let listed = match std::fs::metadata(root) {
Ok(meta) if meta.is_dir() => std::fs::read_dir(root).map(|_| ()),
Ok(_) => std::fs::File::open(root).map(|_| ()),
Err(e) => Err(e),
};
match listed {
Ok(()) => readable.push(root.clone()),
Err(e) => refused.push((root.clone(), e.to_string())),
}
}
(readable, refused)
}
pub const IGNORE_FILE: &str = ".semlithignore";
fn walk_error_path(e: &ignore::Error) -> Option<PathBuf> {
match e {
ignore::Error::WithPath { path, .. } => Some(path.clone()),
ignore::Error::WithDepth { err, .. } | ignore::Error::WithLineNumber { err, .. } => {
walk_error_path(err)
}
ignore::Error::Loop { child, .. } => Some(child.clone()),
_ => None,
}
}
pub(crate) enum Handed {
Walked(Option<std::time::Instant>),
Rest {
budget: std::time::Duration,
bytes_total: Option<u64>,
},
}
pub(crate) enum Walk {
Done(Walked),
Roots(Vec<PathBuf>),
}
pub(crate) struct Walked {
pub walk_ms: u64,
pub files: Vec<PathBuf>,
pub named: Vec<PathBuf>,
pub unreadable: Vec<(PathBuf, String)>,
pub generated: Vec<PathBuf>,
pub credentials: Vec<PathBuf>,
pub excluded: Vec<(PathBuf, String)>,
}
const ALWAYS: &[&str] = &[
"node_modules",
"bower_components",
".next",
".nuxt",
".svelte-kit",
".turbo",
".parcel-cache",
".yarn",
"__pycache__",
".venv",
"site-packages",
".tox",
".mypy_cache",
".pytest_cache",
".ruff_cache",
".eggs",
".gradle",
"DerivedData",
".swiftpm",
".dart_tool",
"_build",
".terraform",
".serverless",
".gradle-cache",
];
const GENERATED: &[(&str, &[&str])] = &[
("target", &["Cargo.toml"]),
("vendor", &["go.mod", "composer.json", "Gemfile"]),
("deps", &["mix.exs"]),
(
"build",
&[
"build.gradle",
"build.gradle.kts",
"pom.xml",
"CMakeLists.txt",
"meson.build",
"package.json",
],
),
(
"dist",
&[
"package.json",
"pyproject.toml",
"setup.py",
"rollup.config.js",
],
),
("out", &["package.json", "tsconfig.json"]),
("bin", &["*.csproj", "*.sln", "*.fsproj"]),
("obj", &["*.csproj", "*.sln", "*.fsproj"]),
("coverage", &["package.json", "pyproject.toml"]),
("Pods", &["Podfile"]),
];
pub fn is_generated_dir(path: &Path) -> bool {
if !default_ignores_on(std::env::var("SEMLITH_DEFAULT_IGNORES").ok().as_deref()) {
return false;
}
let Some(name) = path.file_name().and_then(|n| n.to_str()) else {
return false;
};
if ALWAYS.contains(&name) {
return true;
}
let Some(parent) = path.parent() else {
return false;
};
GENERATED
.iter()
.filter(|(dir, _)| *dir == name)
.any(|(_, manifests)| manifests.iter().any(|m| sibling_exists(parent, m)))
}
pub fn default_ignores_on(value: Option<&str>) -> bool {
!matches!(value, Some("0") | Some("off") | Some("false"))
}
fn sibling_exists(parent: &Path, manifest: &str) -> bool {
let Some(extension) = manifest.strip_prefix("*.") else {
return parent.join(manifest).exists();
};
std::fs::read_dir(parent)
.map(|entries| {
entries.flatten().any(|entry| {
entry
.path()
.extension()
.is_some_and(|e| e.eq_ignore_ascii_case(extension))
})
})
.unwrap_or(false)
}
fn walk(roots: &[PathBuf]) -> Walked {
walk_allowing(roots, &[])
}
fn bytes_of(paths: &[PathBuf]) -> u64 {
paths
.iter()
.map(|p| embeddable_bytes(p.metadata().ok()))
.sum()
}
const HEAD_FILES: usize = 48;
fn recent_head(
roots: &[PathBuf],
boundary: &ResolvedBoundary,
accepted: &[PathBuf],
) -> Vec<PathBuf> {
let mut out: Vec<PathBuf> = Vec::new();
let mut repos: Vec<PathBuf> = Vec::new();
for root in roots {
let root = canonical(root);
if !root.is_dir() {
continue;
}
let mut level = vec![root];
for _ in 0..3 {
let mut next = Vec::new();
for dir in level {
if dir.join(".git").exists() {
repos.push(dir);
continue;
}
for entry in std::fs::read_dir(&dir).into_iter().flatten().flatten() {
let path = entry.path();
if entry.file_type().is_ok_and(|t| t.is_dir())
&& !entry.file_name().to_string_lossy().starts_with('.')
&& !is_generated_dir(&path)
{
next.push(path);
}
}
}
level = next;
}
}
let each = (HEAD_FILES / repos.len().max(1)).max(4);
for root in repos {
let mut taken = 0;
let Ok(log) = std::process::Command::new("git")
.arg("-C")
.arg(&root)
.args([
"log",
"-n",
"60",
"--name-only",
"--format=%x00",
"--relative",
"--",
".",
])
.stderr(std::process::Stdio::null())
.output()
else {
continue;
};
let text = String::from_utf8_lossy(&log.stdout);
let mut lines: Vec<&str> = Vec::new();
for commit in text.split('\0') {
let mut names: Vec<&str> = commit.lines().collect();
names.sort_by_key(|name| !filter::is_code(name.trim()));
lines.extend(names);
}
for line in lines {
if out.len() >= HEAD_FILES {
return out;
}
if taken >= each {
break;
}
let line = line.trim();
if line.is_empty() {
continue;
}
let rel = Path::new(line);
let hidden = rel
.components()
.any(|c| c.as_os_str().to_string_lossy().starts_with('.'));
let path = root.join(rel);
if hidden || out.contains(&path) || !path.is_file() {
continue;
}
let generated = path
.ancestors()
.skip(1)
.take_while(|a| a.starts_with(&root) && *a != root)
.any(|a| is_generated_dir(a) && !accepted.iter().any(|x| x.as_path() == a));
let ignored = path
.parent()
.and_then(semlithignore_for)
.is_some_and(|m| m.matched_path_or_any_parents(&path, false).is_ignore());
if generated || ignored || boundary.refuses(&path, true).is_some() {
continue;
}
out.push(path);
taken += 1;
}
}
out
}
fn recent_first(paths: &mut [PathBuf]) {
let committed = git_recency(paths);
let when = |path: &Path| -> i64 {
let mtime = path
.metadata()
.ok()
.and_then(|m| m.modified().ok())
.and_then(|t| t.duration_since(std::time::UNIX_EPOCH).ok())
.map_or(0, |d| d.as_secs() as i64);
mtime.max(committed.get(path).copied().unwrap_or(0))
};
let mut keyed: Vec<(i64, PathBuf)> = paths.iter().map(|p| (when(p), p.clone())).collect();
keyed.sort_by_key(|(at, _)| std::cmp::Reverse(*at));
for (slot, (_, path)) in paths.iter_mut().zip(keyed) {
*slot = path;
}
}
fn git_recency(paths: &[PathBuf]) -> std::collections::HashMap<PathBuf, i64> {
const COMMITS: &str = "400";
let mut tops: std::collections::BTreeSet<PathBuf> = Default::default();
let mut seen: std::collections::HashSet<PathBuf> = Default::default();
for path in paths {
let mut dir = path.parent();
while let Some(d) = dir {
if !seen.insert(d.to_path_buf()) {
break;
}
if d.join(".git").exists() {
tops.insert(d.to_path_buf());
break;
}
dir = d.parent();
}
}
let mut out = std::collections::HashMap::new();
for top in tops {
let Ok(log) = std::process::Command::new("git")
.arg("-C")
.arg(&top)
.args(["log", "-n", COMMITS, "--name-only", "--format=%x00%ct"])
.stderr(std::process::Stdio::null())
.output()
else {
continue;
};
let mut at = 0i64;
for line in String::from_utf8_lossy(&log.stdout).lines() {
if let Some(stamp) = line.strip_prefix('\0') {
at = stamp.trim().parse().unwrap_or(0);
} else if !line.is_empty() {
out.entry(top.join(line)).or_insert(at);
}
}
}
out
}
fn walk_allowing(roots: &[PathBuf], allowed: &[PathBuf]) -> Walked {
let mut out = Vec::new();
let mut named = Vec::new();
let mut unreadable = Vec::new();
let mut seen = std::collections::HashSet::new();
let mut dirs: Vec<PathBuf> = Vec::new();
let mut yielded: std::collections::HashSet<PathBuf> = std::collections::HashSet::new();
let generated: std::sync::Arc<std::sync::Mutex<Vec<PathBuf>>> = Default::default();
for root in roots {
if root.is_file() {
let path = canonical(root);
if seen.insert(path.clone()) {
named.push(path);
}
continue;
}
let canonical_root = canonical(root);
let mut builder = ignore::WalkBuilder::new(root);
builder
.hidden(true)
.git_ignore(true)
.git_global(true)
.git_exclude(true)
.parents(true)
.require_git(false)
.add_custom_ignore_filename(IGNORE_FILE)
.filter_entry({
let generated = generated.clone();
let allowed = allowed.to_vec();
move |e| {
if e.file_name() == ".semlith" {
return false;
}
if e.file_type().is_some_and(|t| t.is_dir())
&& is_generated_dir(e.path())
&& !allowed.iter().any(|a| canonical(e.path()) == *a)
{
if let Ok(mut seen) = generated.lock() {
seen.push(e.path().to_path_buf());
}
return false;
}
true
}
});
for result in builder.build() {
let entry = match result {
Ok(e) => e,
Err(e) => {
let at = walk_error_path(&e).unwrap_or_else(|| root.clone());
unreadable.push((at, format!("the walk could not read it: {e}")));
continue;
}
};
if entry.file_type().is_some_and(|t| t.is_dir()) {
dirs.push(entry.path().to_path_buf());
}
if !entry.file_type().is_some_and(|t| t.is_file()) {
continue;
}
yielded.insert(entry.path().to_path_buf());
let path = match entry.path().strip_prefix(root) {
Ok(rest) if !entry.path_is_symlink() => canonical_root.join(rest),
_ => canonical(entry.path()),
};
if seen.insert(path.clone()) {
out.push(path);
}
}
}
let mut credentials = Vec::new();
let mut excluded: Vec<(PathBuf, String)> = Vec::new();
let walked_dirs: std::collections::HashSet<&PathBuf> = dirs.iter().collect();
let generated_now: Vec<PathBuf> = generated.lock().map(|g| g.clone()).unwrap_or_default();
let home = crate::home::user_home().ok().map(|h| canonical(&h));
for dir in &dirs {
let Ok(entries) = std::fs::read_dir(dir) else {
continue;
};
let semlithignore = semlithignore_for(dir);
let mut resolved: Option<PathBuf> = None;
let mut key = |name: &std::ffi::OsStr| -> PathBuf {
resolved.get_or_insert_with(|| canonical(dir)).join(name)
};
for entry in entries.flatten() {
let name = entry.file_name();
let name = name.to_string_lossy();
let Ok(kind) = entry.file_type() else {
continue;
};
let path = entry.path();
if name.starts_with('.') {
if kind.is_file()
&& matches!(
filter::denied_against(&path, home.as_deref()),
Some(filter::Denied::Name(_))
)
{
credentials.push(key(&entry.file_name()));
}
continue;
}
let held = if kind.is_dir() {
walked_dirs.contains(&path) || generated_now.iter().any(|g| g == &path)
} else if kind.is_file() {
yielded.contains(&path)
} else {
true
};
if held {
continue;
}
let rule = if semlithignore.as_ref().is_some_and(|m| {
m.matched_path_or_any_parents(&path, kind.is_dir())
.is_ignore()
}) {
IGNORE_FILE
} else {
".gitignore"
};
excluded.push((key(&entry.file_name()), rule.to_string()));
}
}
credentials.sort();
credentials.dedup();
excluded.sort();
out.sort();
named.sort();
unreadable.sort();
let mut generated = generated.lock().map(|g| g.clone()).unwrap_or_default();
generated.sort();
generated.dedup();
Walked {
walk_ms: 0,
files: out,
named,
unreadable,
generated,
credentials,
excluded,
}
}
fn semlithignore_for(dir: &Path) -> Option<ignore::gitignore::Gitignore> {
let mut at = Some(dir);
while let Some(d) = at {
let file = d.join(IGNORE_FILE);
if file.is_file() {
return Some(ignore::gitignore::Gitignore::new(&file).0);
}
at = d.parent();
}
None
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn the_head_takes_code_first_from_repositories_below_the_root() {
let root = tempfile::tempdir().unwrap();
let repo = root.path().join("group").join("app");
std::fs::create_dir_all(repo.join("src")).unwrap();
for (name, text) in [
("CHANGELOG.md", "# Changes\n"),
("LICENSE", "MIT\n"),
("src/main.rs", "fn main() {}\n"),
] {
std::fs::write(repo.join(name), text).unwrap();
}
let git = |args: &[&str]| {
std::process::Command::new("git")
.arg("-C")
.arg(&repo)
.args(["-c", "user.name=t", "-c", "user.email=t@t"])
.args(args)
.output()
.unwrap()
};
if !git(&["init", "-q"]).status.success() {
return; }
git(&["add", "."]);
assert!(git(&["commit", "-q", "-m", "one"]).status.success());
let head = recent_head(
&[root.path().to_path_buf()],
&Boundary::default().resolved(Some(Path::new("/nonexistent-home"))),
&[],
);
let names: Vec<String> = head
.iter()
.map(|p| {
p.strip_prefix(canonical(&repo))
.unwrap()
.display()
.to_string()
})
.collect();
let main = Path::new("src").join("main.rs").display().to_string();
assert_eq!(names.first(), Some(&main), "{names:?}");
assert_eq!(names.len(), 3, "{names:?}");
}
#[test]
fn product_code_comes_before_tests_and_prose_stays_put() {
let hit = |path: &str| {
(
Hit {
score: 1.0,
path: path.to_string(),
start_line: 1,
end_line: 2,
text: path.to_string(),
store: None,
lists: Vec::new(),
image: None,
fresh: true,
symbol: None,
symbol_kind: None,
symbol_line: None,
provenance: None,
copies: Vec::new(),
},
1.0,
)
};
let mut hits = vec![
hit("/r/README.md"),
hit("/r/tests/watch.rs"),
hit("/r/tests/fixtures/x.rs"),
hit("/r/src/lib.rs"),
hit("/r/src/mcp.rs"),
];
product_first(&mut hits);
let order: Vec<&str> = hits.iter().map(|(h, _)| h.path.as_str()).collect();
assert_eq!(
order,
[
"/r/README.md",
"/r/src/lib.rs",
"/r/tests/watch.rs",
"/r/tests/fixtures/x.rs",
"/r/src/mcp.rs"
]
);
}
#[test]
fn semlithignore_leaves_a_copy_out_and_a_root_inside_it_whole() {
let dir = tempfile::tempdir().unwrap();
let root = canonical(dir.path());
std::fs::create_dir_all(root.join("src")).unwrap();
std::fs::create_dir_all(root.join("tests/fixtures/corpus/src")).unwrap();
std::fs::write(root.join("src/lib.rs"), "fn a() {}").unwrap();
std::fs::write(root.join("tests/fixtures/corpus/src/lib.rs"), "fn a() {}").unwrap();
std::fs::write(root.join(IGNORE_FILE), "tests/fixtures/corpus/\n").unwrap();
let whole = walk(std::slice::from_ref(&root)).files;
assert_eq!(whole, vec![root.join("src/lib.rs")], "{whole:?}");
let corpus = root.join("tests/fixtures/corpus");
let inside = walk(std::slice::from_ref(&corpus)).files;
assert_eq!(inside, vec![corpus.join("src/lib.rs")], "{inside:?}");
}
#[test]
fn the_graph_list_adds_candidates_and_does_not_amplify_them() {
let both = 10u64;
let only_graph = 99u64;
let lists: [CandidateList; 3] = [
("vector", false, vec![(both, 1.0), (11, 1.0)]),
("keyword", false, vec![(both, 1.0)]),
("graph", false, vec![(both, 1.0), (only_graph, 1.0)]),
];
let text_only: [CandidateList; 2] = [lists[0].clone(), lists[1].clone()];
let (with_graph, badges) = fuse(&lists, |_| RRF_K);
let (without_graph, _) = fuse(&text_only, |_| RRF_K);
let score = |fused: &[Scored], id: u64| {
fused
.iter()
.find(|(key, _)| *key == (false, id))
.map(|(_, score)| *score)
.expect("the candidate is in the fused set")
};
assert_eq!(
score(&with_graph, both),
score(&without_graph, both),
"a chunk both text lists found was scored a third time by the graph"
);
let slot = with_graph
.iter()
.position(|(key, _)| *key == (false, both))
.unwrap();
assert_eq!(badges[slot], vec!["vector", "keyword", "graph"]);
assert!(
with_graph
.iter()
.any(|(key, _)| *key == (false, only_graph)),
"the graph list stopped adding candidates of its own"
);
}
#[test]
fn a_query_is_identifier_shaped_only_when_it_is_one_token_of_identifier() {
for identifier in [
"record_retrieval",
"store::edges_out",
"Semlith",
"self.index.search",
"format!",
"k8s",
] {
assert_eq!(
shape_of(identifier),
Shape::Identifier,
"{identifier:?} is an identifier"
);
}
for question in [
"where does a retrieval get written down",
"record retrieval",
"how does indexing work?",
"",
" ",
"42",
] {
assert_eq!(shape_of(question), Shape::Question, "{question:?}");
}
}
#[test]
fn the_shape_decides_what_the_keyword_list_is_worth() {
assert_eq!(Shape::Identifier.keyword_weight(), 2.0);
assert_eq!(Shape::Question.keyword_weight(), 1.0);
assert_eq!(Shape::Identifier.weighting(), "keyword weighted 2×");
assert_eq!(Shape::Question.weighting(), "vector and keyword equal");
}
#[test]
fn prefer_lifts_one_side_and_never_removes_the_other() {
assert_eq!(Prefer::Any.multiplier(true), 1.0);
assert_eq!(Prefer::Any.multiplier(false), 1.0);
assert!(Prefer::Code.multiplier(true) > Prefer::Code.multiplier(false));
assert!(Prefer::Docs.multiplier(false) > Prefer::Docs.multiplier(true));
assert!(Prefer::Code.multiplier(false) > 0.0);
assert!(Prefer::Docs.multiplier(true) > 0.0);
}
#[test]
fn an_unknown_preference_is_an_error() {
assert_eq!(Prefer::parse("code").unwrap(), Prefer::Code);
assert_eq!(Prefer::parse("DOCS").unwrap(), Prefer::Docs);
assert_eq!(Prefer::parse("").unwrap(), Prefer::Any);
assert!(Prefer::parse("source").is_err());
}
#[test]
fn every_rerank_factor_is_a_tiebreak_rather_than_a_ranking() {
let strongest = 1.0 + GRAPH_PROXIMITY;
for penalty in [STALE_PENALTY, TEST_PENALTY] {
let weakest = 1.0 - penalty;
assert!(
strongest / weakest < 1.6,
"the rerank spans {strongest}/{weakest}, which is a ranking rather than a \
tiebreak"
);
}
const { assert!(GRAPH_PROXIMITY > 0.0 && STALE_PENALTY > 0.0 && TEST_PENALTY > 0.0) };
const { assert!(STALE_PENALTY < 1.0 && TEST_PENALTY < 1.0) };
}
#[test]
fn a_read_target_is_a_span_or_a_name_by_its_shape() {
assert_eq!(
Target::parse("src/store.rs:1041-1080"),
Target::Span {
path: "src/store.rs".to_string(),
start: 1041,
end: 1080
}
);
assert_eq!(
Target::parse("src/store.rs:12"),
Target::Span {
path: "src/store.rs".to_string(),
start: 12,
end: 12
}
);
assert_eq!(
Target::parse("a.rs:80-40"),
Target::Span {
path: "a.rs".to_string(),
start: 40,
end: 80
}
);
for name in [
"record_retrieval",
"store::edges_out",
"src/store.rs",
"a.rs:notanumber",
] {
assert!(
matches!(Target::parse(name), Target::Symbol(_)),
"{name:?} is a name"
);
}
}
#[test]
fn a_hit_that_straddles_a_definitions_start_is_labelled_with_it() {
let symbols = [
(100u32, 180u32, "outer".to_string(), "function".to_string()),
(120u32, 140u32, "inner".to_string(), "function".to_string()),
];
let pick = |start: u32, end: u32| {
let contains = symbols
.iter()
.filter(|(s, e, _, _)| *s <= start && *e >= start)
.min_by_key(|(s, e, _, _)| e.saturating_sub(*s));
contains
.or_else(|| {
symbols
.iter()
.filter(|(s, e, _, _)| *s <= end && *e >= start)
.min_by_key(|(s, e, _, _)| e.saturating_sub(*s))
})
.map(|(_, _, name, _)| name.as_str())
};
assert_eq!(pick(115, 130), Some("outer"));
assert_eq!(pick(90, 105), Some("outer"));
assert_eq!(pick(125, 135), Some("inner"));
assert_eq!(pick(200, 210), None);
}
#[test]
fn normalize_gives_unit_length() {
let mut v = vec![3.0, 4.0];
normalize(&mut v);
assert!((v.iter().map(|x| x * x).sum::<f32>() - 1.0).abs() < 1e-6);
let mut z = vec![0.0, 0.0];
normalize(&mut z);
assert_eq!(z, vec![0.0, 0.0]);
}
#[test]
fn one_large_file_does_not_queue_more_than_a_batch() {
let text = "a line of perfectly ordinary text\n".repeat(4000);
let chunks = chunk::chunk_text(&text);
assert!(
chunks.len() > EMBED_BATCH * 4,
"test file is too small to prove anything: {} chunks",
chunks.len()
);
let mut queued = 0usize;
let mut high_water = 0usize;
for _ in &chunks {
queued += 1;
high_water = high_water.max(queued);
if queued >= EMBED_BATCH {
queued = 0;
}
}
assert_eq!(
high_water, EMBED_BATCH,
"queue grew past one batch within a single file"
);
}
#[test]
fn a_new_store_records_the_default_model_and_reopens_on_it() {
let dir = tempdir();
let created = Semlith::open(&dir, None).unwrap();
assert_eq!(*created.model(), default_model());
assert_eq!(created.dim(), 384);
drop(created);
let reopened = Semlith::open(&dir, None).unwrap();
assert_eq!(*reopened.model(), default_model());
}
#[test]
fn an_existing_store_keeps_its_own_model() {
use fastembed::EmbeddingModel;
let dir = tempdir2();
let old = Model::Builtin(EmbeddingModel::BGESmallENV15);
let created = Semlith::open(&dir, Some(old.clone())).unwrap();
assert_eq!(*created.model(), old);
drop(created);
let reopened = Semlith::open(&dir, None).unwrap();
assert_eq!(*reopened.model(), old);
assert!(Semlith::open(&dir, Some(Model::Granite)).is_err());
}
fn continuation(n: usize) -> (tempfile::TempDir, tempfile::TempDir, PathBuf, Semlith) {
let corpus = tempfile::tempdir().unwrap();
let root = canonical(corpus.path());
for i in 0..n {
std::fs::write(root.join(format!("{i:04}.txt")), b"").unwrap();
}
let dir = tempfile::tempdir().unwrap();
let mut s = Semlith::open(dir.path(), None).unwrap();
s.boundary = Boundary::within(vec![root.clone()]);
(corpus, dir, root, s)
}
#[test]
fn a_continuation_slice_checks_only_the_files_it_reaches() {
const N: usize = 300;
const K: usize = 10;
let (_corpus, _dir, root, mut s) = continuation(N);
let files: Vec<PathBuf> = (0..N).map(|i| root.join(format!("{i:04}.txt"))).collect();
let asked = std::cell::Cell::new(0usize);
let control = || {
asked.set(asked.get() + 1);
if asked.get() > K {
Flow::Yield
} else {
Flow::Run
}
};
let before = canonical_calls();
let report = s
.index_rest_held_under(
files,
std::time::Duration::from_secs(3600),
Some(0),
&control,
|_, _| {},
)
.unwrap();
let spent = canonical_calls() - before;
assert_eq!(report.scanned, K);
assert_eq!(report.pending.len(), N - K);
assert!(
spent <= K as u64 + 4,
"{spent} canonicalisations for a slice that reached {K} of {N} files"
);
}
#[test]
fn a_continuation_slice_refuses_what_a_first_slice_refuses() {
let (_corpus, _dir, _root, mut s) = continuation(0);
let elsewhere = tempfile::tempdir().unwrap();
let outside = canonical(elsewhere.path()).join("outside.txt");
std::fs::write(&outside, b"not this store's").unwrap();
let first = s
.index_paths_within(
std::slice::from_ref(&outside),
std::time::Duration::from_secs(3600),
|_, _| {},
)
.unwrap();
let rest = s
.index_rest_held_under(
vec![outside.clone()],
std::time::Duration::from_secs(3600),
None,
&|| Flow::Run,
|_, _| {},
)
.unwrap();
assert_eq!(first.refused.len(), 1, "{:?}", first.refused);
assert!(first.refused[0].1.contains("outside this store's roots"));
assert_eq!(rest.refused, first.refused);
assert_eq!(rest.indexed, 0);
}
#[test]
fn a_cut_slice_counts_only_what_it_did_not_reach_as_remaining() {
const N: usize = 20;
const K: usize = 5;
let (_corpus, _dir, root, mut s) = continuation(N);
let elsewhere = tempfile::tempdir().unwrap();
let outside = canonical(elsewhere.path()).join("outside.txt");
std::fs::write(&outside, b"not this store's").unwrap();
let asked = std::cell::Cell::new(0usize);
let control = || {
asked.set(asked.get() + 1);
if asked.get() > K {
Flow::Yield
} else {
Flow::Run
}
};
let report = s
.index_within_held_under(
&[root, outside],
std::time::Duration::from_secs(3600),
&control,
|_, _| {},
)
.unwrap();
assert_eq!(report.refused.len(), 1);
assert_eq!(report.pending.len(), N - K);
assert_eq!(report.remaining, N - K);
assert_eq!(report.scanned + report.remaining, N + 1);
}
fn tempdir() -> PathBuf {
let d = std::env::temp_dir().join(format!("semlith-test-{}", std::process::id()));
let _ = std::fs::remove_dir_all(&d);
d
}
fn tempdir2() -> PathBuf {
let d = std::env::temp_dir().join(format!("semlith-test-b-{}", std::process::id()));
let _ = std::fs::remove_dir_all(&d);
d
}
}
#[cfg(test)]
mod plain_path_tests {
use super::plain;
#[test]
fn a_verbatim_prefix_is_not_what_anybody_pastes_back() {
assert_eq!(
plain(r"\\?\C:\work\api\src\lock.rs"),
r"C:\work\api\src\lock.rs"
);
assert_eq!(plain(r"\\?\c:\work"), r"c:\work");
assert_eq!(
plain(r"\\?\UNC\server\share\notes.md"),
r"\\server\share\notes.md"
);
let volume = r"\\?\Volume{9f8a7b6c-0000-0000-0000-000000000000}\data";
assert_eq!(plain(volume), volume);
assert_eq!(
plain("/home/someone/api/src/lock.rs"),
"/home/someone/api/src/lock.rs"
);
assert_eq!(plain(r"C:\work\api"), r"C:\work\api");
assert_eq!(
plain(r"\\server\share\notes.md"),
r"\\server\share\notes.md"
);
assert_eq!(plain(""), "");
}
}