From bc468c1aef5495975c713c2b4882ccf785f94eb0 Mon Sep 17 00:00:00 2001 From: Dmitriy Kovalenko Date: Tue, 8 Sep 2026 11:00:26 -0700 Subject: [PATCH] =?UTF-8?q?perf(core):=20cut=20index=20RSS=20172=E2=86=925?= =?UTF-8?q?5MB=20and=20peak=20255=E2=86=92175MB=20on=20linux=20tree?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- README.md | 1 + crates/fff-core/src/file_picker.rs | 19 +- crates/fff-core/src/index/bigram_filter.rs | 506 ++++++++++++++++++--- crates/fff-core/src/index/bigram_query.rs | 29 +- crates/fff-core/src/index/column_slab.rs | 191 ++++++++ crates/fff-core/src/index/constraints.rs | 8 +- crates/fff-core/src/index/mod.rs | 3 + crates/fff-core/src/scan.rs | 1 + crates/fff-core/src/simd_path.rs | 161 ++++--- crates/fff-core/src/types.rs | 59 ++- crates/fff-mcp/Cargo.toml | 2 +- crates/fff-mcp/src/main.rs | 13 + crates/fff-nvim/src/lib.rs | 13 + packages/fff-python/tests/test_finder.py | 11 + 14 files changed, 828 insertions(+), 189 deletions(-) create mode 100644 crates/fff-core/src/index/column_slab.rs diff --git a/README.md b/README.md index c0109fc65..9ffe3d74c 100644 --- a/README.md +++ b/README.md @@ -575,6 +575,7 @@ Run `:FFFScan` to force a rescan. - `:FFFOpenLog` opens the current session's log file. - Historical log files are stored near the main log file `/log/fff++.log` (up to 20 files) - For a crash backtrace, run `lldb -- nvim` or `gdb -- nvim` and reproduce +- fff keeps its allocator off transparent huge pages to keep the index RSS low (~15-20% on large repos); set `MIMALLOC_ALLOW_LARGE_OS_PAGES=2` before starting Neovim to trade that memory back for a slightly faster initial index build diff --git a/crates/fff-core/src/file_picker.rs b/crates/fff-core/src/file_picker.rs index 346504074..1c03e360a 100644 --- a/crates/fff-core/src/file_picker.rs +++ b/crates/fff-core/src/file_picker.rs @@ -278,7 +278,7 @@ impl FileSync { fn tombstone_files_with_arena(&mut self, mut predicate: F, mut on_tombstone: T) -> usize where F: FnMut(&FileItem, ArenaPtr) -> bool, - T: FnMut(&FileItem, ArenaPtr), + T: FnMut(&mut FileItem, ArenaPtr), { let base_arena = self.arena_base_ptr(); let overflow_arena = self.arena_overflow_ptr(); @@ -1816,6 +1816,7 @@ impl FilePicker { } let base_path = self.base_path.clone(); + let cache_budget = &self.cache_budget; let mut path_buf = [0u8; crate::simd_path::PATH_BUF_SIZE]; let tombstoned = self.sync_data.tombstone_files_with_arena( |file, arena| { @@ -1824,6 +1825,7 @@ impl FilePicker { .any(|prefix| file.relative_path_starts_with(arena, prefix)) }, |file, arena| { + file.invalidate_mmap(cache_budget); if let Some(callback) = callback.as_mut() { callback(file.write_absolute_path(arena, &base_path, &mut path_buf)); } @@ -2385,6 +2387,21 @@ fn common_dir_prefix_len(a: &str, b: &str) -> usize { last_sep } +/// Keep mimalloc off 2 MiB huge pages: with THP the arena is resident at +/// 2 MiB granularity and idle index memory inflates RSS by ~2x. Env overrides win. +/// Must run before the first allocation (see `fff_nvim`'s init-array hook). +#[cfg(feature = "mimalloc-collect")] +pub extern "C" fn tune_mimalloc() { + // SAFETY: getenv/mi_option_set touch static tables only; no allocation happens here. + unsafe { + let user_set = !libc::getenv(c"MIMALLOC_ALLOW_LARGE_OS_PAGES".as_ptr()).is_null() + || !libc::getenv(c"MIMALLOC_LARGE_OS_PAGES".as_ptr()).is_null(); + if !user_set { + libmimalloc_sys::mi_option_set(libmimalloc_sys::mi_option_large_os_pages, 0); + } + } +} + /// Ask the global allocator to return freed pages to the OS. /// Enabled via the `mimalloc-collect` feature (set by fff-nvim). /// No-op when the feature is off (tests, system allocator). diff --git a/crates/fff-core/src/index/bigram_filter.rs b/crates/fff-core/src/index/bigram_filter.rs index cd6499ced..8999bb404 100644 --- a/crates/fff-core/src/index/bigram_filter.rs +++ b/crates/fff-core/src/index/bigram_filter.rs @@ -1,11 +1,12 @@ use crate::constants::MAX_INDEXABLE_FILE_SIZE; use ahash::AHashMap; -use rayon::iter::{IndexedParallelIterator, ParallelIterator}; +use rayon::iter::{IndexedParallelIterator, IntoParallelRefIterator, ParallelIterator}; use rayon::slice::ParallelSlice; use std::cell::UnsafeCell; use std::sync::OnceLock; use std::sync::atomic::{AtomicU16, AtomicUsize, Ordering}; +use crate::index::ColumnSlab; use crate::{FileItem, constants}; /// Maximum number of distinct bigrams tracked in the inverted index. @@ -17,9 +18,31 @@ const MAX_BIGRAM_COLUMNS: usize = 5000; /// Sentinel value: bigram has no allocated column. const NO_COLUMN: u16 = u16::MAX; +/// Bigram keys only ever pair printable bytes (32..=126) which is 95 ^ 2 +pub const BIGRAM_KEY_SLOTS: usize = 95 * 95; + +// Slot in the compact lookup for a printable bigram key (`hi << 8 | lo`). +#[inline(always)] +fn key_slot(key: u16) -> usize { + let hi = (key >> 8) as usize; + let lo = (key & 0xFF) as usize; + debug_assert!((32..=126).contains(&hi) && (32..=126).contains(&lo)); + (hi - 32) * 95 + (lo - 32) +} + /// 1024 × u64 = 8 KB covers all 65536 possible bigram keys. const SEEN_WORDS: usize = 1024; +/// Below this many bitset words (512 KB) sparse encoding runs inline instead of on the pool. +const PARALLEL_ENCODE_MIN_WORDS: usize = 1 << 16; + +// Slab base pointers of the consecutive and skip-1 builders for one file insert. +#[derive(Clone, Copy)] +struct SlabPtrs { + consec: *mut u64, + skip: *mut u64, +} + /// Content size where the branchless two-pass `add_long_content` overtakes /// the single-pass `add_short_content`: ~-35% on 4 KB files, but its fixed /// flush scan dominates files under ~1 KB. See bigram_bench `bigram_build`. @@ -36,8 +59,9 @@ pub struct BigramIndexBuilder { // we use lookup as atomics only in the builder because it is filled by the rayon threads // the actual index uses pure u16 for the allocations lookup: Vec, - /// Flat bitset data, materialised on first use. - col_data: OnceLock>>, + /// Flat bitset data, materialised on first use. Stays `None` when the OS + /// refuses the mapping: the index then compresses to nothing (no prefilter). + col_data: OnceLock>>, next_column: AtomicU16, words: usize, file_count: usize, @@ -52,8 +76,8 @@ unsafe impl Sync for BigramIndexBuilder {} impl BigramIndexBuilder { pub fn new(file_count: usize) -> Self { let words = file_count.div_ceil(64); - let mut lookup = Vec::with_capacity(65536); - lookup.resize_with(65536, || AtomicU16::new(NO_COLUMN)); + let mut lookup = Vec::with_capacity(BIGRAM_KEY_SLOTS); + lookup.resize_with(BIGRAM_KEY_SLOTS, || AtomicU16::new(NO_COLUMN)); Self { lookup, col_data: OnceLock::new(), @@ -67,23 +91,33 @@ impl BigramIndexBuilder { /// Lazily materialise the full `MAX_BIGRAM_COLUMNS * words` bitset /// on first access. #[inline(always)] - fn col_data_cell(&self) -> &UnsafeCell> { - self.col_data.get_or_init(|| { - let total = MAX_BIGRAM_COLUMNS * self.words; - UnsafeCell::new(vec![0u64; total].into_boxed_slice()) - }) + fn col_data_cell(&self) -> Option<&UnsafeCell> { + self.col_data + .get_or_init(|| { + let words = MAX_BIGRAM_COLUMNS.checked_mul(self.words)?; + let slab = ColumnSlab::new(words); + if slab.is_none() { + tracing::warn!( + bytes = words.saturating_mul(8), + "bigram slab allocation refused by the OS; content index disabled" + ); + } + slab.map(UnsafeCell::new) + }) + .as_ref() } /// Raw pointer to the start of the bitset slab. Used for in-place /// `|=` writes under the partitioning invariant. #[inline(always)] - fn col_data_ptr(&self) -> *mut u64 { - unsafe { (*self.col_data_cell().get()).as_mut_ptr() } + fn col_data_ptr(&self) -> Option<*mut u64> { + Some(unsafe { (*self.col_data_cell()?.get()).as_mut_ptr() }) } #[inline] fn get_or_alloc_column(&self, key: u16) -> u16 { - let current = self.lookup[key as usize].load(Ordering::Relaxed); + let slot = key_slot(key); + let current = self.lookup[slot].load(Ordering::Relaxed); if current != NO_COLUMN { return current; } @@ -92,7 +126,7 @@ impl BigramIndexBuilder { return NO_COLUMN; } - match self.lookup[key as usize].compare_exchange( + match self.lookup[slot].compare_exchange( NO_COLUMN, new_col, Ordering::Relaxed, @@ -108,7 +142,7 @@ impl BigramIndexBuilder { #[cfg(test)] fn column_bitset(&self, col: u16) -> &[u64] { let start = col as usize * self.words; - let slab = unsafe { &*self.col_data_cell().get() }; + let slab = unsafe { &*self.col_data_cell().expect("slab").get() }; &slab[start..start + self.words] } @@ -122,6 +156,17 @@ impl BigramIndexBuilder { let word_idx = file_idx / 64; let bit_mask = 1u64 << (file_idx % 64); + // No slab means the OS refused the mapping: leave the index empty. + let (Some(consec_base), Some(skip_base)) = + (self.col_data_ptr(), skip_builder.col_data_ptr()) + else { + return; + }; + let bases = SlabPtrs { + consec: consec_base, + skip: skip_base, + }; + NORM_BUF.with_borrow_mut(|buf| { let len = content.len(); if buf.len() < len { @@ -134,9 +179,9 @@ impl BigramIndexBuilder { // Both paths record the identical bigram set; the split exists // purely for speed (see LONG_CONTENT_MIN_LEN). if len >= LONG_CONTENT_MIN_LEN { - self.add_long_content(skip_builder, n, word_idx, bit_mask); + self.add_long_content(skip_builder, n, word_idx, bit_mask, bases); } else { - self.add_short_content(skip_builder, n, word_idx, bit_mask); + self.add_short_content(skip_builder, n, word_idx, bit_mask, bases); } }); @@ -148,7 +193,14 @@ impl BigramIndexBuilder { // pairs touching the 0 sentinel — flush_seen masks those out. ~-35% vs // the single pass on 4 KB files. #[inline(always)] - fn add_long_content(&self, skip_builder: &Self, n: &[u8], word_idx: usize, bit_mask: u64) { + fn add_long_content( + &self, + skip_builder: &Self, + n: &[u8], + word_idx: usize, + bit_mask: u64, + bases: SlabPtrs, + ) { // Stack-local dedup bitsets: 1024 × u64 = 8 KB each, covers all 65536 // bigram keys. Has to fit in L1 cache. let mut seen_consec = [0u64; SEEN_WORDS]; @@ -173,18 +225,27 @@ impl BigramIndexBuilder { n1 = cur; } - self.flush_seen(&seen_consec, word_idx, bit_mask); - skip_builder.flush_seen(&seen_skip, word_idx, bit_mask); + self.flush_seen(&seen_consec, word_idx, bit_mask, bases.consec); + skip_builder.flush_seen(&seen_skip, word_idx, bit_mask, bases.skip); } #[inline(always)] - fn add_short_content(&self, skip_builder: &Self, n: &[u8], word_idx: usize, bit_mask: u64) { + fn add_short_content( + &self, + skip_builder: &Self, + n: &[u8], + word_idx: usize, + bit_mask: u64, + bases: SlabPtrs, + ) { let mut seen_consec = [0u64; SEEN_WORDS]; let mut seen_skip = [0u64; SEEN_WORDS]; - let consec_base = self.col_data_ptr(); + let SlabPtrs { + consec: consec_base, + skip: skip_base, + } = bases; let consec_words = self.words; - let skip_base = skip_builder.col_data_ptr(); let skip_words = skip_builder.words; let mut n0 = n[0]; @@ -261,8 +322,13 @@ impl BigramIndexBuilder { } } - fn flush_seen(&self, seen: &[u64; SEEN_WORDS], word_idx: usize, bit_mask: u64) { - let col_base = self.col_data_ptr(); + fn flush_seen( + &self, + seen: &[u64; SEEN_WORDS], + word_idx: usize, + bit_mask: u64, + col_base: *mut u64, + ) { let words = self.words; // SEEN_WORDS is a multiple of 8, so the remainder is always empty. for (blk, block) in seen.as_chunks::<8>().0.iter().enumerate() { @@ -319,27 +385,25 @@ impl BigramIndexBuilder { let old_lookup = self.lookup; // If no file ever populated content, col_data was never // materialised. Treat as empty — every column falls through. - let col_data: Option> = self.col_data.into_inner().map(UnsafeCell::into_inner); - - let mut lookup: Vec = vec![NO_COLUMN; 65536]; - let mut dense_data: Vec = Vec::with_capacity(cols * words); - let mut dense_count: usize = 0; - + let mut col_data: Option = self + .col_data + .into_inner() + .flatten() + .map(UnsafeCell::into_inner); + + // Pass 1: pick the columns worth keeping, in slab order so the in-place + // compaction below only ever moves a column towards the front. + let mut kept: Vec<(usize, u16, u32)> = Vec::new(); if let Some(col_data) = col_data.as_deref() { - for key in 0..65536usize { - let old_col = old_lookup[key].load(Ordering::Relaxed); + for (slot, old_col) in old_lookup.iter().enumerate() { + let old_col = old_col.load(Ordering::Relaxed); if old_col == NO_COLUMN || old_col as usize >= cols { continue; } let col_start = old_col as usize * words; let bitset = &col_data[col_start..col_start + words]; - - // count set bits to decide if this column is worth keeping. - let mut popcount = 0u32; - for &word in bitset.iter().take(words) { - popcount += word.count_ones(); - } + let popcount: u32 = bitset.iter().map(|w| w.count_ones()).sum(); // drop bigrams appearing in too few files let not_to_rare = if let Some(min_pct) = min_density_pct { @@ -349,7 +413,6 @@ impl BigramIndexBuilder { // Default: popcount ≥ words × 2 (~3.1% of files). (popcount as usize * 4) >= dense_bytes }; - if !not_to_rare { continue; } @@ -360,18 +423,79 @@ impl BigramIndexBuilder { continue; } - let dense_idx = dense_count as u16; - lookup[key] = dense_idx; - dense_count += 1; - - dense_data.extend_from_slice(bitset); + kept.push((slot, old_col, popcount)); + } + } + kept.sort_unstable_by_key(|&(_, old_col, _)| old_col); + let column = |old_col: u16| -> &[u64] { + let start = old_col as usize * words; + &col_data.as_deref().expect("kept columns imply a slab")[start..start + words] + }; + + // Low-density columns become gap lists when that encodes smaller than a dense bitset. + let encode = |&(_, old_col, popcount): &(usize, u16, u32)| -> Option> { + if (popcount as usize) >= dense_bytes { + return None; + } + let mut out = Vec::with_capacity(popcount as usize + 8); + encode_sparse_column(column(old_col), &mut out); + (out.len() < dense_bytes).then_some(out) + }; + // Small indexes finish faster inline than waking the pool; this keeps the + // window between "scan done" and "index installed" tight for tiny repos. + let encoded: Vec>> = if kept.len() * words < PARALLEL_ENCODE_MIN_WORDS { + kept.iter().map(encode).collect() + } else { + crate::parallelism::BACKGROUND_THREAD_POOL + .install(|| kept.par_iter().map(encode).collect()) + }; + + // Pass 3: compact dense columns to the front of the builder slab (no + // copy into a fresh buffer) and number sparse ones after them. + let mut lookup: Vec = vec![NO_COLUMN; BIGRAM_KEY_SLOTS]; + let mut dense_count: usize = 0; + let mut sparse_slots: Vec = Vec::new(); + let mut sparse_offsets: Vec = vec![0]; + let mut sparse_data: Vec = Vec::new(); + + for ((slot, old_col, _), sparse) in kept.iter().zip(encoded) { + match sparse { + Some(bytes) => { + sparse_slots.push(*slot); + sparse_data.extend_from_slice(&bytes); + sparse_offsets.push(sparse_data.len() as u32); + } + None => { + let src = *old_col as usize * words; + let dst = dense_count * words; + if src != dst { + let slab = col_data.as_mut().expect("kept columns imply a slab"); + slab.as_mut_slice().copy_within(src..src + words, dst); + } + lookup[*slot] = dense_count as u16; + dense_count += 1; + } + } + } + let dense_data = match col_data { + Some(mut slab) => { + slab.truncate(dense_count * words); + slab } + None => ColumnSlab::empty(), + }; + + for (i, slot) in sparse_slots.into_iter().enumerate() { + lookup[slot] = (dense_count + i) as u16; } + sparse_data.shrink_to_fit(); BigramFilter { lookup, dense_data, dense_count, + sparse_offsets, + sparse_data, words, file_count, populated, @@ -380,6 +504,69 @@ impl BigramIndexBuilder { } } +// Append the set bit positions of `bitset` as LEB128 gaps. +fn encode_sparse_column(bitset: &[u64], out: &mut Vec) { + let mut prev = 0usize; + for (w, &word) in bitset.iter().enumerate() { + let mut bits = word; + while bits != 0 { + let pos = w * 64 + bits.trailing_zeros() as usize; + bits &= bits - 1; + let mut gap = pos - prev; + prev = pos; + while gap >= 0x80 { + out.push((gap as u8) | 0x80); + gap >>= 7; + } + out.push(gap as u8); + } + } +} + +// `result &= column` for a sparse column: sorted positions are merged in one +// pass, zeroing every word the column leaves untouched. +fn and_sparse_column(result: &mut [u64], data: &[u8]) { + let mut word = 0usize; + let mut mask = 0u64; + let mut pos = 0usize; + let mut i = 0usize; + while i < data.len() { + let mut gap = 0usize; + let mut shift = 0; + loop { + let b = data[i]; + i += 1; + gap |= ((b & 0x7F) as usize) << shift; + if b & 0x80 == 0 { + break; + } + shift += 7; + } + pos += gap; + let w = pos >> 6; + if w != word { + if word < result.len() { + result[word] &= mask; + } + let end = w.min(result.len()); + result[(word + 1).min(end)..end].fill(0); + word = w; + mask = 0; + } + mask |= 1u64 << (pos & 63); + } + if word < result.len() { + result[word] &= mask; + result[word + 1..].fill(0); + } +} + +/// One bigram column: a stride-`words` dense bitset or a varint gap list. +pub(crate) enum ColumnRef<'a> { + Dense(&'a [u64]), + Sparse(&'a [u8]), +} + unsafe impl Send for BigramIndexBuilder {} /// Inverted bigram index with optional "skip-1" extension @@ -389,8 +576,12 @@ pub struct BigramFilter { lookup: Vec, /// Flat buffer of all dense column data laid out at fixed stride `words`. /// Column `i` starts at `i * words`. - dense_data: Vec, // do not try to change this to u8 it has to be wordsize + dense_data: ColumnSlab, // do not try to change this to u8 it has to be wordsize dense_count: usize, + /// Sparse columns are numbered after the dense ones: column `dense_count + i` + /// lives at `sparse_data[sparse_offsets[i]..sparse_offsets[i + 1]]`. + sparse_offsets: Vec, + sparse_data: Vec, words: usize, file_count: usize, populated: usize, @@ -425,19 +616,14 @@ impl BigramFilter { result[last] = (1u64 << (self.file_count % 64)) - 1; } - let words = self.words; let mut has_filter = false; let mut prev = pattern[0]; for &b in &pattern[1..] { if (32..=126).contains(&prev) && (32..=126).contains(&b) { let key = (prev.to_ascii_lowercase() as u16) << 8 | b.to_ascii_lowercase() as u16; - let col = self.lookup[key as usize]; - if col != NO_COLUMN { - let offset = col as usize * words; - // SAFETY: compress() guarantees offset + words <= dense_data.len() - let slice = unsafe { self.dense_data.get_unchecked(offset..offset + words) }; - bitset_and(&mut result, slice); + if let Some(col) = self.column_ref(key) { + Self::and_column(&mut result, col); has_filter = true; } } @@ -465,7 +651,6 @@ impl BigramFilter { result[last] = (1u64 << (self.file_count % 64)) - 1; } - let words = self.words; let mut has_filter = false; for i in 0..pattern.len().saturating_sub(2) { @@ -473,11 +658,8 @@ impl BigramFilter { let b = pattern[i + 2]; if (32..=126).contains(&a) && (32..=126).contains(&b) { let key = (a.to_ascii_lowercase() as u16) << 8 | b.to_ascii_lowercase() as u16; - let col = self.lookup[key as usize]; - if col != NO_COLUMN { - let offset = col as usize * words; - let slice = unsafe { self.dense_data.get_unchecked(offset..offset + words) }; - bitset_and(&mut result, slice); + if let Some(col) = self.column_ref(key) { + Self::and_column(&mut result, col); has_filter = true; } } @@ -486,6 +668,47 @@ impl BigramFilter { has_filter.then_some(result) } + /// Resolve a bigram key to its stored column, if the index kept one. + #[inline] + pub(crate) fn column_ref(&self, key: u16) -> Option> { + let col = self.column(key); + if col == NO_COLUMN { + return None; + } + let col = col as usize; + if col < self.dense_count { + let offset = col * self.words; + self.dense_data + .get(offset..offset + self.words) + .map(ColumnRef::Dense) + } else { + let i = col - self.dense_count; + let start = *self.sparse_offsets.get(i)? as usize; + let end = *self.sparse_offsets.get(i + 1)? as usize; + self.sparse_data.get(start..end).map(ColumnRef::Sparse) + } + } + + #[inline] + pub(crate) fn and_column(result: &mut [u64], col: ColumnRef<'_>) { + match col { + ColumnRef::Dense(bits) => bitset_and(result, bits), + ColumnRef::Sparse(data) => and_sparse_column(result, data), + } + } + + /// Column as a materialized bitset (borrowed for dense, decoded for sparse). + pub(crate) fn column_bitset(&self, key: u16) -> Option> { + Some(match self.column_ref(key)? { + ColumnRef::Dense(bits) => std::borrow::Cow::Borrowed(bits), + ColumnRef::Sparse(data) => { + let mut bits = vec![u64::MAX; self.words]; + and_sparse_column(&mut bits, data); + std::borrow::Cow::Owned(bits) + } + }) + } + /// Attach a skip-1 bigram index for tighter candidate filtering. pub fn set_skip_index(&mut self, skip: BigramFilter) { self.skip_index = Some(Box::new(skip)); @@ -511,23 +734,45 @@ impl BigramFilter { } pub fn columns_used(&self) -> usize { - self.dense_count + self.dense_count + self.sparse_count() + } + + /// Number of columns stored as varint gap lists. + pub fn sparse_count(&self) -> usize { + self.sparse_offsets.len().saturating_sub(1) + } + + /// Bytes held by the sparse (gap-encoded) columns. + pub fn sparse_bytes(&self) -> usize { + self.sparse_data.len() + self.sparse_offsets.len() * std::mem::size_of::() } - /// Total heap bytes used by this index (lookup + dense data + skip). + /// Total heap bytes used by this index (lookup + dense + sparse + skip). pub fn heap_bytes(&self) -> usize { let lookup_bytes = self.lookup.len() * std::mem::size_of::(); let dense_bytes = self.dense_data.len() * std::mem::size_of::(); let skip_bytes = self.skip_index.as_ref().map_or(0, |s| s.heap_bytes()); - lookup_bytes + dense_bytes + skip_bytes + lookup_bytes + dense_bytes + self.sparse_bytes() + skip_bytes } /// Check whether a bigram key is present in this index. pub fn has_key(&self, key: u16) -> bool { - self.lookup[key as usize] != NO_COLUMN + self.column(key) != NO_COLUMN + } + + /// Dense column for a printable bigram key, or `u16::MAX` when absent. + #[inline] + pub fn column(&self, key: u16) -> u16 { + let hi = key >> 8; + let lo = key & 0xFF; + if !(32..=126).contains(&hi) || !(32..=126).contains(&lo) { + return NO_COLUMN; + } + self.lookup[key_slot(key)] } - /// Raw lookup table (65536 entries mapping bigram key → column index). + /// Compact lookup table: [`BIGRAM_KEY_SLOTS`] entries mapping printable + /// bigram slots (see [`Self::column`]) to column index. pub fn lookup(&self) -> &[u16] { &self.lookup } @@ -566,10 +811,17 @@ impl BigramFilter { file_count: usize, populated: usize, ) -> Self { + // Without memory for the columns every lookup misses: no prefilter, still correct. + let (dense_data, dense_count) = match ColumnSlab::from_vec(dense_data) { + Some(slab) => (slab, dense_count), + None => (ColumnSlab::empty(), 0), + }; Self { lookup, dense_data, dense_count, + sparse_offsets: vec![0], + sparse_data: Vec::new(), words, file_count, populated, @@ -797,9 +1049,8 @@ const BIGRAM_CHUNK_FILES: usize = 4 * 64; const SKIP_INDEX_MIN_DENSITY_PCT: u32 = 12; thread_local! { - /// Reusable read buffer that is allocated per thread and used for reading files - static READ_BUF: std::cell::RefCell> = - std::cell::RefCell::new(vec![0u8; MAX_INDEXABLE_FILE_SIZE].into_boxed_slice()); + /// Per-thread file read buffer, grown on demand and released after the build. + static READ_BUF: std::cell::RefCell> = const { std::cell::RefCell::new(Vec::new()) }; } /// Reads bigram chunk, we *SHOULD NOT* use mmap cache here because bigram is built off-lock @@ -860,6 +1111,10 @@ pub(crate) fn build_bigram_index( READ_BUF.with(|read_cell| { let mut buf = read_cell.borrow_mut(); + let want = (file.size as usize).min(MAX_INDEXABLE_FILE_SIZE); + if buf.len() < want { + buf.resize(want, 0); + } let mut path_buf = [0u8; crate::simd_path::PATH_BUF_SIZE]; if let Some(content) = read_bigram_chunk( @@ -867,7 +1122,7 @@ pub(crate) fn build_bigram_index( base_fd, base_path, arena, - &mut buf[..], + &mut buf[..want], &mut path_buf, ) { // we have to manually ensure that every byte is a valid text byte to @@ -931,6 +1186,18 @@ pub(crate) fn sniff_binary_for_non_indexable( } } +/// Drop the per-thread read/normalize buffers (up to 2 x 2 MiB per pool thread) +/// so idle workers don't pin them in RSS. Call after the index is installed: +/// the broadcast waits for every worker and must not delay searchability. +pub(crate) fn release_thread_buffers() { + fn release() { + READ_BUF.with_borrow_mut(|buf| *buf = Vec::new()); + NORM_BUF.with_borrow_mut(|buf| *buf = Vec::new()); + } + crate::parallelism::BACKGROUND_THREAD_POOL.broadcast(|_| release()); + release(); +} + /// Open the base directory for the `openat` fast path. Returns `-1` on /// failure — callers interpret a negative fd as "fall back to absolute /// paths". @@ -957,6 +1224,96 @@ fn open_base_dir_fd(base_path: &std::path::Path) -> libc::c_int { mod tests { use super::*; + fn sparse_roundtrip(bitset: &[u64]) { + let mut data = Vec::new(); + encode_sparse_column(bitset, &mut data); + let mut got = vec![u64::MAX; bitset.len()]; + and_sparse_column(&mut got, &data); + assert_eq!(got, bitset); + + // AND semantics against an arbitrary partner bitset + let partner: Vec = (0..bitset.len() as u64) + .map(|i| i.wrapping_mul(0x9E37_79B9_7F4A_7C15) ^ 0x5555_5555_5555_5555) + .collect(); + let mut anded = partner.clone(); + and_sparse_column(&mut anded, &data); + let expected: Vec = partner.iter().zip(bitset).map(|(a, b)| a & b).collect(); + assert_eq!(anded, expected); + } + + #[test] + fn sparse_column_roundtrip_variants() { + sparse_roundtrip(&[]); + sparse_roundtrip(&[0]); + sparse_roundtrip(&[1]); + sparse_roundtrip(&[1 << 63]); + sparse_roundtrip(&[0, 0, 0, 1 << 5, 0, 0]); + sparse_roundtrip(&[u64::MAX, u64::MAX]); + // gaps above 127 exercise multi-byte varints + let mut wide = vec![0u64; 64]; + wide[0] = 1; + wide[10] = 1 << 3; + wide[63] = 1 << 63; + sparse_roundtrip(&wide); + let pseudo: Vec = (0..37u64) + .map(|i| i.wrapping_mul(0xD1B5_4A32_D192_ED03)) + .collect(); + sparse_roundtrip(&pseudo); + } + + #[test] + fn unmappable_slab_degrades_to_no_prefilter() { + // 2^40 files -> ~700 TB slab: the OS refuses, the index must stay empty. + let n = 1usize << 40; + let consec = BigramIndexBuilder::new(n); + let skip = BigramIndexBuilder::new(n); + consec.add_file_content(&skip, 0, b"hello world"); + assert!(!consec.is_ready()); + + let index = consec.compress(None); + assert_eq!(index.columns_used(), 0); + assert!(!index.has_key(key(b'h', b'e'))); + assert_eq!(skip.compress(Some(1)).columns_used(), 0); + } + + #[test] + fn compress_picks_sparse_for_rare_bigrams_and_queries_agree() { + // 4096 files: "zq" in 5% of them (sparse), "ab" in 50% (dense) + let n = 4096; + let consec = BigramIndexBuilder::new(n); + let skip = BigramIndexBuilder::new(n); + for i in 0..n { + let mut content = String::from("padding text "); + if i % 20 == 0 { + content.push_str("zq"); + } + if i % 2 == 0 { + content.push_str(" ab"); + } + consec.add_file_content(&skip, i, content.as_bytes()); + } + let index = consec.compress(Some(1)); + assert!(index.sparse_count() >= 1, "rare column should be sparse"); + assert!(index.dense_count() >= 1, "common column should stay dense"); + + let zq = index.query(b"zq").expect("zq tracked"); + for i in 0..n { + assert_eq!(BigramFilter::is_candidate(&zq, i), i % 20 == 0, "file {i}"); + } + let ab = index.query(b"ab").expect("ab tracked"); + for i in 0..n { + assert_eq!(BigramFilter::is_candidate(&ab, i), i % 2 == 0, "file {i}"); + } + let both = index.query(b"zq ab").expect("tracked"); + for i in 0..n { + assert_eq!( + BigramFilter::is_candidate(&both, i), + i % 20 == 0, + "file {i}" + ); + } + } + /// Build a key the same way `add_file_content` does: two printable-ASCII /// bytes, lowercased, packed as `(hi << 8) | lo`. fn key(a: u8, b: u8) -> u16 { @@ -988,7 +1345,10 @@ mod tests { /// Query: does the builder record file 0 as having this bigram set? fn builder_has_key_for_file_0(b: &BigramIndexBuilder, k: u16) -> bool { - let col = b.lookup[k as usize].load(Ordering::Relaxed); + if (k >> 8) < 32 || (k >> 8) > 126 || (k & 0xFF) < 32 || (k & 0xFF) > 126 { + return false; + } + let col = b.lookup[key_slot(k)].load(Ordering::Relaxed); if col == NO_COLUMN { return false; } @@ -1141,8 +1501,8 @@ mod tests { let key_zw = key(b'z', b'w'); // file 0 has "xy" but not "zw" - let col_xy = consec.lookup[key_xy as usize].load(Ordering::Relaxed); - let col_zw = consec.lookup[key_zw as usize].load(Ordering::Relaxed); + let col_xy = consec.lookup[key_slot(key_xy)].load(Ordering::Relaxed); + let col_zw = consec.lookup[key_slot(key_zw)].load(Ordering::Relaxed); let bitset_xy = consec.column_bitset(col_xy)[0]; let bitset_zw = consec.column_bitset(col_zw)[0]; assert_eq!(bitset_xy & 0b01, 0b01, "file 0 should have xy"); @@ -1213,8 +1573,8 @@ mod tests { let kab = key(b'a', b'b'); let kcd = key(b'c', b'd'); - let col_ab = consec.lookup[kab as usize].load(Ordering::Relaxed); - let col_cd = consec.lookup[kcd as usize].load(Ordering::Relaxed); + let col_ab = consec.lookup[key_slot(kab)].load(Ordering::Relaxed); + let col_cd = consec.lookup[key_slot(kcd)].load(Ordering::Relaxed); let ab_bitset = consec.column_bitset(col_ab); let cd_bitset = consec.column_bitset(col_cd); diff --git a/crates/fff-core/src/index/bigram_query.rs b/crates/fff-core/src/index/bigram_query.rs index b834f9035..d3963b2fa 100644 --- a/crates/fff-core/src/index/bigram_query.rs +++ b/crates/fff-core/src/index/bigram_query.rs @@ -92,34 +92,9 @@ impl BigramQuery { match self { BigramQuery::Any => None, - BigramQuery::Consec(key) => { - let col = index.lookup()[*key as usize]; - if col == u16::MAX { - return None; - } - let words = index.words(); - let offset = col as usize * words; - let data = index.dense_data(); - if offset + words > data.len() { - return None; - } - Some(Cow::Borrowed(&data[offset..offset + words])) - } + BigramQuery::Consec(key) => index.column_bitset(*key), - BigramQuery::Skip1(key) => { - let skip = index.skip_index()?; - let col = skip.lookup()[*key as usize]; - if col == u16::MAX { - return None; - } - let words = skip.words(); - let offset = col as usize * words; - let data = skip.dense_data(); - if offset + words > data.len() { - return None; - } - Some(Cow::Borrowed(&data[offset..offset + words])) - } + BigramQuery::Skip1(key) => index.skip_index()?.column_bitset(*key), BigramQuery::And(children) => { let mut result: Option> = None; diff --git a/crates/fff-core/src/index/column_slab.rs b/crates/fff-core/src/index/column_slab.rs new file mode 100644 index 000000000..f03da4dbf --- /dev/null +++ b/crates/fff-core/src/index/column_slab.rs @@ -0,0 +1,191 @@ +use std::ops::Deref; +use std::ptr::NonNull; + +pub struct ColumnSlab { + ptr: *mut u64, + len: usize, + // unix: anonymous mmap so `truncate` can hand trailing pages back to the OS + #[cfg(unix)] + mapped_bytes: usize, + #[cfg(not(unix))] + vec: Vec, +} + +unsafe impl Send for ColumnSlab {} +unsafe impl Sync for ColumnSlab {} + +impl ColumnSlab { + pub fn empty() -> Self { + Self { + ptr: NonNull::dangling().as_ptr(), + len: 0, + #[cfg(unix)] + mapped_bytes: 0, + #[cfg(not(unix))] + vec: Vec::new(), + } + } + + /// `None` when the OS refuses the allocation; callers degrade instead of aborting. + pub fn new(len: usize) -> Option { + if len == 0 { + return Some(Self::empty()); + } + #[cfg(unix)] + { + let mapped_bytes = len.checked_mul(8)?.checked_next_multiple_of(page_size())?; + // SAFETY: anonymous private mapping + let ptr = unsafe { + libc::mmap( + std::ptr::null_mut(), + mapped_bytes, + libc::PROT_READ | libc::PROT_WRITE, + libc::MAP_PRIVATE | libc::MAP_ANONYMOUS, + -1, + 0, + ) + }; + if ptr == libc::MAP_FAILED { + return None; + } + Some(Self { + ptr: ptr as *mut u64, + len, + mapped_bytes, + }) + } + #[cfg(not(unix))] + { + let mut vec = Vec::new(); + vec.try_reserve_exact(len).ok()?; + vec.resize(len, 0); + Self::from_vec(vec) + } + } + + pub fn from_vec(vec: Vec) -> Option { + #[cfg(unix)] + { + let mut slab = Self::new(vec.len())?; + slab.as_mut_slice().copy_from_slice(&vec); + Some(slab) + } + #[cfg(not(unix))] + { + let mut vec = vec; + Some(Self { + ptr: vec.as_mut_ptr(), + len: vec.len(), + vec, + }) + } + } + + #[inline] + pub fn as_mut_ptr(&mut self) -> *mut u64 { + self.ptr + } + + #[inline] + pub fn as_mut_slice(&mut self) -> &mut [u64] { + unsafe { std::slice::from_raw_parts_mut(self.ptr, self.len) } + } + + /// Shrink to the first `len` words, returning the trailing pages to the OS. + pub fn truncate(&mut self, len: usize) { + if len >= self.len { + return; + } + self.len = len; + #[cfg(unix)] + { + let keep = if len == 0 { + 0 + } else { + (len * 8).next_multiple_of(page_size()) + }; + if keep < self.mapped_bytes { + // SAFETY: unmapping a page-aligned tail of our own mapping. + unsafe { + libc::munmap( + (self.ptr as *mut u8).add(keep).cast(), + self.mapped_bytes - keep, + ); + } + self.mapped_bytes = keep; + if keep == 0 { + self.ptr = NonNull::dangling().as_ptr(); + } + } + } + #[cfg(not(unix))] + { + self.vec.truncate(len); + self.vec.shrink_to_fit(); + self.ptr = self.vec.as_mut_ptr(); + } + } +} + +impl Deref for ColumnSlab { + type Target = [u64]; + + #[inline] + fn deref(&self) -> &[u64] { + unsafe { std::slice::from_raw_parts(self.ptr, self.len) } + } +} + +#[cfg(unix)] +impl Drop for ColumnSlab { + fn drop(&mut self) { + if self.mapped_bytes > 0 { + unsafe { + libc::munmap(self.ptr.cast(), self.mapped_bytes); + } + } + } +} + +impl std::fmt::Debug for ColumnSlab { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ColumnSlab") + .field("words", &self.len) + .finish() + } +} + +#[cfg(unix)] +fn page_size() -> usize { + unsafe { libc::sysconf(libc::_SC_PAGESIZE) as usize } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn zeroed_truncate_and_roundtrip() { + let mut slab = ColumnSlab::new(10_000).unwrap(); + assert!(slab.iter().all(|&w| w == 0)); + slab.as_mut_slice()[9_999] = 7; + slab.as_mut_slice()[3] = 5; + slab.truncate(4); + assert_eq!(&slab[..], &[0, 0, 0, 5]); + slab.truncate(100); + assert_eq!(slab.len(), 4); + slab.truncate(0); + assert!(slab.is_empty()); + + let v = ColumnSlab::from_vec(vec![1, 2, 3]).unwrap(); + assert_eq!(&v[..], &[1, 2, 3]); + let empty = ColumnSlab::new(0).unwrap(); + assert!(empty.is_empty()); + } + + #[test] + fn unmappable_size_is_none_not_a_panic() { + assert!(ColumnSlab::new(usize::MAX / 16).is_none()); + assert!(ColumnSlab::new(usize::MAX).is_none()); + } +} diff --git a/crates/fff-core/src/index/constraints.rs b/crates/fff-core/src/index/constraints.rs index fefd1fadb..cdcab3b05 100644 --- a/crates/fff-core/src/index/constraints.rs +++ b/crates/fff-core/src/index/constraints.rs @@ -796,7 +796,7 @@ mod tests { #[test] fn test_apply_constraints_file_path_with_unicode_suffix() { - let arena_ptr = ArenaPtr(std::ptr::null()); + let arena_ptr = ArenaPtr::null(); let item = TestItem { relative_path: "data/유니코드_파일_테스트.csv", @@ -846,7 +846,7 @@ mod tests { #[test] fn test_negated_glob_excludes_matching_files() { - let arena_ptr = ArenaPtr(std::ptr::null()); + let arena_ptr = ArenaPtr::null(); let items = vec![ TestItem { @@ -882,7 +882,7 @@ mod tests { fn test_inline_glob_path_matches_prepass() { // Mixed (extensions + glob) takes the inline-compiled path. // Pure glob takes the prepass bitmap path. Both must give identical results. - let arena_ptr = ArenaPtr(std::ptr::null()); + let arena_ptr = ArenaPtr::null(); let items = vec![ TestItem { relative_path: "src/main.rs", @@ -925,7 +925,7 @@ mod tests { fn test_inline_negated_glob_with_extension() { // Mixed Not(Glob) on inline path — exercise the negate=true branch in // glob_matches_inline through the Not->Glob recursion. - let arena_ptr = ArenaPtr(std::ptr::null()); + let arena_ptr = ArenaPtr::null(); let items = vec![ TestItem { relative_path: "src/main.rs", diff --git a/crates/fff-core/src/index/mod.rs b/crates/fff-core/src/index/mod.rs index 1e19a09d1..bab2a5b65 100644 --- a/crates/fff-core/src/index/mod.rs +++ b/crates/fff-core/src/index/mod.rs @@ -5,6 +5,9 @@ pub(crate) use bigram_filter::*; mod bigram_query; pub use bigram_query::*; +mod column_slab; +pub(crate) use column_slab::ColumnSlab; + mod candidates; pub(crate) use candidates::*; diff --git a/crates/fff-core/src/scan.rs b/crates/fff-core/src/scan.rs index 29f1bb32c..ee5f6a7ec 100644 --- a/crates/fff-core/src/scan.rs +++ b/crates/fff-core/src/scan.rs @@ -329,6 +329,7 @@ impl ScanJob { { picker.set_bigram_index(index); } + crate::index::release_thread_buffers(); // Bigram only sniffs files <= MAX_INDEXABLE_FILE_SIZE; large // unknown-extension binaries slip past it and would otherwise be diff --git a/crates/fff-core/src/simd_path.rs b/crates/fff-core/src/simd_path.rs index 0795e042e..a2bd66d53 100644 --- a/crates/fff-core/src/simd_path.rs +++ b/crates/fff-core/src/simd_path.rs @@ -1,18 +1,17 @@ use ahash::AHashMap; -use smallvec::SmallVec; use std::borrow::Cow; /// SIMD chunk size in bytes (matches NEON/SSE2 register width). /// This must stay in sync with neo_frizbee's internal chunk size. pub(crate) const SIMD_CHUNK_BYTES: usize = 16; -/// 4 chunks = 64 bytes inline, covers ~85% of paths without heap fallback. -const INLINE_CHUNKS: usize = 4; - -pub(crate) type ChunkIndices = SmallVec<[u32; INLINE_CHUNKS]>; - +/// Read-only view of a path store: 16-byte chunk arena plus the flat table of +/// per-path chunk indices. Both point into the owning store's Vecs. #[derive(Clone, Copy)] -pub struct ArenaPtr(pub(crate) *const u8); +pub struct ArenaPtr { + chunks: *const u8, + indices: *const u32, +} // SAFETY: The arena is a read-only immutable part of file sync unsafe impl Send for ArenaPtr {} @@ -20,24 +19,32 @@ unsafe impl Sync for ArenaPtr {} impl ArenaPtr { #[inline] - pub fn new(ptr: *const u8) -> Self { - Self(ptr) + pub fn new(chunks: *const u8, indices: *const u32) -> Self { + Self { chunks, indices } } #[inline] pub fn null() -> Self { - Self(std::ptr::null()) + Self { + chunks: std::ptr::null(), + indices: std::ptr::null(), + } } #[inline] pub fn as_ptr(self) -> *const u8 { - self.0 + self.chunks + } + + #[inline] + fn chunk_ptr(self, idx: u32) -> *const u8 { + unsafe { self.chunks.add(idx as usize * SIMD_CHUNK_BYTES) } } } impl std::fmt::Debug for ArenaPtr { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "--arena-raw-pointer-0x({:?})", self.0) + write!(f, "--arena-raw-pointer-0x({:?})", self.chunks) } } @@ -65,13 +72,11 @@ pub use crate::constants::PATH_BUF_SIZE; /// Chunk pointer capacity needed for the longest path the platform allows. pub(crate) const MAX_PATH_CHUNKS: usize = PATH_BUF_SIZE.div_ceil(SIMD_CHUNK_BYTES); -/// Indices into a shared `SimdChunk` arena representing a file path. -/// -/// All read methods require an explicit `arena_base` pointer from the owning -/// `ChunkedPathStore`. The struct itself contains no raw pointers to the arena +/// A path stored as a run of chunk indices in the owning store's flat index +/// table. 8 bytes; every read needs the store's `ArenaPtr`. #[derive(Clone)] pub(crate) struct ChunkedString { - indices: ChunkIndices, + index_offset: u32, pub byte_len: u16, /// Byte offset where the filename begins. 0 for root-level files. pub filename_offset: u16, @@ -80,32 +85,41 @@ pub(crate) struct ChunkedString { impl ChunkedString { pub fn empty() -> Self { Self { - indices: SmallVec::new(), + index_offset: 0, byte_len: 0, filename_offset: 0, } } #[inline] - pub fn new(indices: ChunkIndices, byte_len: u16, filename_offset: u16) -> Self { + pub fn new(index_offset: u32, byte_len: u16, filename_offset: u16) -> Self { Self { - indices, + index_offset, byte_len, filename_offset, } } - #[cfg(test)] + #[inline] pub fn chunk_count(&self) -> usize { - self.indices.len() + chunks_needed(self.byte_len as usize) + } + + #[inline] + fn indices<'a>(&self, arena: ArenaPtr) -> &'a [u32] { + let count = self.chunk_count(); + if count == 0 { + return &[]; + } + unsafe { core::slice::from_raw_parts(arena.indices.add(self.index_offset as usize), count) } } #[inline] pub fn resolve_ptrs<'a>(&self, arena: ArenaPtr, buf: &'a mut [*const u8]) -> &'a [*const u8] { - let count = self.indices.len().min(buf.len()); - let base = arena.as_ptr(); - for (i, &idx) in self.indices[..count].iter().enumerate() { - buf[i] = unsafe { base.add(idx as usize * SIMD_CHUNK_BYTES) }; + let indices = self.indices(arena); + let count = indices.len().min(buf.len()); + for (slot, &idx) in buf.iter_mut().zip(&indices[..count]) { + *slot = arena.chunk_ptr(idx); } &buf[..count] } @@ -113,15 +127,15 @@ impl ChunkedString { #[inline] fn write_slice_to_vec( indices: &[u32], - base: *const u8, + arena: ArenaPtr, offset_in_chunk: usize, len: usize, vec: &mut Vec, ) { let mut written = 0usize; for (i, &idx) in indices.iter().enumerate() { - let src = unsafe { base.add(idx as usize * SIMD_CHUNK_BYTES) }; - let chunk_bytes = unsafe { core::slice::from_raw_parts(src, SIMD_CHUNK_BYTES) }; + let chunk_bytes = + unsafe { core::slice::from_raw_parts(arena.chunk_ptr(idx), SIMD_CHUNK_BYTES) }; let start = if i == 0 { offset_in_chunk } else { 0 }; let end = SIMD_CHUNK_BYTES.min(start + (len - written)); vec.extend_from_slice(&chunk_bytes[start..end]); @@ -142,12 +156,12 @@ impl ChunkedString { return Cow::Borrowed(""); } - let base = arena.as_ptr(); + let indices = self.indices(arena); let start_chunk = fname_offset / SIMD_CHUNK_BYTES; let offset_in_chunk = fname_offset % SIMD_CHUNK_BYTES; if offset_in_chunk == 0 && fname_len <= SIMD_CHUNK_BYTES { - let ptr = unsafe { base.add(self.indices[start_chunk] as usize * SIMD_CHUNK_BYTES) }; + let ptr = arena.chunk_ptr(indices[start_chunk]); let slice = unsafe { core::slice::from_raw_parts(ptr, fname_len) }; return Cow::Borrowed(unsafe { core::str::from_utf8_unchecked(slice) }); } @@ -155,8 +169,8 @@ impl ChunkedString { let mut out = String::with_capacity(fname_len); let needed_chunks = chunks_needed(offset_in_chunk + fname_len); Self::write_slice_to_vec( - &self.indices[start_chunk..start_chunk + needed_chunks], - base, + &indices[start_chunk..start_chunk + needed_chunks], + arena, offset_in_chunk, fname_len, unsafe { out.as_mut_vec() }, @@ -168,17 +182,19 @@ impl ChunkedString { #[inline] pub fn read_to_buf<'a>(&self, arena: ArenaPtr, buf: &'a mut [u8]) -> &'a str { let total = (self.byte_len as usize).min(buf.len()); - let usable_chunks = total.div_ceil(SIMD_CHUNK_BYTES); - let chunks_to_copy = usable_chunks.min(self.indices.len()); - let base = arena.as_ptr(); + let indices = self.indices(arena); + let chunks_to_copy = total.div_ceil(SIMD_CHUNK_BYTES).min(indices.len()); - for (i, &idx) in self.indices[..chunks_to_copy].iter().enumerate() { - let src = unsafe { base.add(idx as usize * SIMD_CHUNK_BYTES) }; + for (i, &idx) in indices[..chunks_to_copy].iter().enumerate() { let dst_offset = i * SIMD_CHUNK_BYTES; let take = SIMD_CHUNK_BYTES.min(total - dst_offset); unsafe { - core::ptr::copy_nonoverlapping(src, buf.as_mut_ptr().add(dst_offset), take); + core::ptr::copy_nonoverlapping( + arena.chunk_ptr(idx), + buf.as_mut_ptr().add(dst_offset), + take, + ); } } @@ -191,13 +207,14 @@ impl ChunkedString { let dir_len = self.filename_offset as usize; out.reserve(dir_len); - let dir_chunks = chunks_needed(dir_len).min(self.indices.len()); - let base = arena.as_ptr(); + let indices = self.indices(arena); + let dir_chunks = chunks_needed(dir_len).min(indices.len()); let vec = unsafe { out.as_mut_vec() }; - for (i, &idx) in self.indices[..dir_chunks].iter().enumerate() { - let src = unsafe { base.add(idx as usize * SIMD_CHUNK_BYTES) }; + for (i, &idx) in indices[..dir_chunks].iter().enumerate() { let take = SIMD_CHUNK_BYTES.min(dir_len - i * SIMD_CHUNK_BYTES); - vec.extend_from_slice(unsafe { core::slice::from_raw_parts(src, take) }); + vec.extend_from_slice(unsafe { + core::slice::from_raw_parts(arena.chunk_ptr(idx), take) + }); } } @@ -208,12 +225,13 @@ impl ChunkedString { let fname_offset = self.filename_offset as usize; let fname_len = self.byte_len as usize - fname_offset; out.reserve(fname_len); + let indices = self.indices(arena); let start_chunk = fname_offset / SIMD_CHUNK_BYTES; let offset_in_chunk = fname_offset % SIMD_CHUNK_BYTES; let needed_chunks = chunks_needed(offset_in_chunk + fname_len); Self::write_slice_to_vec( - &self.indices[start_chunk..start_chunk + needed_chunks], - arena.as_ptr(), + &indices[start_chunk..start_chunk + needed_chunks], + arena, offset_in_chunk, fname_len, unsafe { out.as_mut_vec() }, @@ -229,12 +247,12 @@ impl ChunkedString { return; } out.reserve(total); - let base = arena.as_ptr(); let vec = unsafe { out.as_mut_vec() }; - for (i, &idx) in self.indices.iter().enumerate() { - let src = unsafe { base.add(idx as usize * SIMD_CHUNK_BYTES) }; + for (i, &idx) in self.indices(arena).iter().enumerate() { let take = SIMD_CHUNK_BYTES.min(total - i * SIMD_CHUNK_BYTES); - vec.extend_from_slice(unsafe { core::slice::from_raw_parts(src, take) }); + vec.extend_from_slice(unsafe { + core::slice::from_raw_parts(arena.chunk_ptr(idx), take) + }); } } } @@ -242,8 +260,8 @@ impl ChunkedString { impl std::fmt::Debug for ChunkedString { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("ChunkedString") - .field("indices", &self.indices.as_slice()) - .field("chunks", &self.indices.len()) + .field("index_offset", &self.index_offset) + .field("chunks", &self.chunk_count()) .field("byte_len", &self.byte_len) .field("filename_offset", &self.filename_offset) .finish() @@ -262,6 +280,7 @@ const fn chunks_needed(byte_len: usize) -> usize { #[derive(Clone, Debug)] pub(crate) struct ChunkedPathStore { arena: Vec, + indices: Vec, } // SAFETY: arena is immutable after construction. Pointers derived from it are @@ -271,7 +290,7 @@ unsafe impl Sync for ChunkedPathStore {} impl ChunkedPathStore { pub fn heap_bytes(&self) -> usize { - self.arena.len() * SIMD_CHUNK_BYTES + self.arena.len() * SIMD_CHUNK_BYTES + self.indices.len() * std::mem::size_of::() } #[cfg(test)] @@ -281,7 +300,7 @@ impl ChunkedPathStore { #[inline] pub fn as_arena_ptr(&self) -> ArenaPtr { - ArenaPtr::new(self.arena.as_ptr() as *const u8) + ArenaPtr::new(self.arena.as_ptr() as *const u8, self.indices.as_ptr()) } } @@ -289,25 +308,34 @@ impl ChunkedPathStore { #[derive(Clone, Debug)] pub(crate) struct ChunkedPathStoreBuilder { arena: Vec, + indices: Vec, chunk_dedup: AHashMap<[u8; SIMD_CHUNK_BYTES], u32>, } impl ChunkedPathStoreBuilder { pub fn new(estimated_files: usize) -> Self { - let est_chunks = estimated_files * INLINE_CHUNKS; // we know that most of repos will fit - // most paths into 64 = 16 * INLINE_CHUNKS + // most paths fit into 64 bytes = 4 chunks; dedup keeps the arena well below that + let est_indices = estimated_files * 4; Self { - arena: Vec::with_capacity(est_chunks), - chunk_dedup: AHashMap::with_capacity(est_chunks), + arena: Vec::with_capacity(est_indices / 2), + indices: Vec::with_capacity(est_indices), + chunk_dedup: AHashMap::with_capacity(est_indices / 2), } } pub fn finish(self) -> ChunkedPathStore { - ChunkedPathStore { arena: self.arena } + let Self { + mut arena, + mut indices, + .. + } = self; + arena.shrink_to_fit(); + indices.shrink_to_fit(); + ChunkedPathStore { arena, indices } } pub fn as_arena_ptr(&self) -> ArenaPtr { - ArenaPtr::new(self.arena.as_ptr() as *const u8) + ArenaPtr::new(self.arena.as_ptr() as *const u8, self.indices.as_ptr()) } /// Like [`add_file_immediate`] but for directory paths where the entire @@ -318,8 +346,7 @@ impl ChunkedPathStoreBuilder { pub fn add_file_immediate(&mut self, rel_path: &str, filename_offset: u16) -> ChunkedString { let path_bytes = rel_path.as_bytes(); - let byte_len = rel_path.len(); - let mut indices = ChunkIndices::with_capacity(chunks_needed(byte_len)); + let index_offset = self.indices.len() as u32; for chunk in path_bytes.chunks(SIMD_CHUNK_BYTES) { let mut chunk_bytes = [0u8; SIMD_CHUNK_BYTES]; @@ -335,10 +362,10 @@ impl ChunkedPathStoreBuilder { } }; - indices.push(arena_idx); + self.indices.push(arena_idx); } - ChunkedString::new(indices, byte_len as u16, filename_offset) + ChunkedString::new(index_offset, path_bytes.len() as u16, filename_offset) } } @@ -571,11 +598,7 @@ mod tests { let mut buf = [0u8; 512]; assert_eq!(cs.read_to_buf(arena, &mut buf), path); - assert!( - cs.chunk_count() <= 6, - "should fit inline in ChunkIndices (INLINE_CHUNKS={})", - INLINE_CHUNKS - ); + assert_eq!(cs.chunk_count(), path.len().div_ceil(SIMD_CHUNK_BYTES)); } #[test] diff --git a/crates/fff-core/src/types.rs b/crates/fff-core/src/types.rs index 72661c967..8056e1b44 100644 --- a/crates/fff-core/src/types.rs +++ b/crates/fff-core/src/types.rs @@ -1,7 +1,7 @@ use std::io::Read; use std::path::{Path, PathBuf}; #[cfg(not(target_os = "windows"))] -use std::sync::OnceLock; +use std::sync::atomic::AtomicPtr; use std::sync::atomic::{AtomicI32, AtomicU8, AtomicU64, AtomicUsize, Ordering}; #[cfg(not(target_os = "windows"))] @@ -254,9 +254,17 @@ pub struct FileItem { pub(crate) path: crate::simd_path::ChunkedString, pub(crate) parent_dir_index: u32, flags: AtomicU8, - /// Lazy mmap cache. Only populated by the actual file read, controlled by the budget. + /// Lazy mmap cache (boxed, set once). Only populated by the actual file + /// read, controlled by the budget. Null while empty. #[cfg(not(target_os = "windows"))] - content: OnceLock, + content: AtomicPtr, +} + +#[cfg(not(target_os = "windows"))] +impl Drop for FileItem { + fn drop(&mut self) { + self.take_content(); + } } impl Clone for FileItem { @@ -270,9 +278,9 @@ impl Clone for FileItem { modification_frecency_score: self.modification_frecency_score, git_status: self.git_status, flags: AtomicU8::new(self.flags.load(Ordering::Relaxed)), - // on clone we have to reset the content lock + // on clone we have to reset the content cache #[cfg(not(target_os = "windows"))] - content: OnceLock::new(), + content: AtomicPtr::new(std::ptr::null_mut()), } } } @@ -315,7 +323,7 @@ impl FileItem { git_status, flags: AtomicU8::new(flags), #[cfg(not(target_os = "windows"))] - content: OnceLock::new(), + content: AtomicPtr::new(std::ptr::null_mut()), } } @@ -632,12 +640,25 @@ impl FileItem { /// invalidating ensures a fresh read on the next access. #[cfg(not(target_os = "windows"))] pub fn invalidate_mmap(&mut self, budget: &ContentCacheBudget) { - if self.content.get().is_some() { + if self.take_content().is_some() { budget.cached_count.fetch_sub(1, Ordering::Relaxed); budget.cached_bytes.fetch_sub(self.size, Ordering::Relaxed); } + } - self.content = OnceLock::new(); + #[cfg(not(target_os = "windows"))] + #[inline] + fn cached_content(&self) -> Option<&[u8]> { + let ptr = self.content.load(Ordering::Acquire); + // SAFETY: a non-null pointer is a leaked Box owned by this item until `take_content`. + (!ptr.is_null()).then(|| unsafe { (&*ptr).as_ref() }) + } + + #[cfg(not(target_os = "windows"))] + fn take_content(&mut self) -> Option> { + let ptr = self.content.swap(std::ptr::null_mut(), Ordering::AcqRel); + // SAFETY: non-null pointers always come from `Box::into_raw` in `get_cached_content`. + (!ptr.is_null()).then(|| unsafe { Box::from_raw(ptr) }) } #[cfg(target_os = "windows")] @@ -694,7 +715,7 @@ impl FileItem { base_path: &Path, budget: &ContentCacheBudget, ) -> Option<&[u8]> { - if let Some(content) = self.content.get() { + if let Some(content) = self.cached_content() { return Some(content); } @@ -717,12 +738,22 @@ impl FileItem { // file updates; the only risk is SIGBUS on a concurrent truncate, // which the watcher mitigates by invalidating on modification. let mmap = unsafe { memmap2::Mmap::map(&file) }.ok()?; - let result = self.content.get_or_init(|| mmap); - - budget.cached_count.fetch_add(1, Ordering::Relaxed); - budget.cached_bytes.fetch_add(self.size, Ordering::Relaxed); + let fresh = Box::into_raw(Box::new(mmap)); + match self.content.compare_exchange( + std::ptr::null_mut(), + fresh, + Ordering::AcqRel, + Ordering::Acquire, + ) { + Ok(_) => { + budget.cached_count.fetch_add(1, Ordering::Relaxed); + budget.cached_bytes.fetch_add(self.size, Ordering::Relaxed); + } + // Lost the race: another thread cached first, keep theirs. + Err(_) => drop(unsafe { Box::from_raw(fresh) }), + } - Some(result) + self.cached_content() } /// Get file content for searching — **always returns content** for eligible diff --git a/crates/fff-mcp/Cargo.toml b/crates/fff-mcp/Cargo.toml index 535a489a4..02a580306 100644 --- a/crates/fff-mcp/Cargo.toml +++ b/crates/fff-mcp/Cargo.toml @@ -19,7 +19,7 @@ ripgrep = ["fff/ripgrep", "fff-query-parser/ripgrep"] zlob = ["fff/zlob", "fff-query-parser/zlob"] [dependencies] -fff = { package = "fff-search", path = "../fff-core", default-features = false , version = "0.10.6", features = ["definitions"] } +fff = { package = "fff-search", path = "../fff-core", default-features = false , version = "0.10.6", features = ["definitions", "mimalloc-collect"] } fff-query-parser = { path = "../fff-query-parser", default-features = false , version = "0.10.6" } mimalloc = { workspace = true } rmcp = { version = "1.7.0", features = ["server", "transport-io"] } diff --git a/crates/fff-mcp/src/main.rs b/crates/fff-mcp/src/main.rs index 3b13f9658..da13171b3 100644 --- a/crates/fff-mcp/src/main.rs +++ b/crates/fff-mcp/src/main.rs @@ -23,6 +23,19 @@ use server::FffServer; #[global_allocator] static GLOBAL: MiMalloc = MiMalloc; +// Runs at load time, before mimalloc maps its first arena. +#[used] +#[cfg_attr( + any(target_os = "linux", target_os = "android"), + unsafe(link_section = ".init_array") +)] +#[cfg_attr( + target_vendor = "apple", + unsafe(link_section = "__DATA,__mod_init_func") +)] +#[cfg_attr(windows, unsafe(link_section = ".CRT$XCU"))] +static TUNE_MIMALLOC: extern "C" fn() = fff::tune_mimalloc; + pub const MCP_INSTRUCTIONS: &str = concat!( "FFF is a fast file finder with frecency-ranked results (frequent/recent files first, git-dirty files boosted).\n", "\n", diff --git a/crates/fff-nvim/src/lib.rs b/crates/fff-nvim/src/lib.rs index 4ea4208af..8e3605a98 100644 --- a/crates/fff-nvim/src/lib.rs +++ b/crates/fff-nvim/src/lib.rs @@ -27,6 +27,19 @@ mod user_config; #[global_allocator] static GLOBAL: MiMalloc = MiMalloc; +// Runs at load time, before mimalloc maps its first arena. +#[used] +#[cfg_attr( + any(target_os = "linux", target_os = "android"), + unsafe(link_section = ".init_array") +)] +#[cfg_attr( + target_vendor = "apple", + unsafe(link_section = "__DATA,__mod_init_func") +)] +#[cfg_attr(windows, unsafe(link_section = ".CRT$XCU"))] +static TUNE_MIMALLOC: extern "C" fn() = fff::tune_mimalloc; + // the global state for neovim lives here for efficiency // lua ffi is pretty bad with the overhead of converting raw pointer into tables pub static FILE_PICKER: Lazy = Lazy::new(SharedFilePicker::default); diff --git a/packages/fff-python/tests/test_finder.py b/packages/fff-python/tests/test_finder.py index c85b4e3bb..564644920 100644 --- a/packages/fff-python/tests/test_finder.py +++ b/packages/fff-python/tests/test_finder.py @@ -4,6 +4,7 @@ import importlib.metadata as metadata import tempfile +import time from pathlib import Path import pytest @@ -16,6 +17,15 @@ def rel(path: str) -> str: return path.replace("\\", "/") +def wait_for_content_index(finder: FileFinder, timeout: float = 10.0) -> None: + # wait_for_scan only covers the file walk; the bigram index lands a bit later + # and changes which files a grep page spans, so cursors must not straddle it. + deadline = time.monotonic() + timeout + while not finder.scan_progress.is_warmup_complete: + assert time.monotonic() < deadline, "content index did not finish" + time.sleep(0.01) + + @pytest.fixture def sample_dir() -> str: with tempfile.TemporaryDirectory() as tmp: @@ -289,6 +299,7 @@ def test_grep_invalid_mode_raises(sample_dir: str) -> None: def test_grep_cursor_paginates_by_file(sample_dir: str) -> None: with FileFinder(sample_dir, watch=False, enable_content_indexing=True) as finder: assert finder.wait_for_scan_blocking(timeout_ms=5000) + wait_for_content_index(finder) first = finder.grep("def", page_limit=1) assert first.total_matched >= 1