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 cloud;
pub mod compact;
pub mod coreml;
pub mod cpucap;
pub mod cuda;
pub mod daemon;
pub mod doctor;
pub mod drop;
pub mod embed;
pub mod filter;
pub mod fleet;
mod formats;
pub mod gpu;
pub mod graph;
pub mod helpers;
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 progress;
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;
pub mod worker;
#[derive(Debug)]
pub struct NotFound(pub String);
impl std::fmt::Display for NotFound {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(&self.0)
}
}
impl std::error::Error for NotFound {}
#[doc(hidden)]
pub mod fuzzing {
pub fn http_head(bytes: &[u8]) {
crate::http::fuzz_head(bytes)
}
}
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";
const MARKUP_KEY: &str = "markup";
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 FIRST_IMPORT_WINDOW: usize = 1024;
const FIRST_IMPORT_ROW_COMMIT: std::time::Duration = std::time::Duration::from_secs(2);
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,
(_, "named") => 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);
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) => {
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("()")
}
pub fn pictured(query: &str) -> bool {
const SEEN: [&str; 46] = [
"picture",
"pictures",
"image",
"images",
"photo",
"photos",
"photograph",
"screenshot",
"screenshots",
"screen",
"diagram",
"diagrams",
"logo",
"icon",
"icons",
"drawing",
"illustration",
"sketch",
"painting",
"png",
"jpg",
"jpeg",
"gif",
"webp",
"red",
"green",
"blue",
"yellow",
"orange",
"purple",
"pink",
"black",
"white",
"grey",
"gray",
"brown",
"circle",
"square",
"triangle",
"rectangle",
"ring",
"stripes",
"striped",
"checkerboard",
"gradient",
"donut",
];
let lower = query.to_lowercase();
lower.contains("look like")
|| lower.contains("looks like")
|| lower
.split(|c: char| !c.is_alphanumeric())
.any(|w| SEEN.contains(&w))
}
fn named_identifiers(query: &str) -> Vec<String> {
let mut out: Vec<String> = Vec::new();
for raw in query.split_whitespace() {
let call = raw.contains("()");
let word = raw.trim_matches(|c: char| !c.is_alphanumeric() && c != '_' && c != ':');
let word = word.rsplit("::").next().unwrap_or(word);
if word.len() < 3
|| !word
.chars()
.next()
.is_some_and(|c| c.is_alphabetic() || c == '_')
{
continue;
}
let inner_capital =
word.chars().skip(1).any(char::is_uppercase) && word.chars().any(char::is_lowercase);
let ticked = raw.starts_with('`');
if (word.contains('_') || inner_capital || ticked || call || raw.contains("::"))
&& !out.iter().any(|w| w == word)
{
out.push(word.to_string());
}
if out.len() == 3 {
break;
}
}
out
}
fn names_releases(query: &str) -> bool {
query.split(|c: char| !c.is_alphanumeric()).any(|w| {
let w = w.to_lowercase();
w.starts_with("changelog") || w.starts_with("release") || w == "version" || w == "news"
})
}
fn is_release_notes(path: &str) -> bool {
let p = Path::new(path);
let stem = p
.file_stem()
.and_then(|s| s.to_str())
.unwrap_or_default()
.to_ascii_lowercase();
(!filter::is_code(path)
&& matches!(
stem.as_str(),
"changelog"
| "changes"
| "history"
| "releases"
| "release-notes"
| "release_notes"
| "news"
))
|| p.components()
.any(|c| c.as_os_str().eq_ignore_ascii_case("changelogs"))
}
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 PROJECT_LIFT: f32 = 1.0;
const NAMED_PROJECT_LIFT: f32 = 3.0;
const PROJECT_HEAD: usize = 20;
fn project_prior(hits: &[(Hit, f32)], query: &str) -> (Vec<f32>, Vec<bool>) {
let mut seen: std::collections::HashMap<PathBuf, Option<PathBuf>> =
std::collections::HashMap::new();
let mut project_of = |path: &str| -> Option<PathBuf> {
let dir = Path::new(path).parent()?.to_path_buf();
if let Some(found) = seen.get(&dir) {
return found.clone();
}
let found = dir
.ancestors()
.find(|d| d.join(".git").exists())
.map(Path::to_path_buf);
seen.insert(dir, found.clone());
found
};
let projects: Vec<Option<PathBuf>> = hits.iter().map(|(h, _)| project_of(&h.path)).collect();
let words: std::collections::HashSet<String> = query
.split(|c: char| !c.is_alphanumeric())
.map(str::to_lowercase)
.collect();
let named: Vec<bool> = projects
.iter()
.map(|p| {
p.as_ref()
.and_then(|p| p.file_name())
.and_then(|n| n.to_str())
.is_some_and(|n| {
n.to_lowercase()
.split(|c: char| !c.is_alphanumeric())
.any(|t| t.len() >= 4 && words.contains(t))
})
})
.collect();
let named = if named.iter().all(|n| *n) {
vec![false; named.len()]
} else {
named
};
let mut mass: std::collections::HashMap<&PathBuf, f32> = std::collections::HashMap::new();
for ((hit, _), project) in hits.iter().zip(&projects).take(PROJECT_HEAD) {
if let Some(project) = project {
*mass.entry(project).or_default() += hit.score.max(0.0);
}
}
let total: f32 = mass.values().sum();
if mass.len() < 2 || total <= 0.0 {
return (vec![0.0; hits.len()], named);
}
let shares = projects
.iter()
.map(|p| {
p.as_ref()
.and_then(|p| mass.get(p))
.map_or(0.0, |m| m / total)
})
.collect();
(shares, named)
}
const STALE_PENALTY: f32 = 0.10;
const TEST_PENALTY: f32 = 0.10;
type Candidate = (((bool, u64), f32), Vec<&'static str>);
const TEXT_LISTS: f32 = 2.0;
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 FUSE_DEPTH: usize = RANK_DEPTH * 8;
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 FILTER_MEMOS: usize = 8;
struct FilterMemo {
groups: Vec<Vec<String>>,
generation: u64,
newest_chunk: i64,
allowlist: Allowlist,
selected: std::sync::Arc<std::collections::HashSet<u64>>,
}
#[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 = "std::ops::Not::not")]
pub replaced: 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 = "Vec::is_empty")]
#[serde(serialize_with = "serialize_plain_list")]
pub copies: Vec<String>,
#[serde(skip_serializing_if = "Vec::is_empty")]
pub agreed: 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,
Image,
Phase,
}
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",
Self::Image => "image",
Self::Phase => "phase",
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, serde::Deserialize)]
pub enum SkipReason {
Empty,
TooLarge,
NotRegular,
Unreadable(String),
Binary,
NoText,
NotDecodableImage,
Encrypted,
ArchiveTooLarge,
NotText(u8),
}
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(),
Self::Encrypted => "encrypted: no text without its password".to_string(),
Self::ArchiveTooLarge => format!(
"decompresses past {} MiB, left out as a zip bomb would be",
crate::formats::MAX_ARCHIVE_TEXT / (1024 * 1024)
),
Self::NotText(share) => {
format!("binary: {share} % of its first 8 KiB is not valid UTF-8 text")
}
}
}
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",
Self::Encrypted => "encrypted",
Self::ArchiveTooLarge => "archive too large",
Self::NotText(_) => "binary",
}
}
}
#[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,
pub phase: progress::Phase,
pub phase_detail: Option<String>,
pub expected_chunks: u64,
pub images: usize,
pub images_total: usize,
pub progress: f64,
pub eta_ms: Option<u64>,
pub eta_range_ms: Option<(u64, u64)>,
pub image_weight: f64,
}
#[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,
pub expected_left: u64,
pub images_left: usize,
pub images_flight: usize,
#[serde(skip)]
pub(crate) current_est: u64,
#[serde(skip)]
pub(crate) current_rows: u64,
#[serde(skip)]
pub(crate) live: std::cell::RefCell<Live>,
}
#[derive(Debug, Clone)]
pub(crate) struct Live {
pub phase: progress::Phase,
pub detail: Option<String>,
pub high: f64,
pub eta: progress::Eta,
pub started: std::time::Instant,
pub image_weight: f64,
pub images_parallel: bool,
}
impl Default for Live {
fn default() -> Self {
Self {
phase: progress::Phase::Read,
detail: None,
high: 0.0,
eta: progress::Eta::default(),
started: std::time::Instant::now(),
image_weight: progress::IMAGE_UNITS,
images_parallel: false,
}
}
}
struct ImageJob {
path: PathBuf,
key: String,
bytes: Vec<u8>,
hash: String,
width: u32,
height: u32,
}
struct ImageDone {
job: ImageJob,
vector: std::result::Result<Vec<f32>, String>,
}
const IMAGE_BYTES_IN_FLIGHT: u64 = 256 * 1024 * 1024;
#[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;
const PLAN_PATHS: usize = 50;
#[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>,
#[serde(flatten)]
pub assessment: serde_json::Map<String, serde_json::Value>,
#[serde(skip_serializing_if = "Option::is_none")]
pub chunks: Option<u64>,
}
impl Review {
fn new(
path: String,
class: &str,
rule: String,
matches: Vec<keyscan::Match>,
confidence: Option<u8>,
) -> Self {
let stored: Vec<serde_json::Value> = matches
.iter()
.filter_map(|m| serde_json::to_value(m).ok())
.collect();
let assessment = match keyscan::assess(class, &path, &stored, confidence) {
serde_json::Value::Object(map) => map,
_ => Default::default(),
};
Self {
path,
class: class.to_string(),
rule,
matches,
confidence,
assessment,
chunks: None,
}
}
}
#[derive(Debug, Clone, Default, Serialize)]
pub struct ScanProgress {
pub phase: &'static str,
pub scanned: usize,
pub total: usize,
pub bytes: u64,
pub bytes_total: u64,
pub flagged: usize,
}
#[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>,
#[serde(serialize_with = "serialize_plain_lists")]
pub not_indexed_paths: std::collections::BTreeMap<String, Vec<String>>,
pub not_indexed_why: std::collections::BTreeMap<String, String>,
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,
pub chunks: u64,
pub images: usize,
pub image_bytes: u64,
#[serde(skip)]
pub counts: std::collections::HashMap<String, u64>,
}
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()
);
let chunks = report.embedded.max(report.batched);
let expected_chunks = report.rows as u64
+ report.expected_left
+ report.current_est.saturating_sub(report.current_rows);
let images_total = report.images + report.images_flight + report.images_left;
let (phase, phase_detail, share, eta, image_weight) = {
let mut live = report.live.borrow_mut();
match accel::waiting_for() {
Some(waiting)
if matches!(
live.phase,
progress::Phase::Read | progress::Phase::Embed | progress::Phase::Lane
) =>
{
live.phase = progress::Phase::Lane;
live.detail = Some(waiting);
}
None if live.phase == progress::Phase::Read && chunks > 0 => {
live.phase = progress::Phase::Embed;
live.detail = None;
}
None if live.phase == progress::Phase::Lane => {
live.phase = progress::Phase::Embed;
live.detail = None;
}
_ => {}
}
let work = progress::Work {
files: (report.scanned as u64, total as u64),
chunks: (chunks as u64, expected_chunks),
images: (report.images as u64, images_total as u64),
image_weight: live.image_weight,
parallel: live.images_parallel,
};
live.high = live.high.max(work.share()).min(0.999);
let at = live.started.elapsed().as_secs_f64();
let hold = matches!(
live.phase,
progress::Phase::Lane | progress::Phase::Save | progress::Phase::Walk
);
let (done, all) = work.eta_units();
live.eta.observe_unless(hold, at, done);
let eta = (all > 0.0)
.then(|| live.eta.left(at, all - done, hold))
.flatten();
(
live.phase,
live.detail.clone(),
live.high,
eta,
live.image_weight,
)
};
on_file(
path,
IndexProgress {
outcome,
scanned: report.scanned,
indexed: report.indexed,
chunks,
total,
symbols: report.symbols,
why,
threads: report.threads,
lanes: report.lanes.clone(),
bytes: report.bytes,
bytes_total: report.bytes_total,
rows: report.rows,
phase,
phase_detail,
expected_chunks,
images: report.images,
images_total,
progress: share,
eta_ms: eta.map(|(ms, _)| ms),
eta_range_ms: eta.map(|(_, range)| range),
image_weight,
},
);
}
fn say_phase(
on_file: &mut dyn FnMut(&Path, IndexProgress),
report: &IndexReport,
total: usize,
phase: progress::Phase,
detail: Option<String>,
) {
{
let mut live = report.live.borrow_mut();
if live.phase == phase && live.detail == detail {
return;
}
live.phase = phase;
live.detail = detail;
}
say_file(
on_file,
report,
total,
Path::new(""),
FileOutcome::Phase,
None,
);
}
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: index::Vectors,
images: index::Vectors,
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>,
write_parts: pipeline::WriteParts,
budget_generation: u64,
lane_chunks: std::collections::BTreeMap<String, usize>,
tokenizer: Option<tokenizers::Tokenizer>,
clip: image::Clip,
generation: u64,
shares_vectors: bool,
filter_memo: std::collections::VecDeque<FilterMemo>,
pub quiet: bool,
pub boundary: Boundary,
pub gitignore: bool,
pub(crate) force: bool,
pub vector_cache: Option<cache::Location>,
pub planned: Option<std::sync::Arc<std::collections::HashMap<String, u64>>>,
pub planned_images: Option<usize>,
pub expect_rest: Option<(u64, usize)>,
awaiting: std::collections::HashMap<i64, (usize, String, bool)>,
file_of: std::collections::HashMap<u64, i64>,
landed: Vec<(i64, String)>,
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 = index::Vectors::own(VectorIndex::open(&dir, dim, BIT_WIDTH, sharded)?);
let images = index::Vectors::own(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,
write_parts: Default::default(),
budget_generation: index::budget_generation(),
lane_chunks: std::collections::BTreeMap::new(),
tokenizer: None,
clip: image::Clip::default(),
generation,
shares_vectors: false,
filter_memo: std::collections::VecDeque::new(),
quiet: false,
boundary: Boundary::default(),
gitignore: true,
force: false,
vector_cache: None,
planned: None,
planned_images: None,
expect_rest: None,
awaiting: Default::default(),
file_of: Default::default(),
landed: Vec::new(),
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;
let (dir, dim) = (self.dir.clone(), self.dim);
let text = move || VectorIndex::open(&dir, dim, BIT_WIDTH, sharded);
let pictures_dir = self.dir.join(image::INDEX_DIR);
let pictures = move || VectorIndex::open(&pictures_dir, image::DIM, BIT_WIDTH, true);
if self.shares_vectors {
let generation = generation(&self.db)?;
self.index = index::shared(&self.dir, "text", generation, text)?;
self.images = index::shared(&self.dir, "images", generation, pictures)?;
} else {
self.index = index::Vectors::own(text()?);
self.images = index::Vectors::own(pictures()?);
}
Ok(())
}
pub fn share_vectors(&mut self) -> Result<()> {
self.shares_vectors = true;
self.reopen_indexes()
}
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");
}
pub(crate) fn holds(&self, path: &Path) -> bool {
self.db
.query_row(
"SELECT 1 FROM files WHERE path = ?1",
[path.to_string_lossy()],
|_| Ok(()),
)
.is_ok()
}
fn accepted_folders(&self) -> Vec<PathBuf> {
store::acceptances(&self.db)
.unwrap_or_default()
.into_iter()
.filter(|a| a.class == store::class::POLICY && a.mode != "refused")
.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, self.gitignore);
walked.walk_ms = started.elapsed().as_millis() as u64;
walked
}
pub fn plan(&mut self, roots: &[PathBuf]) -> Result<Plan> {
self.plan_with(roots, &mut |_| {})
}
pub fn plan_with(
&mut self,
roots: &[PathBuf],
on: &mut dyn FnMut(ScanProgress),
) -> Result<Plan> {
let started = std::time::Instant::now();
on(ScanProgress {
phase: "walk",
..ScanProgress::default()
});
let walked = self.walk(roots);
let total = walked.files.len() + walked.named.len();
let bytes_total: u64 = walked
.files
.iter()
.chain(&walked.named)
.map(|p| embeddable_bytes(p.metadata().ok()))
.sum();
let mut said = ScanProgress {
phase: "read",
total,
bytes_total,
..ScanProgress::default()
};
on(said.clone());
let mut last_said = std::time::Instant::now();
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, path: &Path| {
*plan.not_indexed.entry(class.to_string()).or_insert(0) += 1;
let listed = plan.not_indexed_paths.entry(class.to_string()).or_default();
if listed.len() < PLAN_PATHS {
listed.push(path.to_string_lossy().into_owned());
}
};
for dir in &walked.generated {
plan.review.push(Review::new(
dir.to_string_lossy().into_owned(),
store::class::POLICY,
"a generated or vendored folder the walk steps over".to_string(),
Vec::new(),
None,
));
not(&mut plan, store::class::POLICY, dir);
}
for (path, _) in &walked.unreadable {
not(&mut plan, store::class::UNINDEXABLE, path);
}
for path in &walked.credentials {
plan.credential.push(path.to_string_lossy().into_owned());
not(&mut plan, store::class::CREDENTIAL, path);
}
for (path, _) in &walked.excluded {
not(&mut plan, store::class::EXCLUDED, path);
}
let boundary = self.boundary.resolved(home.as_deref());
for (path, walked) in all {
said.scanned += 1;
said.bytes += embeddable_bytes(path.metadata().ok());
said.flagged = plan.review.len() + plan.credential.len();
if last_said.elapsed() >= std::time::Duration::from_millis(100) {
on(said.clone());
last_said = std::time::Instant::now();
}
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, &path);
} else {
not(&mut plan, store::class::EXCLUDED, &path);
}
continue;
}
let Ok((file, meta)) = open_regular(&path) else {
not(&mut plan, store::class::UNINDEXABLE, &path);
continue;
};
if meta.len() == 0 {
not(&mut plan, store::class::UNINDEXABLE, &path);
continue;
}
let accepted = store::acceptance(&self.db, &key)?;
if let Some(kept) = accepted.as_ref().filter(|a| a.mode == "refused") {
not(&mut plan, &kept.class.clone(), &path);
continue;
}
if meta.len() > chunk::MAX_FILE_BYTES
&& !accepted
.as_ref()
.is_some_and(|a| a.class == store::class::POLICY)
{
plan.review.push(Review::new(
key,
store::class::POLICY,
format!(
"over {} MiB ({} MiB)",
chunk::MAX_FILE_BYTES / (1024 * 1024),
meta.len() / (1024 * 1024)
),
Vec::new(),
None,
));
not(&mut plan, store::class::POLICY, &path);
continue;
}
let Ok(bytes) = read_capped(file, meta.len()) else {
not(&mut plan, store::class::UNINDEXABLE, &path);
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.images += 1;
plan.image_bytes += meta.len();
}
continue;
}
let read = match worker::read(&path, &bytes, true) {
Ok(read) => read,
Err(worker::Stop::Crashed(_)) => {
not(&mut plan, store::class::CRASHED, &path);
continue;
}
Err(worker::Stop::Unavailable(_)) => {
not(&mut plan, store::class::UNINDEXABLE, &path);
continue;
}
};
let (text, symbols) = match read.text {
Ok(Ok((text, _))) => (
text,
read.graph
.and_then(|g| g.ok())
.flatten()
.map(|e| e.symbols)
.unwrap_or_default(),
),
other => {
if let Ok(Err(why)) = &other {
plan.not_indexed_why
.insert(crate::plain(&path.to_string_lossy()), why.as_str());
}
not(&mut plan, store::class::UNINDEXABLE, &path);
continue;
}
};
let count = |plan: &mut Plan, key: &str| {
let n = planned_chunks(&path, &text, &symbols);
plan.chunks += n;
plan.counts.insert(key.to_string(), n);
};
if self.boundary.allow_secrets {
plan.embed += usize::from(!unchanged);
plan.unchanged += usize::from(unchanged);
if !unchanged {
plan.embed_bytes += meta.len();
count(&mut plan, &key);
}
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();
count(&mut plan, &key);
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 n = planned_chunks(&path, &text, &symbols);
let live: Vec<keyscan::Match> =
found.into_iter().filter(|m| m.dummy.is_none()).collect();
let confidence = live.iter().map(|m| m.confidence).max();
let mut review =
Review::new(key, store::class::CONTENT, rule, live, confidence);
review.chunks = Some(n);
plan.review.push(review);
not(&mut plan, store::class::CONTENT, &path);
}
}
}
said.phase = "rules";
said.flagged = plan.review.len() + plan.credential.len();
on(said);
let images = plan.images as f64 * progress::image_units();
let units = if accel::images_beside_text() {
(plan.chunks as f64).max(images)
} else {
plan.chunks as f64 + images
};
plan.eta_ms = match accel::expected_rate() {
Some(rate) => Some((units / rate * 1000.0) as u64),
None => 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 + plan.image_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 unrecorded = || {
let at = Path::new(&key);
if at.is_dir() || std::fs::metadata(at).is_ok_and(|m| m.len() > chunk::MAX_FILE_BYTES) {
store::class::POLICY
} else {
store::class::CONTENT
}
};
let class = row
.as_ref()
.map_or_else(unrecorded, |r| r.class.as_str())
.to_string();
if class == store::class::CRASHED {
if mode != "as-is" {
anyhow::bail!("{key} crashed a parser; it can only be tried again (as-is)");
}
self.writing(|me| {
worker::clear(&me.dir, &key)?;
store::unrefuse(&me.db, &key)
})?;
return Ok(store::Acceptance {
path: key,
class,
mode: mode.to_string(),
fingerprints: Vec::new(),
confidence: None,
at: now(),
source: source.to_string(),
});
}
if !(store::class::reviewable(&class) || class == store::class::DUMMY) {
anyhow::bail!(
"{key} is not indexable ({}); accepting cannot change that",
class
);
}
if mode == "redacted" && class == store::class::POLICY {
anyhow::bail!("{key} holds nothing to redact; index it as it is or keep it out");
}
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(
read_regular(Path::new(&key), u64::MAX - 1),
|| format!("reading {key}"),
)?;
let text = match worker::read(Path::new(&key), &bytes, false)
.map_err(|stop| anyhow::anyhow!("{key}: {}", stop.reason()))?
.text
{
Ok(Ok((text, _))) => text,
Ok(Err(why)) => anyhow::bail!("{key}: {}", why.as_str()),
Err(panic) => anyhow::bail!("{key}: {panic}"),
};
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 fn index_paths_under(
&mut self,
roots: &[PathBuf],
control: &dyn Fn() -> Flow,
on_file: impl FnMut(&Path, IndexProgress),
) -> Result<IndexReport> {
let _lock = lock::StoreLock::acquire(&self.dir)?;
self.index_walk_under(roots, control, on_file)
}
pub fn index_rest_under(
&mut self,
files: Vec<PathBuf>,
control: &dyn Fn() -> Flow,
on_file: impl FnMut(&Path, IndexProgress),
) -> Result<IndexReport> {
let _lock = lock::StoreLock::acquire(&self.dir)?;
let unbounded = std::time::Duration::from_secs(100 * 365 * 24 * 3600);
self.index_rest_held_under(files, unbounded, None, control, on_file)
}
pub fn undo_run(&mut self, written: &[String]) -> Result<usize> {
let _lock = lock::StoreLock::acquire(&self.dir)?;
self.undo_held(written)
}
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 = store::mark_vector_format(&self.db)
.and_then(|()| store::release_automerge(&self.db))
.and_then(|()| body(self));
let _ = self.tx_commit();
let _ = store::release_automerge(&self.db);
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_secret(&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 = self.index.is_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 markup_extracted = {
let settings = crate::home::store_settings(&self.dir);
settings.kind == "docs" || settings.lean == "docs"
};
let markup_rule = if markup_extracted { "extracted" } else { "raw" };
let markup_moved = store::get_meta(&self.db, MARKUP_KEY)?
.as_deref()
.unwrap_or("extracted")
!= markup_rule;
let downloads = add::downloads_dir(&self.dir);
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));
let (chunks, images) = match self.expect_rest.take().filter(|_| budget.is_some()) {
Some(known) => known,
None => expected_of(&paths, self.planned.as_deref()),
};
report.expected_left = chunks;
report.images_left = images;
paths
}
None => {
total = head.len();
report.bytes_total = bytes_of(&head);
let (chunks, images) = expected_of(&head, self.planned.as_deref());
report.expected_left = chunks;
report.images_left = images;
if let Some(planned) = self.planned.as_deref() {
report.expected_left = report.expected_left.max(planned.values().sum());
total = total.max(planned.len());
}
if let Some(images) = self.planned_images {
report.images_left = report.images_left.max(images);
}
head.clone()
}
};
{
let mut live = report.live.borrow_mut();
live.eta.set_prior(accel::expected_rate());
live.image_weight = progress::image_units();
live.images_parallel = accel::images_beside_text();
}
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 cache_at = if self.force {
None
} else if let Some(own) = &self.vector_cache {
(own.cap > 0).then(|| own.clone())
} else if accel::cache_in_use() {
cache::Location::machine()
} else {
None
};
let ctx = pipeline::Context {
rechunk,
rescan,
regraph,
allow_secrets: self.boundary.allow_secrets,
hashes: if self.force {
Default::default()
} else {
let mut hashes = store::all_hashes(&self.db)?;
if markup_moved {
hashes.retain(|path, _| {
let path = Path::new(path);
!chunk::is_markup(path) || path.starts_with(&downloads)
});
}
hashes
},
uncounted: store::uncounted(&self.db)?,
over_cap: store::acceptances(&self.db)?
.into_iter()
.filter(|a| a.class == store::class::POLICY && a.mode != "refused")
.map(|a| a.path)
.collect(),
prehashed,
each: each.cloned(),
tokenizer,
clocks: pipeline::Clocks::default(),
cache_at: cache_at.clone(),
cache: cache_at.as_ref().map(|_| cache::Scope::of(&self.model)),
variants: accel::cache_variants(),
lookups: Default::default(),
hits: Default::default(),
journal: Some(worker::Journal::open(&self.dir)?),
};
let mut cached = CacheWrites::default();
let mut clock = pipeline::WriterClock::start(walk_ms);
self.write_parts.take();
self.awaiting.clear();
self.file_of.clear();
self.landed.clear();
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();
let gitignore = self.gitignore;
let first_import = store::first_import(&self.db, self.is_empty())?;
if first_import {
store::hold_automerge(&self.db)?;
}
let (window_size, row_commit) = if first_import {
(FIRST_IMPORT_WINDOW, FIRST_IMPORT_ROW_COMMIT)
} else {
(SORT_WINDOW, ROW_COMMIT)
};
std::thread::scope(|scope| -> Result<()> {
let prefetch = pipeline::Prefetch::start(scope, first, sealed, &ctx, prepare_threads);
let (image_send, image_jobs) = std::sync::mpsc::channel::<ImageJob>();
let (image_back, image_done) = std::sync::mpsc::channel::<ImageDone>();
let mut image_send = Some(image_send);
let mut image_jobs = Some(image_jobs);
let mut image_lane: Option<std::thread::ScopedJoinHandle<'_, image::Clip>> = None;
let mut clip = Some(std::mem::take(&mut self.clip));
let quiet = self.quiet;
let mut images_out = 0usize;
let mut image_bytes_out = 0u64;
let mut walker = pending_walk.map(|(roots, accepted)| {
scope.spawn(move || {
let started = std::time::Instant::now();
let mut walked = walk_allowing(&roots, &accepted, gitignore);
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);
let (chunks, images) = expected_of(&rest, me.planned.as_deref());
if me.planned.is_none() {
report.expected_left += chunks;
}
if me.planned_images.is_none() {
report.images_left += images;
}
prefetch.extend_and_seal(rest);
if report.live.borrow().phase == progress::Phase::Walk {
say_phase(
on_file,
report,
*total,
progress::Phase::Read,
Some(format!("{} files found", *total)),
);
}
Ok(())
};
let mut next = 0usize;
if walker.is_some() {
say_phase(
&mut on_file,
&report,
total,
progress::Phase::Walk,
Some(if head_set.is_empty() {
"finding the files under the roots".to_string()
} else {
"finding the files; reading the most recently changed first".to_string()
}),
);
}
loop {
if stage.is_some()
&& matches!(
report.live.borrow().phase,
progress::Phase::Read | progress::Phase::Walk
)
&& walker.is_none()
{
say_phase(&mut on_file, &report, total, progress::Phase::Embed, None);
}
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;
}
}
while let Ok(done) = image_done.try_recv() {
images_out -= 1;
report.images_flight = images_out;
image_bytes_out = image_bytes_out.saturating_sub(done.job.bytes.len() as u64);
if let Some((failed_path, why)) =
self.land_image(done, &mut report, &mut written, &mut completed)?
{
say_file(
&mut on_file,
&report,
total,
&failed_path,
FileOutcome::Failed,
Some(why),
);
}
}
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_ticking(seen, &mut || {
let size = embeddable_bytes(path.metadata().ok());
let name = path.file_name().map_or_else(
|| path.display().to_string(),
|n| n.to_string_lossy().into_owned(),
);
report.live.borrow_mut().detail = Some(format!(
"reading and chunking {name} ({})",
human_bytes(size as i64)
));
say_file(
&mut on_file,
&report,
total,
&path,
FileOutcome::Progress,
None,
);
});
clock.waited_on_prepare(waiting);
{
let mut live = report.live.borrow_mut();
if live
.detail
.as_deref()
.is_some_and(|d| d.starts_with("reading and chunking"))
{
live.detail = None;
}
}
let (estimated, image) = expected_one(&path, self.planned.as_deref());
report.expected_left = report.expected_left.saturating_sub(estimated);
report.current_est = estimated;
report.current_rows = 0;
if image {
report.images_left = report.images_left.saturating_sub(1);
} else {
let mut live = report.live.borrow_mut();
if live
.detail
.as_deref()
.is_some_and(|d| d.starts_with("embedding images"))
{
live.detail = None;
}
}
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 read_mtime: Option<i64>;
let (hash, len, text, units, found, ready) = match prepared {
pipeline::Prepared::Refused(_) => unreachable!("handled above"),
pipeline::Prepared::Unchanged {
size,
mtime,
read_mtime,
file_bytes,
} => {
report.bytes += file_bytes;
self.tx_begin()?;
store::heal_stamps(&self.db, &key, size as i64, mtime, read_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()?;
let (chunks, images) = self.evict(&key)?;
if chunks > 0 || images > 0 {
report.removed += 1;
if gone {
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,
read_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_secret(&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,
read_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::Image,
None,
);
if image_lane.is_none()
&& let (Some(jobs), Some(mut model)) = (image_jobs.take(), clip.take())
{
let back = image_back.clone();
image_lane = Some(scope.spawn(move || {
for job in jobs {
crate::cpucap::pace();
let vector = model
.embed_image(&job.path, &job.bytes, quiet)
.map_err(|e| format!("{e:#}"));
if back.send(ImageDone { job, vector }).is_err() {
break;
}
}
model
}));
}
while image_bytes_out > IMAGE_BYTES_IN_FLIGHT && images_out > 0 {
let Ok(done) = image_done.recv() else { break };
images_out -= 1;
report.images_flight = images_out;
image_bytes_out =
image_bytes_out.saturating_sub(done.job.bytes.len() as u64);
if let Some((path, why)) =
self.land_image(done, &mut report, &mut written, &mut completed)?
{
say_file(
&mut on_file,
&report,
total,
&path,
FileOutcome::Failed,
Some(why),
);
}
}
image_bytes_out += bytes.len() as u64;
images_out += 1;
report.images_flight = images_out;
if let Some(send) = &image_send {
let _ = send.send(ImageJob {
path: path.clone(),
key: key.clone(),
bytes,
hash,
width,
height,
});
}
continue;
}
pipeline::Prepared::Crashed { reason, file_bytes } => {
report.bytes += file_bytes;
self.tx_begin()?;
if store::file_hash(&self.db, &key)?.is_some() {
self.evict(&key)?;
}
store::refuse(
&self.db,
&key,
store::class::CRASHED,
&reason,
&[],
1,
now(),
)?;
report
.refused
.push((path.display().to_string(), reason.clone()));
say_file(
&mut on_file,
&report,
total,
&path,
FileOutcome::Skipped,
Some(reason),
);
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,
read_mtime: seen,
text,
units,
found,
ready,
} => {
report.bytes += file_bytes;
read_mtime = seen;
if chunk::is_markup(&path)
&& (markup_extracted || path.starts_with(&downloads))
{
let text = chunk::markup_text(&text);
let found = keyscan::scan(&key, &text);
(hash, len, text, units, found, None)
} else {
(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_secret(&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,
None,
&ctx,
None,
std::time::Instant::now(),
),
(None, Some(ready)) => ready,
(None, None) => pipeline::take_ready(
&path,
&text,
None,
&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,
Some(match chunks.len() {
1 => "1 chunk".to_string(),
n => format!("{n} chunks"),
}),
);
let timed = std::time::Instant::now();
for id in store::delete_file(&self.db, &key, now())? {
self.index.remove(id)?;
}
self.write_parts.add("evict", timed);
let timed = std::time::Instant::now();
let file_id = store::insert_file(&self.db, &key, PENDING, len, now())?;
store::set_read_mtime(&self.db, file_id, read_mtime)?;
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)?;
}
self.write_parts.add("rows", timed);
let mut spans: Vec<(u32, u32, i64)> = Vec::with_capacity(chunks.len());
let mut halted = false;
let count = chunks.len();
self.awaiting.insert(file_id, (0, hash.clone(), false));
for (((ord, c), piece), hash) in chunks.iter().enumerate().zip(pieces).zip(hashes) {
let timed = std::time::Instant::now();
let id = store::insert_chunk(
&self.db,
file_id,
ord,
c.start_line,
c.end_line,
&c.text,
)?;
self.write_parts.add("rows", timed);
spans.push((c.start_line, c.end_line, id));
self.file_of.insert(id as u64, file_id);
if let Some(left) = self.awaiting.get_mut(&file_id) {
left.0 += 1;
}
report.rows += 1;
report.current_rows += 1;
window.ids.push(id as u64);
window.pieces.push(piece);
window.hashes.push(hash);
if window.ids.len() >= window_size {
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;
report.threads = live_threads(report.threads);
say_file(
&mut on_file,
&report,
total,
&path,
FileOutcome::Progress,
None,
)
},
)? {
halted = true;
break;
}
self.tx_begin()?;
}
}
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 timed = std::time::Instant::now();
let (symbols, edges) = self.write_graph(extraction, file_id, &spans)?;
self.write_parts.add("graph", timed);
report.symbols += symbols;
report.edges += edges;
if let Some(entry) = self.awaiting.get_mut(&file_id) {
entry.2 = true;
if entry.0 == 0
&& let Some((_, hash, _)) = self.awaiting.remove(&file_id)
{
self.landed.push((file_id, hash));
}
}
report.indexed += 1;
report.chunks += count;
report.current_est = 0;
report.current_rows = 0;
if committed.elapsed() >= row_commit {
self.tx_commit()?;
committed = std::time::Instant::now();
}
since_checkpoint += 1;
if checkpointing && since_checkpoint >= every {
if !window.ids.is_empty()
&& !self.hand_over(
scope,
&mut stage,
&mut window,
&cpu_back,
&paused,
&ask,
&mut clock,
&mut |n| {
report.batched = n;
report.threads = live_threads(report.threads);
say_file(
&mut on_file,
&report,
total,
&path,
FileOutcome::Progress,
None,
)
},
)?
{
report.remaining = prefetch.len().0 - seen - 1;
report.stopped = true;
break;
}
if let Some(stage) = stage.as_mut() {
while let Some(done) = stage.next(std::time::Duration::ZERO) {
self.land(
done.map_err(anyhow::Error::msg)?,
&mut report,
&mut run_lanes,
&mut cached,
)?;
}
}
say_phase(
&mut on_file,
&report,
total,
progress::Phase::Save,
Some(format!(
"checkpoint: writing the index to disk ({} chunks still embedding carry on)",
progress::grouped(report.rows.saturating_sub(report.embedded) as u64)
)),
);
completed.append(&mut self.landed);
self.checkpoint(&mut completed)?;
since_checkpoint = 0;
say_phase(&mut on_file, &report, total, progress::Phase::Embed, None);
}
}
let last = prefetch
.path(prefetch.len().0.saturating_sub(1))
.unwrap_or_default();
drop(prefetch);
if !report.stopped {
if report.rows > report.embedded || !window.ids.is_empty() {
say_phase(
&mut on_file,
&report,
total,
progress::Phase::Drain,
Some(format!(
"{} chunks still embedding",
progress::grouped(report.rows.saturating_sub(report.embedded) as u64)
)),
);
}
let drained = self.hand_over(
scope,
&mut stage,
&mut window,
&cpu_back,
&paused,
&ask,
&mut clock,
&mut |n| {
report.batched = n;
report.threads = live_threads(report.threads);
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;
}
}
image_send = None;
if images_out > 0 {
say_phase(
&mut on_file,
&report,
total,
progress::Phase::Drain,
Some(format!(
"{} images still with the image model",
progress::grouped(images_out as u64)
)),
);
}
while images_out > 0 {
let Ok(done) = image_done.recv() else { break };
images_out -= 1;
report.images_flight = images_out;
if let Some((failed_path, why)) =
self.land_image(done, &mut report, &mut written, &mut completed)?
{
say_file(
&mut on_file,
&report,
total,
&failed_path,
FileOutcome::Failed,
Some(why),
);
}
}
drop(image_send);
drop(image_back);
match image_lane.take() {
Some(lane) => {
if let Ok(model) = lane.join() {
self.clip = model;
}
}
None => {
if let Some(model) = clip.take() {
self.clip = model;
}
}
}
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);
let tail = std::time::Instant::now();
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) = ctx.cache_at.as_ref().and_then(cache::Cache::open_in)
{
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 timed = std::time::Instant::now();
let _ = cache.record(&fresh, &hits, report.cache_lookups as u64);
self.write_parts.add("cache", timed);
}
let mut ghosts = 0usize;
for key in &written {
if !Path::new(key).exists() {
let (chunks, images) = self.evict(key)?;
ghosts += usize::from(chunks + images > 0);
}
}
if ghosts > 0 {
report.removed += ghosts;
self.save()?;
}
report.written = written;
if sweep && report.pending.is_empty() && !report.stopped {
say_phase(
&mut on_file,
&report,
total,
progress::Phase::Finalize,
Some("dropping files gone from disk".to_string()),
);
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;
}
}
let accepted = self.accepted_folders();
let roots: Option<Vec<PathBuf>> = boundary
.boundary
.roots
.as_ref()
.map(|roots| roots.iter().map(|r| canonical(r)).collect());
for row in store::refusals(&self.db)? {
let path = Path::new(&row.path);
if !path.exists()
|| (under_skipped_dir(path, roots.as_deref())
&& !accepted.iter().any(|a| path.starts_with(a)))
{
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 markup_moved {
store::set_meta(&self.db, MARKUP_KEY, markup_rule)?;
}
}
if report.indexed > 0 || report.removed > 0 || !self.index.exists() {
say_phase(
&mut on_file,
&report,
total,
progress::Phase::Save,
Some("writing the index to disk".to_string()),
);
let timed = std::time::Instant::now();
self.save()?;
self.write_parts.add("save", timed);
}
completed.append(&mut self.landed);
self.commit_hashes(&mut completed)?;
let tail = tail.elapsed().as_millis() as u64;
report.stages.wall_ms += tail;
report.stages.write_ms += tail;
report.stages.write_parts_ms = self.write_parts.take();
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 first_import && report.pending.is_empty() && !report.stopped {
store::end_first_import(&self.db)?;
}
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(());
}
let timed = std::time::Instant::now();
self.save()?;
self.write_parts.add("save", timed);
self.commit_hashes(completed)
}
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 write_graph(
&self,
extraction: Option<graph::Extraction>,
file_id: i64,
spans: &[(u32, u32, i64)],
) -> Result<(usize, usize)> {
match extraction {
Some(extraction) => store::write_extraction(&self.db, &extraction, file_id, spans),
None => Ok((0, 0)),
}
}
fn commit_hashes(&mut self, completed: &mut Vec<(i64, String)>) -> Result<()> {
if completed.is_empty() {
return Ok(());
}
let _charged = pipeline::Charge(&self.write_parts, "hashes", std::time::Instant::now());
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;
report.threads = live_threads(report.threads);
tick(report);
}
}
}
}
Ok(true)
}
fn land_image(
&mut self,
done: ImageDone,
report: &mut IndexReport,
written: &mut Vec<String>,
completed: &mut Vec<(i64, String)>,
) -> Result<Option<(PathBuf, String)>> {
let ImageDone { job, vector } = done;
let _charged = pipeline::Charge(&self.write_parts, "images", std::time::Instant::now());
let vector = match vector {
Ok(vector) => vector,
Err(why) => {
failed(report, &job.path, &anyhow::anyhow!("{why}"));
return Ok(Some((job.path, why)));
}
};
self.tx_begin()?;
for id in store::image_ids_of(&self.db, &job.key)? {
self.images.remove(id as u64)?;
}
for id in store::delete_file(&self.db, &job.key, now())? {
self.index.remove(id)?;
}
let file_id =
store::insert_file(&self.db, &job.key, PENDING, job.bytes.len() as u64, now())?;
written.push(job.key.clone());
let image_id = store::insert_image(&self.db, file_id, job.width, job.height)?;
self.images.add(&vector, &[image_id as u64])?;
completed.push((file_id, job.hash));
report.indexed += 1;
report.images += 1;
Ok(None)
}
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()
);
crate::cpucap::pace();
self.follow_budget();
let timed = std::time::Instant::now();
let (index, exact) = (&self.index, &mut self.exact);
crate::cpucap::within(|| -> Result<()> {
index.add(&flat, &done.ids)?;
exact.append(&flat, &done.ids)
})?;
self.write_parts.add("vectors", timed);
for id in &done.ids {
if let Some(file) = self.file_of.remove(id)
&& let Some(left) = self.awaiting.get_mut(&file)
{
left.0 = left.0.saturating_sub(1);
if left.0 == 0
&& left.2
&& let Some((_, hash, _)) = self.awaiting.remove(&file)
{
self.landed.push((file, hash));
}
}
}
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")?;
store::drop_insert_trigger(&self.db)?;
}
Ok(())
}
fn tx_commit(&self) -> Result<()> {
let _charged = pipeline::Charge(&self.write_parts, "commit", std::time::Instant::now());
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_as(key, false)?;
}
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)> {
self.evict_as(key, true)
}
fn evict_secret(&mut self, key: &str) -> Result<(usize, usize)> {
let ids: Vec<u64> = store::chunk_ids_of(&self.db, key)?;
let vectors: Vec<Vec<f32>> = ids.iter().filter_map(|id| self.exact.get(*id)).collect();
store::secure_delete_fts(&self.db)?;
let evicted = self.evict(key)?;
self.exact.zero(&ids)?;
let mut caches: Vec<cache::Location> = cache::Location::machine().into_iter().collect();
if let Some(own) = &self.vector_cache {
caches.push(own.clone());
}
for at in caches {
if let Some(mut cache) = cache::Cache::open_in(&at) {
let _ = cache.forget_vectors(&vectors);
}
}
self.scrub |= evicted.0 + evicted.1 > 0;
Ok(evicted)
}
fn evict_as(&mut self, key: &str, retire: bool) -> Result<(usize, usize)> {
let images = store::image_ids_of(&self.db, key)?;
let ids = if retire {
store::delete_file(&self.db, key, now())?
} else {
store::delete_file_unretired(&self.db, key)?
};
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,
Option<std::sync::Arc<std::collections::HashSet<u64>>>,
)> {
if filter.is_empty() {
return Ok((Allowlist::All, None));
}
let newest_chunk = store::newest_chunk(&self.db)?;
let (generation, groups) = (self.generation, filter.groups());
self.filter_memo
.retain(|m| m.generation == generation && m.newest_chunk == newest_chunk);
if let Some(at) = self.filter_memo.iter().position(|m| m.groups == groups) {
let memo = self.filter_memo.remove(at).expect("position is in range");
let out = (memo.allowlist.clone(), Some(memo.selected.clone()));
self.filter_memo.push_front(memo);
return Ok(out);
}
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);
}
}
let allowlist = if ids.is_empty() {
Allowlist::Empty
} else if ids.len() == self.len() {
Allowlist::All
} else {
Allowlist::subset(ids)
};
let selected = std::sync::Arc::new(candidates.into_iter().collect());
self.filter_memo.push_front(FilterMemo {
groups: filter.groups().to_vec(),
generation: self.generation,
newest_chunk,
allowlist: allowlist.clone(),
selected: std::sync::Arc::clone(&selected),
});
self.filter_memo.truncate(FILTER_MEMOS);
Ok((allowlist, Some(selected)))
}
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 = FUSE_DEPTH.max(k);
let shape = shape_of(query);
let (allowlist, selected) = 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 named: Vec<u64> = if shape != Shape::Identifier {
let names = named_identifiers(query);
let mut ids = Vec::new();
for name in &names {
if let Some(chunk) =
store::symbols_by_names(&self.db, std::slice::from_ref(name), filter.groups())?
.into_iter()
.find_map(|row| row.chunk_id)
{
let id = chunk as u64;
if !ids.contains(&id) {
ids.push(id);
}
}
}
ids
} else {
Vec::new()
};
let (dense_scores, dense_ids) = self.search_vectors(vector, depth, &allowlist)?;
let (dense_scores, dense_ids) = self.rescored(vector, dense_scores, dense_ids);
let keyword_ids = match &selected {
Some(set) => store::keyword_search_within(&self.db, query, depth, set)?,
None => store::keyword_search(&self.db, query, depth, &[])?,
};
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, truncated) = if pictured(query) {
self.image_search(query, depth, filter)?
} else {
(Vec::new(), false)
};
let image_list: Vec<(u64, f32)> = if truncated {
Vec::new()
} else {
images
.iter()
.filter(|(_, similarity)| *similarity >= image_floor())
.map(|(id, _)| (*id, TEXT_LISTS))
.collect()
};
let candidate_lists = [
(
"definition",
false,
definitions.iter().map(|id| (*id, 0.0)).collect(),
),
("image", true, image_list),
("named", false, named.iter().map(|id| (*id, 1.0)).collect()),
(
"vector",
false,
dense_ids.iter().map(|id| (*id, 1.0)).collect(),
),
(
"keyword",
false,
keyword_ids
.iter()
.map(|id| (*id, shape.keyword_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());
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,
copies: Vec::new(),
agreed: Vec::new(),
},
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,
copies: Vec::new(),
agreed: Vec::new(),
},
similarity,
));
}
}
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);
let (prior, named) = project_prior(&hits, query);
for (((hit, _), share), named) in hits.iter_mut().zip(prior).zip(named) {
hit.score *= 1.0 + PROJECT_LIFT * share;
if named {
hit.score *= 1.0 + NAMED_PROJECT_LIFT;
}
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);
}
if !names_releases(query) {
hits.sort_by_key(|(h, _)| is_release_notes(&h.path));
}
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.split('\n').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,
replaced: 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 at = Path::new(&span.path);
if canonical(at) != at || filter::denied(at).is_some_and(|d| d != filter::Denied::Hidden) {
span.replaced = at.exists();
return Ok(());
}
let file = match open_regular(at) {
Ok((file, _)) => file,
Err(_) => {
span.replaced = at.symlink_metadata().is_ok();
return Ok(());
}
};
let mut bytes = Vec::new();
{
use std::io::Read;
if file
.take(chunk::MAX_FILE_BYTES + 1)
.read_to_end(&mut bytes)
.is_err()
{
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(stamp) = stamps.get(&span.path) else {
return Ok(());
};
span.fresh = std::fs::metadata(&span.path).is_ok_and(|meta| stamp.fresh(&meta));
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(stamp) = stamps.get(path) else {
fresh.insert(path.as_str(), true);
continue;
};
let state = match std::fs::metadata(path) {
Ok(meta) => stamp.fresh(&meta),
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)>, bool)> {
if store::image_count(&self.db)? == 0 {
return Ok((Vec::new(), false));
}
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(), false));
}
index::Allowlist::subset(ids)
};
let (mut vector, truncated) = 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(),
truncated,
))
}
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 * KB {
format!("{:.1} GB", b / (KB * KB * KB))
} else 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_lists<S: serde::Serializer>(
groups: &std::collections::BTreeMap<String, Vec<String>>,
s: S,
) -> Result<S::Ok, S::Error> {
s.collect_map(
groups
.iter()
.map(|(k, paths)| (k, paths.iter().map(|p| plain(p)).collect::<Vec<String>>())),
)
}
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) };
}
fn live_threads(known: usize) -> usize {
match pipeline::CPU_THREADS.load(std::sync::atomic::Ordering::Relaxed) {
0 => known,
live => live,
}
}
pub fn open_regular(path: &Path) -> std::io::Result<(std::fs::File, std::fs::Metadata)> {
let mut options = std::fs::OpenOptions::new();
options.read(true);
#[cfg(unix)]
{
use std::os::unix::fs::OpenOptionsExt;
options.custom_flags(libc::O_NOFOLLOW | libc::O_NONBLOCK);
}
#[cfg(windows)]
{
use std::os::windows::fs::OpenOptionsExt;
options.custom_flags(0x0020_0000);
}
let file = options.open(path);
#[cfg(unix)]
let file = file.map_err(|e| {
if e.raw_os_error() == Some(libc::ELOOP) {
std::io::Error::other("it is a symbolic link, not a file")
} else {
e
}
});
let file = file?;
let meta = file.metadata()?;
#[cfg(windows)]
{
use std::os::windows::fs::MetadataExt;
if meta.file_attributes() & 0x400 != 0 {
return Err(std::io::Error::other("it is a link, not a file"));
}
}
if !meta.is_file() {
return Err(std::io::Error::other("it is not a regular file"));
}
Ok((file, meta))
}
pub fn read_regular(path: &Path, cap: u64) -> std::io::Result<Vec<u8>> {
let (file, _) = open_regular(path)?;
read_capped(file, cap)
}
fn read_capped(file: std::fs::File, cap: u64) -> std::io::Result<Vec<u8>> {
use std::io::Read;
let mut bytes = Vec::new();
file.take(cap.saturating_add(1)).read_to_end(&mut bytes)?;
if bytes.len() as u64 > cap {
return Err(std::io::Error::new(
std::io::ErrorKind::FileTooLarge,
format!("it is over {} MiB", cap >> 20),
));
}
Ok(bytes)
}
pub fn git(dir: &Path) -> std::process::Command {
let (null, never) = if cfg!(windows) {
("NUL", "NUL")
} else {
("/dev/null", "/bin/false")
};
let mut cmd = std::process::Command::new("git");
cmd.env("GIT_CONFIG_NOSYSTEM", "1")
.arg("-C")
.arg(dir)
.args(["-c", "core.fsmonitor=false"])
.arg("-c")
.arg(format!("core.hooksPath={null}"))
.arg("-c")
.arg(format!("gpg.program={never}"))
.args(["-c", "log.showSignature=false", "-c", "core.pager=cat"])
.stdin(std::process::Stdio::null());
cmd
}
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 covered =
|(s, e, _, _): &&(u32, u32, String, String)| (*e).min(end).saturating_sub((*s).max(start));
let contains = || {
symbols
.iter()
.filter(|(s, e, _, _)| *s <= start && *e >= start)
.min_by_key(|(s, e, _, _)| e.saturating_sub(*s))
};
let starts_inside = symbols
.iter()
.filter(|(s, _, _, _)| *s >= start && *s <= end)
.min_by_key(|(s, e, _, _)| (*s, std::cmp::Reverse(e.saturating_sub(*s))))
.filter(|inside| {
!contains()
.is_some_and(|before| before.1 < inside.0 && covered(&before) > covered(inside))
});
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";
pub const HIDDEN_RULE: &str = "hidden file or folder";
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",
".pnpm-store",
"jspm_packages",
"storybook-static",
".astro",
".vite",
".output",
".vercel",
".netlify",
".angular",
".docusaurus",
".expo",
".wrangler",
"__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",
"vite.config.js",
"vite.config.ts",
"vite.config.mjs",
"astro.config.mjs",
"astro.config.ts",
"deno.json",
],
),
(
"out",
&[
"package.json",
"tsconfig.json",
"next.config.js",
"next.config.mjs",
"next.config.ts",
],
),
("bin", &["*.csproj", "*.sln", "*.fsproj"]),
("obj", &["*.csproj", "*.sln", "*.fsproj"]),
("coverage", &["package.json", "pyproject.toml"]),
("Pods", &["Podfile"]),
];
fn under_skipped_dir(path: &Path, roots: Option<&[PathBuf]>) -> bool {
path.ancestors()
.skip(1)
.take_while(|a| roots.is_none_or(|roots| !roots.iter().any(|r| r == a)))
.any(|a| is_generated_dir(a) || filter::is_denied_dir(a))
}
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)
}
pub fn push_files(root: &Path, cap: u64) -> (Vec<PathBuf>, Vec<(PathBuf, String)>) {
let walked = walk_allowing(&[root.to_path_buf()], &[], true);
let home = crate::home::user_home().ok().map(|h| canonical(&h));
let boundary = Boundary {
roots: None,
allow_secrets: false,
}
.resolved(home.as_deref());
let mut ok = Vec::new();
let mut refused: Vec<(PathBuf, String)> = walked.unreadable;
for path in walked.credentials {
refused.push((path, "a credential file by its name".to_string()));
}
for path in walked.files.into_iter().chain(walked.named) {
if let Some(refusal) = boundary.refuses(&path, true) {
refused.push((path, refusal.why));
continue;
}
let bytes = match read_regular(&path, cap) {
Ok(bytes) => bytes,
Err(e) if e.kind() == std::io::ErrorKind::FileTooLarge => {
refused.push((path, format!("over {} MB", cap / (1 << 20))));
continue;
}
Err(_) => {
refused.push((path, "could not be read".to_string()));
continue;
}
};
let text = String::from_utf8_lossy(&bytes);
if let Some(live) = keyscan::scan(&path.to_string_lossy(), &text)
.into_iter()
.find(|m| m.dummy.is_none())
{
refused.push((
path,
format!("holds what looks like {} at line {}", live.kind, live.line),
));
continue;
}
ok.push(path);
}
refused.sort();
(ok, refused)
}
#[cfg(test)]
fn walk(roots: &[PathBuf]) -> Walked {
walk_allowing(roots, &[], true)
}
fn bytes_of(paths: &[PathBuf]) -> u64 {
paths
.iter()
.map(|p| embeddable_bytes(p.metadata().ok()))
.sum()
}
fn planned_chunks(path: &Path, text: &str, symbols: &[graph::Symbol]) -> u64 {
chunk::chunk_file(path, text, symbols).len() as u64
}
fn expected_of(
paths: &[PathBuf],
planned: Option<&std::collections::HashMap<String, u64>>,
) -> (u64, usize) {
let mut chunks = 0;
let mut images = 0;
for path in paths {
let (n, image) = expected_one(path, planned);
chunks += n;
images += usize::from(image);
}
(chunks, images)
}
fn expected_one(
path: &Path,
planned: Option<&std::collections::HashMap<String, u64>>,
) -> (u64, bool) {
let bytes = embeddable_bytes(path.metadata().ok());
let planned = planned.map(|p| p.get(path.to_string_lossy().as_ref()).copied().unwrap_or(0));
progress::estimate(path, bytes, planned)
}
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) = git(&root)
.args([
"log",
"--no-show-signature",
"-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) = git(&top)
.args([
"log",
"--no-show-signature",
"-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], gitignore: bool) -> 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(gitignore)
.git_global(gitignore)
.git_exclude(gitignore)
.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.depth() > 0
&& e.file_type().is_some_and(|t| t.is_dir())
&& filter::is_denied_dir(e.path())
{
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 yielded_here = if kind.is_dir() {
walked_dirs.contains(&path)
} else {
yielded.contains(&path)
};
if !yielded_here && name != ".git" && (kind.is_file() || kind.is_dir()) {
excluded.push((key(&entry.file_name()), HIDDEN_RULE.to_string()));
}
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 only_a_query_describing_something_seen_asks_for_pictures() {
for seen in [
"a red circle",
"a black and white checkerboard pattern",
"yellow horizontal stripes",
"an orange ring, like a donut",
"a grey gradient fading from dark to light",
"a screenshot of the flaskr login page",
"what does the debugger look like",
] {
assert!(pictured(seen), "{seen}");
}
for text in [
"register a blueprint url prefix",
"where does flask register a blueprint",
"how does flask dispatch a request to a view",
"how does ripgrep detect binary files",
"search_preferring",
] {
assert!(!pictured(text), "{text}");
}
}
#[test]
fn a_skipped_dotfile_is_listed_not_silent() {
let tmp = tempfile::tempdir().unwrap();
let root = tmp.path().join("repo");
std::fs::create_dir_all(root.join(".cargo")).unwrap();
std::fs::create_dir_all(root.join(".git")).unwrap();
std::fs::create_dir_all(root.join("src")).unwrap();
std::fs::write(root.join(".cargo/config.toml"), "x").unwrap();
std::fs::write(root.join(".nvim.lua"), "x").unwrap();
std::fs::write(root.join(".ignore"), "!/.github/\n").unwrap();
std::fs::create_dir_all(root.join(".github")).unwrap();
std::fs::write(root.join(".github/ci.yml"), "x").unwrap();
std::fs::write(root.join("src/a.rs"), "fn a() {}").unwrap();
let walked = walk_allowing(std::slice::from_ref(&root), &[], true);
let hidden: Vec<String> = walked
.excluded
.iter()
.filter(|(_, rule)| rule == HIDDEN_RULE)
.map(|(p, _)| p.file_name().unwrap().to_string_lossy().into_owned())
.collect();
assert!(hidden.contains(&".cargo".to_string()), "{hidden:?}");
assert!(hidden.contains(&".nvim.lua".to_string()), "{hidden:?}");
assert!(hidden.contains(&".ignore".to_string()), "{hidden:?}");
assert!(!hidden.contains(&".git".to_string()), "{hidden:?}");
assert!(
!hidden.contains(&".github".to_string()),
"a whitelisted folder: {hidden:?}"
);
}
fn def(start: u32, end: u32, name: &str) -> (u32, u32, String, String) {
(start, end, name.to_string(), "function".to_string())
}
#[test]
fn a_chunk_is_named_for_the_definition_it_is_mostly_about() {
let symbols = [def(5, 50, "main"), def(52, 55, "DirEntry")];
let (_, name, _) = enclosing_definition(&symbols, 25, 53).unwrap();
assert_eq!(name, "main");
}
#[test]
fn a_chunk_that_starts_on_a_tail_belongs_to_the_definition_below() {
let symbols = [def(1, 12, "before"), def(14, 40, "after")];
let (_, name, _) = enclosing_definition(&symbols, 11, 30).unwrap();
assert_eq!(name, "after");
let (_, name, _) = enclosing_definition(&symbols, 13, 30).unwrap();
assert_eq!(name, "after");
}
#[test]
fn a_nested_definition_starting_inside_the_chunk_wins_over_its_parent() {
let symbols = [def(1, 100, "Impl"), def(20, 60, "method")];
let (_, name, _) = enclosing_definition(&symbols, 18, 50).unwrap();
assert_eq!(name, "method");
}
#[test]
#[ignore = "downloads an embedding model on first run"]
fn alternating_scopes_both_stay_resolved() {
let corpus = tempfile::tempdir().unwrap();
let store = tempfile::tempdir().unwrap();
let others: Vec<String> = (0..FILTER_MEMOS).map(|i| format!("other{i}")).collect();
for dir in ["one", "two"]
.into_iter()
.chain(others.iter().map(String::as_str))
{
std::fs::create_dir(corpus.path().join(dir)).unwrap();
std::fs::write(
corpus.path().join(dir).join("retry.rs"),
format!("fn backoff_{dir}() {{ let delay = base * 2u32.pow(attempt); }}\n"),
)
.unwrap();
}
let mut s = Semlith::open(store.path(), None).unwrap();
s.quiet = true;
s.index_paths(&[corpus.path().to_path_buf()], |_, _| {})
.unwrap();
let one = Filter::new(&["one/**".into()], &[], &[]).unwrap();
let two = Filter::new(&["two/**".into()], &[], &[]).unwrap();
for filter in [&one, &two, &one, &two] {
s.search_filtered("retry backoff", 4, filter).unwrap();
}
assert_eq!(s.filter_memo.len(), 2, "each scope resolved once");
assert_eq!(s.filter_memo[0].groups, two.groups());
for dir in &others {
let f = Filter::new(&[format!("{dir}/**")], &[], &[]).unwrap();
s.search_filtered("retry backoff", 4, &f).unwrap();
}
assert_eq!(s.filter_memo.len(), FILTER_MEMOS);
assert!(s.filter_memo.iter().all(|m| m.groups != one.groups()));
}
#[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,
copies: Vec::new(),
agreed: 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 release_notes_are_named_files_and_changelog_directories() {
for path in [
"/w/repo/CHANGELOG.md",
"/w/repo/CHANGES.rst",
"/w/repo/HISTORY.md",
"/w/repo/NEWS",
"/w/repo/docs/changelogs/v1.2.md",
] {
assert!(is_release_notes(path), "{path}");
}
for path in [
"/w/repo/README.md",
"/w/repo/src/history.rs",
"/w/repo/src/news_feed.rs",
] {
assert!(!is_release_notes(path), "{path}");
}
assert!(names_releases("what changed in the 1.2 release"));
assert!(!names_releases("how does the parser recover from an error"));
}
#[test]
fn a_sentence_names_its_code_shaped_words() {
assert_eq!(
named_identifiers("why does `parse_args` call Config::load and buildRequest() twice"),
["parse_args", "load", "buildRequest"]
);
assert!(named_identifiers("how does the server shut down").is_empty());
assert_eq!(
named_identifiers("a_b c_d e_f g_h").len(),
3,
"capped at three"
);
}
#[test]
fn the_project_prior_follows_the_head_and_ignores_one_project() {
let root = tempfile::tempdir().unwrap();
for repo in ["one", "two"] {
std::fs::create_dir_all(root.path().join(repo).join(".git")).unwrap();
std::fs::create_dir_all(root.path().join(repo).join("src")).unwrap();
}
let hit = |repo: &str, score: f32| {
(
Hit {
score,
path: root
.path()
.join(repo)
.join("src/a.rs")
.display()
.to_string(),
start_line: 1,
end_line: 2,
text: String::new(),
store: None,
lists: vec!["vector"],
image: None,
fresh: true,
symbol: None,
symbol_kind: None,
symbol_line: None,
copies: Vec::new(),
agreed: Vec::new(),
},
0.0,
)
};
let hits = vec![hit("one", 3.0), hit("one", 2.0), hit("two", 1.0)];
let (shares, named) = project_prior(&hits, "how does two handle it");
assert!(shares[0] > shares[2], "{shares:?}");
assert!((shares[0] - 5.0 / 6.0).abs() < 1e-6, "{shares:?}");
assert_eq!(
named,
[false, false, false],
"a word of three letters names nothing"
);
let alone = vec![hit("one", 3.0), hit("one", 2.0)];
let (shares, _) = project_prior(&alone, "anything");
assert_eq!(shares, [0.0, 0.0]);
}
#[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() {
for penalty in [STALE_PENALTY, TEST_PENALTY] {
let weakest = 1.0 - penalty;
assert!(
1.0 / weakest < 1.6,
"the rerank spans 1/{weakest}, which is a ranking rather than a tiebreak"
);
}
const { assert!(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 skipped_dir_tests {
use super::*;
#[test]
fn a_path_below_generated_output_or_git_is_under_a_skipped_dir() {
let dir = tempfile::tempdir().unwrap();
let root = canonical(dir.path());
std::fs::write(root.join("package.json"), "{}").unwrap();
let roots = [root.clone()];
let under = |rel: &str| under_skipped_dir(&root.join(rel), Some(&roots));
assert!(under("node_modules/.bin/esbuild"));
assert!(under("dist/assets/app.js"));
assert!(under(".git/index.lock"));
assert!(under("sub/.aws/notes.md"));
assert!(!under("src/dist.rs"));
assert!(!under("notes.md"));
let named = root.join("node_modules");
assert!(!under_skipped_dir(
&named.join("pkg.json"),
Some(std::slice::from_ref(&named))
));
}
}
#[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(""), "");
}
}