use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::time::UNIX_EPOCH;
use serde::{Deserialize, Serialize};
mod build;
mod carry_over;
mod chunking;
pub use chunking::*;
mod coordinator;
pub mod coverage;
mod evict;
pub(crate) mod external_origin;
pub use coordinator::{SearchIndexBuildProgress, get_or_start_build};
mod provider_origin;
pub use provider_origin::ChunkOrigins;
mod served_ranking;
pub mod source_gate;
#[cfg(test)]
mod tests;
mod walk;
pub use walk::is_code_file;
use walk::{admitted_scoped_source_paths, list_code_files};
#[cfg(test)]
use walk::{list_code_files_with, list_scoped_code_files};
pub(crate) const MAX_BM25_FILES: usize = 5000;
const CHUNK_COUNT_WARNING: usize = 50_000;
const ZSTD_LEVEL: i32 = 9;
const DEFAULT_BM25_IGNORES: &[&str] = &[
"vendor/**",
"dist/**",
"build/**",
"public/vendor/**",
"public/js/**",
"public/css/**",
"public/build/**",
".next/**",
".nuxt/**",
"__pycache__/**",
"*.min.js",
"*.min.css",
"*.bundle.js",
"*.chunk.js",
];
fn max_bm25_cache_bytes() -> u64 {
let mb = std::env::var("LEAN_CTX_BM25_MAX_CACHE_MB")
.ok()
.and_then(|v| v.parse::<u64>().ok())
.unwrap_or_else(|| crate::core::config::Config::load().bm25_max_cache_mb_effective());
mb * 1024 * 1024
}
pub fn persist_ceiling_bytes() -> u64 {
max_bm25_cache_bytes()
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SaveOutcome {
Persisted { compressed_bytes: u64 },
SkippedTooLarge {
compressed_bytes: u64,
limit_bytes: u64,
},
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct CodeChunk {
pub file_path: String,
pub symbol_name: String,
pub kind: ChunkKind,
pub start_line: usize,
pub end_line: usize,
pub content: String,
#[serde(default)]
pub tokens: Vec<String>,
pub token_count: usize,
#[serde(default)]
pub input_origin: Option<crate::core::context_admission::origin::InputOrigin>,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum ChunkKind {
Function,
Struct,
Impl,
Module,
Class,
Method,
Other,
Issue,
PullRequest,
WikiPage,
DbSchema,
ApiEndpoint,
Ticket,
ExternalOther,
}
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
pub struct IndexedFileState {
pub mtime_ms: u64,
pub size_bytes: u64,
}
impl IndexedFileState {
fn from_path(path: &Path) -> Option<Self> {
let meta = path.metadata().ok()?;
let size_bytes = meta.len();
let mtime_ms = meta
.modified()
.ok()
.and_then(|t| t.duration_since(UNIX_EPOCH).ok())
.map(|d| d.as_millis() as u64)?;
Some(Self {
mtime_ms,
size_bytes,
})
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BM25Index {
pub chunks: Vec<CodeChunk>,
pub inverted: HashMap<String, Vec<(usize, f64)>>,
pub avg_doc_len: f64,
pub doc_count: usize,
pub doc_freqs: HashMap<String, usize>,
#[serde(default)]
pub files: HashMap<String, IndexedFileState>,
#[serde(default)]
pub dirs: HashMap<String, u64>,
#[serde(default)]
pub admission_policy: Option<String>,
#[serde(default, skip)]
pub content_truncated: bool,
#[serde(default, skip)]
pub source_root: Option<PathBuf>,
#[serde(default, skip)]
pub provider_origins: std::sync::Arc<ChunkOrigins>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SearchResult {
pub chunk_idx: usize,
pub score: f64,
pub file_path: String,
pub symbol_name: String,
pub kind: ChunkKind,
pub start_line: usize,
pub end_line: usize,
pub snippet: String,
}
const BM25_K1: f64 = 1.2;
const BM25_B: f64 = 0.75;
fn structural_query_boost(query: &str, chunk: &CodeChunk) -> f64 {
let raw_terms: std::collections::HashSet<String> = query
.split(|c: char| !(c.is_alphanumeric() || c == '_'))
.filter(|term| term.len() >= 3)
.map(str::to_ascii_lowercase)
.collect();
let path_terms: std::collections::HashSet<String> = tokenize(&chunk.file_path)
.into_iter()
.map(|term| term.to_ascii_lowercase())
.collect();
let mut boost = 0.0;
if raw_terms.contains(&chunk.symbol_name.to_ascii_lowercase()) {
boost += 32.0;
}
boost += raw_terms.intersection(&path_terms).count() as f64 * 8.0;
let has_identifier = query.split_whitespace().any(|term| {
let has_lower = term.chars().any(char::is_lowercase);
let has_upper = term.chars().any(char::is_uppercase);
(has_lower && has_upper) || term.contains('_') || term.contains('.')
});
if has_identifier {
let extension = std::path::Path::new(&chunk.file_path)
.extension()
.and_then(std::ffi::OsStr::to_str)
.unwrap_or_default()
.to_ascii_lowercase();
if matches!(extension.as_str(), "md" | "mdx" | "rst" | "txt" | "adoc") {
boost -= 8.0;
} else if !extension.is_empty() {
boost += 4.0;
}
}
boost
}
impl Default for BM25Index {
fn default() -> Self {
Self::new()
}
}
const ADMITTED_BATCH: usize = 256;
const FILE_CHUNK_COUNT: usize = 4096;
type ChunkedFile = (Vec<CodeChunk>, usize, usize);
fn extract_sorted_chunks(
path: &str,
content: &str,
chunk_bytes: &mut usize,
chunk_count: &mut usize,
) -> Result<Vec<CodeChunk>, String> {
let mut chunks = extract_admitted_chunks(path, content, chunk_bytes, chunk_count)?;
chunks.sort_by(|a, b| {
(a.start_line, a.end_line, &a.symbol_name).cmp(&(b.start_line, b.end_line, &b.symbol_name))
});
Ok(chunks)
}
fn parallel_admitted_chunks(
admitted: &[(&str, String)],
) -> Vec<(usize, Option<Result<ChunkedFile, String>>)> {
use rayon::prelude::*;
admitted
.par_iter()
.map(|(path, content)| {
let full_allowance =
(64 * 1024 * 1024usize).min(content.len().saturating_mul(5).max(64 * 1024));
if crate::core::memory_guard::abort_requested() {
return (full_allowance, None);
}
let mut bytes = full_allowance;
let mut count = FILE_CHUNK_COUNT;
let result = extract_sorted_chunks(path, content, &mut bytes, &mut count)
.map(|chunks| (chunks, full_allowance - bytes, FILE_CHUNK_COUNT - count));
(full_allowance, Some(result))
})
.collect()
}
static PERSIST: std::sync::Mutex<()> = std::sync::Mutex::new(());
pub(crate) fn persist_lock() -> std::sync::MutexGuard<'static, ()> {
PERSIST
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
}
impl BM25Index {
pub fn new() -> Self {
Self {
chunks: Vec::new(),
inverted: HashMap::new(),
avg_doc_len: 0.0,
doc_count: 0,
doc_freqs: HashMap::new(),
files: HashMap::new(),
dirs: HashMap::new(),
admission_policy: None,
content_truncated: false,
source_root: None,
provider_origins: std::sync::Arc::default(),
}
}
pub fn memory_usage_bytes(&self) -> usize {
let chunks_size: usize = self
.chunks
.iter()
.map(|c| {
c.content.len()
+ c.file_path.len()
+ c.symbol_name.len()
+ c.tokens.iter().map(String::len).sum::<usize>()
+ 64
})
.sum();
let inverted_size: usize = self
.inverted
.iter()
.map(|(k, v)| k.len() + v.len() * 16 + 32)
.sum();
let files_size: usize = self.files.keys().map(|k| k.len() + 24).sum();
let dirs_size: usize = self.dirs.keys().map(|k| k.len() + 16).sum();
let freqs_size: usize = self.doc_freqs.keys().map(|k| k.len() + 16).sum();
chunks_size + inverted_size + files_size + dirs_size + freqs_size
}
pub fn unload(&mut self) {
let usage = self.memory_usage_bytes();
self.chunks = Vec::new();
self.inverted = HashMap::new();
self.doc_freqs = HashMap::new();
self.files = HashMap::new();
self.dirs = HashMap::new();
self.provider_origins = std::sync::Arc::default();
self.avg_doc_len = 0.0;
self.doc_count = 0;
tracing::info!(
"[bm25] unloaded index, freed ~{:.1}MB",
usage as f64 / 1_048_576.0
);
}
pub fn shrink_resident_content_to_snippet(&mut self, keep_lines: usize) {
let before = self.memory_usage_bytes();
for chunk in &mut self.chunks {
if chunk.content.lines().nth(keep_lines).is_some() {
let trimmed: String = chunk
.content
.lines()
.take(keep_lines)
.collect::<Vec<_>>()
.join("\n");
chunk.content = trimmed;
chunk.content.shrink_to_fit();
}
}
self.content_truncated = true;
let after = self.memory_usage_bytes();
tracing::debug!(
"[bm25] shrank resident content to {keep_lines} lines/chunk, freed ~{:.1}MB",
before.saturating_sub(after) as f64 / 1_048_576.0
);
}
#[cfg(any(test, feature = "test-support"))]
pub fn from_chunks_for_test(chunks: Vec<CodeChunk>) -> Self {
let mut index = Self::new();
for mut chunk in chunks {
if chunk.token_count == 0 {
chunk.token_count = tokenize(&chunk.content).len();
}
index.add_chunk(chunk);
}
index.finalize();
index
}
pub fn build_from_directory(root: &Path) -> Self {
Self::build_from_directory_inner(root, &HashMap::new())
}
pub fn build_from_admitted_sources(
root: &Path,
subdir: Option<&str>,
in_scope: &dyn Fn(&str) -> bool,
admit: impl FnMut(&str) -> Result<Option<String>, String>,
) -> Result<Self, String> {
let rules = CorpusRules::load();
Self::build_scoped_from_admitted_sources(root, &rules, subdir, in_scope, admit)
}
fn build_scoped_from_admitted_sources(
root: &Path,
rules: &CorpusRules,
subdir: Option<&str>,
in_scope: &dyn Fn(&str) -> bool,
mut admit: impl FnMut(&str) -> Result<Option<String>, String>,
) -> Result<Self, String> {
if !is_safe_bm25_root(root) {
return Err("source root is not safe to scan".into());
}
let mut index = Self::new();
let mut chunk_bytes = 64 * 1024 * 1024;
let mut chunk_count = 100_000;
if crate::core::memory_guard::is_under_pressure() {
crate::core::memory_guard::reclaim_now();
}
let paths = admitted_scoped_source_paths(root, rules, subdir, in_scope)?;
for batch in paths.chunks(ADMITTED_BATCH) {
let mut admitted = Vec::with_capacity(batch.len());
for path in batch {
if crate::core::memory_guard::abort_requested() {
return Err("source search withheld under memory pressure".into());
}
let Some(content) = admit(path)? else {
continue;
};
if looks_minified(&content) {
continue;
}
admitted.push((path.as_str(), content));
}
let chunked = parallel_admitted_chunks(&admitted);
for ((path, content), (full_allowance, unbounded)) in admitted.iter().zip(chunked) {
let allowance = chunk_bytes.min(full_allowance);
if !crate::core::index_admission::admit(
crate::core::index_admission::BuildKind::Bm25,
allowance as u64,
)
.parallel_ok
{
return Err("source search exceeds available memory headroom".into());
}
let count_allowance = chunk_count.min(FILE_CHUNK_COUNT);
let (chunks, used_bytes, used_count) = match unbounded {
Some(Ok(result))
if allowance == full_allowance && count_allowance == FILE_CHUNK_COUNT =>
{
result
}
Some(Err(error))
if allowance == full_allowance && count_allowance == FILE_CHUNK_COUNT =>
{
return Err(error);
}
_ => {
let mut file_chunk_bytes = allowance;
let mut file_chunk_count = count_allowance;
let chunks = extract_sorted_chunks(
path,
content,
&mut file_chunk_bytes,
&mut file_chunk_count,
)?;
(
chunks,
allowance - file_chunk_bytes,
count_allowance - file_chunk_count,
)
}
};
chunk_bytes -= used_bytes;
chunk_count -= used_count;
for chunk in chunks {
if crate::core::memory_guard::abort_requested() {
return Err("source search withheld under memory pressure".into());
}
index.add_chunk_with_content_limit(chunk, None);
}
}
}
index.finalize();
Ok(index)
}
pub fn build_with_content_hint(root: &Path, content_hint: &HashMap<String, String>) -> Self {
Self::build_from_directory_inner(root, content_hint)
}
fn build_from_directory_inner(root: &Path, content_hint: &HashMap<String, String>) -> Self {
let root_str = root.to_string_lossy();
if !super::graph_index::is_safe_scan_root_public(&root_str) {
tracing::warn!("[bm25: scan aborted for unsafe root {root_str}]");
return Self::new();
}
let files = list_code_files(root);
let dirs = dir_states(root, &files);
let admission = crate::core::context_admission::stores::StoreAdmission::current();
if files.len() >= build::PARALLEL_MIN_FILES
&& !crate::core::memory_guard::is_under_pressure()
&& !crate::core::memory_guard::abort_requested()
&& crate::core::index_admission::admit_files(
crate::core::index_admission::BuildKind::Bm25,
root,
&files,
)
.parallel_ok
{
let mut index = Self::build_parallel(root, content_hint, &files, &admission);
index.dirs = dirs;
index.admission_policy = Some(admission.digest().to_owned());
return index;
}
let mut index = Self::build_sequential(root, content_hint, &files, &admission);
index.dirs = dirs;
index.admission_policy = Some(admission.digest().to_owned());
index
}
fn group_prev_chunks_by_file(prev: &BM25Index) -> HashMap<String, Vec<CodeChunk>> {
let mut old_by_file: HashMap<String, Vec<CodeChunk>> = HashMap::new();
for c in &prev.chunks {
old_by_file
.entry(c.file_path.clone())
.or_default()
.push(c.clone());
}
for v in old_by_file.values_mut() {
v.sort_by(|a, b| {
a.start_line
.cmp(&b.start_line)
.then_with(|| a.end_line.cmp(&b.end_line))
.then_with(|| a.symbol_name.cmp(&b.symbol_name))
});
}
old_by_file
}
pub fn rebuild_incremental(root: &Path, prev: &BM25Index) -> Self {
let admission = crate::core::context_admission::stores::StoreAdmission::current();
if prev.admission_policy.as_deref() != Some(admission.digest()) {
return Self::build_from_directory(root);
}
let old_by_file = Self::group_prev_chunks_by_file(prev);
let files = list_code_files(root);
let dirs = dir_states(root, &files);
if files.len() >= build::PARALLEL_MIN_FILES
&& !crate::core::memory_guard::is_under_pressure()
&& !crate::core::memory_guard::abort_requested()
&& crate::core::index_admission::admit_files(
crate::core::index_admission::BuildKind::Bm25,
root,
&files,
)
.parallel_ok
{
let mut index =
Self::rebuild_incremental_parallel(root, prev, &old_by_file, &files, &admission);
index.dirs = dirs;
index.admission_policy = Some(admission.digest().to_owned());
return index;
}
let mut index =
Self::rebuild_incremental_sequential(root, prev, &old_by_file, &files, &admission);
index.dirs = dirs;
index.admission_policy = Some(admission.digest().to_owned());
index
}
pub(crate) fn rebuild_incremental_sequential(
root: &Path,
prev: &BM25Index,
old_by_file: &HashMap<String, Vec<CodeChunk>>,
files: &[String],
admission: &crate::core::context_admission::stores::StoreAdmission,
) -> Self {
let mut index = Self::new();
const MAX_FILE_SIZE_BYTES: u64 = 2 * 1024 * 1024;
let root_key = root.to_string_lossy().to_string();
let total = files.len() as u64;
crate::core::index_progress::report_bm25(&root_key, 0, total);
for (i, rel) in files.iter().enumerate() {
if i.is_multiple_of(16) || i + 1 == files.len() {
crate::core::index_progress::report_bm25(&root_key, (i + 1) as u64, total);
}
if i.is_multiple_of(500) && crate::core::memory_guard::is_under_pressure() {
tracing::warn!(
"[bm25: stopping incremental rebuild at file {i}/{} due to memory pressure]",
files.len()
);
break;
}
let abs = root.join(rel);
let Some(state) = IndexedFileState::from_path(&abs) else {
continue;
};
let unchanged = prev.files.get(rel).is_some_and(|old| *old == state);
if unchanged
&& let Some(chunks) = old_by_file.get(rel)
&& chunks.first().is_some_and(|c| !c.content.is_empty())
{
for chunk in chunks {
index.add_chunk(chunk.clone());
}
index.files.insert(rel.clone(), state);
continue;
}
if state.size_bytes > MAX_FILE_SIZE_BYTES {
continue;
}
let content = if crate::core::extractors::is_binary_document(&abs) {
match std::fs::read(&abs) {
Ok(bytes) => crate::core::extractors::extract(&abs, &bytes).text,
Err(_) => continue,
}
} else {
match crate::core::text_decode::read_text_checked(&abs) {
Ok(decoded) if decoded.complete() => decoded.text,
_ => continue,
}
};
if content.is_empty() {
continue;
}
let Some(content) = admission.admit(&content, &abs, None) else {
continue;
};
if looks_minified(&content) {
tracing::debug!("[bm25: skipping minified payload {rel}]");
continue;
}
let mut chunks = extract_chunks(rel, &content);
chunks.sort_by(|a, b| {
a.start_line
.cmp(&b.start_line)
.then_with(|| a.end_line.cmp(&b.end_line))
.then_with(|| a.symbol_name.cmp(&b.symbol_name))
});
for chunk in chunks {
index.add_chunk(chunk);
}
index.files.insert(rel.clone(), state);
}
index.finalize();
index
}
fn add_chunk(&mut self, chunk: CodeChunk) {
self.add_chunk_with_content_limit(chunk, Some(10));
}
fn add_chunk_with_content_limit(&mut self, mut chunk: CodeChunk, keep_lines: Option<usize>) {
let idx = self.chunks.len();
let enriched = enrich_for_bm25(&chunk);
let tokens = chunking::index_terms(&enriched);
for token in &tokens {
let lower = token.to_lowercase();
let postings = self.inverted.entry(lower.clone()).or_default();
if postings.last().map(|(last_idx, _)| *last_idx) != Some(idx) {
*self.doc_freqs.entry(lower).or_insert(0) += 1;
}
postings.push((idx, 1.0));
}
if let Some(keep_lines) = keep_lines.filter(|&n| chunk.content.lines().nth(n).is_some()) {
chunk.content = chunk
.content
.lines()
.take(keep_lines)
.collect::<Vec<_>>()
.join("\n");
chunk.content.shrink_to_fit();
}
self.chunks.push(CodeChunk {
token_count: tokens.len(),
tokens: Vec::new(),
..chunk
});
}
fn finalize(&mut self) {
self.doc_count = self.chunks.len();
if self.doc_count == 0 {
return;
}
let total_len: usize = self.chunks.iter().map(|c| c.token_count).sum();
self.avg_doc_len = total_len as f64 / self.doc_count as f64;
}
fn rank(&self, query: &str, excluded: Option<&[bool]>) -> Vec<SearchResult> {
let query_tokens = tokenize(query);
if query_tokens.is_empty() || self.doc_count == 0 {
return Vec::new();
}
let (doc_count, avg_doc_len) = match excluded {
None => (self.doc_count, self.avg_doc_len),
Some(excluded) => self.served_corpus(excluded),
};
if doc_count == 0 {
return Vec::new();
}
let n = self.chunks.len();
let mut scores = vec![0.0f64; n];
let mut touched = Vec::with_capacity(n.min(256));
for token in &query_tokens {
let lower = token.to_lowercase();
let Some(postings) = self.inverted.get(&lower) else {
continue;
};
let df = match excluded {
None => *self.doc_freqs.get(&lower).unwrap_or(&0),
Some(excluded) => served_ranking::doc_freq(postings, excluded),
};
let df = df as f64;
if df == 0.0 {
continue;
}
let idf = ((doc_count as f64 - df + 0.5) / (df + 0.5) + 1.0).ln();
for &(idx, weight) in postings {
if excluded.is_some_and(|excluded| served_ranking::is_excluded(excluded, idx)) {
continue;
}
let doc_len = self.chunks[idx].token_count as f64;
let norm_len = doc_len / avg_doc_len.max(1.0);
let bm25 = idf * (weight * (BM25_K1 + 1.0))
/ (weight + BM25_K1 * (1.0 - BM25_B + BM25_B * norm_len));
if scores[idx] == 0.0 {
touched.push(idx);
}
scores[idx] += bm25;
}
}
for &idx in &touched {
scores[idx] += structural_query_boost(query, &self.chunks[idx]);
}
let mut results: Vec<SearchResult> = touched
.iter()
.filter(|&&idx| scores[idx] > 0.0)
.map(|&idx| {
let chunk = &self.chunks[idx];
let snippet = query_snippet(&chunk.content, &query_tokens);
SearchResult {
chunk_idx: idx,
score: scores[idx],
file_path: chunk.file_path.clone(),
symbol_name: chunk.symbol_name.clone(),
kind: chunk.kind.clone(),
start_line: chunk.start_line,
end_line: chunk.end_line,
snippet,
}
})
.collect();
results.sort_by(|a, b| {
b.score
.partial_cmp(&a.score)
.unwrap_or(std::cmp::Ordering::Equal)
.then_with(|| a.file_path.cmp(&b.file_path))
.then_with(|| a.symbol_name.cmp(&b.symbol_name))
.then_with(|| a.start_line.cmp(&b.start_line))
.then_with(|| a.end_line.cmp(&b.end_line))
});
results
}
pub fn save(&self, root: &Path) -> std::io::Result<SaveOutcome> {
if self.chunks.len() > CHUNK_COUNT_WARNING {
tracing::warn!(
"[bm25] index has {} chunks (threshold {}), consider adding extra_ignore_patterns",
self.chunks.len(),
CHUNK_COUNT_WARNING
);
}
let dir = index_dir(root);
std::fs::create_dir_all(&dir)?;
let target = dir.join("bm25_index.bin.zst");
let tmp = dir.join("bm25_index.bin.zst.tmp");
{
let file = std::fs::File::create(&tmp)?;
let buf_writer = std::io::BufWriter::new(file);
let mut encoder = zstd::Encoder::new(buf_writer, ZSTD_LEVEL)
.map_err(|e| std::io::Error::other(format!("zstd encoder init: {e}")))?;
postcard::to_io(self, &mut encoder)
.map_err(|e| std::io::Error::other(e.to_string()))?;
encoder.finish()?;
}
let compressed_bytes = std::fs::metadata(&tmp)?.len();
let max_bytes = max_bm25_cache_bytes();
if compressed_bytes > max_bytes {
let _ = std::fs::remove_file(&tmp);
tracing::warn!(
"[bm25] compressed index too large ({:.1} MB, limit {:.0} MB), refusing to persist: {}",
compressed_bytes as f64 / 1_048_576.0,
max_bytes / (1024 * 1024),
dir.display()
);
return Ok(SaveOutcome::SkippedTooLarge {
compressed_bytes,
limit_bytes: max_bytes,
});
}
tracing::info!(
"[bm25] index: {:.1} MB zstd compressed",
compressed_bytes as f64 / 1_048_576.0,
);
std::fs::rename(&tmp, &target)?;
let _ = std::fs::remove_file(dir.join("bm25_index.bin"));
let _ = std::fs::remove_file(dir.join("bm25_index.json"));
let _ = std::fs::write(
dir.join("project_root.txt"),
root.to_string_lossy().as_bytes(),
);
self.save_provider_origins(&dir);
Ok(SaveOutcome::Persisted { compressed_bytes })
}
pub fn load(root: &Path) -> Option<Self> {
let admission = crate::core::context_admission::stores::StoreAdmission::current();
let canonical = |path: &Path| crate::core::pathutil::safe_canonicalize_bounded(path, 2000);
if std::fs::read_to_string(index_dir(root).join("project_root.txt"))
.is_ok_and(|stored| canonical(Path::new(stored.trim())) != canonical(root))
{
return None;
}
Self::load_persisted(root)
.filter(|idx| idx.admission_policy.as_deref() == Some(admission.digest()))
.map(|idx| idx.rooted_at(root).with_provider_origins(&index_dir(root)))
}
fn load_persisted(root: &Path) -> Option<Self> {
let dir = index_dir(root);
let max_bytes = max_bm25_cache_bytes();
let zst_path = dir.join("bm25_index.bin.zst");
if zst_path.exists() {
let meta = std::fs::metadata(&zst_path).ok()?;
if meta.len() > max_bytes {
tracing::warn!(
"[bm25] compressed index too large ({:.1} GB, limit {:.0} MB), quarantining: {}",
meta.len() as f64 / 1_073_741_824.0,
max_bytes / (1024 * 1024),
zst_path.display()
);
let quarantined = zst_path.with_extension("zst.quarantined");
let _ = std::fs::rename(&zst_path, &quarantined);
return None;
}
let compressed = std::fs::read(&zst_path).ok()?;
let max_decompressed = max_decompressed_bytes(meta.len());
let data = bounded_zstd_decode(&compressed, max_decompressed)?;
let idx: Self = postcard::from_bytes(&data).ok()?;
return Some(idx);
}
let bin_path = dir.join("bm25_index.bin");
if bin_path.exists() {
let meta = std::fs::metadata(&bin_path).ok()?;
if meta.len() > max_bytes {
tracing::warn!(
"[bm25] index too large ({:.1} GB, limit {:.0} MB), quarantining: {}",
meta.len() as f64 / 1_073_741_824.0,
max_bytes / (1024 * 1024),
bin_path.display()
);
let quarantined = bin_path.with_extension("bin.quarantined");
let _ = std::fs::rename(&bin_path, &quarantined);
return None;
}
let data = std::fs::read(&bin_path).ok()?;
let idx: Self = postcard::from_bytes(&data).ok()?;
if let Ok(compressed) = zstd::encode_all(data.as_slice(), ZSTD_LEVEL) {
let zst_tmp = zst_path.with_extension("zst.tmp");
if std::fs::write(&zst_tmp, &compressed).is_ok()
&& std::fs::rename(&zst_tmp, &zst_path).is_ok()
{
tracing::info!(
"[bm25] migrated {:.1} MB → {:.1} MB zstd",
data.len() as f64 / 1_048_576.0,
compressed.len() as f64 / 1_048_576.0
);
let _ = std::fs::remove_file(&bin_path);
}
}
return Some(idx);
}
let json_path = dir.join("bm25_index.json");
if json_path.exists() {
let meta = std::fs::metadata(&json_path).ok()?;
if meta.len() > max_bytes {
tracing::warn!(
"[bm25] index too large ({:.1} GB, limit {:.0} MB), quarantining: {}",
meta.len() as f64 / 1_073_741_824.0,
max_bytes / (1024 * 1024),
json_path.display()
);
let quarantined = json_path.with_extension("json.quarantined");
let _ = std::fs::rename(&json_path, &quarantined);
return None;
}
let data = std::fs::read_to_string(&json_path).ok()?;
return serde_json::from_str(&data).ok();
}
None
}
pub fn load_or_build(root: &Path) -> Self {
Self::load_or_build_inner(root, false)
}
pub fn load_or_build_fast(root: &Path) -> Self {
Self::load_or_build_inner(root, true)
}
fn load_or_build_inner(root: &Path, fast_stale: bool) -> Self {
if !is_safe_bm25_root(root) {
return Self::default();
}
if let Some(idx) = Self::load(root) {
let stale = if fast_stale {
bm25_index_looks_stale_fast(&idx, root)
} else {
bm25_index_looks_stale(&idx, root)
};
if !stale {
return idx;
}
tracing::debug!(
"[bm25_index: stale index detected for {}; rebuilding]",
root.display()
);
let rebuilt = if idx.files.is_empty() {
Self::build_from_directory(root)
} else {
Self::rebuild_incremental(root, &idx)
}
.carrying_external_from(&idx);
let _ = rebuilt.save(root);
return rebuilt.rooted_at(root);
}
let built = Self::build_from_directory(root);
let _ = built.save(root);
built.rooted_at(root)
}
pub fn index_file_path(root: &Path) -> PathBuf {
let dir = index_dir(root);
let zst = dir.join("bm25_index.bin.zst");
if zst.exists() {
return zst;
}
let bin = dir.join("bm25_index.bin");
if bin.exists() {
return bin;
}
dir.join("bm25_index.json")
}
pub fn persisted_chunk_count(root: &Path) -> Option<u64> {
use std::io::Read;
let path = Self::index_file_path(root);
let file = std::fs::File::open(&path).ok()?;
let mut head: Box<dyn Read> = match path.extension().and_then(|e| e.to_str()) {
Some("zst") => Box::new(zstd::Decoder::new(file).ok()?),
Some("bin") => Box::new(file),
_ => return None,
};
let mut count = 0u64;
for shift in (0..64).step_by(7) {
let mut byte = [0u8; 1];
head.read_exact(&mut byte).ok()?;
count |= u64::from(byte[0] & 0x7f) << shift;
if byte[0] & 0x80 == 0 {
return Some(count);
}
}
None
}
}
fn is_safe_bm25_root(root: &Path) -> bool {
super::graph_index::is_safe_scan_root_public(&root.to_string_lossy())
}
fn bm25_index_looks_stale(index: &BM25Index, root: &Path) -> bool {
bm25_index_looks_stale_inner(index, root, false)
}
pub fn bm25_index_looks_stale_fast(index: &BM25Index, root: &Path) -> bool {
bm25_index_looks_stale_inner(index, root, true)
}
fn bm25_index_looks_stale_inner(index: &BM25Index, root: &Path, fast: bool) -> bool {
if index.chunks.is_empty() {
return false;
}
if index
.files
.keys()
.any(|rel| crate::core::io_boundary::is_secret_like(Path::new(rel)).is_some())
{
return true;
}
if index.files.is_empty() {
let mut seen = std::collections::HashSet::<&str>::new();
for chunk in &index.chunks {
let rel = chunk.file_path.trim_start_matches(['/', '\\']);
if rel.is_empty() {
continue;
}
if !seen.insert(rel) {
continue;
}
if !root.join(rel).exists() {
return true;
}
}
return false;
}
if fast {
let sample_size = index.files.len().min(SENTINEL_SAMPLE_SIZE);
let step = if index.files.len() > sample_size {
index.files.len() / sample_size
} else {
1
};
for (i, (rel, old_state)) in index.files.iter().enumerate() {
if i % step != 0 {
continue;
}
let abs = root.join(rel);
if !abs.exists() {
return true;
}
let Some(cur) = IndexedFileState::from_path(&abs) else {
return true;
};
if &cur != old_state {
return true;
}
}
return dirs_look_stale(index, root);
}
for (rel, old_state) in &index.files {
let abs = root.join(rel);
if !abs.exists() {
return true;
}
let Some(cur) = IndexedFileState::from_path(&abs) else {
return true;
};
if &cur != old_state {
return true;
}
}
for rel in list_code_files(root) {
if !index.files.contains_key(&rel) {
return true;
}
}
false
}
const SENTINEL_SAMPLE_SIZE: usize = 10;
fn dirs_look_stale(index: &BM25Index, root: &Path) -> bool {
if index.dirs.is_empty() {
return true;
}
for (rel, recorded) in &index.dirs {
let abs = if rel.is_empty() {
root.to_path_buf()
} else {
root.join(rel)
};
if dir_fingerprint(&abs) != Some(*recorded) {
return true;
}
}
false
}
fn dir_fingerprint(path: &Path) -> Option<u64> {
use std::hash::{Hash, Hasher};
let meta = std::fs::metadata(path).ok()?;
let modified = meta.modified().ok()?.duration_since(UNIX_EPOCH).ok()?;
let mut entries = std::fs::read_dir(path)
.ok()?
.filter_map(Result::ok)
.map(|entry| entry.file_name().to_string_lossy().into_owned())
.collect::<Vec<_>>();
entries.sort_unstable();
let mut hasher = std::collections::hash_map::DefaultHasher::new();
modified.as_nanos().hash(&mut hasher);
entries.hash(&mut hasher);
Some(hasher.finish())
}
fn dir_states(root: &Path, files: &[String]) -> HashMap<String, u64> {
let mut rels: std::collections::HashSet<String> = std::collections::HashSet::new();
rels.insert(String::new());
for rel in files {
let mut cur = Path::new(rel);
while let Some(parent) = cur.parent() {
let key = parent.to_string_lossy().to_string();
if !rels.insert(key) {
break;
}
cur = parent;
}
}
rels.into_iter()
.filter_map(|rel| {
let abs = if rel.is_empty() {
root.to_path_buf()
} else {
root.join(&rel)
};
dir_fingerprint(&abs).map(|fingerprint| (rel, fingerprint))
})
.collect()
}
fn bounded_zstd_decode(compressed: &[u8], max_bytes: u64) -> Option<Vec<u8>> {
use std::io::Read;
let mut decoder = zstd::Decoder::new(compressed).ok()?;
let mut buf = Vec::new();
let mut chunk = vec![0u8; 65536];
let mut total = 0u64;
loop {
let n = decoder.read(&mut chunk).ok()?;
if n == 0 {
break;
}
total += n as u64;
if total > max_bytes {
tracing::warn!(
"[bm25] decompressed index exceeds limit ({:.0} MB > {:.0} MB), aborting load",
total as f64 / (1024.0 * 1024.0),
max_bytes as f64 / (1024.0 * 1024.0)
);
return None;
}
buf.extend_from_slice(&chunk[..n]);
}
Some(buf)
}
fn max_decompressed_bytes(compressed_bytes: u64) -> u64 {
const MAX_EXPANSION_RATIO: u64 = 20;
const FLOOR: u64 = 64 * 1024 * 1024;
let file_derived = compressed_bytes
.saturating_mul(MAX_EXPANSION_RATIO)
.max(FLOOR);
match crate::core::memory_guard::rss_limit_bytes() {
Some(limit) => file_derived.min(limit.max(FLOOR)),
None => file_derived,
}
}
fn index_dir(root: &Path) -> PathBuf {
crate::core::index_namespace::vectors_dir(root)
}
const CORPUS_MAX_DEPTH: usize = 20;
pub(crate) struct CorpusRules {
filter: crate::core::index_filter::IndexFileFilter,
ignore_patterns: Vec<(String, glob::Pattern)>,
max_files: usize,
}
impl CorpusRules {
pub(crate) fn load() -> Self {
Self::from_config(&crate::core::config::Config::load())
}
fn from_config(cfg: &crate::core::config::Config) -> Self {
let filter = crate::core::index_filter::IndexFileFilter::resolve(cfg);
let ignore_patterns = DEFAULT_BM25_IGNORES
.iter()
.map(|p| (*p).to_string())
.chain(cfg.extra_ignore_patterns.iter().cloned())
.filter_map(|p| glob::Pattern::new(&p).ok().map(|g| (p, g)))
.collect();
let max_files = if cfg.bm25_max_files == 0 {
usize::MAX
} else {
cfg.bm25_max_files as usize
};
Self {
filter,
ignore_patterns,
max_files,
}
}
fn walker(&self, root: &Path) -> ignore::WalkBuilder {
let mut builder = ignore::WalkBuilder::new(root);
builder
.hidden(crate::core::walk_filter::SKIP_HIDDEN_IN_CONTENT_WALK)
.git_ignore(self.filter.respect_gitignore)
.git_global(self.filter.respect_gitignore)
.git_exclude(self.filter.respect_gitignore)
.require_git(false)
.max_depth(Some(CORPUS_MAX_DEPTH));
builder
}
fn ignore_pattern_for(&self, rel: &str) -> Option<&str> {
self.ignore_patterns
.iter()
.find(|(_, p)| p.matches(rel))
.map(|(src, _)| src.as_str())
}
fn filter_excludes(&self, rel: &str) -> bool {
self.filter.is_excluded(&rel.replace('\\', "/"))
}
}