Track already queried chunks in the editor

Kirill Bulatov created

Change summary

crates/editor/src/bracket_colorization.rs        | 24 +++++
crates/editor/src/editor.rs                      |  7 +
crates/language/src/buffer.rs                    | 73 +++++++++++++----
crates/language/src/buffer/row_chunk.rs          | 27 ++++++
crates/multi_buffer/src/multi_buffer.rs          | 11 +-
crates/project/src/lsp_store.rs                  | 32 ++-----
crates/project/src/lsp_store/inlay_hint_cache.rs |  9 +
7 files changed, 127 insertions(+), 56 deletions(-)

Detailed changes

crates/editor/src/bracket_colorization.rs 🔗

@@ -12,10 +12,14 @@ impl Editor {
             return;
         }
 
+        if invalidate {
+            self.fetched_tree_sitter_chunks.clear();
+        }
+
         let multi_buffer_snapshot = self.buffer().read(cx).snapshot(cx);
         let bracket_matches = self.visible_excerpts(cx).into_iter().fold(
             HashMap::default(),
-            |mut acc, (excerpt_id, (buffer, _, buffer_range))| {
+            |mut acc, (excerpt_id, (buffer, buffer_version, buffer_range))| {
                 let buffer_snapshot = buffer.read(cx).snapshot();
                 if language_settings::language_settings(
                     buffer_snapshot.language().map(|language| language.name()),
@@ -24,9 +28,24 @@ impl Editor {
                 )
                 .colorize_brackets
                 {
+                    let fetched_chunks = self
+                        .fetched_tree_sitter_chunks
+                        .entry(excerpt_id)
+                        .or_default();
+
                     for (depth, open_range, close_range) in buffer_snapshot
-                        .bracket_ranges(buffer_range.start..buffer_range.end)
+                        .fetch_bracket_ranges(
+                            buffer_range.start..buffer_range.end,
+                            Some((&buffer_version, fetched_chunks)),
+                        )
                         .into_iter()
+                        .flat_map(|(chunk_range, pairs)| {
+                            if fetched_chunks.insert(chunk_range) {
+                                pairs
+                            } else {
+                                Vec::new()
+                            }
+                        })
                         .filter_map(|pair| {
                             let buffer_open_range = buffer_snapshot
                                 .anchor_before(pair.open_range.start)
@@ -63,7 +82,6 @@ impl Editor {
             self.clear_highlights::<RainbowBracketHighlight>(cx);
         }
 
-        // todo! can we skip the re-highlighting entirely, if it's not adding anything on top?
         let editor_background = cx.theme().colors().editor_background;
         for (depth, bracket_highlights) in bracket_matches {
             let bracket_color = cx.theme().accents().color_for_index(depth as u32);

crates/editor/src/editor.rs 🔗

@@ -1199,6 +1199,7 @@ pub struct Editor {
     folding_newlines: Task<()>,
     pub lookup_key: Option<Box<dyn Any + Send + Sync>>,
     applicable_language_settings: HashMap<Option<LanguageName>, LanguageSettings>,
+    fetched_tree_sitter_chunks: HashMap<ExcerptId, HashSet<Range<BufferRow>>>,
 }
 
 fn debounce_value(debounce_ms: u64) -> Option<Duration> {
@@ -2297,6 +2298,7 @@ impl Editor {
             folding_newlines: Task::ready(()),
             lookup_key: None,
             applicable_language_settings: HashMap::default(),
+            fetched_tree_sitter_chunks: HashMap::default(),
         };
 
         if is_minimap {
@@ -3248,7 +3250,6 @@ impl Editor {
             refresh_linked_ranges(self, window, cx);
 
             self.refresh_selected_text_highlights(false, window, cx);
-            self.colorize_brackets(false, cx);
             self.refresh_matching_bracket_highlights(window, cx);
             self.update_visible_edit_prediction(window, cx);
             self.edit_prediction_requires_modifier_in_indent_conflict = true;
@@ -21128,6 +21129,9 @@ impl Editor {
             multi_buffer::Event::ExcerptsExpanded { ids } => {
                 self.refresh_inlay_hints(InlayHintRefreshReason::NewLinesShown, cx);
                 self.refresh_document_highlights(cx);
+                for id in ids {
+                    self.fetched_tree_sitter_chunks.remove(id);
+                }
                 self.colorize_brackets(false, cx);
                 cx.emit(EditorEvent::ExcerptsExpanded { ids: ids.clone() })
             }
@@ -22429,7 +22433,6 @@ fn insert_extra_newline_tree_sitter(buffer: &MultiBufferSnapshot, range: Range<u
 
         for pair in buffer
             .all_bracket_ranges(range.clone())
-            .into_iter()
             .filter(move |pair| {
                 pair.open_range.start <= range.start && pair.close_range.end >= range.end
             })

crates/language/src/buffer.rs 🔗

@@ -21,9 +21,9 @@ pub use crate::{
     proto,
 };
 use anyhow::{Context as _, Result};
-use clock::Lamport;
 pub use clock::ReplicaId;
-use collections::HashMap;
+use clock::{Global, Lamport};
+use collections::{HashMap, HashSet};
 use fs::MTime;
 use futures::channel::oneshot;
 use gpui::{
@@ -4145,33 +4145,65 @@ impl BufferSnapshot {
         self.syntax.matches(range, self, query)
     }
 
-    /// Returns all bracket pairs that intersect with the range given.
+    /// Finds all [`RowChunks`] applicable to the given range, then returns all bracket pairs that intersect with those chunks.
+    /// Hence, may return more bracket pairs than the range contains.
     ///
-    /// The resulting collection is not ordered.
-    fn fetch_bracket_ranges(&self, range: Range<usize>) -> Vec<BracketMatch> {
+    /// Will omit known chunks.
+    /// The resulting bracket match collections are not ordered.
+    pub fn fetch_bracket_ranges(
+        &self,
+        range: Range<usize>,
+        known_chunks: Option<(&Global, &HashSet<Range<BufferRow>>)>,
+    ) -> HashMap<Range<BufferRow>, Vec<BracketMatch>> {
         let mut tree_sitter_data = self.latest_tree_sitter_data().clone();
+
+        let known_chunks = match known_chunks {
+            Some((known_version, known_chunks)) => {
+                if !tree_sitter_data
+                    .chunks
+                    .version()
+                    .changed_since(known_version)
+                {
+                    known_chunks.clone()
+                } else {
+                    HashSet::default()
+                }
+            }
+            None => HashSet::default(),
+        };
+
         let mut new_bracket_matches = HashMap::default();
-        let mut all_bracket_matches = Vec::new();
+        let mut all_bracket_matches = HashMap::default();
+
         for chunk in tree_sitter_data
             .chunks
             .applicable_chunks(&[self.anchor_before(range.start)..self.anchor_after(range.end)])
         {
-            let chunk_brackets = tree_sitter_data.brackets_by_chunks.remove(chunk.id);
-            let bracket_matches = match chunk_brackets {
+            if known_chunks.contains(&chunk.row_range()) {
+                continue;
+            }
+            let Some(chunk_range) = tree_sitter_data.chunks.chunk_range(chunk) else {
+                continue;
+            };
+            let chunk_range = chunk_range.to_offset(&tree_sitter_data.chunks.snapshot);
+
+            let bracket_matches = match tree_sitter_data.brackets_by_chunks[chunk.id].take() {
                 Some(cached_brackets) => cached_brackets,
                 None => {
-                    let mut matches = self.syntax.matches(range.clone(), &self.text, |grammar| {
-                        grammar.brackets_config.as_ref().map(|c| &c.query)
-                    });
+                    let mut matches =
+                        self.syntax
+                            .matches(chunk_range.clone(), &self.text, |grammar| {
+                                grammar.brackets_config.as_ref().map(|c| &c.query)
+                            });
                     let configs = matches
                         .grammars()
                         .iter()
                         .map(|grammar| grammar.brackets_config.as_ref().unwrap())
                         .collect::<Vec<_>>();
 
-                    // todo!
+                    // todo! this seems like a wrong parameter: instead, use chunk range, `Range<BufferRow>`, as a key part + add bracket_id that will be used for each bracket
                     let mut depth = 0;
-                    let range = range.clone();
+                    let chunk_range = chunk_range.clone();
                     let new_matches = iter::from_fn(move || {
                         while let Some(mat) = matches.peek() {
                             let mut open = None;
@@ -4193,7 +4225,7 @@ impl BufferSnapshot {
                             };
 
                             let bracket_range = open_range.start..=close_range.end;
-                            if !bracket_range.overlaps(&range) {
+                            if !bracket_range.overlaps(&chunk_range) {
                                 continue;
                             }
 
@@ -4214,7 +4246,7 @@ impl BufferSnapshot {
                     new_matches
                 }
             };
-            all_bracket_matches.extend(bracket_matches);
+            all_bracket_matches.insert(chunk.row_range(), bracket_matches);
         }
 
         let mut latest_tree_sitter_data = self.latest_tree_sitter_data();
@@ -4241,8 +4273,14 @@ impl BufferSnapshot {
         tree_sitter_data
     }
 
-    pub fn all_bracket_ranges(&self, range: Range<usize>) -> Vec<BracketMatch> {
-        self.fetch_bracket_ranges(range)
+    pub fn all_bracket_ranges(&self, range: Range<usize>) -> impl Iterator<Item = BracketMatch> {
+        self.fetch_bracket_ranges(range.clone(), None)
+            .into_values()
+            .flatten()
+            .filter(move |bracket_match| {
+                let bracket_range = bracket_match.open_range.start..=bracket_match.close_range.end;
+                bracket_range.overlaps(&range)
+            })
     }
 
     /// Returns bracket range pairs overlapping or adjacent to `range`
@@ -4253,7 +4291,6 @@ impl BufferSnapshot {
         // Find bracket pairs that *inclusively* contain the given range.
         let range = range.start.to_previous_offset(self)..range.end.to_next_offset(self);
         self.all_bracket_ranges(range)
-            .into_iter()
             .filter(|pair| !pair.newline_only)
     }
 

crates/language/src/buffer/row_chunk.rs 🔗

@@ -4,7 +4,7 @@
 use std::{ops::Range, sync::Arc};
 
 use clock::Global;
-use text::OffsetRangeExt as _;
+use text::{Anchor, OffsetRangeExt as _, Point};
 
 use crate::BufferRow;
 
@@ -72,7 +72,7 @@ impl RowChunks {
             .filter(move |chunk| -> bool {
                 // Be lenient and yield multiple chunks if they "touch" the exclusive part of the range.
                 // This will result in LSP hints [re-]queried for more ranges, but also more hints already visible when scrolling around.
-                let chunk_range = chunk.start..=chunk.end_exclusive;
+                let chunk_range = chunk.row_range();
                 row_ranges.iter().any(|row_range| {
                     chunk_range.contains(&row_range.start())
                         || chunk_range.contains(&row_range.end())
@@ -80,6 +80,23 @@ impl RowChunks {
             })
             .copied()
     }
+
+    pub fn chunk_range(&self, chunk: RowChunk) -> Option<Range<Anchor>> {
+        if !self.chunks.contains(&chunk) {
+            return None;
+        }
+
+        let start = Point::new(chunk.start, 0);
+        let end = if self.chunks.last() == Some(&chunk) {
+            Point::new(
+                chunk.end_exclusive,
+                self.snapshot.line_len(chunk.end_exclusive),
+            )
+        } else {
+            Point::new(chunk.end_exclusive, 0)
+        };
+        Some(self.snapshot.anchor_before(start)..self.snapshot.anchor_after(end))
+    }
 }
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
@@ -88,3 +105,9 @@ pub struct RowChunk {
     pub start: BufferRow,
     pub end_exclusive: BufferRow,
 }
+
+impl RowChunk {
+    pub fn row_range(&self) -> Range<BufferRow> {
+        self.start..self.end_exclusive
+    }
+}

crates/multi_buffer/src/multi_buffer.rs 🔗

@@ -21,11 +21,10 @@ use gpui::{App, Context, Entity, EntityId, EventEmitter};
 use itertools::Itertools;
 use language::{
     AutoindentMode, BracketMatch, Buffer, BufferChunks, BufferRow, BufferSnapshot, Capability,
-    CharClassifier, CharKind, CharScopeContext, Chunk, CursorShape, DiagnosticEntryRef, DiskState, File,
-    IndentGuideSettings, IndentSize,
-    Language, LanguageScope, OffsetRangeExt, OffsetUtf16, Outline, OutlineItem, Point, PointUtf16,
-    Selection, TextDimension, TextObject, ToOffset as _, ToPoint as _, TransactionId,
-    TreeSitterOptions, Unclipped,
+    CharClassifier, CharKind, CharScopeContext, Chunk, CursorShape, DiagnosticEntryRef, DiskState,
+    File, IndentGuideSettings, IndentSize, Language, LanguageScope, OffsetRangeExt, OffsetUtf16,
+    Outline, OutlineItem, Point, PointUtf16, Selection, TextDimension, TextObject, ToOffset as _,
+    ToPoint as _, TransactionId, TreeSitterOptions, Unclipped,
     language_settings::{LanguageSettings, language_settings},
 };
 
@@ -4896,7 +4895,7 @@ impl MultiBufferSnapshot {
             .map(|matches_iter| matches_iter.map(BracketMatch::bracket_ranges))
     }
 
-    pub fn bracket_matches<T: ToOffset>(
+    fn bracket_matches<T: ToOffset>(
         &self,
         range: Range<T>,
     ) -> Option<impl Iterator<Item = BracketMatch> + '_> {

crates/project/src/lsp_store.rs 🔗

@@ -117,7 +117,7 @@ use std::{
     time::{Duration, Instant},
 };
 use sum_tree::Dimensions;
-use text::{Anchor, BufferId, LineEnding, OffsetRangeExt, Point, ToPoint as _};
+use text::{Anchor, BufferId, LineEnding, OffsetRangeExt, ToPoint as _};
 
 use util::{
     ConnectionResult, ResultExt as _, debug_panic, defer, maybe, merge_json_value_into,
@@ -6624,7 +6624,7 @@ impl LspStore {
         self.latest_lsp_data(buffer, cx)
             .inlay_hints
             .applicable_chunks(ranges)
-            .map(|chunk| chunk.start..chunk.end_exclusive)
+            .map(|chunk| chunk.row_range())
             .collect()
     }
 
@@ -6647,7 +6647,6 @@ impl LspStore {
         known_chunks: Option<(clock::Global, HashSet<Range<BufferRow>>)>,
         cx: &mut Context<Self>,
     ) -> HashMap<Range<BufferRow>, Task<Result<CacheInlayHints>>> {
-        let buffer_snapshot = buffer.read(cx).snapshot();
         let next_hint_id = self.next_hint_id.clone();
         let lsp_data = self.latest_lsp_data(&buffer, cx);
         let mut lsp_refresh_requested = false;
@@ -6675,14 +6674,12 @@ impl LspStore {
         let mut ranges_to_query = None;
         let applicable_chunks = existing_inlay_hints
             .applicable_chunks(ranges.as_slice())
-            .filter(|chunk| !known_chunks.contains(&(chunk.start..chunk.end_exclusive)))
+            .filter(|chunk| !known_chunks.contains(&chunk.row_range()))
             .collect::<Vec<_>>();
         if applicable_chunks.is_empty() {
             return HashMap::default();
         }
 
-        let last_chunk_number = existing_inlay_hints.buffer_chunks_len() - 1;
-
         for row_chunk in applicable_chunks {
             match (
                 existing_inlay_hints
@@ -6696,19 +6693,12 @@ impl LspStore {
                     .cloned(),
             ) {
                 (None, None) => {
-                    let end = if last_chunk_number == row_chunk.id {
-                        Point::new(
-                            row_chunk.end_exclusive,
-                            buffer_snapshot.line_len(row_chunk.end_exclusive),
-                        )
-                    } else {
-                        Point::new(row_chunk.end_exclusive, 0)
+                    let Some(chunk_range) = existing_inlay_hints.chunk_range(row_chunk) else {
+                        continue;
                     };
-                    ranges_to_query.get_or_insert_with(Vec::new).push((
-                        row_chunk,
-                        buffer_snapshot.anchor_before(Point::new(row_chunk.start, 0))
-                            ..buffer_snapshot.anchor_after(end),
-                    ));
+                    ranges_to_query
+                        .get_or_insert_with(Vec::new)
+                        .push((row_chunk, chunk_range));
                 }
                 (None, Some(fetched_hints)) => hint_fetch_tasks.push((row_chunk, fetched_hints)),
                 (Some(cached_hints), None) => {
@@ -6716,7 +6706,7 @@ impl LspStore {
                         if for_server.is_none_or(|for_server| for_server == server_id) {
                             cached_inlay_hints
                                 .get_or_insert_with(HashMap::default)
-                                .entry(row_chunk.start..row_chunk.end_exclusive)
+                                .entry(row_chunk.row_range())
                                 .or_insert_with(HashMap::default)
                                 .entry(server_id)
                                 .or_insert_with(Vec::new)
@@ -6730,7 +6720,7 @@ impl LspStore {
                         if for_server.is_none_or(|for_server| for_server == server_id) {
                             cached_inlay_hints
                                 .get_or_insert_with(HashMap::default)
-                                .entry(row_chunk.start..row_chunk.end_exclusive)
+                                .entry(row_chunk.row_range())
                                 .or_insert_with(HashMap::default)
                                 .entry(server_id)
                                 .or_insert_with(Vec::new)
@@ -6817,7 +6807,7 @@ impl LspStore {
                 .map(|(row_chunk, hints)| (row_chunk, Task::ready(Ok(hints))))
                 .chain(hint_fetch_tasks.into_iter().map(|(chunk, hints_fetch)| {
                     (
-                        chunk.start..chunk.end_exclusive,
+                        chunk.row_range(),
                         cx.spawn(async move |_, _| {
                             hints_fetch.await.map_err(|e| {
                                 if e.error_code() != ErrorCode::Internal {

crates/project/src/lsp_store/inlay_hint_cache.rs 🔗

@@ -8,6 +8,7 @@ use language::{
     row_chunk::{RowChunk, RowChunks},
 };
 use lsp::LanguageServerId;
+use text::Anchor;
 
 use crate::{InlayHint, InlayId};
 
@@ -182,10 +183,6 @@ impl BufferInlayHints {
         Some(hint)
     }
 
-    pub fn buffer_chunks_len(&self) -> usize {
-        self.chunks.len()
-    }
-
     pub(crate) fn invalidate_for_server_refresh(
         &mut self,
         for_server: LanguageServerId,
@@ -229,4 +226,8 @@ impl BufferInlayHints {
             }
         }
     }
+
+    pub fn chunk_range(&self, chunk: RowChunk) -> Option<Range<Anchor>> {
+        self.chunks.chunk_range(chunk)
+    }
 }