perf(grep): rank the first file batch before searching

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);

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;
 

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,

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<usize>) -> &[&'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,15 +112,15 @@ 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 * <word> 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]>,
     base_count: usize,
     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::<Vec<_>>(),
+                        result
+                            .batch(0..result.len())
+                            .iter()
+                            .map(|file| file.modified)
+                            .collect::<Vec<_>>(),
                         expected[offset.min(expected.len())..],
                     );
                 }