Introduce MultiBuffer::stream_excerpts_with_context_lines

This allows us to push excerpts in a streaming fashion without blocking
the main thread.

Cargo.lock

@@ -5510,6 +5510,7 @@ dependencies = [
  "anyhow",
  "collections",
  "editor",
+ "futures 0.3.25",
  "gpui",
  "language",
  "log",

crates/editor/src/multi_buffer.rs

@@ -4,6 +4,7 @@ pub use anchor::{Anchor, AnchorRangeExt};
 use anyhow::Result;
 use clock::ReplicaId;
 use collections::{BTreeMap, Bound, HashMap, HashSet};
+use futures::{channel::mpsc, SinkExt};
 use git::diff::DiffHunk;
 use gpui::{AppContext, Entity, ModelContext, ModelHandle, Task};
 pub use language::Completion;
@@ -764,6 +765,63 @@ impl MultiBuffer {
         None
     }
 
+    pub fn stream_excerpts_with_context_lines(
+        &mut self,
+        excerpts: Vec<(ModelHandle<Buffer>, Vec<Range<text::Anchor>>)>,
+        context_line_count: u32,
+        cx: &mut ModelContext<Self>,
+    ) -> (Task<()>, mpsc::Receiver<Range<Anchor>>) {
+        let (mut tx, rx) = mpsc::channel(256);
+        let task = cx.spawn(|this, mut cx| async move {
+            for (buffer, ranges) in excerpts {
+                let buffer_id = buffer.id();
+                let buffer_snapshot = buffer.read_with(&cx, |buffer, _| buffer.snapshot());
+
+                let mut excerpt_ranges = Vec::new();
+                let mut range_counts = Vec::new();
+                cx.background()
+                    .scoped(|scope| {
+                        scope.spawn(async {
+                            let (ranges, counts) =
+                                build_excerpt_ranges(&buffer_snapshot, &ranges, context_line_count);
+                            excerpt_ranges = ranges;
+                            range_counts = counts;
+                        });
+                    })
+                    .await;
+
+                let mut ranges = ranges.into_iter();
+                let mut range_counts = range_counts.into_iter();
+                for excerpt_ranges in excerpt_ranges.chunks(100) {
+                    let excerpt_ids = this.update(&mut cx, |this, cx| {
+                        this.push_excerpts(buffer.clone(), excerpt_ranges.iter().cloned(), cx)
+                    });
+
+                    for (excerpt_id, range_count) in
+                        excerpt_ids.into_iter().zip(range_counts.by_ref())
+                    {
+                        for range in ranges.by_ref().take(range_count) {
+                            let start = Anchor {
+                                buffer_id: Some(buffer_id),
+                                excerpt_id: excerpt_id.clone(),
+                                text_anchor: range.start,
+                            };
+                            let end = Anchor {
+                                buffer_id: Some(buffer_id),
+                                excerpt_id: excerpt_id.clone(),
+                                text_anchor: range.end,
+                            };
+                            if tx.send(start..end).await.is_err() {
+                                break;
+                            }
+                        }
+                    }
+                }
+            }
+        });
+        (task, rx)
+    }
+
     pub fn push_excerpts<O>(
         &mut self,
         buffer: ModelHandle<Buffer>,
@@ -788,39 +846,8 @@ impl MultiBuffer {
     {
         let buffer_id = buffer.id();
         let buffer_snapshot = buffer.read(cx).snapshot();
-        let max_point = buffer_snapshot.max_point();
-
-        let mut range_counts = Vec::new();
-        let mut excerpt_ranges = Vec::new();
-        let mut range_iter = ranges
-            .iter()
-            .map(|range| {
-                range.start.to_point(&buffer_snapshot)..range.end.to_point(&buffer_snapshot)
-            })
-            .peekable();
-        while let Some(range) = range_iter.next() {
-            let excerpt_start = Point::new(range.start.row.saturating_sub(context_line_count), 0);
-            let mut excerpt_end =
-                Point::new(range.end.row + 1 + context_line_count, 0).min(max_point);
-            let mut ranges_in_excerpt = 1;
-
-            while let Some(next_range) = range_iter.peek() {
-                if next_range.start.row <= excerpt_end.row + context_line_count {
-                    excerpt_end =
-                        Point::new(next_range.end.row + 1 + context_line_count, 0).min(max_point);
-                    ranges_in_excerpt += 1;
-                    range_iter.next();
-                } else {
-                    break;
-                }
-            }
-
-            excerpt_ranges.push(ExcerptRange {
-                context: excerpt_start..excerpt_end,
-                primary: Some(range),
-            });
-            range_counts.push(ranges_in_excerpt);
-        }
+        let (excerpt_ranges, range_counts) =
+            build_excerpt_ranges(&buffer_snapshot, &ranges, context_line_count);
 
         let excerpt_ids = self.push_excerpts(buffer, excerpt_ranges, cx);
 
@@ -3605,6 +3632,47 @@ impl ToPointUtf16 for PointUtf16 {
     }
 }
 
+fn build_excerpt_ranges<T>(
+    buffer: &BufferSnapshot,
+    ranges: &[Range<T>],
+    context_line_count: u32,
+) -> (Vec<ExcerptRange<Point>>, Vec<usize>)
+where
+    T: text::ToPoint,
+{
+    let max_point = buffer.max_point();
+    let mut range_counts = Vec::new();
+    let mut excerpt_ranges = Vec::new();
+    let mut range_iter = ranges
+        .iter()
+        .map(|range| range.start.to_point(buffer)..range.end.to_point(buffer))
+        .peekable();
+    while let Some(range) = range_iter.next() {
+        let excerpt_start = Point::new(range.start.row.saturating_sub(context_line_count), 0);
+        let mut excerpt_end = Point::new(range.end.row + 1 + context_line_count, 0).min(max_point);
+        let mut ranges_in_excerpt = 1;
+
+        while let Some(next_range) = range_iter.peek() {
+            if next_range.start.row <= excerpt_end.row + context_line_count {
+                excerpt_end =
+                    Point::new(next_range.end.row + 1 + context_line_count, 0).min(max_point);
+                ranges_in_excerpt += 1;
+                range_iter.next();
+            } else {
+                break;
+            }
+        }
+
+        excerpt_ranges.push(ExcerptRange {
+            context: excerpt_start..excerpt_end,
+            primary: Some(range),
+        });
+        range_counts.push(ranges_in_excerpt);
+    }
+
+    (excerpt_ranges, range_counts)
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

crates/search/Cargo.toml

@@ -19,6 +19,7 @@ theme = { path = "../theme" }
 util = { path = "../util" }
 workspace = { path = "../workspace" }
 anyhow = "1.0"
+futures = "0.3"
 log = { version = "0.4.16", features = ["kv_unstable_serde"] }
 postage = { version = "0.4.1", features = ["futures-traits"] }
 serde = { version = "1.0", features = ["derive", "rc"] }

crates/search/src/project_search.rs

@@ -7,6 +7,7 @@ use editor::{
     items::active_match_index, scroll::autoscroll::Autoscroll, Anchor, Editor, MultiBuffer,
     SelectAll, MAX_TAB_TITLE_LEN,
 };
+use futures::StreamExt;
 use gpui::{
     actions, elements::*, platform::CursorStyle, Action, AnyViewHandle, AppContext, ElementBox,
     Entity, ModelContext, ModelHandle, MouseButton, MutableAppContext, RenderContext, Subscription,
@@ -129,31 +130,23 @@ impl ProjectSearch {
             let matches = search.await.log_err()?;
             let this = this.upgrade(&cx)?;
             let mut matches = matches.into_iter().collect::<Vec<_>>();
-            this.update(&mut cx, |this, cx| {
+            let (_rebuild, mut match_ranges) = this.update(&mut cx, |this, cx| {
                 this.match_ranges.clear();
                 matches.sort_by_key(|(buffer, _)| buffer.read(cx).file().map(|file| file.path()));
-                this.excerpts.update(cx, |excerpts, cx| excerpts.clear(cx));
+                this.excerpts.update(cx, |excerpts, cx| {
+                    excerpts.clear(cx);
+                    excerpts.stream_excerpts_with_context_lines(matches, 1, cx)
+                })
             });
 
-            for matches_chunk in matches.chunks(100) {
+            while let Some(match_range) = match_ranges.next().await {
                 this.update(&mut cx, |this, cx| {
-                    this.excerpts.update(cx, |excerpts, cx| {
-                        for (buffer, buffer_matches) in matches_chunk {
-                            let ranges_to_highlight = excerpts.push_excerpts_with_context_lines(
-                                buffer.clone(),
-                                buffer_matches.clone(),
-                                1,
-                                cx,
-                            );
-                            this.match_ranges.extend(ranges_to_highlight);
-                        }
-                    });
-
+                    this.match_ranges.push(match_range);
+                    while let Ok(Some(match_range)) = match_ranges.try_next() {
+                        this.match_ranges.push(match_range);
+                    }
                     cx.notify();
                 });
-
-                // Don't starve the main thread when adding lots of excerpts.
-                smol::future::yield_now().await;
             }
 
             this.update(&mut cx, |this, cx| {