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