Fix soft-wrapping with fold creases (#28029)

Release Notes:

- Fixed a rendering bug that caused context in the agent to not wrap
properly.

---------

Co-authored-by: Conrad Irwin <conrad.irwin@gmail.com>
Co-authored-by: Zed AI <ai+claude-3.7@zed.dev>

crates/editor/src/display_map/block_map.rs

@@ -1,11 +1,12 @@
 use super::{
     Highlights,
+    fold_map::Chunk,
     wrap_map::{self, WrapEdit, WrapPoint, WrapSnapshot},
 };
 use crate::{EditorStyle, GutterDimensions};
 use collections::{Bound, HashMap, HashSet};
 use gpui::{AnyElement, App, EntityId, Pixels, Window};
-use language::{Chunk, Patch, Point};
+use language::{Patch, Point};
 use multi_buffer::{
     Anchor, ExcerptId, ExcerptInfo, MultiBuffer, MultiBufferRow, MultiBufferSnapshot, RowInfo,
     ToOffset, ToPoint as _,

crates/editor/src/display_map/fold_map.rs

@@ -2,8 +2,9 @@ use super::{
     Highlights,
     inlay_map::{InlayBufferRows, InlayChunks, InlayEdit, InlayOffset, InlayPoint, InlaySnapshot},
 };
-use gpui::{AnyElement, App, ElementId};
-use language::{Chunk, ChunkRenderer, Edit, Point, TextSummary};
+use gpui::{AnyElement, App, ElementId, HighlightStyle, Pixels, Window};
+use language::{Edit, HighlightId, Point, TextSummary};
+use lsp::DiagnosticSeverity;
 use multi_buffer::{
     Anchor, AnchorRangeExt, MultiBufferRow, MultiBufferSnapshot, RowInfo, ToOffset,
 };
@@ -14,7 +15,7 @@ use std::{
     ops::{Add, AddAssign, Deref, DerefMut, Range, Sub},
     sync::Arc,
 };
-use sum_tree::{Bias, Cursor, FilterCursor, SumTree, Summary};
+use sum_tree::{Bias, Cursor, FilterCursor, SumTree, Summary, TreeMap};
 use ui::IntoElement as _;
 use util::post_inc;
 
@@ -177,6 +178,13 @@ impl FoldMapWriter<'_> {
             let mut new_tree = SumTree::new(buffer);
             let mut cursor = self.0.snapshot.folds.cursor::<FoldRange>(buffer);
             for fold in folds {
+                self.0.snapshot.fold_metadata_by_id.insert(
+                    fold.id,
+                    FoldMetadata {
+                        range: fold.range.clone(),
+                        width: None,
+                    },
+                );
                 new_tree.append(cursor.slice(&fold.range, Bias::Right, buffer), buffer);
                 new_tree.push(fold, buffer);
             }
@@ -240,6 +248,7 @@ impl FoldMapWriter<'_> {
                         });
                     }
                     fold_ixs_to_delete.push(*folds_cursor.start());
+                    self.0.snapshot.fold_metadata_by_id.remove(&fold.id);
                 }
                 folds_cursor.next(buffer);
             }
@@ -263,6 +272,42 @@ impl FoldMapWriter<'_> {
         let edits = self.0.sync(snapshot.clone(), edits);
         (self.0.snapshot.clone(), edits)
     }
+
+    pub(crate) fn update_fold_widths(
+        &mut self,
+        new_widths: impl IntoIterator<Item = (FoldId, Pixels)>,
+    ) -> (FoldSnapshot, Vec<FoldEdit>) {
+        let mut edits = Vec::new();
+        let inlay_snapshot = self.0.snapshot.inlay_snapshot.clone();
+        let buffer = &inlay_snapshot.buffer;
+
+        for (id, new_width) in new_widths {
+            if let Some(metadata) = self.0.snapshot.fold_metadata_by_id.get(&id).cloned() {
+                if Some(new_width) != metadata.width {
+                    let buffer_start = metadata.range.start.to_offset(buffer);
+                    let buffer_end = metadata.range.end.to_offset(buffer);
+                    let inlay_range = inlay_snapshot.to_inlay_offset(buffer_start)
+                        ..inlay_snapshot.to_inlay_offset(buffer_end);
+                    edits.push(InlayEdit {
+                        old: inlay_range.clone(),
+                        new: inlay_range.clone(),
+                    });
+
+                    self.0.snapshot.fold_metadata_by_id.insert(
+                        id,
+                        FoldMetadata {
+                            range: metadata.range,
+                            width: Some(new_width),
+                        },
+                    );
+                }
+            }
+        }
+
+        let edits = consolidate_inlay_edits(edits);
+        let edits = self.0.sync(inlay_snapshot, edits);
+        (self.0.snapshot.clone(), edits)
+    }
 }
 
 /// Decides where the fold indicators should be; also tracks parts of a source file that are currently folded.
@@ -290,6 +335,7 @@ impl FoldMap {
                 ),
                 inlay_snapshot: inlay_snapshot.clone(),
                 version: 0,
+                fold_metadata_by_id: TreeMap::default(),
             },
             next_fold_id: FoldId::default(),
         };
@@ -481,6 +527,7 @@ impl FoldMap {
                                 placeholder: Some(TransformPlaceholder {
                                     text: ELLIPSIS,
                                     renderer: ChunkRenderer {
+                                        id: fold.id,
                                         render: Arc::new(move |cx| {
                                             (fold.placeholder.render)(
                                                 fold_id,
@@ -489,6 +536,7 @@ impl FoldMap {
                                             )
                                         }),
                                         constrain_width: fold.placeholder.constrain_width,
+                                        measured_width: self.snapshot.fold_width(&fold_id),
                                     },
                                 }),
                             },
@@ -573,6 +621,7 @@ impl FoldMap {
 pub struct FoldSnapshot {
     transforms: SumTree<Transform>,
     folds: SumTree<Fold>,
+    fold_metadata_by_id: TreeMap<FoldId, FoldMetadata>,
     pub inlay_snapshot: InlaySnapshot,
     pub version: usize,
 }
@@ -582,6 +631,10 @@ impl FoldSnapshot {
         &self.inlay_snapshot.buffer
     }
 
+    fn fold_width(&self, fold_id: &FoldId) -> Option<Pixels> {
+        self.fold_metadata_by_id.get(fold_id)?.width
+    }
+
     #[cfg(test)]
     pub fn text(&self) -> String {
         self.chunks(FoldOffset(0)..self.len(), false, Highlights::default())
@@ -1006,7 +1059,7 @@ impl sum_tree::Summary for TransformSummary {
     }
 }
 
-#[derive(Copy, Clone, Eq, PartialEq, Debug, Default)]
+#[derive(Copy, Clone, Eq, PartialEq, Debug, Default, Ord, PartialOrd, Hash)]
 pub struct FoldId(usize);
 
 impl From<FoldId> for ElementId {
@@ -1045,6 +1098,12 @@ impl Default for FoldRange {
     }
 }
 
+#[derive(Clone, Debug)]
+struct FoldMetadata {
+    range: FoldRange,
+    width: Option<Pixels>,
+}
+
 impl sum_tree::Item for Fold {
     type Summary = FoldSummary;
 
@@ -1181,10 +1240,74 @@ impl Iterator for FoldRows<'_> {
     }
 }
 
+/// A chunk of a buffer's text, along with its syntax highlight and
+/// diagnostic status.
+#[derive(Clone, Debug, Default)]
+pub struct Chunk<'a> {
+    /// The text of the chunk.
+    pub text: &'a str,
+    /// The syntax highlighting style of the chunk.
+    pub syntax_highlight_id: Option<HighlightId>,
+    /// The highlight style that has been applied to this chunk in
+    /// the editor.
+    pub highlight_style: Option<HighlightStyle>,
+    /// The severity of diagnostic associated with this chunk, if any.
+    pub diagnostic_severity: Option<DiagnosticSeverity>,
+    /// Whether this chunk of text is marked as unnecessary.
+    pub is_unnecessary: bool,
+    /// Whether this chunk of text was originally a tab character.
+    pub is_tab: bool,
+    /// An optional recipe for how the chunk should be presented.
+    pub renderer: Option<ChunkRenderer>,
+}
+
+/// A recipe for how the chunk should be presented.
+#[derive(Clone)]
+pub struct ChunkRenderer {
+    /// The id of the fold associated with this chunk.
+    pub id: FoldId,
+    /// Creates a custom element to represent this chunk.
+    pub render: Arc<dyn Send + Sync + Fn(&mut ChunkRendererContext) -> AnyElement>,
+    /// If true, the element is constrained to the shaped width of the text.
+    pub constrain_width: bool,
+    /// The width of the element, as measured during the last layout pass.
+    ///
+    /// This is None if the element has not been laid out yet.
+    pub measured_width: Option<Pixels>,
+}
+
+pub struct ChunkRendererContext<'a, 'b> {
+    pub window: &'a mut Window,
+    pub context: &'b mut App,
+    pub max_width: Pixels,
+}
+
+impl fmt::Debug for ChunkRenderer {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        f.debug_struct("ChunkRenderer")
+            .field("constrain_width", &self.constrain_width)
+            .finish()
+    }
+}
+
+impl Deref for ChunkRendererContext<'_, '_> {
+    type Target = App;
+
+    fn deref(&self) -> &Self::Target {
+        self.context
+    }
+}
+
+impl DerefMut for ChunkRendererContext<'_, '_> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        self.context
+    }
+}
+
 pub struct FoldChunks<'a> {
     transform_cursor: Cursor<'a, Transform, (FoldOffset, InlayOffset)>,
     inlay_chunks: InlayChunks<'a>,
-    inlay_chunk: Option<(InlayOffset, Chunk<'a>)>,
+    inlay_chunk: Option<(InlayOffset, language::Chunk<'a>)>,
     inlay_offset: InlayOffset,
     output_offset: FoldOffset,
     max_output_offset: FoldOffset,
@@ -1292,7 +1415,15 @@ impl<'a> Iterator for FoldChunks<'a> {
 
             self.inlay_offset = chunk_end;
             self.output_offset.0 += chunk.text.len();
-            return Some(chunk);
+            return Some(Chunk {
+                text: chunk.text,
+                syntax_highlight_id: chunk.syntax_highlight_id,
+                highlight_style: chunk.highlight_style,
+                diagnostic_severity: chunk.diagnostic_severity,
+                is_unnecessary: chunk.is_unnecessary,
+                is_tab: chunk.is_tab,
+                renderer: None,
+            });
         }
 
         None

crates/editor/src/display_map/tab_map.rs

@@ -1,8 +1,8 @@
 use super::{
     Highlights,
-    fold_map::{self, FoldChunks, FoldEdit, FoldPoint, FoldSnapshot},
+    fold_map::{self, Chunk, FoldChunks, FoldEdit, FoldPoint, FoldSnapshot},
 };
-use language::{Chunk, Point};
+use language::Point;
 use multi_buffer::MultiBufferSnapshot;
 use std::{cmp, mem, num::NonZeroU32, ops::Range};
 use sum_tree::Bias;

crates/editor/src/display_map/wrap_map.rs

@@ -1,10 +1,10 @@
 use super::{
     Highlights,
-    fold_map::FoldRows,
+    fold_map::{Chunk, FoldRows},
     tab_map::{self, TabEdit, TabPoint, TabSnapshot},
 };
 use gpui::{App, AppContext as _, Context, Entity, Font, LineWrapper, Pixels, Task};
-use language::{Chunk, Point};
+use language::Point;
 use multi_buffer::{MultiBufferSnapshot, RowInfo};
 use smol::future::yield_now;
 use std::sync::LazyLock;
@@ -454,6 +454,7 @@ impl WrapSnapshot {
                 }
 
                 let mut line = String::new();
+                let mut line_fragments = Vec::new();
                 let mut remaining = None;
                 let mut chunks = new_tab_snapshot.chunks(
                     TabPoint::new(edit.new_rows.start, 0)..new_tab_snapshot.max_point(),
@@ -462,15 +463,26 @@ impl WrapSnapshot {
                 );
                 let mut edit_transforms = Vec::<Transform>::new();
                 for _ in edit.new_rows.start..edit.new_rows.end {
-                    while let Some(chunk) =
-                        remaining.take().or_else(|| chunks.next().map(|c| c.text))
-                    {
-                        if let Some(ix) = chunk.find('\n') {
-                            line.push_str(&chunk[..ix + 1]);
-                            remaining = Some(&chunk[ix + 1..]);
+                    while let Some(chunk) = remaining.take().or_else(|| chunks.next()) {
+                        if let Some(ix) = chunk.text.find('\n') {
+                            let (prefix, suffix) = chunk.text.split_at(ix + 1);
+                            line_fragments.push(gpui::LineFragment::text(prefix));
+                            line.push_str(prefix);
+                            remaining = Some(Chunk {
+                                text: suffix,
+                                ..chunk
+                            });
                             break;
                         } else {
-                            line.push_str(chunk)
+                            if let Some(width) =
+                                chunk.renderer.as_ref().and_then(|r| r.measured_width)
+                            {
+                                line_fragments
+                                    .push(gpui::LineFragment::element(width, chunk.text.len()));
+                            } else {
+                                line_fragments.push(gpui::LineFragment::text(chunk.text));
+                            }
+                            line.push_str(chunk.text);
                         }
                     }
 
@@ -479,7 +491,7 @@ impl WrapSnapshot {
                     }
 
                     let mut prev_boundary_ix = 0;
-                    for boundary in line_wrapper.wrap_line(&line, wrap_width) {
+                    for boundary in line_wrapper.wrap_line(&line_fragments, wrap_width) {
                         let wrapped = &line[prev_boundary_ix..boundary.ix];
                         push_isomorphic(&mut edit_transforms, TextSummary::from(wrapped));
                         edit_transforms.push(Transform::wrap(boundary.next_indent));
@@ -494,6 +506,7 @@ impl WrapSnapshot {
                     }
 
                     line.clear();
+                    line_fragments.clear();
                     yield_now().await;
                 }
 
@@ -1173,7 +1186,7 @@ mod tests {
         display_map::{fold_map::FoldMap, inlay_map::InlayMap, tab_map::TabMap},
         test::test_font,
     };
-    use gpui::{px, test::observe};
+    use gpui::{LineFragment, px, test::observe};
     use rand::prelude::*;
     use settings::SettingsStore;
     use smol::stream::StreamExt;
@@ -1228,8 +1241,7 @@ mod tests {
         log::info!("TabMap text: {:?}", tabs_snapshot.text());
 
         let mut line_wrapper = text_system.line_wrapper(font.clone(), font_size);
-        let unwrapped_text = tabs_snapshot.text();
-        let expected_text = wrap_text(&unwrapped_text, wrap_width, &mut line_wrapper);
+        let expected_text = wrap_text(&tabs_snapshot, wrap_width, &mut line_wrapper);
 
         let (wrap_map, _) =
             cx.update(|cx| WrapMap::new(tabs_snapshot.clone(), font, font_size, wrap_width, cx));
@@ -1246,9 +1258,10 @@ mod tests {
 
         let actual_text = initial_snapshot.text();
         assert_eq!(
-            actual_text, expected_text,
+            actual_text,
+            expected_text,
             "unwrapped text is: {:?}",
-            unwrapped_text
+            tabs_snapshot.text()
         );
         log::info!("Wrapped text: {:?}", actual_text);
 
@@ -1311,8 +1324,7 @@ mod tests {
             let (tabs_snapshot, tab_edits) = tab_map.sync(fold_snapshot, fold_edits, tab_size);
             log::info!("TabMap text: {:?}", tabs_snapshot.text());
 
-            let unwrapped_text = tabs_snapshot.text();
-            let expected_text = wrap_text(&unwrapped_text, wrap_width, &mut line_wrapper);
+            let expected_text = wrap_text(&tabs_snapshot, wrap_width, &mut line_wrapper);
             let (mut snapshot, wrap_edits) =
                 wrap_map.update(cx, |map, cx| map.sync(tabs_snapshot.clone(), tab_edits, cx));
             snapshot.check_invariants();
@@ -1328,8 +1340,9 @@ mod tests {
             }
 
             if !wrap_map.read_with(cx, |map, _| map.is_rewrapping()) {
-                let (mut wrapped_snapshot, wrap_edits) =
-                    wrap_map.update(cx, |map, cx| map.sync(tabs_snapshot, Vec::new(), cx));
+                let (mut wrapped_snapshot, wrap_edits) = wrap_map.update(cx, |map, cx| {
+                    map.sync(tabs_snapshot.clone(), Vec::new(), cx)
+                });
                 let actual_text = wrapped_snapshot.text();
                 let actual_longest_row = wrapped_snapshot.longest_row();
                 log::info!("Wrapping finished: {:?}", actual_text);
@@ -1337,9 +1350,10 @@ mod tests {
                 wrapped_snapshot.verify_chunks(&mut rng);
                 edits.push((wrapped_snapshot.clone(), wrap_edits));
                 assert_eq!(
-                    actual_text, expected_text,
+                    actual_text,
+                    expected_text,
                     "unwrapped text is: {:?}",
-                    unwrapped_text
+                    tabs_snapshot.text()
                 );
 
                 let mut summary = TextSummary::default();
@@ -1425,19 +1439,19 @@ mod tests {
     }
 
     fn wrap_text(
-        unwrapped_text: &str,
+        tab_snapshot: &TabSnapshot,
         wrap_width: Option<Pixels>,
         line_wrapper: &mut LineWrapper,
     ) -> String {
         if let Some(wrap_width) = wrap_width {
             let mut wrapped_text = String::new();
-            for (row, line) in unwrapped_text.split('\n').enumerate() {
+            for (row, line) in tab_snapshot.text().split('\n').enumerate() {
                 if row > 0 {
-                    wrapped_text.push('\n')
+                    wrapped_text.push('\n');
                 }
 
                 let mut prev_ix = 0;
-                for boundary in line_wrapper.wrap_line(line, wrap_width) {
+                for boundary in line_wrapper.wrap_line(&[LineFragment::text(&line)], wrap_width) {
                     wrapped_text.push_str(&line[prev_ix..boundary.ix]);
                     wrapped_text.push('\n');
                     wrapped_text.push_str(&" ".repeat(boundary.next_indent as usize));
@@ -1445,9 +1459,10 @@ mod tests {
                 }
                 wrapped_text.push_str(&line[prev_ix..]);
             }
+
             wrapped_text
         } else {
-            unwrapped_text.to_string()
+            tab_snapshot.text()
         }
     }