diff --git a/crates/fff-core/src/grep/fuzzy_grep.rs b/crates/fff-core/src/grep/fuzzy_grep.rs index 566a233cb..8e67a2932 100644 --- a/crates/fff-core/src/grep/fuzzy_grep.rs +++ b/crates/fff-core/src/grep/fuzzy_grep.rs @@ -12,7 +12,7 @@ use super::types::{GrepMatch, GrepResult, GrepSearchOptions}; #[allow(clippy::too_many_arguments)] pub(super) fn fuzzy_grep_search<'a>( grep_text: &str, - files_to_search: &[&'a FileItem], + files_to_search: &mut super::prefilter::RankedFiles<'a>, options: &GrepSearchOptions, total_files: usize, filtered_file_count: usize, @@ -136,7 +136,7 @@ pub(super) fn fuzzy_grep_search<'a>( while chunk_start < files_to_search.len() { let chunk_end = (chunk_start + chunk_size).min(files_to_search.len()); - let chunk = &files_to_search[chunk_start..chunk_end]; + let chunk = files_to_search.batch(chunk_start..chunk_end); let chunk_offset = chunk_start; chunk_start = chunk_end; chunk_size = (chunk_size * growth).min(max_chunk); diff --git a/crates/fff-core/src/grep/grep.rs b/crates/fff-core/src/grep/grep.rs index bfacd3989..2ede1f227 100644 --- a/crates/fff-core/src/grep/grep.rs +++ b/crates/fff-core/src/grep/grep.rs @@ -336,7 +336,7 @@ fn grep_search_parsed<'a>( GrepMode::Fuzzy => { let bigram_candidates = fuzzy_candidates(bigram_index, bigram_overlay, &grep_text); - let (files_to_search, filtered_file_count) = prefilter_with_filepath_retry( + let (mut files_to_search, filtered_file_count) = prefilter_with_filepath_retry( files, constraints_from_query, bigram_candidates.as_deref(), @@ -352,7 +352,7 @@ fn grep_search_parsed<'a>( return super::fuzzy_grep::fuzzy_grep_search( &grep_text, - &files_to_search, + &mut files_to_search, options, total_files, filtered_file_count, @@ -403,7 +403,7 @@ fn grep_search_parsed<'a>( regex_candidates(bigram_index, bigram_overlay, &effective_pattern) }; - let (files_to_search, filtered_file_count) = prefilter_with_filepath_retry( + let (mut files_to_search, filtered_file_count) = prefilter_with_filepath_retry( files, constraints_from_query, bigram_candidates.as_deref(), @@ -433,7 +433,7 @@ fn grep_search_parsed<'a>( let should_prefilter = regex.is_none(); let mut result = perform_grep( - &files_to_search, + &mut files_to_search, options, &GrepContext { total_files, @@ -583,7 +583,7 @@ impl GrepContext<'_, '_> { fields(prefiltered_count = files_to_search.len()) )] pub(super) fn perform_grep<'a, F>( - files_to_search: &[&'a FileItem], + files_to_search: &mut super::prefilter::RankedFiles<'a>, options: &GrepSearchOptions, ctx: &GrepContext<'_, '_>, search_file: F, @@ -627,7 +627,7 @@ where while chunk_start < files_to_search.len() { let chunk_end = (chunk_start + chunk_size).min(files_to_search.len()); - let chunk = &files_to_search[chunk_start..chunk_end]; + let chunk = files_to_search.batch(chunk_start..chunk_end); chunk_start = chunk_end; let chunk_offset = files_consumed; diff --git a/crates/fff-core/src/grep/multi_pattern.rs b/crates/fff-core/src/grep/multi_pattern.rs index 7913d4f69..1ad3e8983 100644 --- a/crates/fff-core/src/grep/multi_pattern.rs +++ b/crates/fff-core/src/grep/multi_pattern.rs @@ -126,7 +126,7 @@ pub(crate) fn multi_grep_search<'a>( let base_file_count = bigram_boundary(bigram_overlay, files.len()); // Constraints are separate from patterns, so a miss must not broaden the search. - let (files_to_search, filtered_file_count) = prefilter_files( + let (mut files_to_search, filtered_file_count) = prefilter_files( files, constraints, bigram_candidates.as_deref(), @@ -157,7 +157,7 @@ pub(crate) fn multi_grep_search<'a>( let ac_matcher = AhoCorasickMatcher { ac: &ac }; perform_grep( - &files_to_search, + &mut files_to_search, options, &GrepContext { total_files, diff --git a/crates/fff-core/src/grep/prefilter.rs b/crates/fff-core/src/grep/prefilter.rs index ea7cab180..f0eb60218 100644 --- a/crates/fff-core/src/grep/prefilter.rs +++ b/crates/fff-core/src/grep/prefilter.rs @@ -4,6 +4,66 @@ use crate::index::constraints::{ConstraintPlan, ConstraintsBuffers}; use crate::sort_buffer::sort_with_buffer; use crate::types::FileItem; use fff_query_parser::Constraint; +use std::cmp::Ordering; +use std::ops::Range; + +const FIRST_RANKED: usize = 8192; + +pub(super) struct RankedFiles<'a> { + files: Vec<&'a FileItem>, + ranked: usize, + offset: usize, +} + +fn rank(a: &&FileItem, b: &&FileItem) -> Ordering { + b.total_frecency_score() + .cmp(&a.total_frecency_score()) + .then_with(|| b.modified.cmp(&a.modified)) +} + +fn rank_with_ties(a: &&FileItem, b: &&FileItem) -> Ordering { + // All references point into one input slice, whose order breaks ties. + rank(a, b).then_with(|| std::ptr::from_ref(*a).cmp(&std::ptr::from_ref(*b))) +} + +impl<'a> RankedFiles<'a> { + fn new(mut files: Vec<&'a FileItem>, options: &GrepSearchOptions) -> Self { + let offset = options.file_offset.min(files.len()); + let prefix = offset.saturating_add(FIRST_RANKED); + let ranked = if offset == files.len() { + files.len() + } else if prefix <= files.len() / 8 && options.page_limit <= FIRST_RANKED { + files.select_nth_unstable_by(prefix, rank_with_ties); + files[..prefix].sort_unstable_by(rank_with_ties); + prefix + } else { + sort_with_buffer(&mut files, rank); + files.len() + }; + Self { + files, + ranked, + offset, + } + } + + pub(super) fn len(&self) -> usize { + self.files.len() - self.offset + } + + pub(super) fn is_empty(&self) -> bool { + self.len() == 0 + } + + pub(super) fn batch(&mut self, range: Range) -> &[&'a FileItem] { + let range = self.offset + range.start..self.offset + range.end; + if range.end > self.ranked { + self.files[self.ranked..].sort_unstable_by(rank_with_ties); + self.ranked = self.files.len(); + } + &self.files[range] + } +} /// Prefilter with a FilePath-constraint fallback: when constraints yield 0 /// files and the query had FilePath constraints, retry without them (the path @@ -17,7 +77,7 @@ pub(super) fn prefilter_with_filepath_retry<'a>( options: &GrepSearchOptions, arena: crate::simd_path::ArenaPtr, overflow_arena: crate::simd_path::ArenaPtr, -) -> (Vec<&'a FileItem>, usize) { +) -> (RankedFiles<'a>, usize) { let (files_to_search, filtered_file_count) = prefilter_files( files, constraints, @@ -52,7 +112,7 @@ pub(super) fn prefilter_with_filepath_retry<'a>( /// Single pass prefilter that doesn't involve file reading /// allocates only amount of memory required for storing references of the FileItems have to be /// opened for grepping unaviodably, in the worst case allocates N * memory if no prefilter needed -pub(crate) fn prefilter_files<'a>( +pub(super) fn prefilter_files<'a>( files: &'a [FileItem], constraints: &[Constraint<'_>], bigram_candidates: Option<&[u64]>, @@ -60,7 +120,7 @@ pub(crate) fn prefilter_files<'a>( options: &GrepSearchOptions, arena: crate::simd_path::ArenaPtr, overflow_arena: crate::simd_path::ArenaPtr, -) -> (Vec<&'a FileItem>, usize) { +) -> (RankedFiles<'a>, usize) { let max_file_size = options.max_file_size; let definition_only = options.classify_definitions == super::DefinitionMode::DefinitionOnly; let candidate_count = bigram_candidates.map(BigramFilter::count_candidates); @@ -90,7 +150,7 @@ pub(crate) fn prefilter_files<'a>( } // squeeze as much prefilters into a single loop as possible - let mut prefiltered: Vec<&FileItem> = match bigram_candidates { + let prefiltered: Vec<&FileItem> = match bigram_candidates { Some(candidates) => { let boundary = base_count.min(files.len()); let (indexed, tail) = files.split_at(boundary); @@ -179,20 +239,7 @@ pub(crate) fn prefilter_files<'a>( let total_count = prefiltered.len(); - sort_with_buffer(&mut prefiltered, |a, b| { - b.total_frecency_score() - .cmp(&a.total_frecency_score()) - .then(b.modified.cmp(&a.modified)) - }); - - if options.file_offset > 0 && options.file_offset < total_count { - let paginated = prefiltered.split_off(options.file_offset); - (paginated, total_count) - } else if options.file_offset >= total_count { - (Vec::new(), total_count) - } else { - (prefiltered, total_count) - } + (RankedFiles::new(prefiltered, options), total_count) } fn strip_file_path_constraint_if_present<'a>( @@ -220,6 +267,49 @@ mod tests { use crate::grep::DefinitionMode; use crate::simd_path::ArenaPtr; + #[test] + fn partial_ranking_matches_stable_sort_across_batch_boundaries() { + let files: Vec<_> = (0..100_000) + .map(|i| { + let mut file = FileItem::new_raw(0, 1, (i * 17 % 31) as u64, None, false); + file.access_frecency_score = (i * 13 % 7) as i16; + file + }) + .collect(); + let candidates: Vec<_> = files.iter().filter(|file| file.modified != 8).collect(); + let mut expected = candidates.clone(); + expected.sort_by(rank); + for (offset, limit) in [ + (0, 50), + (971, 50), + (50_000, 50), + (0, 10_000), + (usize::MAX, 50), + ] { + let mut list = RankedFiles::new( + candidates.clone(), + &GrepSearchOptions { + file_offset: offset, + page_limit: limit, + ..Default::default() + }, + ); + let expected = &expected[offset.min(expected.len())..]; + assert_eq!(list.len(), expected.len()); + let mut start = 0; + for end in [1, FIRST_RANKED - 1, FIRST_RANKED + 3, list.len()] { + let end = end.min(list.len()); + assert!( + list.batch(start..end) + .iter() + .zip(&expected[start..end]) + .all(|(a, b)| std::ptr::eq(*a, *b)) + ); + start = end; + } + } + } + #[test] fn definition_only_prefilters_before_sorting_and_offsets() { let files: Vec<_> = (0..70) @@ -261,7 +351,7 @@ mod tests { file_offset: offset, ..Default::default() }; - let (result, count) = prefilter_files( + let (mut result, count) = prefilter_files( &files, &[], candidates, @@ -272,7 +362,11 @@ mod tests { ); assert_eq!(count, expected.len()); assert_eq!( - result.iter().map(|file| file.modified).collect::>(), + result + .batch(0..result.len()) + .iter() + .map(|file| file.modified) + .collect::>(), expected[offset.min(expected.len())..], ); }