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. Both paths produce an identical index (see
291        // `build` module + determinism tests).
292        if files.len() >= build::PARALLEL_MIN_FILES
293            && !crate::core::memory_guard::is_under_pressure()
294            && !crate::core::memory_guard::abort_requested()
295        {
296            return Self::build_parallel(root, content_hint, &files);
297        }
298        Self::build_sequential(root, content_hint, &files)
299    }
300
301    /// Group a previous index's chunks by file, each file's list sorted by
302    /// (start_line, end_line, symbol_name) — the deterministic reuse order shared
303    /// by both incremental rebuild paths (and the equivalence test, so it feeds
304    /// them identical inputs).
305    fn group_prev_chunks_by_file(prev: &BM25Index) -> HashMap<String, Vec<CodeChunk>> {
306        let mut old_by_file: HashMap<String, Vec<CodeChunk>> = HashMap::new();
307        for c in &prev.chunks {
308            old_by_file
309                .entry(c.file_path.clone())
310                .or_default()
311                .push(c.clone());
312        }
313        for v in old_by_file.values_mut() {
314            v.sort_by(|a, b| {
315                a.start_line
316                    .cmp(&b.start_line)
317                    .then_with(|| a.end_line.cmp(&b.end_line))
318                    .then_with(|| a.symbol_name.cmp(&b.symbol_name))
319            });
320        }
321        old_by_file
322    }
323
324    pub fn rebuild_incremental(root: &Path, prev: &BM25Index) -> Self {
325        let old_by_file = Self::group_prev_chunks_by_file(prev);
326        let files = list_code_files(root);
327
328        // #581: mirror `build()`'s dispatch. The edit loop is the hottest path in
329        // daily use, and its serial cost is dominated by re-tokenizing the *many
330        // unchanged* files' reused chunks — not the few changed ones. So the
331        // parallel path fans the entire tokenization (changed files via
332        // `prepare_file`, unchanged via re-`prepare_chunk`) across the rayon pool
333        // and merges sequentially in file order. The sequential path keeps the
334        // per-file memory-pressure early-break and serves small corpora / pressure.
335        // Both produce an identical index (see determinism tests).
336        if files.len() >= build::PARALLEL_MIN_FILES
337            && !crate::core::memory_guard::is_under_pressure()
338            && !crate::core::memory_guard::abort_requested()
339        {
340            return Self::rebuild_incremental_parallel(root, prev, &old_by_file, &files);
341        }
342        Self::rebuild_incremental_sequential(root, prev, &old_by_file, &files)
343    }
344
345    /// Sequential incremental rebuild with per-file memory-pressure guards. Reuses
346    /// unchanged files' chunks and re-extracts changed ones. Used for small
347    /// corpora and as the safe fallback under memory pressure.
348    pub(crate) fn rebuild_incremental_sequential(
349        root: &Path,
350        prev: &BM25Index,
351        old_by_file: &HashMap<String, Vec<CodeChunk>>,
352        files: &[String],
353    ) -> Self {
354        let mut index = Self::new();
355        const MAX_FILE_SIZE_BYTES: u64 = 2 * 1024 * 1024;
356
357        for (i, rel) in files.iter().enumerate() {
358            if i.is_multiple_of(500) && crate::core::memory_guard::is_under_pressure() {
359                tracing::warn!(
360                    "[bm25: stopping incremental rebuild at file {i}/{} due to memory pressure]",
361                    files.len()
362                );
363                break;
364            }
365
366            let abs = root.join(rel);
367            let Some(state) = IndexedFileState::from_path(&abs) else {
368                continue;
369            };
370
371            let unchanged = prev.files.get(rel).is_some_and(|old| *old == state);
372            if unchanged
373                && let Some(chunks) = old_by_file.get(rel)
374                && chunks.first().is_some_and(|c| !c.content.is_empty())
375            {
376                for chunk in chunks {
377                    index.add_chunk(chunk.clone());
378                }
379                index.files.insert(rel.clone(), state);
380                continue;
381            }
382
383            if state.size_bytes > MAX_FILE_SIZE_BYTES {
384                continue;
385            }
386            let content = if crate::core::extractors::is_binary_document(&abs) {
387                match std::fs::read(&abs) {
388                    Ok(bytes) => crate::core::extractors::extract(&abs, &bytes).text,
389                    Err(_) => continue,
390                }
391            } else {
392                match std::fs::read_to_string(&abs) {
393                    Ok(c) => c,
394                    Err(_) => continue,
395                }
396            };
397            if content.is_empty() {
398                continue;
399            }
400            let mut chunks = extract_chunks(rel, &content);
401            chunks.sort_by(|a, b| {
402                a.start_line
403                    .cmp(&b.start_line)
404                    .then_with(|| a.end_line.cmp(&b.end_line))
405                    .then_with(|| a.symbol_name.cmp(&b.symbol_name))
406            });
407            for chunk in chunks {
408                index.add_chunk(chunk);
409            }
410            index.files.insert(rel.clone(), state);
411        }
412
413        index.finalize();
414        index
415    }
416
417    fn add_chunk(&mut self, chunk: CodeChunk) {
418        let idx = self.chunks.len();
419
420        let enriched = enrich_for_bm25(&chunk);
421        let tokens = tokenize(&enriched);
422        for token in &tokens {
423            let lower = token.to_lowercase();
424            let postings = self.inverted.entry(lower.clone()).or_default();
425            if postings.last().map(|(last_idx, _)| *last_idx) != Some(idx) {
426                *self.doc_freqs.entry(lower).or_insert(0) += 1;
427            }
428            postings.push((idx, 1.0));
429        }
430
431        self.chunks.push(CodeChunk {
432            token_count: tokens.len(),
433            tokens: Vec::new(),
434            ..chunk
435        });
436    }
437
438    fn finalize(&mut self) {
439        self.doc_count = self.chunks.len();
440        if self.doc_count == 0 {
441            return;
442        }
443
444        let total_len: usize = self.chunks.iter().map(|c| c.token_count).sum();
445        self.avg_doc_len = total_len as f64 / self.doc_count as f64;
446    }
447
448    pub fn search(&self, query: &str, top_k: usize) -> Vec<SearchResult> {
449        let query_tokens = tokenize(query);
450        if query_tokens.is_empty() || self.doc_count == 0 {
451            return Vec::new();
452        }
453
454        // Pre-allocated score array: O(1) per-access vs HashMap overhead.
455        // Kolmogorov-optimal: minimal allocation for the scoring operation.
456        let n = self.chunks.len();
457        let mut scores = vec![0.0f64; n];
458        let mut touched = Vec::with_capacity(n.min(256));
459
460        for token in &query_tokens {
461            let lower = token.to_lowercase();
462            let df = *self.doc_freqs.get(&lower).unwrap_or(&0) as f64;
463            if df == 0.0 {
464                continue;
465            }
466
467            let idf = ((self.doc_count as f64 - df + 0.5) / (df + 0.5) + 1.0).ln();
468
469            if let Some(postings) = self.inverted.get(&lower) {
470                for &(idx, weight) in postings {
471                    let doc_len = self.chunks[idx].token_count as f64;
472                    let norm_len = doc_len / self.avg_doc_len.max(1.0);
473                    let bm25 = idf * (weight * (BM25_K1 + 1.0))
474                        / (weight + BM25_K1 * (1.0 - BM25_B + BM25_B * norm_len));
475
476                    if scores[idx] == 0.0 {
477                        touched.push(idx);
478                    }
479                    scores[idx] += bm25;
480                }
481            }
482        }
483
484        let mut results: Vec<SearchResult> = touched
485            .iter()
486            .filter(|&&idx| scores[idx] > 0.0)
487            .map(|&idx| {
488                let chunk = &self.chunks[idx];
489                let snippet = chunk.content.lines().take(5).collect::<Vec<_>>().join("\n");
490                SearchResult {
491                    chunk_idx: idx,
492                    score: scores[idx],
493                    file_path: chunk.file_path.clone(),
494                    symbol_name: chunk.symbol_name.clone(),
495                    kind: chunk.kind.clone(),
496                    start_line: chunk.start_line,
497                    end_line: chunk.end_line,
498                    snippet,
499                }
500            })
501            .collect();
502
503        results.sort_by(|a, b| {
504            b.score
505                .partial_cmp(&a.score)
506                .unwrap_or(std::cmp::Ordering::Equal)
507                .then_with(|| a.file_path.cmp(&b.file_path))
508                .then_with(|| a.symbol_name.cmp(&b.symbol_name))
509                .then_with(|| a.start_line.cmp(&b.start_line))
510                .then_with(|| a.end_line.cmp(&b.end_line))
511        });
512        results.truncate(top_k);
513        results
514    }
515
516    pub fn save(&self, root: &Path) -> std::io::Result<SaveOutcome> {
517        if self.chunks.len() > CHUNK_COUNT_WARNING {
518            tracing::warn!(
519                "[bm25] index has {} chunks (threshold {}), consider adding extra_ignore_patterns",
520                self.chunks.len(),
521                CHUNK_COUNT_WARNING
522            );
523        }
524
525        let dir = index_dir(root);
526        std::fs::create_dir_all(&dir)?;
527        let data = postcard::to_allocvec(self).map_err(|e| std::io::Error::other(e.to_string()))?;
528
529        let compressed = zstd::encode_all(data.as_slice(), ZSTD_LEVEL)
530            .map_err(|e| std::io::Error::other(format!("zstd compress: {e}")))?;
531        let compressed_bytes = compressed.len() as u64;
532
533        let max_bytes = max_bm25_cache_bytes();
534        if compressed_bytes > max_bytes {
535            // Do NOT pretend success: a silent `Ok(())` here made `load` return
536            // `None` forever and the index rebuild on every call (issue #249).
537            // Report the refusal so the orchestrator can record an actionable
538            // note and the agent-facing tools can stop claiming the index will
539            // be "ready next call".
540            tracing::warn!(
541                "[bm25] compressed index too large ({:.1} MB, limit {:.0} MB), refusing to persist: {}",
542                compressed_bytes as f64 / 1_048_576.0,
543                max_bytes / (1024 * 1024),
544                dir.display()
545            );
546            return Ok(SaveOutcome::SkippedTooLarge {
547                compressed_bytes,
548                limit_bytes: max_bytes,
549            });
550        }
551
552        tracing::info!(
553            "[bm25] index: {:.1} MB postcard → {:.1} MB zstd ({:.0}% saved)",
554            data.len() as f64 / 1_048_576.0,
555            compressed_bytes as f64 / 1_048_576.0,
556            (1.0 - compressed_bytes as f64 / data.len().max(1) as f64) * 100.0
557        );
558
559        let target = dir.join("bm25_index.bin.zst");
560        let tmp = dir.join("bm25_index.bin.zst.tmp");
561        std::fs::write(&tmp, &compressed)?;
562        std::fs::rename(&tmp, &target)?;
563
564        let _ = std::fs::remove_file(dir.join("bm25_index.bin"));
565        let _ = std::fs::remove_file(dir.join("bm25_index.json"));
566
567        let _ = std::fs::write(
568            dir.join("project_root.txt"),
569            root.to_string_lossy().as_bytes(),
570        );
571
572        Ok(SaveOutcome::Persisted { compressed_bytes })
573    }
574
575    pub fn load(root: &Path) -> Option<Self> {
576        let dir = index_dir(root);
577        let max_bytes = max_bm25_cache_bytes();
578
579        let zst_path = dir.join("bm25_index.bin.zst");
580        if zst_path.exists() {
581            let meta = std::fs::metadata(&zst_path).ok()?;
582            if meta.len() > max_bytes {
583                tracing::warn!(
584                    "[bm25] compressed index too large ({:.1} GB, limit {:.0} MB), quarantining: {}",
585                    meta.len() as f64 / 1_073_741_824.0,
586                    max_bytes / (1024 * 1024),
587                    zst_path.display()
588                );
589                let quarantined = zst_path.with_extension("zst.quarantined");
590                let _ = std::fs::rename(&zst_path, &quarantined);
591                return None;
592            }
593            let compressed = std::fs::read(&zst_path).ok()?;
594            let max_decompressed = max_bytes * 20; // allow 20x expansion ratio
595            let data = bounded_zstd_decode(&compressed, max_decompressed)?;
596            let idx: Self = postcard::from_bytes(&data).ok()?;
597            return Some(idx);
598        }
599
600        let bin_path = dir.join("bm25_index.bin");
601        if bin_path.exists() {
602            let meta = std::fs::metadata(&bin_path).ok()?;
603            if meta.len() > max_bytes {
604                tracing::warn!(
605                    "[bm25] index too large ({:.1} GB, limit {:.0} MB), quarantining: {}",
606                    meta.len() as f64 / 1_073_741_824.0,
607                    max_bytes / (1024 * 1024),
608                    bin_path.display()
609                );
610                let quarantined = bin_path.with_extension("bin.quarantined");
611                let _ = std::fs::rename(&bin_path, &quarantined);
612                return None;
613            }
614            let data = std::fs::read(&bin_path).ok()?;
615            let idx: Self = postcard::from_bytes(&data).ok()?;
616            // Auto-migrate: compress legacy .bin to .bin.zst
617            if let Ok(compressed) = zstd::encode_all(data.as_slice(), ZSTD_LEVEL) {
618                let zst_tmp = zst_path.with_extension("zst.tmp");
619                if std::fs::write(&zst_tmp, &compressed).is_ok()
620                    && std::fs::rename(&zst_tmp, &zst_path).is_ok()
621                {
622                    tracing::info!(
623                        "[bm25] migrated {:.1} MB → {:.1} MB zstd",
624                        data.len() as f64 / 1_048_576.0,
625                        compressed.len() as f64 / 1_048_576.0
626                    );
627                    let _ = std::fs::remove_file(&bin_path);
628                }
629            }
630            return Some(idx);
631        }
632
633        let json_path = dir.join("bm25_index.json");
634        if json_path.exists() {
635            let meta = std::fs::metadata(&json_path).ok()?;
636            if meta.len() > max_bytes {
637                tracing::warn!(
638                    "[bm25] index too large ({:.1} GB, limit {:.0} MB), quarantining: {}",
639                    meta.len() as f64 / 1_073_741_824.0,
640                    max_bytes / (1024 * 1024),
641                    json_path.display()
642                );
643                let quarantined = json_path.with_extension("json.quarantined");
644                let _ = std::fs::rename(&json_path, &quarantined);
645                return None;
646            }
647            let data = std::fs::read_to_string(&json_path).ok()?;
648            return serde_json::from_str(&data).ok();
649        }
650
651        None
652    }
653
654    pub fn load_or_build(root: &Path) -> Self {
655        Self::load_or_build_inner(root, false)
656    }
657
658    /// Like `load_or_build` but uses a fast sentinel-sampling staleness check
659    /// that skips the expensive full directory walk for new-file detection.
660    pub fn load_or_build_fast(root: &Path) -> Self {
661        Self::load_or_build_inner(root, true)
662    }
663
664    fn load_or_build_inner(root: &Path, fast_stale: bool) -> Self {
665        if !is_safe_bm25_root(root) {
666            return Self::default();
667        }
668        if let Some(idx) = Self::load(root) {
669            let stale = if fast_stale {
670                bm25_index_looks_stale_fast(&idx, root)
671            } else {
672                bm25_index_looks_stale(&idx, root)
673            };
674            if !stale {
675                return idx;
676            }
677            tracing::debug!(
678                "[bm25_index: stale index detected for {}; rebuilding]",
679                root.display()
680            );
681            let rebuilt = if idx.files.is_empty() {
682                Self::build_from_directory(root)
683            } else {
684                Self::rebuild_incremental(root, &idx)
685            };
686            let _ = rebuilt.save(root);
687            return rebuilt;
688        }
689
690        let built = Self::build_from_directory(root);
691        let _ = built.save(root);
692        built
693    }
694
695    pub fn index_file_path(root: &Path) -> PathBuf {
696        let dir = index_dir(root);
697        let zst = dir.join("bm25_index.bin.zst");
698        if zst.exists() {
699            return zst;
700        }
701        let bin = dir.join("bm25_index.bin");
702        if bin.exists() {
703            return bin;
704        }
705        dir.join("bm25_index.json")
706    }
707
708    /// Ingest external `ContentChunk`s into the BM25 index.
709    /// Converts each chunk to a `CodeChunk` (backward-compatible) and
710    /// rebuilds the inverted index. Returns the number of chunks ingested.
711    pub fn ingest_content_chunks(
712        &mut self,
713        chunks: impl IntoIterator<Item = super::content_chunk::ContentChunk>,
714    ) -> usize {
715        let mut count = 0usize;
716        for cc in chunks {
717            self.add_chunk(cc.into());
718            count += 1;
719        }
720        if count > 0 {
721            self.finalize();
722        }
723        count
724    }
725
726    /// Number of chunks originating from external providers.
727    pub fn external_chunk_count(&self) -> usize {
728        self.chunks
729            .iter()
730            .filter(|c| c.file_path.contains("://"))
731            .count()
732    }
733
734    /// Remove every chunk whose `file_path` starts with `prefix` (e.g.
735    /// `health://`) and rebuild the inverted index. Lets a recomputed source
736    /// (like the code-health fabric) evict its prior pass so stale entries never
737    /// linger in search. Returns the number of chunks removed.
738    pub fn remove_chunks_with_prefix(&mut self, prefix: &str) -> usize {
739        let before = self.chunks.len();
740        self.chunks.retain(|c| !c.file_path.starts_with(prefix));
741        let removed = before - self.chunks.len();
742        if removed > 0 {
743            self.finalize();
744        }
745        removed
746    }
747}
748
749fn is_safe_bm25_root(root: &Path) -> bool {
750    super::graph_index::is_safe_scan_root_public(&root.to_string_lossy())
751}
752
753fn bm25_index_looks_stale(index: &BM25Index, root: &Path) -> bool {
754    bm25_index_looks_stale_inner(index, root, false)
755}
756
757/// Fast staleness check: samples a subset of tracked files and skips the
758/// expensive `list_code_files()` walk for new-file detection.
759pub fn bm25_index_looks_stale_fast(index: &BM25Index, root: &Path) -> bool {
760    bm25_index_looks_stale_inner(index, root, true)
761}
762
763fn bm25_index_looks_stale_inner(index: &BM25Index, root: &Path, fast: bool) -> bool {
764    if index.chunks.is_empty() {
765        return false;
766    }
767
768    if index.files.is_empty() {
769        let mut seen = std::collections::HashSet::<&str>::new();
770        for chunk in &index.chunks {
771            let rel = chunk.file_path.trim_start_matches(['/', '\\']);
772            if rel.is_empty() {
773                continue;
774            }
775            if !seen.insert(rel) {
776                continue;
777            }
778            if !root.join(rel).exists() {
779                return true;
780            }
781        }
782        return false;
783    }
784
785    if fast {
786        let sample_size = index.files.len().min(SENTINEL_SAMPLE_SIZE);
787        let step = if index.files.len() > sample_size {
788            index.files.len() / sample_size
789        } else {
790            1
791        };
792        for (i, (rel, old_state)) in index.files.iter().enumerate() {
793            if i % step != 0 {
794                continue;
795            }
796            let abs = root.join(rel);
797            if !abs.exists() {
798                return true;
799            }
800            let Some(cur) = IndexedFileState::from_path(&abs) else {
801                return true;
802            };
803            if &cur != old_state {
804                return true;
805            }
806        }
807        return false;
808    }
809
810    for (rel, old_state) in &index.files {
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
823    for rel in list_code_files(root) {
824        if !index.files.contains_key(&rel) {
825            return true;
826        }
827    }
828
829    false
830}
831
832const SENTINEL_SAMPLE_SIZE: usize = 10;
833
834fn bounded_zstd_decode(compressed: &[u8], max_bytes: u64) -> Option<Vec<u8>> {
835    use std::io::Read;
836    let mut decoder = zstd::Decoder::new(compressed).ok()?;
837    let mut buf = Vec::new();
838    let mut chunk = vec![0u8; 65536];
839    let mut total = 0u64;
840    loop {
841        let n = decoder.read(&mut chunk).ok()?;
842        if n == 0 {
843            break;
844        }
845        total += n as u64;
846        if total > max_bytes {
847            tracing::warn!(
848                "[bm25] decompressed index exceeds limit ({:.0} MB > {:.0} MB), aborting load",
849                total as f64 / (1024.0 * 1024.0),
850                max_bytes as f64 / (1024.0 * 1024.0)
851            );
852            return None;
853        }
854        buf.extend_from_slice(&chunk[..n]);
855    }
856    Some(buf)
857}
858
859fn index_dir(root: &Path) -> PathBuf {
860    crate::core::index_namespace::vectors_dir(root)
861}
862
863fn list_code_files(root: &Path) -> Vec<String> {
864    let walker = ignore::WalkBuilder::new(root)
865        .hidden(true)
866        .git_ignore(true)
867        .git_global(true)
868        .git_exclude(true)
869        .require_git(false)
870        .max_depth(Some(20))
871        .filter_entry(crate::core::walk_filter::keep_entry)
872        .build();
873
874    let cfg = crate::core::config::Config::load();
875    let mut ignore_patterns: Vec<glob::Pattern> = DEFAULT_BM25_IGNORES
876        .iter()
877        .filter_map(|p| glob::Pattern::new(p).ok())
878        .collect();
879    ignore_patterns.extend(
880        cfg.extra_ignore_patterns
881            .iter()
882            .filter_map(|p| glob::Pattern::new(p).ok()),
883    );
884
885    let mut files: Vec<String> = Vec::new();
886    for entry in walker.flatten() {
887        let path = entry.path();
888        if !path.is_file() {
889            continue;
890        }
891        if !crate::core::ingestion::is_ingestible(path) {
892            continue;
893        }
894        let rel = path
895            .strip_prefix(root)
896            .unwrap_or(path)
897            .to_string_lossy()
898            .to_string();
899        if rel.is_empty() {
900            continue;
901        }
902        if ignore_patterns.iter().any(|p| p.matches(&rel)) {
903            continue;
904        }
905        if files.len() >= MAX_BM25_FILES {
906            tracing::warn!(
907                "[bm25] file cap reached ({MAX_BM25_FILES}), skipping remaining files in {}",
908                root.display()
909            );
910            break;
911        }
912        files.push(rel);
913    }
914
915    files.sort();
916    files.dedup();
917    files
918}
919
920pub fn is_code_file(path: &Path) -> bool {
921    let ext = path
922        .extension()
923        .and_then(|e| e.to_str())
924        .unwrap_or("")
925        .to_lowercase();
926    matches!(
927        ext.as_str(),
928        "rs" | "ts"
929            | "tsx"
930            | "js"
931            | "jsx"
932            | "py"
933            | "go"
934            | "java"
935            | "c"
936            | "cc"
937            | "cpp"
938            | "h"
939            | "hpp"
940            | "rb"
941            | "cs"
942            | "kt"
943            | "swift"
944            | "php"
945            | "scala"
946            | "sql"
947            | "ex"
948            | "exs"
949            | "zig"
950            | "lua"
951            | "dart"
952            | "vue"
953            | "svelte"
954    )
955}