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    /// Further shrinks resident content (chunks are already truncated to 10 lines
221    /// during `add_chunk`; this method can tighten that further). Sets
222    /// `content_truncated` so a later `ensure_embeddings` against this same
223    /// resident index skips re-embedding (which would otherwise feed truncated
224    /// 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, mut 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        // #790: truncate content to snippet AFTER tokenization — full text was
447        // used for BM25 scoring above; stored content is only for result display.
448        const SNIPPET_LINES: usize = 10;
449        if chunk.content.lines().nth(SNIPPET_LINES).is_some() {
450            chunk.content = chunk
451                .content
452                .lines()
453                .take(SNIPPET_LINES)
454                .collect::<Vec<_>>()
455                .join("\n");
456            chunk.content.shrink_to_fit();
457        }
458
459        self.chunks.push(CodeChunk {
460            token_count: tokens.len(),
461            tokens: Vec::new(),
462            ..chunk
463        });
464    }
465
466    fn finalize(&mut self) {
467        self.doc_count = self.chunks.len();
468        if self.doc_count == 0 {
469            return;
470        }
471
472        let total_len: usize = self.chunks.iter().map(|c| c.token_count).sum();
473        self.avg_doc_len = total_len as f64 / self.doc_count as f64;
474    }
475
476    pub fn search(&self, query: &str, top_k: usize) -> Vec<SearchResult> {
477        let query_tokens = tokenize(query);
478        if query_tokens.is_empty() || self.doc_count == 0 {
479            return Vec::new();
480        }
481
482        // Pre-allocated score array: O(1) per-access vs HashMap overhead.
483        // Kolmogorov-optimal: minimal allocation for the scoring operation.
484        let n = self.chunks.len();
485        let mut scores = vec![0.0f64; n];
486        let mut touched = Vec::with_capacity(n.min(256));
487
488        for token in &query_tokens {
489            let lower = token.to_lowercase();
490            let df = *self.doc_freqs.get(&lower).unwrap_or(&0) as f64;
491            if df == 0.0 {
492                continue;
493            }
494
495            let idf = ((self.doc_count as f64 - df + 0.5) / (df + 0.5) + 1.0).ln();
496
497            if let Some(postings) = self.inverted.get(&lower) {
498                for &(idx, weight) in postings {
499                    let doc_len = self.chunks[idx].token_count as f64;
500                    let norm_len = doc_len / self.avg_doc_len.max(1.0);
501                    let bm25 = idf * (weight * (BM25_K1 + 1.0))
502                        / (weight + BM25_K1 * (1.0 - BM25_B + BM25_B * norm_len));
503
504                    if scores[idx] == 0.0 {
505                        touched.push(idx);
506                    }
507                    scores[idx] += bm25;
508                }
509            }
510        }
511
512        let mut results: Vec<SearchResult> = touched
513            .iter()
514            .filter(|&&idx| scores[idx] > 0.0)
515            .map(|&idx| {
516                let chunk = &self.chunks[idx];
517                let snippet = chunk.content.lines().take(5).collect::<Vec<_>>().join("\n");
518                SearchResult {
519                    chunk_idx: idx,
520                    score: scores[idx],
521                    file_path: chunk.file_path.clone(),
522                    symbol_name: chunk.symbol_name.clone(),
523                    kind: chunk.kind.clone(),
524                    start_line: chunk.start_line,
525                    end_line: chunk.end_line,
526                    snippet,
527                }
528            })
529            .collect();
530
531        results.sort_by(|a, b| {
532            b.score
533                .partial_cmp(&a.score)
534                .unwrap_or(std::cmp::Ordering::Equal)
535                .then_with(|| a.file_path.cmp(&b.file_path))
536                .then_with(|| a.symbol_name.cmp(&b.symbol_name))
537                .then_with(|| a.start_line.cmp(&b.start_line))
538                .then_with(|| a.end_line.cmp(&b.end_line))
539        });
540        results.truncate(top_k);
541        results
542    }
543
544    pub fn save(&self, root: &Path) -> std::io::Result<SaveOutcome> {
545        if self.chunks.len() > CHUNK_COUNT_WARNING {
546            tracing::warn!(
547                "[bm25] index has {} chunks (threshold {}), consider adding extra_ignore_patterns",
548                self.chunks.len(),
549                CHUNK_COUNT_WARNING
550            );
551        }
552
553        let dir = index_dir(root);
554        std::fs::create_dir_all(&dir)?;
555
556        // #790: stream postcard → zstd → file to avoid holding the full
557        // serialized buffer in memory. Temp-file + rename for atomicity.
558        let target = dir.join("bm25_index.bin.zst");
559        let tmp = dir.join("bm25_index.bin.zst.tmp");
560        {
561            let file = std::fs::File::create(&tmp)?;
562            let buf_writer = std::io::BufWriter::new(file);
563            let mut encoder = zstd::Encoder::new(buf_writer, ZSTD_LEVEL)
564                .map_err(|e| std::io::Error::other(format!("zstd encoder init: {e}")))?;
565            postcard::to_io(self, &mut encoder)
566                .map_err(|e| std::io::Error::other(e.to_string()))?;
567            encoder.finish()?;
568        }
569
570        let compressed_bytes = std::fs::metadata(&tmp)?.len();
571        let max_bytes = max_bm25_cache_bytes();
572        if compressed_bytes > max_bytes {
573            let _ = std::fs::remove_file(&tmp);
574            tracing::warn!(
575                "[bm25] compressed index too large ({:.1} MB, limit {:.0} MB), refusing to persist: {}",
576                compressed_bytes as f64 / 1_048_576.0,
577                max_bytes / (1024 * 1024),
578                dir.display()
579            );
580            return Ok(SaveOutcome::SkippedTooLarge {
581                compressed_bytes,
582                limit_bytes: max_bytes,
583            });
584        }
585
586        tracing::info!(
587            "[bm25] index: {:.1} MB zstd compressed",
588            compressed_bytes as f64 / 1_048_576.0,
589        );
590
591        std::fs::rename(&tmp, &target)?;
592
593        let _ = std::fs::remove_file(dir.join("bm25_index.bin"));
594        let _ = std::fs::remove_file(dir.join("bm25_index.json"));
595
596        let _ = std::fs::write(
597            dir.join("project_root.txt"),
598            root.to_string_lossy().as_bytes(),
599        );
600
601        Ok(SaveOutcome::Persisted { compressed_bytes })
602    }
603
604    pub fn load(root: &Path) -> Option<Self> {
605        let dir = index_dir(root);
606        let max_bytes = max_bm25_cache_bytes();
607
608        let zst_path = dir.join("bm25_index.bin.zst");
609        if zst_path.exists() {
610            let meta = std::fs::metadata(&zst_path).ok()?;
611            if meta.len() > max_bytes {
612                tracing::warn!(
613                    "[bm25] compressed index too large ({:.1} GB, limit {:.0} MB), quarantining: {}",
614                    meta.len() as f64 / 1_073_741_824.0,
615                    max_bytes / (1024 * 1024),
616                    zst_path.display()
617                );
618                let quarantined = zst_path.with_extension("zst.quarantined");
619                let _ = std::fs::rename(&zst_path, &quarantined);
620                return None;
621            }
622            let compressed = std::fs::read(&zst_path).ok()?;
623            let max_decompressed = max_bytes * 20; // allow 20x expansion ratio
624            let data = bounded_zstd_decode(&compressed, max_decompressed)?;
625            let idx: Self = postcard::from_bytes(&data).ok()?;
626            return Some(idx);
627        }
628
629        let bin_path = dir.join("bm25_index.bin");
630        if bin_path.exists() {
631            let meta = std::fs::metadata(&bin_path).ok()?;
632            if meta.len() > max_bytes {
633                tracing::warn!(
634                    "[bm25] index too large ({:.1} GB, limit {:.0} MB), quarantining: {}",
635                    meta.len() as f64 / 1_073_741_824.0,
636                    max_bytes / (1024 * 1024),
637                    bin_path.display()
638                );
639                let quarantined = bin_path.with_extension("bin.quarantined");
640                let _ = std::fs::rename(&bin_path, &quarantined);
641                return None;
642            }
643            let data = std::fs::read(&bin_path).ok()?;
644            let idx: Self = postcard::from_bytes(&data).ok()?;
645            // Auto-migrate: compress legacy .bin to .bin.zst
646            if let Ok(compressed) = zstd::encode_all(data.as_slice(), ZSTD_LEVEL) {
647                let zst_tmp = zst_path.with_extension("zst.tmp");
648                if std::fs::write(&zst_tmp, &compressed).is_ok()
649                    && std::fs::rename(&zst_tmp, &zst_path).is_ok()
650                {
651                    tracing::info!(
652                        "[bm25] migrated {:.1} MB → {:.1} MB zstd",
653                        data.len() as f64 / 1_048_576.0,
654                        compressed.len() as f64 / 1_048_576.0
655                    );
656                    let _ = std::fs::remove_file(&bin_path);
657                }
658            }
659            return Some(idx);
660        }
661
662        let json_path = dir.join("bm25_index.json");
663        if json_path.exists() {
664            let meta = std::fs::metadata(&json_path).ok()?;
665            if meta.len() > max_bytes {
666                tracing::warn!(
667                    "[bm25] index too large ({:.1} GB, limit {:.0} MB), quarantining: {}",
668                    meta.len() as f64 / 1_073_741_824.0,
669                    max_bytes / (1024 * 1024),
670                    json_path.display()
671                );
672                let quarantined = json_path.with_extension("json.quarantined");
673                let _ = std::fs::rename(&json_path, &quarantined);
674                return None;
675            }
676            let data = std::fs::read_to_string(&json_path).ok()?;
677            return serde_json::from_str(&data).ok();
678        }
679
680        None
681    }
682
683    pub fn load_or_build(root: &Path) -> Self {
684        Self::load_or_build_inner(root, false)
685    }
686
687    /// Like `load_or_build` but uses a fast sentinel-sampling staleness check
688    /// that skips the expensive full directory walk for new-file detection.
689    pub fn load_or_build_fast(root: &Path) -> Self {
690        Self::load_or_build_inner(root, true)
691    }
692
693    fn load_or_build_inner(root: &Path, fast_stale: bool) -> Self {
694        if !is_safe_bm25_root(root) {
695            return Self::default();
696        }
697        if let Some(idx) = Self::load(root) {
698            let stale = if fast_stale {
699                bm25_index_looks_stale_fast(&idx, root)
700            } else {
701                bm25_index_looks_stale(&idx, root)
702            };
703            if !stale {
704                return idx;
705            }
706            tracing::debug!(
707                "[bm25_index: stale index detected for {}; rebuilding]",
708                root.display()
709            );
710            let rebuilt = if idx.files.is_empty() {
711                Self::build_from_directory(root)
712            } else {
713                Self::rebuild_incremental(root, &idx)
714            };
715            let _ = rebuilt.save(root);
716            return rebuilt;
717        }
718
719        let built = Self::build_from_directory(root);
720        let _ = built.save(root);
721        built
722    }
723
724    pub fn index_file_path(root: &Path) -> PathBuf {
725        let dir = index_dir(root);
726        let zst = dir.join("bm25_index.bin.zst");
727        if zst.exists() {
728            return zst;
729        }
730        let bin = dir.join("bm25_index.bin");
731        if bin.exists() {
732            return bin;
733        }
734        dir.join("bm25_index.json")
735    }
736
737    /// Ingest external `ContentChunk`s into the BM25 index.
738    /// Converts each chunk to a `CodeChunk` (backward-compatible) and
739    /// rebuilds the inverted index. Returns the number of chunks ingested.
740    pub fn ingest_content_chunks(
741        &mut self,
742        chunks: impl IntoIterator<Item = super::content_chunk::ContentChunk>,
743    ) -> usize {
744        let mut count = 0usize;
745        for cc in chunks {
746            self.add_chunk(cc.into());
747            count += 1;
748        }
749        if count > 0 {
750            self.finalize();
751        }
752        count
753    }
754
755    /// Number of chunks originating from external providers.
756    pub fn external_chunk_count(&self) -> usize {
757        self.chunks
758            .iter()
759            .filter(|c| c.file_path.contains("://"))
760            .count()
761    }
762
763    /// Remove every chunk whose `file_path` starts with `prefix` (e.g.
764    /// `health://`) and rebuild the inverted index. Lets a recomputed source
765    /// (like the code-health fabric) evict its prior pass so stale entries never
766    /// linger in search. Returns the number of chunks removed.
767    pub fn remove_chunks_with_prefix(&mut self, prefix: &str) -> usize {
768        let before = self.chunks.len();
769        self.chunks.retain(|c| !c.file_path.starts_with(prefix));
770        let removed = before - self.chunks.len();
771        if removed > 0 {
772            self.finalize();
773        }
774        removed
775    }
776}
777
778fn is_safe_bm25_root(root: &Path) -> bool {
779    super::graph_index::is_safe_scan_root_public(&root.to_string_lossy())
780}
781
782fn bm25_index_looks_stale(index: &BM25Index, root: &Path) -> bool {
783    bm25_index_looks_stale_inner(index, root, false)
784}
785
786/// Fast staleness check: samples a subset of tracked files and skips the
787/// expensive `list_code_files()` walk for new-file detection.
788pub fn bm25_index_looks_stale_fast(index: &BM25Index, root: &Path) -> bool {
789    bm25_index_looks_stale_inner(index, root, true)
790}
791
792fn bm25_index_looks_stale_inner(index: &BM25Index, root: &Path, fast: bool) -> bool {
793    if index.chunks.is_empty() {
794        return false;
795    }
796
797    if index.files.is_empty() {
798        let mut seen = std::collections::HashSet::<&str>::new();
799        for chunk in &index.chunks {
800            let rel = chunk.file_path.trim_start_matches(['/', '\\']);
801            if rel.is_empty() {
802                continue;
803            }
804            if !seen.insert(rel) {
805                continue;
806            }
807            if !root.join(rel).exists() {
808                return true;
809            }
810        }
811        return false;
812    }
813
814    if fast {
815        let sample_size = index.files.len().min(SENTINEL_SAMPLE_SIZE);
816        let step = if index.files.len() > sample_size {
817            index.files.len() / sample_size
818        } else {
819            1
820        };
821        for (i, (rel, old_state)) in index.files.iter().enumerate() {
822            if i % step != 0 {
823                continue;
824            }
825            let abs = root.join(rel);
826            if !abs.exists() {
827                return true;
828            }
829            let Some(cur) = IndexedFileState::from_path(&abs) else {
830                return true;
831            };
832            if &cur != old_state {
833                return true;
834            }
835        }
836        return false;
837    }
838
839    for (rel, old_state) in &index.files {
840        let abs = root.join(rel);
841        if !abs.exists() {
842            return true;
843        }
844        let Some(cur) = IndexedFileState::from_path(&abs) else {
845            return true;
846        };
847        if &cur != old_state {
848            return true;
849        }
850    }
851
852    for rel in list_code_files(root) {
853        if !index.files.contains_key(&rel) {
854            return true;
855        }
856    }
857
858    false
859}
860
861const SENTINEL_SAMPLE_SIZE: usize = 10;
862
863fn bounded_zstd_decode(compressed: &[u8], max_bytes: u64) -> Option<Vec<u8>> {
864    use std::io::Read;
865    let mut decoder = zstd::Decoder::new(compressed).ok()?;
866    let mut buf = Vec::new();
867    let mut chunk = vec![0u8; 65536];
868    let mut total = 0u64;
869    loop {
870        let n = decoder.read(&mut chunk).ok()?;
871        if n == 0 {
872            break;
873        }
874        total += n as u64;
875        if total > max_bytes {
876            tracing::warn!(
877                "[bm25] decompressed index exceeds limit ({:.0} MB > {:.0} MB), aborting load",
878                total as f64 / (1024.0 * 1024.0),
879                max_bytes as f64 / (1024.0 * 1024.0)
880            );
881            return None;
882        }
883        buf.extend_from_slice(&chunk[..n]);
884    }
885    Some(buf)
886}
887
888fn index_dir(root: &Path) -> PathBuf {
889    crate::core::index_namespace::vectors_dir(root)
890}
891
892fn list_code_files(root: &Path) -> Vec<String> {
893    let cfg = crate::core::config::Config::load();
894    // #735: the declared corpus filter ([index] config + CLI overlay) decides
895    // membership before anything is chunked; the semantic index chunks this
896    // corpus, so it inherits the same universe.
897    let filter = crate::core::index_filter::IndexFileFilter::resolve(&cfg);
898
899    let walker = ignore::WalkBuilder::new(root)
900        .hidden(true)
901        .git_ignore(filter.respect_gitignore)
902        .git_global(filter.respect_gitignore)
903        .git_exclude(filter.respect_gitignore)
904        .require_git(false)
905        .max_depth(Some(20))
906        .filter_entry(crate::core::walk_filter::keep_entry)
907        .build();
908
909    let mut ignore_patterns: Vec<glob::Pattern> = DEFAULT_BM25_IGNORES
910        .iter()
911        .filter_map(|p| glob::Pattern::new(p).ok())
912        .collect();
913    ignore_patterns.extend(
914        cfg.extra_ignore_patterns
915            .iter()
916            .filter_map(|p| glob::Pattern::new(p).ok()),
917    );
918
919    let mut files: Vec<String> = Vec::new();
920    let mut filtered_out = 0usize;
921    for entry in walker.flatten() {
922        let path = entry.path();
923        if !path.is_file() {
924            continue;
925        }
926        if !crate::core::ingestion::is_ingestible(path) {
927            continue;
928        }
929        let rel = path
930            .strip_prefix(root)
931            .unwrap_or(path)
932            .to_string_lossy()
933            .to_string();
934        if rel.is_empty() {
935            continue;
936        }
937        if ignore_patterns.iter().any(|p| p.matches(&rel)) {
938            continue;
939        }
940        // Match on forward slashes so globs behave identically on Windows;
941        // the index key keeps the platform separator (existing indexes stay
942        // valid).
943        if filter.is_excluded(&rel.replace('\\', "/")) {
944            filtered_out += 1;
945            continue;
946        }
947        if files.len() >= MAX_BM25_FILES {
948            tracing::warn!(
949                "[bm25] file cap reached ({MAX_BM25_FILES}), skipping remaining files in {}",
950                root.display()
951            );
952            break;
953        }
954        files.push(rel);
955    }
956
957    if filtered_out > 0 {
958        tracing::info!(
959            "[bm25] index filter excluded {filtered_out} files ({})",
960            filter.summary().unwrap_or_default()
961        );
962    }
963
964    files.sort();
965    files.dedup();
966    files
967}
968
969pub fn is_code_file(path: &Path) -> bool {
970    let ext = path
971        .extension()
972        .and_then(|e| e.to_str())
973        .unwrap_or("")
974        .to_lowercase();
975    matches!(
976        ext.as_str(),
977        "rs" | "ts"
978            | "tsx"
979            | "js"
980            | "jsx"
981            | "py"
982            | "go"
983            | "java"
984            | "c"
985            | "cc"
986            | "cpp"
987            | "h"
988            | "hpp"
989            | "rb"
990            | "cs"
991            | "kt"
992            | "swift"
993            | "php"
994            | "scala"
995            | "sql"
996            | "ex"
997            | "exs"
998            | "zig"
999            | "lua"
1000            | "dart"
1001            | "vue"
1002            | "svelte"
1003    )
1004}