1use std::collections::HashMap;
2use std::path::{Path, PathBuf};
3use std::time::UNIX_EPOCH;
4
5use serde::{Deserialize, Serialize};
6mod build;
7mod carry_over;
8mod chunking;
9pub use chunking::*;
10mod coordinator;
11pub mod coverage;
12mod evict;
13pub(crate) mod external_origin;
14pub use coordinator::{SearchIndexBuildProgress, get_or_start_build};
15mod provider_origin;
16pub use provider_origin::ChunkOrigins;
17mod served_ranking;
18pub mod source_gate;
19#[cfg(test)]
20mod tests;
21mod walk;
22pub use walk::is_code_file;
23use walk::{admitted_scoped_source_paths, list_code_files};
24#[cfg(test)]
25use walk::{list_code_files_with, list_scoped_code_files};
26
27pub(crate) const MAX_BM25_FILES: usize = 5000;
28const CHUNK_COUNT_WARNING: usize = 50_000;
29const ZSTD_LEVEL: i32 = 9;
30
31const DEFAULT_BM25_IGNORES: &[&str] = &[
32 "vendor/**",
33 "dist/**",
34 "build/**",
35 "public/vendor/**",
36 "public/js/**",
37 "public/css/**",
38 "public/build/**",
39 ".next/**",
40 ".nuxt/**",
41 "__pycache__/**",
42 "*.min.js",
43 "*.min.css",
44 "*.bundle.js",
45 "*.chunk.js",
46];
47
48fn max_bm25_cache_bytes() -> u64 {
49 let mb = std::env::var("LEAN_CTX_BM25_MAX_CACHE_MB")
53 .ok()
54 .and_then(|v| v.parse::<u64>().ok())
55 .unwrap_or_else(|| crate::core::config::Config::load().bm25_max_cache_mb_effective());
56 mb * 1024 * 1024
57}
58
59pub fn persist_ceiling_bytes() -> u64 {
63 max_bm25_cache_bytes()
64}
65
66#[derive(Debug, Clone, Copy, PartialEq, Eq)]
70pub enum SaveOutcome {
71 Persisted { compressed_bytes: u64 },
73 SkippedTooLarge {
78 compressed_bytes: u64,
79 limit_bytes: u64,
80 },
81}
82
83#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
84pub struct CodeChunk {
85 pub file_path: String,
86 pub symbol_name: String,
87 pub kind: ChunkKind,
88 pub start_line: usize,
89 pub end_line: usize,
90 pub content: String,
91 #[serde(default)]
92 pub tokens: Vec<String>,
93 pub token_count: usize,
94 #[serde(default)]
97 pub input_origin: Option<crate::core::context_admission::origin::InputOrigin>,
98}
99
100#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
101pub enum ChunkKind {
102 Function,
103 Struct,
104 Impl,
105 Module,
106 Class,
107 Method,
108 Other,
109 Issue,
111 PullRequest,
112 WikiPage,
113 DbSchema,
114 ApiEndpoint,
115 Ticket,
116 ExternalOther,
117}
118
119#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
120pub struct IndexedFileState {
121 pub mtime_ms: u64,
122 pub size_bytes: u64,
123}
124
125impl IndexedFileState {
126 fn from_path(path: &Path) -> Option<Self> {
127 let meta = path.metadata().ok()?;
128 let size_bytes = meta.len();
129 let mtime_ms = meta
130 .modified()
131 .ok()
132 .and_then(|t| t.duration_since(UNIX_EPOCH).ok())
133 .map(|d| d.as_millis() as u64)?;
134 Some(Self {
135 mtime_ms,
136 size_bytes,
137 })
138 }
139}
140
141#[derive(Debug, Clone, Serialize, Deserialize)]
142pub struct BM25Index {
143 pub chunks: Vec<CodeChunk>,
144 pub inverted: HashMap<String, Vec<(usize, f64)>>,
145 pub avg_doc_len: f64,
146 pub doc_count: usize,
147 pub doc_freqs: HashMap<String, usize>,
148 #[serde(default)]
149 pub files: HashMap<String, IndexedFileState>,
150 #[serde(default)]
159 pub dirs: HashMap<String, u64>,
160 #[serde(default)]
165 pub admission_policy: Option<String>,
166 #[serde(default, skip)]
172 pub content_truncated: bool,
173 #[serde(default, skip)]
178 pub source_root: Option<PathBuf>,
179 #[serde(default, skip)]
183 pub provider_origins: std::sync::Arc<ChunkOrigins>,
184}
185
186#[derive(Debug, Clone, Serialize, Deserialize)]
187pub struct SearchResult {
188 pub chunk_idx: usize,
189 pub score: f64,
190 pub file_path: String,
191 pub symbol_name: String,
192 pub kind: ChunkKind,
193 pub start_line: usize,
194 pub end_line: usize,
195 pub snippet: String,
196}
197
198const BM25_K1: f64 = 1.2;
199const BM25_B: f64 = 0.75;
200
201fn structural_query_boost(query: &str, chunk: &CodeChunk) -> f64 {
202 let raw_terms: std::collections::HashSet<String> = query
203 .split(|c: char| !(c.is_alphanumeric() || c == '_'))
204 .filter(|term| term.len() >= 3)
205 .map(str::to_ascii_lowercase)
206 .collect();
207 let path_terms: std::collections::HashSet<String> = tokenize(&chunk.file_path)
208 .into_iter()
209 .map(|term| term.to_ascii_lowercase())
210 .collect();
211
212 let mut boost = 0.0;
213 if raw_terms.contains(&chunk.symbol_name.to_ascii_lowercase()) {
214 boost += 32.0;
215 }
216 boost += raw_terms.intersection(&path_terms).count() as f64 * 8.0;
217
218 let has_identifier = query.split_whitespace().any(|term| {
219 let has_lower = term.chars().any(char::is_lowercase);
220 let has_upper = term.chars().any(char::is_uppercase);
221 (has_lower && has_upper) || term.contains('_') || term.contains('.')
222 });
223 if has_identifier {
224 let extension = std::path::Path::new(&chunk.file_path)
225 .extension()
226 .and_then(std::ffi::OsStr::to_str)
227 .unwrap_or_default()
228 .to_ascii_lowercase();
229 if matches!(extension.as_str(), "md" | "mdx" | "rst" | "txt" | "adoc") {
230 boost -= 8.0;
231 } else if !extension.is_empty() {
232 boost += 4.0;
233 }
234 }
235 boost
236}
237
238impl Default for BM25Index {
239 fn default() -> Self {
240 Self::new()
241 }
242}
243
244const ADMITTED_BATCH: usize = 256;
246const FILE_CHUNK_COUNT: usize = 4096;
248
249type ChunkedFile = (Vec<CodeChunk>, usize, usize);
250
251fn extract_sorted_chunks(
254 path: &str,
255 content: &str,
256 chunk_bytes: &mut usize,
257 chunk_count: &mut usize,
258) -> Result<Vec<CodeChunk>, String> {
259 let mut chunks = extract_admitted_chunks(path, content, chunk_bytes, chunk_count)?;
260 chunks.sort_by(|a, b| {
261 (a.start_line, a.end_line, &a.symbol_name).cmp(&(b.start_line, b.end_line, &b.symbol_name))
262 });
263 Ok(chunks)
264}
265
266fn parallel_admitted_chunks(
271 admitted: &[(&str, String)],
272) -> Vec<(usize, Option<Result<ChunkedFile, String>>)> {
273 use rayon::prelude::*;
274 admitted
275 .par_iter()
276 .map(|(path, content)| {
277 let full_allowance =
278 (64 * 1024 * 1024usize).min(content.len().saturating_mul(5).max(64 * 1024));
279 if crate::core::memory_guard::abort_requested() {
280 return (full_allowance, None);
281 }
282 let mut bytes = full_allowance;
283 let mut count = FILE_CHUNK_COUNT;
284 let result = extract_sorted_chunks(path, content, &mut bytes, &mut count)
285 .map(|chunks| (chunks, full_allowance - bytes, FILE_CHUNK_COUNT - count));
286 (full_allowance, Some(result))
287 })
288 .collect()
289}
290
291static PERSIST: std::sync::Mutex<()> = std::sync::Mutex::new(());
296
297pub(crate) fn persist_lock() -> std::sync::MutexGuard<'static, ()> {
298 PERSIST
299 .lock()
300 .unwrap_or_else(std::sync::PoisonError::into_inner)
301}
302
303impl BM25Index {
304 pub fn new() -> Self {
305 Self {
306 chunks: Vec::new(),
307 inverted: HashMap::new(),
308 avg_doc_len: 0.0,
309 doc_count: 0,
310 doc_freqs: HashMap::new(),
311 files: HashMap::new(),
312 dirs: HashMap::new(),
313 admission_policy: None,
314 content_truncated: false,
315 source_root: None,
316 provider_origins: std::sync::Arc::default(),
317 }
318 }
319
320 pub fn memory_usage_bytes(&self) -> usize {
322 let chunks_size: usize = self
323 .chunks
324 .iter()
325 .map(|c| {
326 c.content.len()
327 + c.file_path.len()
328 + c.symbol_name.len()
329 + c.tokens.iter().map(String::len).sum::<usize>()
330 + 64
331 })
332 .sum();
333 let inverted_size: usize = self
334 .inverted
335 .iter()
336 .map(|(k, v)| k.len() + v.len() * 16 + 32)
337 .sum();
338 let files_size: usize = self.files.keys().map(|k| k.len() + 24).sum();
339 let dirs_size: usize = self.dirs.keys().map(|k| k.len() + 16).sum();
340 let freqs_size: usize = self.doc_freqs.keys().map(|k| k.len() + 16).sum();
341 chunks_size + inverted_size + files_size + dirs_size + freqs_size
342 }
343
344 pub fn unload(&mut self) {
346 let usage = self.memory_usage_bytes();
347 self.chunks = Vec::new();
348 self.inverted = HashMap::new();
349 self.doc_freqs = HashMap::new();
350 self.files = HashMap::new();
351 self.dirs = HashMap::new();
352 self.provider_origins = std::sync::Arc::default();
353 self.avg_doc_len = 0.0;
354 self.doc_count = 0;
355 tracing::info!(
356 "[bm25] unloaded index, freed ~{:.1}MB",
357 usage as f64 / 1_048_576.0
358 );
359 }
360
361 pub fn shrink_resident_content_to_snippet(&mut self, keep_lines: usize) {
376 let before = self.memory_usage_bytes();
377 for chunk in &mut self.chunks {
378 if chunk.content.lines().nth(keep_lines).is_some() {
381 let trimmed: String = chunk
382 .content
383 .lines()
384 .take(keep_lines)
385 .collect::<Vec<_>>()
386 .join("\n");
387 chunk.content = trimmed;
388 chunk.content.shrink_to_fit();
390 }
391 }
392 self.content_truncated = true;
393 let after = self.memory_usage_bytes();
394 tracing::debug!(
395 "[bm25] shrank resident content to {keep_lines} lines/chunk, freed ~{:.1}MB",
396 before.saturating_sub(after) as f64 / 1_048_576.0
397 );
398 }
399
400 #[cfg(any(test, feature = "test-support"))]
402 pub fn from_chunks_for_test(chunks: Vec<CodeChunk>) -> Self {
403 let mut index = Self::new();
404 for mut chunk in chunks {
405 if chunk.token_count == 0 {
406 chunk.token_count = tokenize(&chunk.content).len();
407 }
408 index.add_chunk(chunk);
409 }
410 index.finalize();
411 index
412 }
413
414 pub fn build_from_directory(root: &Path) -> Self {
415 Self::build_from_directory_inner(root, &HashMap::new())
416 }
417
418 pub fn build_from_admitted_sources(
425 root: &Path,
426 subdir: Option<&str>,
427 in_scope: &dyn Fn(&str) -> bool,
428 admit: impl FnMut(&str) -> Result<Option<String>, String>,
429 ) -> Result<Self, String> {
430 let rules = CorpusRules::load();
431 Self::build_scoped_from_admitted_sources(root, &rules, subdir, in_scope, admit)
432 }
433
434 fn build_scoped_from_admitted_sources(
435 root: &Path,
436 rules: &CorpusRules,
437 subdir: Option<&str>,
438 in_scope: &dyn Fn(&str) -> bool,
439 mut admit: impl FnMut(&str) -> Result<Option<String>, String>,
440 ) -> Result<Self, String> {
441 if !is_safe_bm25_root(root) {
442 return Err("source root is not safe to scan".into());
443 }
444 let mut index = Self::new();
445 let mut chunk_bytes = 64 * 1024 * 1024;
446 let mut chunk_count = 100_000;
447 if crate::core::memory_guard::is_under_pressure() {
449 crate::core::memory_guard::reclaim_now();
450 }
451 let paths = admitted_scoped_source_paths(root, rules, subdir, in_scope)?;
452 for batch in paths.chunks(ADMITTED_BATCH) {
458 let mut admitted = Vec::with_capacity(batch.len());
459 for path in batch {
460 if crate::core::memory_guard::abort_requested() {
464 return Err("source search withheld under memory pressure".into());
465 }
466 let Some(content) = admit(path)? else {
467 continue;
468 };
469 if looks_minified(&content) {
470 continue;
471 }
472 admitted.push((path.as_str(), content));
473 }
474 let chunked = parallel_admitted_chunks(&admitted);
475 for ((path, content), (full_allowance, unbounded)) in admitted.iter().zip(chunked) {
476 let allowance = chunk_bytes.min(full_allowance);
477 if !crate::core::index_admission::admit(
478 crate::core::index_admission::BuildKind::Bm25,
479 allowance as u64,
480 )
481 .parallel_ok
482 {
483 return Err("source search exceeds available memory headroom".into());
484 }
485 let count_allowance = chunk_count.min(FILE_CHUNK_COUNT);
486 let (chunks, used_bytes, used_count) = match unbounded {
490 Some(Ok(result))
491 if allowance == full_allowance && count_allowance == FILE_CHUNK_COUNT =>
492 {
493 result
494 }
495 Some(Err(error))
496 if allowance == full_allowance && count_allowance == FILE_CHUNK_COUNT =>
497 {
498 return Err(error);
499 }
500 _ => {
501 let mut file_chunk_bytes = allowance;
502 let mut file_chunk_count = count_allowance;
503 let chunks = extract_sorted_chunks(
504 path,
505 content,
506 &mut file_chunk_bytes,
507 &mut file_chunk_count,
508 )?;
509 (
510 chunks,
511 allowance - file_chunk_bytes,
512 count_allowance - file_chunk_count,
513 )
514 }
515 };
516 chunk_bytes -= used_bytes;
517 chunk_count -= used_count;
518 for chunk in chunks {
519 if crate::core::memory_guard::abort_requested() {
520 return Err("source search withheld under memory pressure".into());
521 }
522 index.add_chunk_with_content_limit(chunk, None);
525 }
526 }
527 }
528 index.finalize();
529 Ok(index)
530 }
531
532 pub fn build_with_content_hint(root: &Path, content_hint: &HashMap<String, String>) -> Self {
535 Self::build_from_directory_inner(root, content_hint)
536 }
537
538 fn build_from_directory_inner(root: &Path, content_hint: &HashMap<String, String>) -> Self {
539 let root_str = root.to_string_lossy();
540 if !super::graph_index::is_safe_scan_root_public(&root_str) {
541 tracing::warn!("[bm25: scan aborted for unsafe root {root_str}]");
542 return Self::new();
543 }
544 let files = list_code_files(root);
545 let dirs = dir_states(root, &files);
550 let admission = crate::core::context_admission::stores::StoreAdmission::current();
552
553 if files.len() >= build::PARALLEL_MIN_FILES
563 && !crate::core::memory_guard::is_under_pressure()
564 && !crate::core::memory_guard::abort_requested()
565 && crate::core::index_admission::admit_files(
566 crate::core::index_admission::BuildKind::Bm25,
567 root,
568 &files,
569 )
570 .parallel_ok
571 {
572 let mut index = Self::build_parallel(root, content_hint, &files, &admission);
573 index.dirs = dirs;
574 index.admission_policy = Some(admission.digest().to_owned());
575 return index;
576 }
577 let mut index = Self::build_sequential(root, content_hint, &files, &admission);
578 index.dirs = dirs;
579 index.admission_policy = Some(admission.digest().to_owned());
580 index
581 }
582
583 fn group_prev_chunks_by_file(prev: &BM25Index) -> HashMap<String, Vec<CodeChunk>> {
588 let mut old_by_file: HashMap<String, Vec<CodeChunk>> = HashMap::new();
589 for c in &prev.chunks {
590 old_by_file
591 .entry(c.file_path.clone())
592 .or_default()
593 .push(c.clone());
594 }
595 for v in old_by_file.values_mut() {
596 v.sort_by(|a, b| {
597 a.start_line
598 .cmp(&b.start_line)
599 .then_with(|| a.end_line.cmp(&b.end_line))
600 .then_with(|| a.symbol_name.cmp(&b.symbol_name))
601 });
602 }
603 old_by_file
604 }
605
606 pub fn rebuild_incremental(root: &Path, prev: &BM25Index) -> Self {
607 let admission = crate::core::context_admission::stores::StoreAdmission::current();
609 if prev.admission_policy.as_deref() != Some(admission.digest()) {
610 return Self::build_from_directory(root);
611 }
612 let old_by_file = Self::group_prev_chunks_by_file(prev);
613 let files = list_code_files(root);
614 let dirs = dir_states(root, &files);
615
616 if files.len() >= build::PARALLEL_MIN_FILES
626 && !crate::core::memory_guard::is_under_pressure()
627 && !crate::core::memory_guard::abort_requested()
628 && crate::core::index_admission::admit_files(
629 crate::core::index_admission::BuildKind::Bm25,
630 root,
631 &files,
632 )
633 .parallel_ok
634 {
635 let mut index =
636 Self::rebuild_incremental_parallel(root, prev, &old_by_file, &files, &admission);
637 index.dirs = dirs;
638 index.admission_policy = Some(admission.digest().to_owned());
639 return index;
640 }
641 let mut index =
642 Self::rebuild_incremental_sequential(root, prev, &old_by_file, &files, &admission);
643 index.dirs = dirs;
644 index.admission_policy = Some(admission.digest().to_owned());
645 index
646 }
647
648 pub(crate) fn rebuild_incremental_sequential(
652 root: &Path,
653 prev: &BM25Index,
654 old_by_file: &HashMap<String, Vec<CodeChunk>>,
655 files: &[String],
656 admission: &crate::core::context_admission::stores::StoreAdmission,
657 ) -> Self {
658 let mut index = Self::new();
659 const MAX_FILE_SIZE_BYTES: u64 = 2 * 1024 * 1024;
660 let root_key = root.to_string_lossy().to_string();
661 let total = files.len() as u64;
662 crate::core::index_progress::report_bm25(&root_key, 0, total);
663
664 for (i, rel) in files.iter().enumerate() {
665 if i.is_multiple_of(16) || i + 1 == files.len() {
666 crate::core::index_progress::report_bm25(&root_key, (i + 1) as u64, total);
667 }
668 if i.is_multiple_of(500) && crate::core::memory_guard::is_under_pressure() {
669 tracing::warn!(
670 "[bm25: stopping incremental rebuild at file {i}/{} due to memory pressure]",
671 files.len()
672 );
673 break;
674 }
675
676 let abs = root.join(rel);
677 let Some(state) = IndexedFileState::from_path(&abs) else {
678 continue;
679 };
680
681 let unchanged = prev.files.get(rel).is_some_and(|old| *old == state);
682 if unchanged
683 && let Some(chunks) = old_by_file.get(rel)
684 && chunks.first().is_some_and(|c| !c.content.is_empty())
685 {
686 for chunk in chunks {
687 index.add_chunk(chunk.clone());
688 }
689 index.files.insert(rel.clone(), state);
690 continue;
691 }
692
693 if state.size_bytes > MAX_FILE_SIZE_BYTES {
694 continue;
695 }
696 let content = if crate::core::extractors::is_binary_document(&abs) {
697 match std::fs::read(&abs) {
698 Ok(bytes) => crate::core::extractors::extract(&abs, &bytes).text,
699 Err(_) => continue,
700 }
701 } else {
702 match crate::core::text_decode::read_text_checked(&abs) {
705 Ok(decoded) if decoded.complete() => decoded.text,
706 _ => continue,
707 }
708 };
709 if content.is_empty() {
710 continue;
711 }
712 let Some(content) = admission.admit(&content, &abs, None) else {
714 continue;
715 };
716 if looks_minified(&content) {
718 tracing::debug!("[bm25: skipping minified payload {rel}]");
719 continue;
720 }
721 let mut chunks = extract_chunks(rel, &content);
722 chunks.sort_by(|a, b| {
723 a.start_line
724 .cmp(&b.start_line)
725 .then_with(|| a.end_line.cmp(&b.end_line))
726 .then_with(|| a.symbol_name.cmp(&b.symbol_name))
727 });
728 for chunk in chunks {
729 index.add_chunk(chunk);
730 }
731 index.files.insert(rel.clone(), state);
732 }
733
734 index.finalize();
735 index
736 }
737
738 fn add_chunk(&mut self, chunk: CodeChunk) {
739 self.add_chunk_with_content_limit(chunk, Some(10));
740 }
741
742 fn add_chunk_with_content_limit(&mut self, mut chunk: CodeChunk, keep_lines: Option<usize>) {
743 let idx = self.chunks.len();
744
745 let enriched = enrich_for_bm25(&chunk);
746 let tokens = chunking::index_terms(&enriched);
747 for token in &tokens {
748 let lower = token.to_lowercase();
749 let postings = self.inverted.entry(lower.clone()).or_default();
750 if postings.last().map(|(last_idx, _)| *last_idx) != Some(idx) {
751 *self.doc_freqs.entry(lower).or_insert(0) += 1;
752 }
753 postings.push((idx, 1.0));
754 }
755
756 if let Some(keep_lines) = keep_lines.filter(|&n| chunk.content.lines().nth(n).is_some()) {
759 chunk.content = chunk
760 .content
761 .lines()
762 .take(keep_lines)
763 .collect::<Vec<_>>()
764 .join("\n");
765 chunk.content.shrink_to_fit();
766 }
767
768 self.chunks.push(CodeChunk {
769 token_count: tokens.len(),
770 tokens: Vec::new(),
771 ..chunk
772 });
773 }
774
775 fn finalize(&mut self) {
776 self.doc_count = self.chunks.len();
777 if self.doc_count == 0 {
778 return;
779 }
780
781 let total_len: usize = self.chunks.iter().map(|c| c.token_count).sum();
782 self.avg_doc_len = total_len as f64 / self.doc_count as f64;
783 }
784
785 fn rank(&self, query: &str, excluded: Option<&[bool]>) -> Vec<SearchResult> {
791 let query_tokens = tokenize(query);
792 if query_tokens.is_empty() || self.doc_count == 0 {
793 return Vec::new();
794 }
795 let (doc_count, avg_doc_len) = match excluded {
796 None => (self.doc_count, self.avg_doc_len),
797 Some(excluded) => self.served_corpus(excluded),
798 };
799 if doc_count == 0 {
800 return Vec::new();
801 }
802
803 let n = self.chunks.len();
806 let mut scores = vec![0.0f64; n];
807 let mut touched = Vec::with_capacity(n.min(256));
808
809 for token in &query_tokens {
810 let lower = token.to_lowercase();
811 let Some(postings) = self.inverted.get(&lower) else {
812 continue;
813 };
814 let df = match excluded {
815 None => *self.doc_freqs.get(&lower).unwrap_or(&0),
816 Some(excluded) => served_ranking::doc_freq(postings, excluded),
817 };
818 let df = df as f64;
819 if df == 0.0 {
820 continue;
821 }
822
823 let idf = ((doc_count as f64 - df + 0.5) / (df + 0.5) + 1.0).ln();
824
825 for &(idx, weight) in postings {
826 if excluded.is_some_and(|excluded| served_ranking::is_excluded(excluded, idx)) {
827 continue;
828 }
829 let doc_len = self.chunks[idx].token_count as f64;
830 let norm_len = doc_len / avg_doc_len.max(1.0);
831 let bm25 = idf * (weight * (BM25_K1 + 1.0))
832 / (weight + BM25_K1 * (1.0 - BM25_B + BM25_B * norm_len));
833
834 if scores[idx] == 0.0 {
835 touched.push(idx);
836 }
837 scores[idx] += bm25;
838 }
839 }
840
841 for &idx in &touched {
842 scores[idx] += structural_query_boost(query, &self.chunks[idx]);
843 }
844
845 let mut results: Vec<SearchResult> = touched
846 .iter()
847 .filter(|&&idx| scores[idx] > 0.0)
848 .map(|&idx| {
849 let chunk = &self.chunks[idx];
850 let snippet = query_snippet(&chunk.content, &query_tokens);
851 SearchResult {
852 chunk_idx: idx,
853 score: scores[idx],
854 file_path: chunk.file_path.clone(),
855 symbol_name: chunk.symbol_name.clone(),
856 kind: chunk.kind.clone(),
857 start_line: chunk.start_line,
858 end_line: chunk.end_line,
859 snippet,
860 }
861 })
862 .collect();
863
864 results.sort_by(|a, b| {
865 b.score
866 .partial_cmp(&a.score)
867 .unwrap_or(std::cmp::Ordering::Equal)
868 .then_with(|| a.file_path.cmp(&b.file_path))
869 .then_with(|| a.symbol_name.cmp(&b.symbol_name))
870 .then_with(|| a.start_line.cmp(&b.start_line))
871 .then_with(|| a.end_line.cmp(&b.end_line))
872 });
873 results
874 }
875
876 pub fn save(&self, root: &Path) -> std::io::Result<SaveOutcome> {
877 if self.chunks.len() > CHUNK_COUNT_WARNING {
878 tracing::warn!(
879 "[bm25] index has {} chunks (threshold {}), consider adding extra_ignore_patterns",
880 self.chunks.len(),
881 CHUNK_COUNT_WARNING
882 );
883 }
884
885 let dir = index_dir(root);
886 std::fs::create_dir_all(&dir)?;
887
888 let target = dir.join("bm25_index.bin.zst");
891 let tmp = dir.join("bm25_index.bin.zst.tmp");
892 {
893 let file = std::fs::File::create(&tmp)?;
894 let buf_writer = std::io::BufWriter::new(file);
895 let mut encoder = zstd::Encoder::new(buf_writer, ZSTD_LEVEL)
896 .map_err(|e| std::io::Error::other(format!("zstd encoder init: {e}")))?;
897 postcard::to_io(self, &mut encoder)
898 .map_err(|e| std::io::Error::other(e.to_string()))?;
899 encoder.finish()?;
900 }
901
902 let compressed_bytes = std::fs::metadata(&tmp)?.len();
903 let max_bytes = max_bm25_cache_bytes();
904 if compressed_bytes > max_bytes {
905 let _ = std::fs::remove_file(&tmp);
906 tracing::warn!(
907 "[bm25] compressed index too large ({:.1} MB, limit {:.0} MB), refusing to persist: {}",
908 compressed_bytes as f64 / 1_048_576.0,
909 max_bytes / (1024 * 1024),
910 dir.display()
911 );
912 return Ok(SaveOutcome::SkippedTooLarge {
913 compressed_bytes,
914 limit_bytes: max_bytes,
915 });
916 }
917
918 tracing::info!(
919 "[bm25] index: {:.1} MB zstd compressed",
920 compressed_bytes as f64 / 1_048_576.0,
921 );
922
923 std::fs::rename(&tmp, &target)?;
924
925 let _ = std::fs::remove_file(dir.join("bm25_index.bin"));
926 let _ = std::fs::remove_file(dir.join("bm25_index.json"));
927
928 let _ = std::fs::write(
929 dir.join("project_root.txt"),
930 root.to_string_lossy().as_bytes(),
931 );
932 self.save_provider_origins(&dir);
933
934 Ok(SaveOutcome::Persisted { compressed_bytes })
935 }
936
937 pub fn load(root: &Path) -> Option<Self> {
941 let admission = crate::core::context_admission::stores::StoreAdmission::current();
942 let canonical = |path: &Path| crate::core::pathutil::safe_canonicalize_bounded(path, 2000);
947 if std::fs::read_to_string(index_dir(root).join("project_root.txt"))
948 .is_ok_and(|stored| canonical(Path::new(stored.trim())) != canonical(root))
949 {
950 return None;
951 }
952 Self::load_persisted(root)
953 .filter(|idx| idx.admission_policy.as_deref() == Some(admission.digest()))
954 .map(|idx| idx.rooted_at(root).with_provider_origins(&index_dir(root)))
955 }
956
957 fn load_persisted(root: &Path) -> Option<Self> {
958 let dir = index_dir(root);
959 let max_bytes = max_bm25_cache_bytes();
960
961 let zst_path = dir.join("bm25_index.bin.zst");
962 if zst_path.exists() {
963 let meta = std::fs::metadata(&zst_path).ok()?;
964 if meta.len() > max_bytes {
965 tracing::warn!(
966 "[bm25] compressed index too large ({:.1} GB, limit {:.0} MB), quarantining: {}",
967 meta.len() as f64 / 1_073_741_824.0,
968 max_bytes / (1024 * 1024),
969 zst_path.display()
970 );
971 let quarantined = zst_path.with_extension("zst.quarantined");
972 let _ = std::fs::rename(&zst_path, &quarantined);
973 return None;
974 }
975 let compressed = std::fs::read(&zst_path).ok()?;
976 let max_decompressed = max_decompressed_bytes(meta.len());
977 let data = bounded_zstd_decode(&compressed, max_decompressed)?;
978 let idx: Self = postcard::from_bytes(&data).ok()?;
979 return Some(idx);
980 }
981
982 let bin_path = dir.join("bm25_index.bin");
983 if bin_path.exists() {
984 let meta = std::fs::metadata(&bin_path).ok()?;
985 if meta.len() > max_bytes {
986 tracing::warn!(
987 "[bm25] index too large ({:.1} GB, limit {:.0} MB), quarantining: {}",
988 meta.len() as f64 / 1_073_741_824.0,
989 max_bytes / (1024 * 1024),
990 bin_path.display()
991 );
992 let quarantined = bin_path.with_extension("bin.quarantined");
993 let _ = std::fs::rename(&bin_path, &quarantined);
994 return None;
995 }
996 let data = std::fs::read(&bin_path).ok()?;
997 let idx: Self = postcard::from_bytes(&data).ok()?;
998 if let Ok(compressed) = zstd::encode_all(data.as_slice(), ZSTD_LEVEL) {
1000 let zst_tmp = zst_path.with_extension("zst.tmp");
1001 if std::fs::write(&zst_tmp, &compressed).is_ok()
1002 && std::fs::rename(&zst_tmp, &zst_path).is_ok()
1003 {
1004 tracing::info!(
1005 "[bm25] migrated {:.1} MB → {:.1} MB zstd",
1006 data.len() as f64 / 1_048_576.0,
1007 compressed.len() as f64 / 1_048_576.0
1008 );
1009 let _ = std::fs::remove_file(&bin_path);
1010 }
1011 }
1012 return Some(idx);
1013 }
1014
1015 let json_path = dir.join("bm25_index.json");
1016 if json_path.exists() {
1017 let meta = std::fs::metadata(&json_path).ok()?;
1018 if meta.len() > max_bytes {
1019 tracing::warn!(
1020 "[bm25] index too large ({:.1} GB, limit {:.0} MB), quarantining: {}",
1021 meta.len() as f64 / 1_073_741_824.0,
1022 max_bytes / (1024 * 1024),
1023 json_path.display()
1024 );
1025 let quarantined = json_path.with_extension("json.quarantined");
1026 let _ = std::fs::rename(&json_path, &quarantined);
1027 return None;
1028 }
1029 let data = std::fs::read_to_string(&json_path).ok()?;
1030 return serde_json::from_str(&data).ok();
1031 }
1032
1033 None
1034 }
1035
1036 pub fn load_or_build(root: &Path) -> Self {
1037 Self::load_or_build_inner(root, false)
1038 }
1039
1040 pub fn load_or_build_fast(root: &Path) -> Self {
1043 Self::load_or_build_inner(root, true)
1044 }
1045
1046 fn load_or_build_inner(root: &Path, fast_stale: bool) -> Self {
1047 if !is_safe_bm25_root(root) {
1048 return Self::default();
1049 }
1050 if let Some(idx) = Self::load(root) {
1051 let stale = if fast_stale {
1052 bm25_index_looks_stale_fast(&idx, root)
1053 } else {
1054 bm25_index_looks_stale(&idx, root)
1055 };
1056 if !stale {
1057 return idx;
1058 }
1059 tracing::debug!(
1060 "[bm25_index: stale index detected for {}; rebuilding]",
1061 root.display()
1062 );
1063 let rebuilt = if idx.files.is_empty() {
1064 Self::build_from_directory(root)
1065 } else {
1066 Self::rebuild_incremental(root, &idx)
1067 }
1068 .carrying_external_from(&idx);
1069 let _ = rebuilt.save(root);
1070 return rebuilt.rooted_at(root);
1071 }
1072
1073 let built = Self::build_from_directory(root);
1074 let _ = built.save(root);
1075 built.rooted_at(root)
1076 }
1077
1078 pub fn index_file_path(root: &Path) -> PathBuf {
1079 let dir = index_dir(root);
1080 let zst = dir.join("bm25_index.bin.zst");
1081 if zst.exists() {
1082 return zst;
1083 }
1084 let bin = dir.join("bm25_index.bin");
1085 if bin.exists() {
1086 return bin;
1087 }
1088 dir.join("bm25_index.json")
1089 }
1090
1091 pub fn persisted_chunk_count(root: &Path) -> Option<u64> {
1096 use std::io::Read;
1097 let path = Self::index_file_path(root);
1098 let file = std::fs::File::open(&path).ok()?;
1099 let mut head: Box<dyn Read> = match path.extension().and_then(|e| e.to_str()) {
1100 Some("zst") => Box::new(zstd::Decoder::new(file).ok()?),
1101 Some("bin") => Box::new(file),
1102 _ => return None,
1103 };
1104 let mut count = 0u64;
1105 for shift in (0..64).step_by(7) {
1106 let mut byte = [0u8; 1];
1107 head.read_exact(&mut byte).ok()?;
1108 count |= u64::from(byte[0] & 0x7f) << shift;
1109 if byte[0] & 0x80 == 0 {
1110 return Some(count);
1111 }
1112 }
1113 None
1114 }
1115}
1116
1117fn is_safe_bm25_root(root: &Path) -> bool {
1118 super::graph_index::is_safe_scan_root_public(&root.to_string_lossy())
1119}
1120
1121fn bm25_index_looks_stale(index: &BM25Index, root: &Path) -> bool {
1122 bm25_index_looks_stale_inner(index, root, false)
1123}
1124
1125pub fn bm25_index_looks_stale_fast(index: &BM25Index, root: &Path) -> bool {
1128 bm25_index_looks_stale_inner(index, root, true)
1129}
1130
1131fn bm25_index_looks_stale_inner(index: &BM25Index, root: &Path, fast: bool) -> bool {
1132 if index.chunks.is_empty() {
1133 return false;
1134 }
1135
1136 if index
1139 .files
1140 .keys()
1141 .any(|rel| crate::core::io_boundary::is_secret_like(Path::new(rel)).is_some())
1142 {
1143 return true;
1144 }
1145
1146 if index.files.is_empty() {
1147 let mut seen = std::collections::HashSet::<&str>::new();
1148 for chunk in &index.chunks {
1149 let rel = chunk.file_path.trim_start_matches(['/', '\\']);
1150 if rel.is_empty() {
1151 continue;
1152 }
1153 if !seen.insert(rel) {
1154 continue;
1155 }
1156 if !root.join(rel).exists() {
1157 return true;
1158 }
1159 }
1160 return false;
1161 }
1162
1163 if fast {
1164 let sample_size = index.files.len().min(SENTINEL_SAMPLE_SIZE);
1165 let step = if index.files.len() > sample_size {
1166 index.files.len() / sample_size
1167 } else {
1168 1
1169 };
1170 for (i, (rel, old_state)) in index.files.iter().enumerate() {
1171 if i % step != 0 {
1172 continue;
1173 }
1174 let abs = root.join(rel);
1175 if !abs.exists() {
1176 return true;
1177 }
1178 let Some(cur) = IndexedFileState::from_path(&abs) else {
1179 return true;
1180 };
1181 if &cur != old_state {
1182 return true;
1183 }
1184 }
1185 return dirs_look_stale(index, root);
1190 }
1191
1192 for (rel, old_state) in &index.files {
1193 let abs = root.join(rel);
1194 if !abs.exists() {
1195 return true;
1196 }
1197 let Some(cur) = IndexedFileState::from_path(&abs) else {
1198 return true;
1199 };
1200 if &cur != old_state {
1201 return true;
1202 }
1203 }
1204
1205 for rel in list_code_files(root) {
1206 if !index.files.contains_key(&rel) {
1207 return true;
1208 }
1209 }
1210
1211 false
1212}
1213
1214const SENTINEL_SAMPLE_SIZE: usize = 10;
1215
1216fn dirs_look_stale(index: &BM25Index, root: &Path) -> bool {
1223 if index.dirs.is_empty() {
1224 return true;
1225 }
1226 for (rel, recorded) in &index.dirs {
1227 let abs = if rel.is_empty() {
1228 root.to_path_buf()
1229 } else {
1230 root.join(rel)
1231 };
1232 if dir_fingerprint(&abs) != Some(*recorded) {
1233 return true;
1234 }
1235 }
1236 false
1237}
1238
1239fn dir_fingerprint(path: &Path) -> Option<u64> {
1240 use std::hash::{Hash, Hasher};
1241
1242 let meta = std::fs::metadata(path).ok()?;
1243 let modified = meta.modified().ok()?.duration_since(UNIX_EPOCH).ok()?;
1244 let mut entries = std::fs::read_dir(path)
1245 .ok()?
1246 .filter_map(Result::ok)
1247 .map(|entry| entry.file_name().to_string_lossy().into_owned())
1248 .collect::<Vec<_>>();
1249 entries.sort_unstable();
1250 let mut hasher = std::collections::hash_map::DefaultHasher::new();
1251 modified.as_nanos().hash(&mut hasher);
1252 entries.hash(&mut hasher);
1253 Some(hasher.finish())
1254}
1255
1256fn dir_states(root: &Path, files: &[String]) -> HashMap<String, u64> {
1260 let mut rels: std::collections::HashSet<String> = std::collections::HashSet::new();
1261 rels.insert(String::new());
1262 for rel in files {
1263 let mut cur = Path::new(rel);
1264 while let Some(parent) = cur.parent() {
1265 let key = parent.to_string_lossy().to_string();
1266 if !rels.insert(key) {
1268 break;
1269 }
1270 cur = parent;
1271 }
1272 }
1273 rels.into_iter()
1274 .filter_map(|rel| {
1275 let abs = if rel.is_empty() {
1276 root.to_path_buf()
1277 } else {
1278 root.join(&rel)
1279 };
1280 dir_fingerprint(&abs).map(|fingerprint| (rel, fingerprint))
1281 })
1282 .collect()
1283}
1284
1285fn bounded_zstd_decode(compressed: &[u8], max_bytes: u64) -> Option<Vec<u8>> {
1293 use std::io::Read;
1294 let mut decoder = zstd::Decoder::new(compressed).ok()?;
1295 let mut buf = Vec::new();
1296 let mut chunk = vec![0u8; 65536];
1297 let mut total = 0u64;
1298 loop {
1299 let n = decoder.read(&mut chunk).ok()?;
1300 if n == 0 {
1301 break;
1302 }
1303 total += n as u64;
1304 if total > max_bytes {
1305 tracing::warn!(
1306 "[bm25] decompressed index exceeds limit ({:.0} MB > {:.0} MB), aborting load",
1307 total as f64 / (1024.0 * 1024.0),
1308 max_bytes as f64 / (1024.0 * 1024.0)
1309 );
1310 return None;
1311 }
1312 buf.extend_from_slice(&chunk[..n]);
1313 }
1314 Some(buf)
1315}
1316
1317fn max_decompressed_bytes(compressed_bytes: u64) -> u64 {
1328 const MAX_EXPANSION_RATIO: u64 = 20;
1330 const FLOOR: u64 = 64 * 1024 * 1024;
1332
1333 let file_derived = compressed_bytes
1334 .saturating_mul(MAX_EXPANSION_RATIO)
1335 .max(FLOOR);
1336 match crate::core::memory_guard::rss_limit_bytes() {
1337 Some(limit) => file_derived.min(limit.max(FLOOR)),
1338 None => file_derived,
1339 }
1340}
1341
1342fn index_dir(root: &Path) -> PathBuf {
1343 crate::core::index_namespace::vectors_dir(root)
1344}
1345
1346const CORPUS_MAX_DEPTH: usize = 20;
1348
1349pub(crate) struct CorpusRules {
1353 filter: crate::core::index_filter::IndexFileFilter,
1354 ignore_patterns: Vec<(String, glob::Pattern)>,
1357 max_files: usize,
1358}
1359
1360impl CorpusRules {
1361 pub(crate) fn load() -> Self {
1362 Self::from_config(&crate::core::config::Config::load())
1363 }
1364
1365 fn from_config(cfg: &crate::core::config::Config) -> Self {
1366 let filter = crate::core::index_filter::IndexFileFilter::resolve(cfg);
1370 let ignore_patterns = DEFAULT_BM25_IGNORES
1371 .iter()
1372 .map(|p| (*p).to_string())
1373 .chain(cfg.extra_ignore_patterns.iter().cloned())
1374 .filter_map(|p| glob::Pattern::new(&p).ok().map(|g| (p, g)))
1375 .collect();
1376 let max_files = if cfg.bm25_max_files == 0 {
1378 usize::MAX
1379 } else {
1380 cfg.bm25_max_files as usize
1381 };
1382 Self {
1383 filter,
1384 ignore_patterns,
1385 max_files,
1386 }
1387 }
1388
1389 fn walker(&self, root: &Path) -> ignore::WalkBuilder {
1392 let mut builder = ignore::WalkBuilder::new(root);
1393 builder
1394 .hidden(crate::core::walk_filter::SKIP_HIDDEN_IN_CONTENT_WALK)
1398 .git_ignore(self.filter.respect_gitignore)
1399 .git_global(self.filter.respect_gitignore)
1400 .git_exclude(self.filter.respect_gitignore)
1401 .require_git(false)
1402 .max_depth(Some(CORPUS_MAX_DEPTH));
1403 builder
1404 }
1405
1406 fn ignore_pattern_for(&self, rel: &str) -> Option<&str> {
1408 self.ignore_patterns
1409 .iter()
1410 .find(|(_, p)| p.matches(rel))
1411 .map(|(src, _)| src.as_str())
1412 }
1413
1414 fn filter_excludes(&self, rel: &str) -> bool {
1418 self.filter.is_excluded(&rel.replace('\\', "/"))
1419 }
1420}