fuzzy_nucleo: Fix out of range panic (#54371)

Finn Evers created 1 week ago

Closes ZED-6PK

The issue here was that we could hit cases where the amount of segments
< amount of CPUs, e.g. for 5 candidates and 4 CPUs, we would have 2
candidates per matcher, so for the fourth matcher, we would start
slicing at 3 * 2 = 6 > 5, which is out of bounds.

Instead, make it so that we distribute the candidates across all
matchers so that all matchers have either n or n + 1 candidates.

No release notes since this is only on Nightly.

Release Notes:

- N/A

Change summary

crates/fuzzy_nucleo/src/strings.rs | 109 +++++++++++++++++++++++++++++++
1 file changed, 106 insertions(+), 3 deletions(-)

Detailed changes

crates/fuzzy_nucleo/src/strings.rs

@@ -165,7 +165,8 @@ where
     };
 
     let num_cpus = executor.num_cpus().min(candidates.len());
-    let segment_size = candidates.len().div_ceil(num_cpus);
+    let base_size = candidates.len() / num_cpus;
+    let remainder = candidates.len() % num_cpus;
     let mut segment_results = (0..num_cpus)
         .map(|_| Vec::with_capacity(max_results.min(candidates.len())))
         .collect::<Vec<_>>();
@@ -182,8 +183,9 @@ where
             {
                 let query = &query;
                 scope.spawn(async move {
-                    let segment_start = segment_idx * segment_size;
-                    let segment_end = (segment_start + segment_size).min(candidates.len());
+                    let segment_start = segment_idx * base_size + segment_idx.min(remainder);
+                    let segment_end =
+                        (segment_idx + 1) * base_size + (segment_idx + 1).min(remainder);
 
                     match_string_helper(
                         &candidates[segment_start..segment_end],
@@ -738,4 +740,105 @@ mod tests {
             "no candidate contains 'xyzzy', so nothing should match"
         );
     }
+
+    #[gpui::test]
+    async fn test_segment_size_not_divisible_by_cpus(executor: BackgroundExecutor) {
+        executor.set_num_cpus(4);
+        let cs = candidates(&["alpha", "beta", "gamma", "delta", "epsilon"]);
+        let cancel = AtomicBool::new(false);
+        let results = match_strings_async(
+            &cs,
+            "a",
+            Case::Ignore,
+            LengthPenalty::Off,
+            10,
+            &cancel,
+            executor,
+        )
+        .await;
+        let matched: Vec<&str> = results.iter().map(|m| m.string.as_ref()).collect();
+        assert!(matched.contains(&"alpha"));
+        assert!(matched.contains(&"gamma"));
+        assert!(matched.contains(&"delta"));
+    }
+
+    #[gpui::test]
+    async fn test_segment_size_with_many_cpus_few_candidates(executor: BackgroundExecutor) {
+        executor.set_num_cpus(16);
+        let cs = candidates(&["one", "two", "three"]);
+        let cancel = AtomicBool::new(false);
+        let results = match_strings_async(
+            &cs,
+            "o",
+            Case::Ignore,
+            LengthPenalty::Off,
+            10,
+            &cancel,
+            executor,
+        )
+        .await;
+        let matched: Vec<&str> = results.iter().map(|m| m.string.as_ref()).collect();
+        assert!(matched.contains(&"one"));
+        assert!(matched.contains(&"two"));
+    }
+
+    #[gpui::test]
+    async fn test_segment_size_single_candidate(executor: BackgroundExecutor) {
+        executor.set_num_cpus(8);
+        let cs = candidates(&["lonely"]);
+        let cancel = AtomicBool::new(false);
+        let results = match_strings_async(
+            &cs,
+            "lone",
+            Case::Ignore,
+            LengthPenalty::Off,
+            10,
+            &cancel,
+            executor,
+        )
+        .await;
+        assert_eq!(results.len(), 1);
+        assert_eq!(results[0].string.as_ref(), "lonely");
+    }
+
+    #[gpui::test]
+    async fn test_segment_size_candidates_equal_cpus(executor: BackgroundExecutor) {
+        executor.set_num_cpus(4);
+        let cs = candidates(&["aaa", "bbb", "ccc", "ddd"]);
+        let cancel = AtomicBool::new(false);
+        let results = match_strings_async(
+            &cs,
+            "a",
+            Case::Ignore,
+            LengthPenalty::Off,
+            10,
+            &cancel,
+            executor,
+        )
+        .await;
+        assert_eq!(results.len(), 1);
+        assert_eq!(results[0].string.as_ref(), "aaa");
+    }
+
+    #[gpui::test]
+    async fn test_segment_size_candidates_one_more_than_cpus(executor: BackgroundExecutor) {
+        executor.set_num_cpus(3);
+        let cs = candidates(&["ant", "ape", "dog", "axe"]);
+        let cancel = AtomicBool::new(false);
+        let results = match_strings_async(
+            &cs,
+            "a",
+            Case::Ignore,
+            LengthPenalty::Off,
+            10,
+            &cancel,
+            executor,
+        )
+        .await;
+        let matched: Vec<&str> = results.iter().map(|m| m.string.as_ref()).collect();
+        assert!(matched.contains(&"ant"));
+        assert!(matched.contains(&"ape"));
+        assert!(matched.contains(&"axe"));
+        assert!(!matched.contains(&"dog"));
+    }
 }