Fix randomized test failures on `BlockMap` with excerpt headers

Antonio Scandurra and Nathan Sobo created

Co-Authored-By: Nathan Sobo <nathan@zed.dev>

Change summary

crates/editor/src/display_map/block_map.rs |  86 ++++++------
crates/editor/src/editor.rs                |  12 +
crates/editor/src/multi_buffer.rs          | 156 +++++++++++++++--------
3 files changed, 150 insertions(+), 104 deletions(-)

Detailed changes

crates/editor/src/display_map/block_map.rs 🔗

@@ -1,6 +1,6 @@
 use super::wrap_map::{self, WrapEdit, WrapPoint, WrapSnapshot};
 use crate::{Anchor, ToPoint as _};
-use collections::{HashMap, HashSet};
+use collections::{Bound, HashMap, HashSet};
 use gpui::{AppContext, ElementBox};
 use language::{BufferSnapshot, Chunk};
 use parking_lot::Mutex;
@@ -156,16 +156,22 @@ pub struct BlockBufferRows<'a> {
 
 impl BlockMap {
     pub fn new(wrap_snapshot: WrapSnapshot, excerpt_header_height: u8) -> Self {
-        Self {
+        let row_count = wrap_snapshot.max_point().row() + 1;
+        let map = Self {
             next_block_id: AtomicUsize::new(0),
             blocks: Vec::new(),
-            transforms: Mutex::new(SumTree::from_item(
-                Transform::isomorphic(wrap_snapshot.text_summary().lines.row + 1),
-                &(),
-            )),
-            wrap_snapshot: Mutex::new(wrap_snapshot),
+            transforms: Mutex::new(SumTree::from_item(Transform::isomorphic(row_count), &())),
+            wrap_snapshot: Mutex::new(wrap_snapshot.clone()),
             excerpt_header_height,
-        }
+        };
+        map.sync(
+            &wrap_snapshot,
+            vec![Edit {
+                old: 0..row_count,
+                new: 0..row_count,
+            }],
+        );
+        map
     }
 
     pub fn read(&self, wrap_snapshot: WrapSnapshot, edits: Vec<WrapEdit>) -> BlockSnapshot {
@@ -275,6 +281,7 @@ impl BlockMap {
             let new_buffer_start =
                 wrap_snapshot.to_point(WrapPoint::new(new_start.0, 0), Bias::Left);
             let start_anchor = buffer.anchor_before(new_buffer_start);
+            let start_bound = Bound::Included(start_anchor.clone());
             let start_block_ix = match self.blocks[last_block_ix..].binary_search_by(|probe| {
                 probe
                     .position
@@ -285,14 +292,15 @@ impl BlockMap {
                 Ok(ix) | Err(ix) => last_block_ix + ix,
             };
 
-            let end_anchor;
+            let end_bound;
             let end_block_ix = if new_end.0 > wrap_snapshot.max_point().row() {
-                end_anchor = Anchor::max();
+                end_bound = Bound::Unbounded;
                 self.blocks.len()
             } else {
                 let new_buffer_end =
                     wrap_snapshot.to_point(WrapPoint::new(new_end.0, 0), Bias::Left);
-                end_anchor = buffer.anchor_before(new_buffer_end);
+                let end_anchor = buffer.anchor_before(new_buffer_end);
+                end_bound = Bound::Excluded(end_anchor.clone());
                 match self.blocks[start_block_ix..].binary_search_by(|probe| {
                     probe
                         .position
@@ -330,10 +338,12 @@ impl BlockMap {
             );
             blocks_in_edit.extend(
                 buffer
-                    .excerpt_boundaries_in_range(start_anchor..end_anchor)
+                    .excerpt_boundaries_in_range((start_bound, end_bound))
                     .map(|excerpt_boundary| {
                         (
-                            excerpt_boundary.row,
+                            wrap_snapshot
+                                .from_point(Point::new(excerpt_boundary.row, 0), Bias::Left)
+                                .row(),
                             TransformBlock::ExcerptHeader {
                                 buffer: excerpt_boundary.buffer,
                                 range: excerpt_boundary.range,
@@ -343,8 +353,7 @@ impl BlockMap {
                     }),
             );
 
-            // When multiple blocks are on the same row, newer blocks appear above older
-            // blocks. This is arbitrary, but we currently rely on it in ProjectDiagnosticsEditor.
+            // Place excerpt headers above custom blocks on the same row.
             blocks_in_edit.sort_unstable_by(|(row_a, block_a), (row_b, block_b)| {
                 row_a.cmp(&row_b).then_with(|| match (block_a, block_b) {
                     (
@@ -359,7 +368,7 @@ impl BlockMap {
                     ) => block_a
                         .disposition
                         .cmp(&block_b.disposition)
-                        .then_with(|| block_a.id.cmp(&block_b.id).reverse()),
+                        .then_with(|| block_a.id.cmp(&block_b.id)),
                 })
             });
 
@@ -936,7 +945,6 @@ mod tests {
     use crate::multi_buffer::MultiBuffer;
     use gpui::{elements::Empty, Element};
     use rand::prelude::*;
-    use std::cmp::Reverse;
     use std::env;
     use text::RandomCharIter;
 
@@ -1213,7 +1221,7 @@ mod tests {
         let (tab_map, tabs_snapshot) = TabMap::new(folds_snapshot.clone(), tab_size);
         let (wrap_map, wraps_snapshot) =
             WrapMap::new(tabs_snapshot, font_id, font_size, wrap_width, cx);
-        let mut block_map = BlockMap::new(wraps_snapshot, excerpt_header_height);
+        let mut block_map = BlockMap::new(wraps_snapshot.clone(), excerpt_header_height);
         let mut custom_blocks = Vec::new();
 
         for _ in 0..operations {
@@ -1326,25 +1334,22 @@ mod tests {
                     }
                 };
                 let row = wraps_snapshot.from_point(position, Bias::Left).row();
-                (row, block.disposition, *id, block.height)
+                (row, block.disposition, Some(*id), block.height)
             }));
-            expected_blocks.extend(
-                buffer_snapshot
-                    .excerpt_boundaries_in_range(0..buffer_snapshot.len())
-                    .map(|boundary| {
-                        let position =
-                            wraps_snapshot.from_point(Point::new(boundary.row, 0), Bias::Left);
-                        (
-                            position.row(),
-                            BlockDisposition::Above,
-                            BlockId(usize::MAX),
-                            excerpt_header_height,
-                        )
-                    }),
-            );
-            expected_blocks.sort_unstable_by_key(|(row, disposition, id, _)| {
-                (*row, *disposition, Reverse(*id))
-            });
+            expected_blocks.extend(buffer_snapshot.excerpt_boundaries_in_range(0..).map(
+                |boundary| {
+                    let position =
+                        wraps_snapshot.from_point(Point::new(boundary.row, 0), Bias::Left);
+                    (
+                        position.row(),
+                        BlockDisposition::Above,
+                        None,
+                        excerpt_header_height,
+                    )
+                },
+            ));
+            expected_blocks
+                .sort_unstable_by_key(|(row, disposition, id, _)| (*row, *disposition, *id));
             let mut sorted_blocks_iter = expected_blocks.iter().peekable();
 
             let input_buffer_rows = buffer_snapshot.buffer_rows(0).collect::<Vec<_>>();
@@ -1421,14 +1426,7 @@ mod tests {
             assert_eq!(
                 blocks_snapshot
                     .blocks_in_range(0..(expected_row_count as u32))
-                    .map(|(row, block)| (
-                        row,
-                        if let Some((block, _)) = block.as_custom() {
-                            block.id
-                        } else {
-                            BlockId(usize::MAX)
-                        }
-                    ))
+                    .map(|(row, block)| (row, block.as_custom().map(|(b, _)| b.id)))
                     .collect::<Vec<_>>(),
                 expected_block_positions
             );

crates/editor/src/editor.rs 🔗

@@ -10,7 +10,7 @@ mod test;
 use aho_corasick::AhoCorasick;
 use anyhow::Result;
 use clock::ReplicaId;
-use collections::{BTreeMap, HashMap, HashSet};
+use collections::{BTreeMap, Bound, HashMap, HashSet};
 pub use display_map::DisplayPoint;
 use display_map::*;
 pub use element::*;
@@ -2605,7 +2605,10 @@ impl Editor {
 
                 // Don't move lines across excerpts
                 if buffer
-                    .excerpt_boundaries_in_range(insertion_point..range_to_move.end)
+                    .excerpt_boundaries_in_range((
+                        Bound::Excluded(insertion_point),
+                        Bound::Included(range_to_move.end),
+                    ))
                     .next()
                     .is_none()
                 {
@@ -2709,7 +2712,10 @@ impl Editor {
 
                 // Don't move lines across excerpt boundaries
                 if buffer
-                    .excerpt_boundaries_in_range(range_to_move.start..insertion_point)
+                    .excerpt_boundaries_in_range((
+                        Bound::Excluded(range_to_move.start),
+                        Bound::Included(insertion_point),
+                    ))
                     .next()
                     .is_none()
                 {

crates/editor/src/multi_buffer.rs 🔗

@@ -3,7 +3,7 @@ mod anchor;
 pub use anchor::{Anchor, AnchorRangeExt};
 use anyhow::Result;
 use clock::ReplicaId;
-use collections::{HashMap, HashSet};
+use collections::{Bound, HashMap, HashSet};
 use gpui::{AppContext, Entity, ModelContext, ModelHandle, Task};
 pub use language::Completion;
 use language::{
@@ -14,7 +14,7 @@ use std::{
     cell::{Ref, RefCell},
     cmp, fmt, io,
     iter::{self, FromIterator},
-    ops::{Range, Sub},
+    ops::{Range, RangeBounds, Sub},
     str,
     sync::Arc,
     time::{Duration, Instant},
@@ -229,11 +229,9 @@ impl MultiBuffer {
                 let buffer = buffer_handle.read(cx);
                 let end_ix = buffer.clip_offset(rng.gen_range(0..=buffer.len()), Bias::Right);
                 let start_ix = buffer.clip_offset(rng.gen_range(0..=end_ix), Bias::Left);
-                let header_height = rng.gen_range(0..=5);
                 log::info!(
-                    "Inserting excerpt from buffer {} with header height {} and range {:?}: {:?}",
+                    "Inserting excerpt from buffer {} and range {:?}: {:?}",
                     buffer_handle.id(),
-                    header_height,
                     start_ix..end_ix,
                     &buffer.text()[start_ix..end_ix]
                 );
@@ -1765,24 +1763,54 @@ impl MultiBufferSnapshot {
         }
     }
 
-    pub fn excerpt_boundaries_in_range<'a, T: ToOffset>(
+    pub fn excerpt_boundaries_in_range<'a, R, T>(
         &'a self,
-        range: Range<T>,
-    ) -> impl Iterator<Item = ExcerptBoundary> + 'a {
-        let start = range.start.to_offset(self);
-        let end = range.end.to_offset(self);
-        let mut cursor = self
-            .excerpts
-            .cursor::<(usize, (Option<&ExcerptId>, Point))>();
-        cursor.seek(&start, Bias::Right, &());
+        range: R,
+    ) -> impl Iterator<Item = ExcerptBoundary> + 'a
+    where
+        R: RangeBounds<T>,
+        T: ToOffset,
+    {
+        let start_offset;
+        let start = match range.start_bound() {
+            Bound::Included(start) => {
+                start_offset = start.to_offset(self);
+                Bound::Included(start_offset)
+            }
+            Bound::Excluded(start) => {
+                start_offset = start.to_offset(self);
+                Bound::Excluded(start_offset)
+            }
+            Bound::Unbounded => {
+                start_offset = 0;
+                Bound::Unbounded
+            }
+        };
+        let end = match range.end_bound() {
+            Bound::Included(end) => Bound::Included(end.to_offset(self)),
+            Bound::Excluded(end) => Bound::Excluded(end.to_offset(self)),
+            Bound::Unbounded => Bound::Unbounded,
+        };
+        let bounds = (start, end);
+
+        let mut cursor = self.excerpts.cursor::<(usize, Point)>();
+        cursor.seek(&start_offset, Bias::Right, &());
+        if cursor.item().is_none() {
+            cursor.prev(&());
+        }
+        if !bounds.contains(&cursor.start().0) {
+            cursor.next(&());
+        }
 
         let mut prev_buffer_id = cursor.prev_item().map(|excerpt| excerpt.buffer_id);
         std::iter::from_fn(move || {
-            if start <= cursor.start().0 && end > cursor.start().0 {
+            if self.singleton {
+                None
+            } else if bounds.contains(&cursor.start().0) {
                 let excerpt = cursor.item()?;
                 let starts_new_buffer = Some(excerpt.buffer_id) != prev_buffer_id;
                 let boundary = ExcerptBoundary {
-                    row: cursor.start().1 .1.row,
+                    row: cursor.start().1.row,
                     buffer: excerpt.buffer.clone(),
                     range: excerpt.range.clone(),
                     starts_new_buffer,
@@ -2649,52 +2677,47 @@ mod tests {
         );
         assert_eq!(snapshot.buffer_rows(4).collect::<Vec<_>>(), [Some(3)]);
         assert_eq!(snapshot.buffer_rows(5).collect::<Vec<_>>(), []);
-        assert!(snapshot
-            .excerpt_boundaries_in_range(Point::new(1, 0)..Point::new(1, 5))
-            .next()
-            .is_none());
-        assert!(snapshot
-            .excerpt_boundaries_in_range(Point::new(1, 0)..Point::new(2, 0))
-            .next()
-            .is_some());
-        assert!(snapshot
-            .excerpt_boundaries_in_range(Point::new(1, 0)..Point::new(4, 0))
-            .next()
-            .is_some());
-        assert!(snapshot
-            .excerpt_boundaries_in_range(Point::new(2, 0)..Point::new(3, 0))
-            .next()
-            .is_none());
-        assert!(snapshot
-            .excerpt_boundaries_in_range(Point::new(4, 0)..Point::new(4, 2))
-            .next()
-            .is_none());
-        assert!(snapshot
-            .excerpt_boundaries_in_range(Point::new(4, 2)..Point::new(4, 2))
-            .next()
-            .is_none());
 
         assert_eq!(
-            snapshot
-                .excerpt_boundaries_in_range(Point::new(0, 0)..Point::new(4, 2))
-                .map(|boundary| (
-                    boundary.row,
-                    boundary
-                        .buffer
-                        .text_for_range(boundary.range)
-                        .collect::<String>(),
-                    boundary.starts_new_buffer
-                ))
-                .collect::<Vec<_>>(),
+            boundaries_in_range(Point::new(0, 0)..Point::new(4, 2), &snapshot),
             &[
-                (0, "".to_string(), true),
-                (0, "".to_string(), true),
-                (0, "".to_string(), true),
-                (0, "".to_string(), true),
-                (0, "".to_string(), true),
-                (0, "".to_string(), true),
+                (0, "bbbb\nccccc".to_string(), true),
+                (2, "ddd\neeee".to_string(), false),
+                (4, "jj".to_string(), true),
             ]
         );
+        assert_eq!(
+            boundaries_in_range(Point::new(0, 0)..Point::new(2, 0), &snapshot),
+            &[(0, "bbbb\nccccc".to_string(), true)]
+        );
+        assert_eq!(
+            boundaries_in_range(Point::new(1, 0)..Point::new(1, 5), &snapshot),
+            &[]
+        );
+        assert_eq!(
+            boundaries_in_range(Point::new(1, 0)..Point::new(2, 0), &snapshot),
+            &[]
+        );
+        assert_eq!(
+            boundaries_in_range(Point::new(1, 0)..Point::new(4, 0), &snapshot),
+            &[(2, "ddd\neeee".to_string(), false)]
+        );
+        assert_eq!(
+            boundaries_in_range(Point::new(1, 0)..Point::new(4, 0), &snapshot),
+            &[(2, "ddd\neeee".to_string(), false)]
+        );
+        assert_eq!(
+            boundaries_in_range(Point::new(2, 0)..Point::new(3, 0), &snapshot),
+            &[(2, "ddd\neeee".to_string(), false)]
+        );
+        assert_eq!(
+            boundaries_in_range(Point::new(4, 0)..Point::new(4, 2), &snapshot),
+            &[(4, "jj".to_string(), true)]
+        );
+        assert_eq!(
+            boundaries_in_range(Point::new(4, 2)..Point::new(4, 2), &snapshot),
+            &[]
+        );
 
         buffer_1.update(cx, |buffer, cx| {
             buffer.edit(
@@ -2766,6 +2789,25 @@ mod tests {
                 "eeee",   //
             )
         );
+
+        fn boundaries_in_range(
+            range: Range<Point>,
+            snapshot: &MultiBufferSnapshot,
+        ) -> Vec<(u32, String, bool)> {
+            snapshot
+                .excerpt_boundaries_in_range(range)
+                .map(|boundary| {
+                    (
+                        boundary.row,
+                        boundary
+                            .buffer
+                            .text_for_range(boundary.range)
+                            .collect::<String>(),
+                        boundary.starts_new_buffer,
+                    )
+                })
+                .collect::<Vec<_>>()
+        }
     }
 
     #[gpui::test]