multibuffer: Speed up anchor resolution by avoiding an excerpts seek (#53340)

Cole Miller and Anthony Eid created

Before, every creation of `AnchorSeekTarget` called
`MultiBufferSnapshot::buffer_for_path`, which seeks the excerpts sum
tree. Now we avoid that by

- looking up the buffer snapshot for the anchor's buffer in the
`buffers` map (keyed by `BufferId`)
- handling the case where the anchor's buffer doesn't exist at the
original path key with a separate `AnchorSeekTarget::Missing` variant

We also added a separate optimization to `AnchorSeekTarget::cmp`, to
skip the full `PathKey` comparison when the cursor's `PathKeyIndex` is
the same as the anchor's. This should help with singleton multibuffers.

Self-Review Checklist:

- [x] I've reviewed my own diff for quality, security, and reliability
- [x] Unsafe blocks (if any) have justifying comments
- [x] The content is consistent with the [UI/UX
checklist](https://github.com/zed-industries/zed/blob/main/CONTRIBUTING.md#uiux-checklist)
- [x] Tests cover the new/changed behavior
- [x] Performance impact has been considered and is acceptable

Release Notes:

- N/A

---------

Co-authored-by: Anthony Eid <hello@anthonyeid.me>

Change summary

crates/multi_buffer/src/anchor.rs       | 77 ++++++++++++++++++--------
crates/multi_buffer/src/multi_buffer.rs | 44 ++++++++-------
2 files changed, 75 insertions(+), 46 deletions(-)

Detailed changes

crates/multi_buffer/src/anchor.rs 🔗

@@ -34,21 +34,28 @@ pub enum Anchor {
     Max,
 }
 
-pub(crate) enum AnchorSeekTarget {
+pub(crate) enum AnchorSeekTarget<'a> {
+    // buffer no longer exists at its original path key in the multibuffer
+    Missing {
+        path_key: &'a PathKey,
+    },
+    // we have excerpts for the buffer at the expected path key
     Excerpt {
-        path_key: PathKey,
-        anchor: ExcerptAnchor,
-        // None when the buffer no longer exists in the multibuffer
-        snapshot: Option<BufferSnapshot>,
+        path_key: &'a PathKey,
+        path_key_index: PathKeyIndex,
+        anchor: text::Anchor,
+        snapshot: &'a BufferSnapshot,
     },
+    // no excerpts and it's a min or max anchor
     Empty,
 }
 
-impl std::fmt::Debug for AnchorSeekTarget {
+impl std::fmt::Debug for AnchorSeekTarget<'_> {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
             Self::Excerpt {
                 path_key,
+                path_key_index: _,
                 anchor,
                 snapshot: _,
             } => f
@@ -56,7 +63,11 @@ impl std::fmt::Debug for AnchorSeekTarget {
                 .field("path_key", path_key)
                 .field("anchor", anchor)
                 .finish(),
-            Self::Empty => write!(f, "Empty"),
+            Self::Missing { path_key } => f
+                .debug_struct("Missing")
+                .field("path_key", path_key)
+                .finish(),
+            Self::Empty => f.debug_struct("Empty").finish(),
         }
     }
 }
@@ -110,15 +121,16 @@ impl ExcerptAnchor {
             return self.text_anchor.buffer_id.cmp(&other.text_anchor.buffer_id);
         }
 
-        let Some(buffer) = snapshot.buffer_for_path(&self_path_key) else {
-            return Ordering::Equal;
-        };
-        // Comparing two anchors into buffer A that formerly existed at path P,
-        // when path P has since been reused for a different buffer B
-        if buffer.remote_id() != self.text_anchor.buffer_id {
+        // two anchors into the same buffer at the same path
+        // TODO(cole) buffer_for_path is slow
+        let Some(buffer) = snapshot
+            .buffer_for_path(&self_path_key)
+            .filter(|buffer| buffer.remote_id() == self.text_anchor.buffer_id)
+        else {
+            // buffer no longer exists at the original path (which may have been reused for a different buffer),
+            // so no way to compare the anchors
             return Ordering::Equal;
         };
-        assert_eq!(self.text_anchor.buffer_id, buffer.remote_id());
         let text_cmp = self.text_anchor().cmp(&other.text_anchor(), buffer);
         if text_cmp != Ordering::Equal {
             return text_cmp;
@@ -234,21 +246,33 @@ impl ExcerptAnchor {
                 .is_ge()
     }
 
-    pub(crate) fn seek_target(&self, snapshot: &MultiBufferSnapshot) -> AnchorSeekTarget {
+    pub(crate) fn seek_target<'a>(
+        &self,
+        snapshot: &'a MultiBufferSnapshot,
+    ) -> AnchorSeekTarget<'a> {
         self.try_seek_target(snapshot)
             .expect("anchor is from different multi-buffer")
     }
 
-    pub(crate) fn try_seek_target(
+    pub(crate) fn try_seek_target<'a>(
         &self,
-        snapshot: &MultiBufferSnapshot,
-    ) -> Option<AnchorSeekTarget> {
+        snapshot: &'a MultiBufferSnapshot,
+    ) -> Option<AnchorSeekTarget<'a>> {
         let path_key = snapshot.try_path_for_anchor(*self)?;
-        let buffer = snapshot.buffer_for_path(&path_key).cloned();
+
+        let Some(state) = snapshot
+            .buffers
+            .get(&self.buffer_id())
+            .filter(|state| &state.path_key == path_key)
+        else {
+            return Some(AnchorSeekTarget::Missing { path_key });
+        };
+
         Some(AnchorSeekTarget::Excerpt {
             path_key,
-            anchor: *self,
-            snapshot: buffer,
+            path_key_index: self.path,
+            anchor: self.text_anchor(),
+            snapshot: &state.buffer_snapshot,
         })
     }
 }
@@ -372,7 +396,10 @@ impl Anchor {
         }
     }
 
-    pub(crate) fn seek_target(&self, snapshot: &MultiBufferSnapshot) -> AnchorSeekTarget {
+    pub(crate) fn seek_target<'a>(
+        &self,
+        snapshot: &'a MultiBufferSnapshot,
+    ) -> AnchorSeekTarget<'a> {
         let Some(excerpt_anchor) = self.to_excerpt_anchor(snapshot) else {
             return AnchorSeekTarget::Empty;
         };
@@ -406,10 +433,10 @@ impl Anchor {
         }
     }
 
-    pub(crate) fn try_seek_target(
+    pub(crate) fn try_seek_target<'a>(
         &self,
-        snapshot: &MultiBufferSnapshot,
-    ) -> Option<AnchorSeekTarget> {
+        snapshot: &'a MultiBufferSnapshot,
+    ) -> Option<AnchorSeekTarget<'a>> {
         let Some(excerpt_anchor) = self.to_excerpt_anchor(snapshot) else {
             return Some(AnchorSeekTarget::Empty);
         };

crates/multi_buffer/src/multi_buffer.rs 🔗

@@ -870,6 +870,7 @@ impl ExcerptRange<text::Anchor> {
 #[derive(Clone, Debug)]
 pub struct ExcerptSummary {
     path_key: PathKey,
+    path_key_index: Option<PathKeyIndex>,
     max_anchor: Option<text::Anchor>,
     widest_line_number: u32,
     text: MBTextSummary,
@@ -880,6 +881,7 @@ impl ExcerptSummary {
     pub fn min() -> Self {
         ExcerptSummary {
             path_key: PathKey::min(),
+            path_key_index: None,
             max_anchor: None,
             widest_line_number: 0,
             text: MBTextSummary::default(),
@@ -6386,11 +6388,11 @@ impl MultiBufferSnapshot {
         self.buffers.get(&id).map(|state| &state.buffer_snapshot)
     }
 
-    fn try_path_for_anchor(&self, anchor: ExcerptAnchor) -> Option<PathKey> {
-        self.path_keys_by_index.get(&anchor.path).cloned()
+    fn try_path_for_anchor(&self, anchor: ExcerptAnchor) -> Option<&PathKey> {
+        self.path_keys_by_index.get(&anchor.path)
     }
 
-    pub fn path_for_anchor(&self, anchor: ExcerptAnchor) -> PathKey {
+    pub fn path_for_anchor(&self, anchor: ExcerptAnchor) -> &PathKey {
         self.try_path_for_anchor(anchor)
             .expect("invalid anchor: path was never added to multibuffer")
     }
@@ -7327,6 +7329,7 @@ impl sum_tree::Item for Excerpt {
         }
         ExcerptSummary {
             path_key: self.path_key.clone(),
+            path_key_index: Some(self.path_key_index),
             max_anchor: Some(self.range.context.end),
             widest_line_number: self.max_buffer_row,
             text: text.into(),
@@ -7425,6 +7428,7 @@ impl sum_tree::ContextLessSummary for ExcerptSummary {
         );
 
         self.path_key = summary.path_key.clone();
+        self.path_key_index = summary.path_key_index;
         self.max_anchor = summary.max_anchor;
         self.text += summary.text;
         self.widest_line_number = cmp::max(self.widest_line_number, summary.widest_line_number);
@@ -7432,38 +7436,36 @@ impl sum_tree::ContextLessSummary for ExcerptSummary {
     }
 }
 
-impl sum_tree::SeekTarget<'_, ExcerptSummary, ExcerptSummary> for AnchorSeekTarget {
+impl sum_tree::SeekTarget<'_, ExcerptSummary, ExcerptSummary> for AnchorSeekTarget<'_> {
     fn cmp(
         &self,
         cursor_location: &ExcerptSummary,
         _cx: <ExcerptSummary as sum_tree::Summary>::Context<'_>,
     ) -> cmp::Ordering {
         match self {
+            AnchorSeekTarget::Missing { path_key } => {
+                // Want to end up after any excerpts for (a different buffer at) the original path
+                match Ord::cmp(*path_key, &cursor_location.path_key) {
+                    Ordering::Less => Ordering::Less,
+                    Ordering::Equal | Ordering::Greater => Ordering::Greater,
+                }
+            }
             AnchorSeekTarget::Excerpt {
                 path_key,
+                path_key_index,
                 anchor,
                 snapshot,
             } => {
-                let path_comparison = Ord::cmp(path_key, &cursor_location.path_key);
-                if path_comparison.is_ne() {
-                    path_comparison
-                } else if let Some(snapshot) = snapshot {
-                    if anchor.text_anchor.buffer_id != snapshot.remote_id() {
-                        Ordering::Greater
-                    } else if let Some(max_anchor) = cursor_location.max_anchor {
-                        debug_assert_eq!(max_anchor.buffer_id, snapshot.remote_id());
-                        anchor.text_anchor().cmp(&max_anchor, snapshot)
-                    } else {
-                        Ordering::Greater
-                    }
+                if Some(*path_key_index) != cursor_location.path_key_index {
+                    Ord::cmp(*path_key, &cursor_location.path_key)
+                } else if let Some(max_anchor) = cursor_location.max_anchor {
+                    debug_assert_eq!(max_anchor.buffer_id, snapshot.remote_id());
+                    anchor.cmp(&max_anchor, snapshot)
                 } else {
-                    // shouldn't happen because we expect this buffer not to have any excerpts
-                    // (otherwise snapshot would have been Some)
-                    Ordering::Equal
+                    Ordering::Greater
                 }
             }
-            // This should be dead code because Empty is only constructed for an empty snapshot
-            AnchorSeekTarget::Empty => Ordering::Equal,
+            AnchorSeekTarget::Empty => Ordering::Greater,
         }
     }
 }