editor: Parallelize find_all_matches

Piotr Osiewicz and Smit Barmase created

Co-authored-by: Smit Barmase <heysmitbarmase@gmail.com>

Change summary

crates/editor/src/items.rs | 136 ++++++++++++++++++++++++++++++---------
1 file changed, 105 insertions(+), 31 deletions(-)

Detailed changes

crates/editor/src/items.rs 🔗

@@ -11,7 +11,7 @@ use anyhow::{Context as _, Result, anyhow};
 use collections::{HashMap, HashSet};
 use file_icons::FileIcons;
 use fs::MTime;
-use futures::future::try_join_all;
+use futures::{channel::oneshot, future::try_join_all};
 use git::status::GitSummary;
 use gpui::{
     AnyElement, App, AsyncWindowContext, Context, Entity, EntityId, EventEmitter, Font,
@@ -22,22 +22,24 @@ use language::{
     SelectionGoal, proto::serialize_anchor as serialize_text_anchor,
 };
 use lsp::DiagnosticSeverity;
-use multi_buffer::{MultiBufferOffset, PathKey};
+use multi_buffer::{BufferOffset, MultiBufferOffset, PathKey};
 use project::{
     File, Project, ProjectItem as _, ProjectPath, lsp_store::FormatTrigger,
     project_settings::ProjectSettings, search::SearchQuery,
 };
+use rope::TextSummary;
 use rpc::proto::{self, update_view};
 use settings::Settings;
 use std::{
     any::{Any, TypeId},
     borrow::Cow,
     cmp::{self, Ordering},
+    num::NonZeroU32,
     ops::Range,
     path::{Path, PathBuf},
     sync::Arc,
 };
-use text::{BufferId, BufferSnapshot, Selection};
+use text::{BufferId, BufferSnapshot, OffsetRangeExt, Selection};
 use ui::{IconDecorationKind, prelude::*};
 use util::{ResultExt, TryFutureExt, paths::PathExt, rel_path::RelPath};
 use workspace::item::{Dedup, ItemSettings, SerializableItem, TabContentParams};
@@ -1836,6 +1838,7 @@ impl SearchableItem for Editor {
                 ranges.iter().cloned().collect::<Vec<_>>()
             });
 
+        let executor = cx.background_executor().clone();
         cx.background_spawn(async move {
             let mut ranges = Vec::new();
 
@@ -1844,38 +1847,71 @@ impl SearchableItem for Editor {
             } else {
                 search_within_ranges
             };
-
+            let num_cpus = executor.num_cpus();
             for range in search_within_ranges {
                 for (search_buffer, search_range, deleted_hunk_anchor) in
                     buffer.range_to_buffer_ranges_with_deleted_hunks(range)
                 {
-                    ranges.extend(
-                        query
-                            .search(
-                                search_buffer,
-                                Some(search_range.start.0..search_range.end.0),
-                            )
-                            .await
-                            .into_iter()
-                            .filter_map(|match_range| {
-                                if let Some(deleted_hunk_anchor) = deleted_hunk_anchor {
-                                    let start = search_buffer
-                                        .anchor_after(search_range.start + match_range.start);
-                                    let end = search_buffer
-                                        .anchor_before(search_range.start + match_range.end);
-                                    Some(
-                                        deleted_hunk_anchor.with_diff_base_anchor(start)
-                                            ..deleted_hunk_anchor.with_diff_base_anchor(end),
-                                    )
-                                } else {
-                                    let start = search_buffer
-                                        .anchor_after(search_range.start + match_range.start);
-                                    let end = search_buffer
-                                        .anchor_before(search_range.start + match_range.end);
-                                    buffer.buffer_anchor_range_to_anchor_range(start..end)
-                                }
-                            }),
-                    );
+                    let query = query.clone();
+
+                    let mut results = Vec::new();
+                    executor
+                        .scoped(|scope| {
+                            for search_range in chunk_search_range(
+                                search_buffer.text.clone(),
+                                &query,
+                                num_cpus as u32,
+                                search_range,
+                            ) {
+                                let query = query.clone();
+                                let buffer = buffer.clone();
+
+                                let (tx, rx) = oneshot::channel();
+                                results.push(rx);
+                                scope.spawn(async move {
+                                    let chunk_result = query
+                                        .search(
+                                            search_buffer,
+                                            Some(search_range.start..search_range.end),
+                                        )
+                                        .await
+                                        .into_iter()
+                                        .filter_map(|match_range| {
+                                            if let Some(deleted_hunk_anchor) = deleted_hunk_anchor {
+                                                let start = search_buffer.anchor_after(
+                                                    search_range.start + match_range.start,
+                                                );
+                                                let end = search_buffer.anchor_before(
+                                                    search_range.start + match_range.end,
+                                                );
+                                                Some(
+                                                    deleted_hunk_anchor.with_diff_base_anchor(start)
+                                                        ..deleted_hunk_anchor
+                                                            .with_diff_base_anchor(end),
+                                                )
+                                            } else {
+                                                let start = search_buffer.anchor_after(
+                                                    search_range.start + match_range.start,
+                                                );
+                                                let end = search_buffer.anchor_before(
+                                                    search_range.start + match_range.end,
+                                                );
+                                                buffer
+                                                    .buffer_anchor_range_to_anchor_range(start..end)
+                                            }
+                                        })
+                                        .collect::<Vec<_>>();
+                                    _ = tx.send(chunk_result);
+                                });
+                            }
+                        })
+                        .await;
+
+                    for rx in results {
+                        if let Ok(results) = rx.await {
+                            ranges.extend(results.into_iter());
+                        }
+                    }
                 }
             }
 
@@ -2074,6 +2110,44 @@ fn deserialize_path_key(path_key: proto::PathKey) -> Option<PathKey> {
     })
 }
 
+fn chunk_search_range(
+    buffer: BufferSnapshot,
+    query: &SearchQuery,
+    num_cpus: u32,
+    initial_range: Range<BufferOffset>,
+) -> Box<dyn Iterator<Item = Range<usize>> + 'static> {
+    let range = initial_range.to_offset(&buffer);
+    let summary: TextSummary = buffer.text_summary_for_range(initial_range);
+    let num_chunks = if !query.is_regex() && !query.as_str().contains('\n') {
+        NonZeroU32::new(summary.lines.row.min(num_cpus))
+    } else {
+        NonZeroU32::new(1)
+    };
+
+    let Some(num_chunks) = num_chunks else {
+        return Box::new(std::iter::empty());
+    };
+
+    let mut chunk_start = range.start;
+    let rope = buffer.as_rope().clone();
+    let total_bytes = summary.len;
+    let average_chunk_length = total_bytes / (num_chunks.get() as usize);
+    Box::new(std::iter::from_fn(move || {
+        if chunk_start >= total_bytes {
+            return None;
+        }
+        let candidate_position = chunk_start + average_chunk_length;
+        let adjusted = rope.ceil_char_boundary(candidate_position);
+        let mut as_point = rope.offset_to_point(adjusted);
+        as_point.row += 1;
+        as_point.column = 0;
+        let end_offset = buffer.point_to_offset(as_point).min(total_bytes);
+        let ret = chunk_start..end_offset;
+        chunk_start = end_offset;
+        Some(ret)
+    }))
+}
+
 #[cfg(test)]
 mod tests {
     use crate::editor_tests::init_test;