Skip to main content

lean_ctx/core/bm25_index/
mod.rs

1use std::collections::HashMap;
2use std::path::{Path, PathBuf};
3use std::time::UNIX_EPOCH;
4
5use serde::{Deserialize, Serialize};
6mod build;
7mod chunking;
8pub use chunking::*;
9mod coordinator;
10pub use coordinator::{SearchIndexBuildProgress, get_or_start_build};
11#[cfg(test)]
12mod tests;
13
14const MAX_BM25_FILES: usize = 5000;
15const CHUNK_COUNT_WARNING: usize = 50_000;
16const ZSTD_LEVEL: i32 = 9;
17
18const DEFAULT_BM25_IGNORES: &[&str] = &[
19    "vendor/**",
20    "dist/**",
21    "build/**",
22    "public/vendor/**",
23    "public/js/**",
24    "public/css/**",
25    "public/build/**",
26    ".next/**",
27    ".nuxt/**",
28    "__pycache__/**",
29    "*.min.js",
30    "*.min.css",
31    "*.bundle.js",
32    "*.chunk.js",
33];
34
35fn max_bm25_cache_bytes() -> u64 {
36    // Single source of truth: `Config::bm25_max_cache_mb_effective` (env override
37    // › explicit config › disk-budget › generous default). Decoupled from the RAM
38    // profile so large repos persist instead of rebuilding forever (issue #249).
39    let mb = std::env::var("LEAN_CTX_BM25_MAX_CACHE_MB")
40        .ok()
41        .and_then(|v| v.parse::<u64>().ok())
42        .unwrap_or_else(|| crate::core::config::Config::load().bm25_max_cache_mb_effective());
43    mb * 1024 * 1024
44}
45
46/// Effective on-disk ceiling (bytes) for the persisted BM25 index. Single source
47/// of truth shared with `doctor` so its "oversized index" warning matches what
48/// `save`/`load` actually enforce.
49pub fn persist_ceiling_bytes() -> u64 {
50    max_bm25_cache_bytes()
51}
52
53/// Outcome of persisting a BM25 index to disk. Distinguishes a real write from a
54/// size-capped refusal so callers never mistake "refused to persist" for
55/// success (the bug behind the perpetual "index warming" report, issue #249).
56#[derive(Debug, Clone, Copy, PartialEq, Eq)]
57pub enum SaveOutcome {
58    /// Written to disk. Carries the compressed (zstd) size in bytes.
59    Persisted { compressed_bytes: u64 },
60    /// Built fine but NOT written — the compressed size exceeds the disk
61    /// ceiling. The in-memory index is still usable for this process; callers
62    /// should surface the remedy (raise the cap / add ignore patterns) instead
63    /// of silently rebuilding on every call.
64    SkippedTooLarge {
65        compressed_bytes: u64,
66        limit_bytes: u64,
67    },
68}
69
70#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
71pub struct CodeChunk {
72    pub file_path: String,
73    pub symbol_name: String,
74    pub kind: ChunkKind,
75    pub start_line: usize,
76    pub end_line: usize,
77    pub content: String,
78    #[serde(default)]
79    pub tokens: Vec<String>,
80    pub token_count: usize,
81}
82
83#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
84pub enum ChunkKind {
85    Function,
86    Struct,
87    Impl,
88    Module,
89    Class,
90    Method,
91    Other,
92    // -- External source kinds (Context Engine) --
93    Issue,
94    PullRequest,
95    WikiPage,
96    DbSchema,
97    ApiEndpoint,
98    Ticket,
99    ExternalOther,
100}
101
102#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
103pub struct IndexedFileState {
104    pub mtime_ms: u64,
105    pub size_bytes: u64,
106}
107
108impl IndexedFileState {
109    fn from_path(path: &Path) -> Option<Self> {
110        let meta = path.metadata().ok()?;
111        let size_bytes = meta.len();
112        let mtime_ms = meta
113            .modified()
114            .ok()
115            .and_then(|t| t.duration_since(UNIX_EPOCH).ok())
116            .map(|d| d.as_millis() as u64)?;
117        Some(Self {
118            mtime_ms,
119            size_bytes,
120        })
121    }
122}
123
124#[derive(Debug, Clone, Serialize, Deserialize)]
125pub struct BM25Index {
126    pub chunks: Vec<CodeChunk>,
127    pub inverted: HashMap<String, Vec<(usize, f64)>>,
128    pub avg_doc_len: f64,
129    pub doc_count: usize,
130    pub doc_freqs: HashMap<String, usize>,
131    #[serde(default)]
132    pub files: HashMap<String, IndexedFileState>,
133    /// True once `shrink_resident_content_to_snippet` has trimmed each chunk's
134    /// `content` down to the snippet lines. Resident-only RAM-saving state: never
135    /// persisted (`skip`) so the on-disk index keeps full content, and a reload
136    /// always starts as `false`. Guards the embedding pass against re-embedding
137    /// truncated bodies (see `ensure_embeddings`).
138    #[serde(default, skip)]
139    pub content_truncated: bool,
140}
141
142#[derive(Debug, Clone, Serialize, Deserialize)]
143pub struct SearchResult {
144    pub chunk_idx: usize,
145    pub score: f64,
146    pub file_path: String,
147    pub symbol_name: String,
148    pub kind: ChunkKind,
149    pub start_line: usize,
150    pub end_line: usize,
151    pub snippet: String,
152}
153
154const BM25_K1: f64 = 1.2;
155const BM25_B: f64 = 0.75;
156
157impl Default for BM25Index {
158    fn default() -> Self {
159        Self::new()
160    }
161}
162
163impl BM25Index {
164    pub fn new() -> Self {
165        Self {
166            chunks: Vec::new(),
167            inverted: HashMap::new(),
168            avg_doc_len: 0.0,
169            doc_count: 0,
170            doc_freqs: HashMap::new(),
171            files: HashMap::new(),
172            content_truncated: false,
173        }
174    }
175
176    /// Approximate heap memory used by this index in bytes.
177    pub fn memory_usage_bytes(&self) -> usize {
178        let chunks_size: usize = self
179            .chunks
180            .iter()
181            .map(|c| {
182                c.content.len()
183                    + c.file_path.len()
184                    + c.symbol_name.len()
185                    + c.tokens.iter().map(String::len).sum::<usize>()
186                    + 64
187            })
188            .sum();
189        let inverted_size: usize = self
190            .inverted
191            .iter()
192            .map(|(k, v)| k.len() + v.len() * 16 + 32)
193            .sum();
194        let files_size: usize = self.files.keys().map(|k| k.len() + 24).sum();
195        let freqs_size: usize = self.doc_freqs.keys().map(|k| k.len() + 16).sum();
196        chunks_size + inverted_size + files_size + freqs_size
197    }
198
199    /// Drops all in-memory data, effectively freeing heap. Index can be re-loaded from disk.
200    pub fn unload(&mut self) {
201        let usage = self.memory_usage_bytes();
202        self.chunks = Vec::new();
203        self.inverted = HashMap::new();
204        self.doc_freqs = HashMap::new();
205        self.files = HashMap::new();
206        self.avg_doc_len = 0.0;
207        self.doc_count = 0;
208        tracing::info!(
209            "[bm25] unloaded index, freed ~{:.1}MB",
210            usage as f64 / 1_048_576.0
211        );
212    }
213
214    /// Shrinks each resident chunk's `content` to its first `keep_lines` lines,
215    /// reclaiming the RAM held by the full source bodies once the embedding pass
216    /// has already consumed them. The search path only ever reads
217    /// `content.lines().take(5)` for snippets, so the trimmed copy is functionally
218    /// complete for BM25/dense/hybrid result rendering.
219    ///
220    /// RESIDENT-ONLY: this mutates the in-memory copy. The persisted `.bin.zst`
221    /// keeps full content (truncation never runs before `save`), so a reload
222    /// restores complete bodies. Sets `content_truncated` so a later
223    /// `ensure_embeddings` against this same resident index skips re-embedding
224    /// (which would otherwise feed truncated bodies to the embedder).
225    ///
226    /// Idempotent and only ever shrinks: chunks shorter than `keep_lines` are
227    /// left untouched.
228    pub fn shrink_resident_content_to_snippet(&mut self, keep_lines: usize) {
229        let before = self.memory_usage_bytes();
230        for chunk in &mut self.chunks {
231            // Cheap line-count gate: only allocate a new string when the body is
232            // actually longer than the snippet window.
233            if chunk.content.lines().nth(keep_lines).is_some() {
234                let trimmed: String = chunk
235                    .content
236                    .lines()
237                    .take(keep_lines)
238                    .collect::<Vec<_>>()
239                    .join("\n");
240                chunk.content = trimmed;
241                // Reclaim the spare capacity left by the larger original body.
242                chunk.content.shrink_to_fit();
243            }
244        }
245        self.content_truncated = true;
246        let after = self.memory_usage_bytes();
247        tracing::debug!(
248            "[bm25] shrank resident content to {keep_lines} lines/chunk, freed ~{:.1}MB",
249            before.saturating_sub(after) as f64 / 1_048_576.0
250        );
251    }
252
253    /// Builds an index from explicit chunks (unit tests; avoids filesystem walking).
254    #[cfg(test)]
255    pub(crate) fn from_chunks_for_test(chunks: Vec<CodeChunk>) -> Self {
256        let mut index = Self::new();
257        for mut chunk in chunks {
258            if chunk.token_count == 0 {
259                chunk.token_count = tokenize(&chunk.content).len();
260            }
261            index.add_chunk(chunk);
262        }
263        index.finalize();
264        index
265    }
266
267    pub fn build_from_directory(root: &Path) -> Self {
268        Self::build_from_directory_inner(root, &HashMap::new())
269    }
270
271    /// Like `build_from_directory` but reuses file content from a prior scan
272    /// (e.g. the graph index walk) to avoid redundant disk reads.
273    pub fn build_with_content_hint(root: &Path, content_hint: &HashMap<String, String>) -> Self {
274        Self::build_from_directory_inner(root, content_hint)
275    }
276
277    fn build_from_directory_inner(root: &Path, content_hint: &HashMap<String, String>) -> Self {
278        let root_str = root.to_string_lossy();
279        if !super::graph_index::is_safe_scan_root_public(&root_str) {
280            tracing::warn!("[bm25: scan aborted for unsafe root {root_str}]");
281            return Self::new();
282        }
283        let files = list_code_files(root);
284
285        // #933: parallel fast path for the common case. The per-file parse +
286        // tokenize work is pure and thread-safe, so we fan it across a rayon pool
287        // and merge sequentially in file order — measured 2.85 s → 0.69 s (~4x) on
288        // the lean-ctx repo. Under memory pressure (or for tiny corpora, where pool
289        // setup is not worth it) we fall back to the sequential build, which carries
290        // the incremental early-break. #685 adds admission control: a corpus whose
291        // estimated peak would blow past the guardian's hard threshold degrades to
292        // the sequential build up front instead of OOMing mid-fan-out. Both paths
293        // produce an identical index (see `build` module + determinism tests).
294        if files.len() >= build::PARALLEL_MIN_FILES
295            && !crate::core::memory_guard::is_under_pressure()
296            && !crate::core::memory_guard::abort_requested()
297            && crate::core::index_admission::admit_files(
298                crate::core::index_admission::BuildKind::Bm25,
299                root,
300                &files,
301            )
302            .parallel_ok
303        {
304            return Self::build_parallel(root, content_hint, &files);
305        }
306        Self::build_sequential(root, content_hint, &files)
307    }
308
309    /// Group a previous index's chunks by file, each file's list sorted by
310    /// (start_line, end_line, symbol_name) — the deterministic reuse order shared
311    /// by both incremental rebuild paths (and the equivalence test, so it feeds
312    /// them identical inputs).
313    fn group_prev_chunks_by_file(prev: &BM25Index) -> HashMap<String, Vec<CodeChunk>> {
314        let mut old_by_file: HashMap<String, Vec<CodeChunk>> = HashMap::new();
315        for c in &prev.chunks {
316            old_by_file
317                .entry(c.file_path.clone())
318                .or_default()
319                .push(c.clone());
320        }
321        for v in old_by_file.values_mut() {
322            v.sort_by(|a, b| {
323                a.start_line
324                    .cmp(&b.start_line)
325                    .then_with(|| a.end_line.cmp(&b.end_line))
326                    .then_with(|| a.symbol_name.cmp(&b.symbol_name))
327            });
328        }
329        old_by_file
330    }
331
332    pub fn rebuild_incremental(root: &Path, prev: &BM25Index) -> Self {
333        let old_by_file = Self::group_prev_chunks_by_file(prev);
334        let files = list_code_files(root);
335
336        // #581: mirror `build()`'s dispatch. The edit loop is the hottest path in
337        // daily use, and its serial cost is dominated by re-tokenizing the *many
338        // unchanged* files' reused chunks — not the few changed ones. So the
339        // parallel path fans the entire tokenization (changed files via
340        // `prepare_file`, unchanged via re-`prepare_chunk`) across the rayon pool
341        // and merges sequentially in file order. The sequential path keeps the
342        // per-file memory-pressure early-break and serves small corpora / pressure.
343        // #685: oversized corpora degrade to sequential via admission control.
344        // Both produce an identical index (see determinism tests).
345        if files.len() >= build::PARALLEL_MIN_FILES
346            && !crate::core::memory_guard::is_under_pressure()
347            && !crate::core::memory_guard::abort_requested()
348            && crate::core::index_admission::admit_files(
349                crate::core::index_admission::BuildKind::Bm25,
350                root,
351                &files,
352            )
353            .parallel_ok
354        {
355            return Self::rebuild_incremental_parallel(root, prev, &old_by_file, &files);
356        }
357        Self::rebuild_incremental_sequential(root, prev, &old_by_file, &files)
358    }
359
360    /// Sequential incremental rebuild with per-file memory-pressure guards. Reuses
361    /// unchanged files' chunks and re-extracts changed ones. Used for small
362    /// corpora and as the safe fallback under memory pressure.
363    pub(crate) fn rebuild_incremental_sequential(
364        root: &Path,
365        prev: &BM25Index,
366        old_by_file: &HashMap<String, Vec<CodeChunk>>,
367        files: &[String],
368    ) -> Self {
369        let mut index = Self::new();
370        const MAX_FILE_SIZE_BYTES: u64 = 2 * 1024 * 1024;
371
372        for (i, rel) in files.iter().enumerate() {
373            if i.is_multiple_of(500) && crate::core::memory_guard::is_under_pressure() {
374                tracing::warn!(
375                    "[bm25: stopping incremental rebuild at file {i}/{} due to memory pressure]",
376                    files.len()
377                );
378                break;
379            }
380
381            let abs = root.join(rel);
382            let Some(state) = IndexedFileState::from_path(&abs) else {
383                continue;
384            };
385
386            let unchanged = prev.files.get(rel).is_some_and(|old| *old == state);
387            if unchanged
388                && let Some(chunks) = old_by_file.get(rel)
389                && chunks.first().is_some_and(|c| !c.content.is_empty())
390            {
391                for chunk in chunks {
392                    index.add_chunk(chunk.clone());
393                }
394                index.files.insert(rel.clone(), state);
395                continue;
396            }
397
398            if state.size_bytes > MAX_FILE_SIZE_BYTES {
399                continue;
400            }
401            let content = if crate::core::extractors::is_binary_document(&abs) {
402                match std::fs::read(&abs) {
403                    Ok(bytes) => crate::core::extractors::extract(&abs, &bytes).text,
404                    Err(_) => continue,
405                }
406            } else {
407                match std::fs::read_to_string(&abs) {
408                    Ok(c) => c,
409                    Err(_) => continue,
410                }
411            };
412            if content.is_empty() {
413                continue;
414            }
415            let mut chunks = extract_chunks(rel, &content);
416            chunks.sort_by(|a, b| {
417                a.start_line
418                    .cmp(&b.start_line)
419                    .then_with(|| a.end_line.cmp(&b.end_line))
420                    .then_with(|| a.symbol_name.cmp(&b.symbol_name))
421            });
422            for chunk in chunks {
423                index.add_chunk(chunk);
424            }
425            index.files.insert(rel.clone(), state);
426        }
427
428        index.finalize();
429        index
430    }
431
432    fn add_chunk(&mut self, chunk: CodeChunk) {
433        let idx = self.chunks.len();
434
435        let enriched = enrich_for_bm25(&chunk);
436        let tokens = tokenize(&enriched);
437        for token in &tokens {
438            let lower = token.to_lowercase();
439            let postings = self.inverted.entry(lower.clone()).or_default();
440            if postings.last().map(|(last_idx, _)| *last_idx) != Some(idx) {
441                *self.doc_freqs.entry(lower).or_insert(0) += 1;
442            }
443            postings.push((idx, 1.0));
444        }
445
446        self.chunks.push(CodeChunk {
447            token_count: tokens.len(),
448            tokens: Vec::new(),
449            ..chunk
450        });
451    }
452
453    fn finalize(&mut self) {
454        self.doc_count = self.chunks.len();
455        if self.doc_count == 0 {
456            return;
457        }
458
459        let total_len: usize = self.chunks.iter().map(|c| c.token_count).sum();
460        self.avg_doc_len = total_len as f64 / self.doc_count as f64;
461    }
462
463    pub fn search(&self, query: &str, top_k: usize) -> Vec<SearchResult> {
464        let query_tokens = tokenize(query);
465        if query_tokens.is_empty() || self.doc_count == 0 {
466            return Vec::new();
467        }
468
469        // Pre-allocated score array: O(1) per-access vs HashMap overhead.
470        // Kolmogorov-optimal: minimal allocation for the scoring operation.
471        let n = self.chunks.len();
472        let mut scores = vec![0.0f64; n];
473        let mut touched = Vec::with_capacity(n.min(256));
474
475        for token in &query_tokens {
476            let lower = token.to_lowercase();
477            let df = *self.doc_freqs.get(&lower).unwrap_or(&0) as f64;
478            if df == 0.0 {
479                continue;
480            }
481
482            let idf = ((self.doc_count as f64 - df + 0.5) / (df + 0.5) + 1.0).ln();
483
484            if let Some(postings) = self.inverted.get(&lower) {
485                for &(idx, weight) in postings {
486                    let doc_len = self.chunks[idx].token_count as f64;
487                    let norm_len = doc_len / self.avg_doc_len.max(1.0);
488                    let bm25 = idf * (weight * (BM25_K1 + 1.0))
489                        / (weight + BM25_K1 * (1.0 - BM25_B + BM25_B * norm_len));
490
491                    if scores[idx] == 0.0 {
492                        touched.push(idx);
493                    }
494                    scores[idx] += bm25;
495                }
496            }
497        }
498
499        let mut results: Vec<SearchResult> = touched
500            .iter()
501            .filter(|&&idx| scores[idx] > 0.0)
502            .map(|&idx| {
503                let chunk = &self.chunks[idx];
504                let snippet = chunk.content.lines().take(5).collect::<Vec<_>>().join("\n");
505                SearchResult {
506                    chunk_idx: idx,
507                    score: scores[idx],
508                    file_path: chunk.file_path.clone(),
509                    symbol_name: chunk.symbol_name.clone(),
510                    kind: chunk.kind.clone(),
511                    start_line: chunk.start_line,
512                    end_line: chunk.end_line,
513                    snippet,
514                }
515            })
516            .collect();
517
518        results.sort_by(|a, b| {
519            b.score
520                .partial_cmp(&a.score)
521                .unwrap_or(std::cmp::Ordering::Equal)
522                .then_with(|| a.file_path.cmp(&b.file_path))
523                .then_with(|| a.symbol_name.cmp(&b.symbol_name))
524                .then_with(|| a.start_line.cmp(&b.start_line))
525                .then_with(|| a.end_line.cmp(&b.end_line))
526        });
527        results.truncate(top_k);
528        results
529    }
530
531    pub fn save(&self, root: &Path) -> std::io::Result<SaveOutcome> {
532        if self.chunks.len() > CHUNK_COUNT_WARNING {
533            tracing::warn!(
534                "[bm25] index has {} chunks (threshold {}), consider adding extra_ignore_patterns",
535                self.chunks.len(),
536                CHUNK_COUNT_WARNING
537            );
538        }
539
540        let dir = index_dir(root);
541        std::fs::create_dir_all(&dir)?;
542        let data = postcard::to_allocvec(self).map_err(|e| std::io::Error::other(e.to_string()))?;
543
544        let compressed = zstd::encode_all(data.as_slice(), ZSTD_LEVEL)
545            .map_err(|e| std::io::Error::other(format!("zstd compress: {e}")))?;
546        let compressed_bytes = compressed.len() as u64;
547
548        let max_bytes = max_bm25_cache_bytes();
549        if compressed_bytes > max_bytes {
550            // Do NOT pretend success: a silent `Ok(())` here made `load` return
551            // `None` forever and the index rebuild on every call (issue #249).
552            // Report the refusal so the orchestrator can record an actionable
553            // note and the agent-facing tools can stop claiming the index will
554            // be "ready next call".
555            tracing::warn!(
556                "[bm25] compressed index too large ({:.1} MB, limit {:.0} MB), refusing to persist: {}",
557                compressed_bytes as f64 / 1_048_576.0,
558                max_bytes / (1024 * 1024),
559                dir.display()
560            );
561            return Ok(SaveOutcome::SkippedTooLarge {
562                compressed_bytes,
563                limit_bytes: max_bytes,
564            });
565        }
566
567        tracing::info!(
568            "[bm25] index: {:.1} MB postcard → {:.1} MB zstd ({:.0}% saved)",
569            data.len() as f64 / 1_048_576.0,
570            compressed_bytes as f64 / 1_048_576.0,
571            (1.0 - compressed_bytes as f64 / data.len().max(1) as f64) * 100.0
572        );
573
574        let target = dir.join("bm25_index.bin.zst");
575        let tmp = dir.join("bm25_index.bin.zst.tmp");
576        std::fs::write(&tmp, &compressed)?;
577        std::fs::rename(&tmp, &target)?;
578
579        let _ = std::fs::remove_file(dir.join("bm25_index.bin"));
580        let _ = std::fs::remove_file(dir.join("bm25_index.json"));
581
582        let _ = std::fs::write(
583            dir.join("project_root.txt"),
584            root.to_string_lossy().as_bytes(),
585        );
586
587        Ok(SaveOutcome::Persisted { compressed_bytes })
588    }
589
590    pub fn load(root: &Path) -> Option<Self> {
591        let dir = index_dir(root);
592        let max_bytes = max_bm25_cache_bytes();
593
594        let zst_path = dir.join("bm25_index.bin.zst");
595        if zst_path.exists() {
596            let meta = std::fs::metadata(&zst_path).ok()?;
597            if meta.len() > max_bytes {
598                tracing::warn!(
599                    "[bm25] compressed index too large ({:.1} GB, limit {:.0} MB), quarantining: {}",
600                    meta.len() as f64 / 1_073_741_824.0,
601                    max_bytes / (1024 * 1024),
602                    zst_path.display()
603                );
604                let quarantined = zst_path.with_extension("zst.quarantined");
605                let _ = std::fs::rename(&zst_path, &quarantined);
606                return None;
607            }
608            let compressed = std::fs::read(&zst_path).ok()?;
609            let max_decompressed = max_bytes * 20; // allow 20x expansion ratio
610            let data = bounded_zstd_decode(&compressed, max_decompressed)?;
611            let idx: Self = postcard::from_bytes(&data).ok()?;
612            return Some(idx);
613        }
614
615        let bin_path = dir.join("bm25_index.bin");
616        if bin_path.exists() {
617            let meta = std::fs::metadata(&bin_path).ok()?;
618            if meta.len() > max_bytes {
619                tracing::warn!(
620                    "[bm25] index too large ({:.1} GB, limit {:.0} MB), quarantining: {}",
621                    meta.len() as f64 / 1_073_741_824.0,
622                    max_bytes / (1024 * 1024),
623                    bin_path.display()
624                );
625                let quarantined = bin_path.with_extension("bin.quarantined");
626                let _ = std::fs::rename(&bin_path, &quarantined);
627                return None;
628            }
629            let data = std::fs::read(&bin_path).ok()?;
630            let idx: Self = postcard::from_bytes(&data).ok()?;
631            // Auto-migrate: compress legacy .bin to .bin.zst
632            if let Ok(compressed) = zstd::encode_all(data.as_slice(), ZSTD_LEVEL) {
633                let zst_tmp = zst_path.with_extension("zst.tmp");
634                if std::fs::write(&zst_tmp, &compressed).is_ok()
635                    && std::fs::rename(&zst_tmp, &zst_path).is_ok()
636                {
637                    tracing::info!(
638                        "[bm25] migrated {:.1} MB → {:.1} MB zstd",
639                        data.len() as f64 / 1_048_576.0,
640                        compressed.len() as f64 / 1_048_576.0
641                    );
642                    let _ = std::fs::remove_file(&bin_path);
643                }
644            }
645            return Some(idx);
646        }
647
648        let json_path = dir.join("bm25_index.json");
649        if json_path.exists() {
650            let meta = std::fs::metadata(&json_path).ok()?;
651            if meta.len() > max_bytes {
652                tracing::warn!(
653                    "[bm25] index too large ({:.1} GB, limit {:.0} MB), quarantining: {}",
654                    meta.len() as f64 / 1_073_741_824.0,
655                    max_bytes / (1024 * 1024),
656                    json_path.display()
657                );
658                let quarantined = json_path.with_extension("json.quarantined");
659                let _ = std::fs::rename(&json_path, &quarantined);
660                return None;
661            }
662            let data = std::fs::read_to_string(&json_path).ok()?;
663            return serde_json::from_str(&data).ok();
664        }
665
666        None
667    }
668
669    pub fn load_or_build(root: &Path) -> Self {
670        Self::load_or_build_inner(root, false)
671    }
672
673    /// Like `load_or_build` but uses a fast sentinel-sampling staleness check
674    /// that skips the expensive full directory walk for new-file detection.
675    pub fn load_or_build_fast(root: &Path) -> Self {
676        Self::load_or_build_inner(root, true)
677    }
678
679    fn load_or_build_inner(root: &Path, fast_stale: bool) -> Self {
680        if !is_safe_bm25_root(root) {
681            return Self::default();
682        }
683        if let Some(idx) = Self::load(root) {
684            let stale = if fast_stale {
685                bm25_index_looks_stale_fast(&idx, root)
686            } else {
687                bm25_index_looks_stale(&idx, root)
688            };
689            if !stale {
690                return idx;
691            }
692            tracing::debug!(
693                "[bm25_index: stale index detected for {}; rebuilding]",
694                root.display()
695            );
696            let rebuilt = if idx.files.is_empty() {
697                Self::build_from_directory(root)
698            } else {
699                Self::rebuild_incremental(root, &idx)
700            };
701            let _ = rebuilt.save(root);
702            return rebuilt;
703        }
704
705        let built = Self::build_from_directory(root);
706        let _ = built.save(root);
707        built
708    }
709
710    pub fn index_file_path(root: &Path) -> PathBuf {
711        let dir = index_dir(root);
712        let zst = dir.join("bm25_index.bin.zst");
713        if zst.exists() {
714            return zst;
715        }
716        let bin = dir.join("bm25_index.bin");
717        if bin.exists() {
718            return bin;
719        }
720        dir.join("bm25_index.json")
721    }
722
723    /// Ingest external `ContentChunk`s into the BM25 index.
724    /// Converts each chunk to a `CodeChunk` (backward-compatible) and
725    /// rebuilds the inverted index. Returns the number of chunks ingested.
726    pub fn ingest_content_chunks(
727        &mut self,
728        chunks: impl IntoIterator<Item = super::content_chunk::ContentChunk>,
729    ) -> usize {
730        let mut count = 0usize;
731        for cc in chunks {
732            self.add_chunk(cc.into());
733            count += 1;
734        }
735        if count > 0 {
736            self.finalize();
737        }
738        count
739    }
740
741    /// Number of chunks originating from external providers.
742    pub fn external_chunk_count(&self) -> usize {
743        self.chunks
744            .iter()
745            .filter(|c| c.file_path.contains("://"))
746            .count()
747    }
748
749    /// Remove every chunk whose `file_path` starts with `prefix` (e.g.
750    /// `health://`) and rebuild the inverted index. Lets a recomputed source
751    /// (like the code-health fabric) evict its prior pass so stale entries never
752    /// linger in search. Returns the number of chunks removed.
753    pub fn remove_chunks_with_prefix(&mut self, prefix: &str) -> usize {
754        let before = self.chunks.len();
755        self.chunks.retain(|c| !c.file_path.starts_with(prefix));
756        let removed = before - self.chunks.len();
757        if removed > 0 {
758            self.finalize();
759        }
760        removed
761    }
762}
763
764fn is_safe_bm25_root(root: &Path) -> bool {
765    super::graph_index::is_safe_scan_root_public(&root.to_string_lossy())
766}
767
768fn bm25_index_looks_stale(index: &BM25Index, root: &Path) -> bool {
769    bm25_index_looks_stale_inner(index, root, false)
770}
771
772/// Fast staleness check: samples a subset of tracked files and skips the
773/// expensive `list_code_files()` walk for new-file detection.
774pub fn bm25_index_looks_stale_fast(index: &BM25Index, root: &Path) -> bool {
775    bm25_index_looks_stale_inner(index, root, true)
776}
777
778fn bm25_index_looks_stale_inner(index: &BM25Index, root: &Path, fast: bool) -> bool {
779    if index.chunks.is_empty() {
780        return false;
781    }
782
783    if index.files.is_empty() {
784        let mut seen = std::collections::HashSet::<&str>::new();
785        for chunk in &index.chunks {
786            let rel = chunk.file_path.trim_start_matches(['/', '\\']);
787            if rel.is_empty() {
788                continue;
789            }
790            if !seen.insert(rel) {
791                continue;
792            }
793            if !root.join(rel).exists() {
794                return true;
795            }
796        }
797        return false;
798    }
799
800    if fast {
801        let sample_size = index.files.len().min(SENTINEL_SAMPLE_SIZE);
802        let step = if index.files.len() > sample_size {
803            index.files.len() / sample_size
804        } else {
805            1
806        };
807        for (i, (rel, old_state)) in index.files.iter().enumerate() {
808            if i % step != 0 {
809                continue;
810            }
811            let abs = root.join(rel);
812            if !abs.exists() {
813                return true;
814            }
815            let Some(cur) = IndexedFileState::from_path(&abs) else {
816                return true;
817            };
818            if &cur != old_state {
819                return true;
820            }
821        }
822        return false;
823    }
824
825    for (rel, old_state) in &index.files {
826        let abs = root.join(rel);
827        if !abs.exists() {
828            return true;
829        }
830        let Some(cur) = IndexedFileState::from_path(&abs) else {
831            return true;
832        };
833        if &cur != old_state {
834            return true;
835        }
836    }
837
838    for rel in list_code_files(root) {
839        if !index.files.contains_key(&rel) {
840            return true;
841        }
842    }
843
844    false
845}
846
847const SENTINEL_SAMPLE_SIZE: usize = 10;
848
849fn bounded_zstd_decode(compressed: &[u8], max_bytes: u64) -> Option<Vec<u8>> {
850    use std::io::Read;
851    let mut decoder = zstd::Decoder::new(compressed).ok()?;
852    let mut buf = Vec::new();
853    let mut chunk = vec![0u8; 65536];
854    let mut total = 0u64;
855    loop {
856        let n = decoder.read(&mut chunk).ok()?;
857        if n == 0 {
858            break;
859        }
860        total += n as u64;
861        if total > max_bytes {
862            tracing::warn!(
863                "[bm25] decompressed index exceeds limit ({:.0} MB > {:.0} MB), aborting load",
864                total as f64 / (1024.0 * 1024.0),
865                max_bytes as f64 / (1024.0 * 1024.0)
866            );
867            return None;
868        }
869        buf.extend_from_slice(&chunk[..n]);
870    }
871    Some(buf)
872}
873
874fn index_dir(root: &Path) -> PathBuf {
875    crate::core::index_namespace::vectors_dir(root)
876}
877
878fn list_code_files(root: &Path) -> Vec<String> {
879    let cfg = crate::core::config::Config::load();
880    // #735: the declared corpus filter ([index] config + CLI overlay) decides
881    // membership before anything is chunked; the semantic index chunks this
882    // corpus, so it inherits the same universe.
883    let filter = crate::core::index_filter::IndexFileFilter::resolve(&cfg);
884
885    let walker = ignore::WalkBuilder::new(root)
886        .hidden(true)
887        .git_ignore(filter.respect_gitignore)
888        .git_global(filter.respect_gitignore)
889        .git_exclude(filter.respect_gitignore)
890        .require_git(false)
891        .max_depth(Some(20))
892        .filter_entry(crate::core::walk_filter::keep_entry)
893        .build();
894
895    let mut ignore_patterns: Vec<glob::Pattern> = DEFAULT_BM25_IGNORES
896        .iter()
897        .filter_map(|p| glob::Pattern::new(p).ok())
898        .collect();
899    ignore_patterns.extend(
900        cfg.extra_ignore_patterns
901            .iter()
902            .filter_map(|p| glob::Pattern::new(p).ok()),
903    );
904
905    let mut files: Vec<String> = Vec::new();
906    let mut filtered_out = 0usize;
907    for entry in walker.flatten() {
908        let path = entry.path();
909        if !path.is_file() {
910            continue;
911        }
912        if !crate::core::ingestion::is_ingestible(path) {
913            continue;
914        }
915        let rel = path
916            .strip_prefix(root)
917            .unwrap_or(path)
918            .to_string_lossy()
919            .to_string();
920        if rel.is_empty() {
921            continue;
922        }
923        if ignore_patterns.iter().any(|p| p.matches(&rel)) {
924            continue;
925        }
926        // Match on forward slashes so globs behave identically on Windows;
927        // the index key keeps the platform separator (existing indexes stay
928        // valid).
929        if filter.is_excluded(&rel.replace('\\', "/")) {
930            filtered_out += 1;
931            continue;
932        }
933        if files.len() >= MAX_BM25_FILES {
934            tracing::warn!(
935                "[bm25] file cap reached ({MAX_BM25_FILES}), skipping remaining files in {}",
936                root.display()
937            );
938            break;
939        }
940        files.push(rel);
941    }
942
943    if filtered_out > 0 {
944        tracing::info!(
945            "[bm25] index filter excluded {filtered_out} files ({})",
946            filter.summary().unwrap_or_default()
947        );
948    }
949
950    files.sort();
951    files.dedup();
952    files
953}
954
955pub fn is_code_file(path: &Path) -> bool {
956    let ext = path
957        .extension()
958        .and_then(|e| e.to_str())
959        .unwrap_or("")
960        .to_lowercase();
961    matches!(
962        ext.as_str(),
963        "rs" | "ts"
964            | "tsx"
965            | "js"
966            | "jsx"
967            | "py"
968            | "go"
969            | "java"
970            | "c"
971            | "cc"
972            | "cpp"
973            | "h"
974            | "hpp"
975            | "rb"
976            | "cs"
977            | "kt"
978            | "swift"
979            | "php"
980            | "scala"
981            | "sql"
982            | "ex"
983            | "exs"
984            | "zig"
985            | "lua"
986            | "dart"
987            | "vue"
988            | "svelte"
989    )
990}