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.rs

@@ -37,7 +37,7 @@ pub use block_map::{
 use block_map::{BlockRow, BlockSnapshot};
 use collections::{HashMap, HashSet};
 pub use crease_map::*;
-pub use fold_map::{Fold, FoldId, FoldPlaceholder, FoldPoint};
+pub use fold_map::{ChunkRenderer, ChunkRendererContext, Fold, FoldId, FoldPlaceholder, FoldPoint};
 use fold_map::{FoldMap, FoldSnapshot};
 use gpui::{App, Context, Entity, Font, HighlightStyle, LineLayout, Pixels, UnderlineStyle};
 pub use inlay_map::Inlay;
@@ -45,8 +45,7 @@ use inlay_map::{InlayMap, InlaySnapshot};
 pub use inlay_map::{InlayOffset, InlayPoint};
 pub use invisibles::{is_invisible, replacement};
 use language::{
-    ChunkRenderer, OffsetUtf16, Point, Subscription as BufferSubscription,
-    language_settings::language_settings,
+    OffsetUtf16, Point, Subscription as BufferSubscription, language_settings::language_settings,
 };
 use lsp::DiagnosticSeverity;
 use multi_buffer::{
@@ -515,6 +514,33 @@ impl DisplayMap {
             .update(cx, |map, cx| map.set_wrap_width(width, cx))
     }
 
+    pub fn update_fold_widths(
+        &mut self,
+        widths: impl IntoIterator<Item = (FoldId, Pixels)>,
+        cx: &mut Context<Self>,
+    ) -> bool {
+        let snapshot = self.buffer.read(cx).snapshot(cx);
+        let edits = self.buffer_subscription.consume().into_inner();
+        let tab_size = Self::tab_size(&self.buffer, cx);
+        let (snapshot, edits) = self.inlay_map.sync(snapshot, edits);
+        let (mut fold_map, snapshot, edits) = self.fold_map.write(snapshot, edits);
+        let (snapshot, edits) = self.tab_map.sync(snapshot, edits, tab_size);
+        let (snapshot, edits) = self
+            .wrap_map
+            .update(cx, |map, cx| map.sync(snapshot, edits, cx));
+        self.block_map.read(snapshot, edits);
+
+        let (snapshot, edits) = fold_map.update_fold_widths(widths);
+        let widths_changed = !edits.is_empty();
+        let (snapshot, edits) = self.tab_map.sync(snapshot, edits, tab_size);
+        let (snapshot, edits) = self
+            .wrap_map
+            .update(cx, |map, cx| map.sync(snapshot, edits, cx));
+        self.block_map.read(snapshot, edits);
+
+        widths_changed
+    }
+
     pub(crate) fn current_inlays(&self) -> impl Iterator<Item = &Inlay> {
         self.inlay_map.current_inlays()
     }

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()
         }
     }
 

crates/editor/src/editor.rs

@@ -58,7 +58,7 @@ use clock::ReplicaId;
 use collections::{BTreeMap, HashMap, HashSet, VecDeque};
 use convert_case::{Case, Casing};
 use display_map::*;
-pub use display_map::{DisplayPoint, FoldPlaceholder};
+pub use display_map::{ChunkRenderer, ChunkRendererContext, DisplayPoint, FoldPlaceholder};
 use editor_settings::GoToDefinitionFallback;
 pub use editor_settings::{
     CurrentLineHighlight, EditorSettings, HideMouseMode, ScrollBeyondLastLine, SearchSettings,
@@ -15045,6 +15045,15 @@ impl Editor {
         self.active_indent_guides_state.dirty = true;
     }
 
+    pub fn update_fold_widths(
+        &mut self,
+        widths: impl IntoIterator<Item = (FoldId, Pixels)>,
+        cx: &mut Context<Self>,
+    ) -> bool {
+        self.display_map
+            .update(cx, |map, cx| map.update_fold_widths(widths, cx))
+    }
+
     pub fn default_fold_placeholder(&self, cx: &App) -> FoldPlaceholder {
         self.display_map.read(cx).fold_placeholder.clone()
     }

crates/editor/src/element.rs

@@ -1,16 +1,16 @@
 use crate::{
-    BlockId, COLUMNAR_SELECTION_MODIFIERS, CURSORS_VISIBLE_FOR, ChunkReplacement,
-    ContextMenuPlacement, CursorShape, CustomBlockId, DisplayDiffHunk, DisplayPoint, DisplayRow,
-    DocumentHighlightRead, DocumentHighlightWrite, EditDisplayMode, Editor, EditorMode,
-    EditorSettings, EditorSnapshot, EditorStyle, FILE_HEADER_HEIGHT, FocusedBlock,
-    GutterDimensions, HalfPageDown, HalfPageUp, HandleInput, HoveredCursor, InlayHintRefreshReason,
-    InlineCompletion, JumpData, LineDown, LineHighlight, LineUp, MAX_LINE_LEN,
-    MIN_LINE_NUMBER_DIGITS, MULTI_BUFFER_EXCERPT_HEADER_HEIGHT, OpenExcerpts, PageDown, PageUp,
-    Point, RowExt, RowRangeExt, SelectPhase, SelectedTextHighlight, Selection, SoftWrap,
-    StickyHeaderExcerpt, ToPoint, ToggleFold,
+    BlockId, COLUMNAR_SELECTION_MODIFIERS, CURSORS_VISIBLE_FOR, ChunkRendererContext,
+    ChunkReplacement, ContextMenuPlacement, CursorShape, CustomBlockId, DisplayDiffHunk,
+    DisplayPoint, DisplayRow, DocumentHighlightRead, DocumentHighlightWrite, EditDisplayMode,
+    Editor, EditorMode, EditorSettings, EditorSnapshot, EditorStyle, FILE_HEADER_HEIGHT,
+    FocusedBlock, GutterDimensions, HalfPageDown, HalfPageUp, HandleInput, HoveredCursor,
+    InlayHintRefreshReason, InlineCompletion, JumpData, LineDown, LineHighlight, LineUp,
+    MAX_LINE_LEN, MIN_LINE_NUMBER_DIGITS, MULTI_BUFFER_EXCERPT_HEADER_HEIGHT, OpenExcerpts,
+    PageDown, PageUp, Point, RowExt, RowRangeExt, SelectPhase, SelectedTextHighlight, Selection,
+    SoftWrap, StickyHeaderExcerpt, ToPoint, ToggleFold,
     code_context_menus::{CodeActionsMenu, MENU_ASIDE_MAX_WIDTH, MENU_ASIDE_MIN_WIDTH, MENU_GAP},
     display_map::{
-        Block, BlockContext, BlockStyle, DisplaySnapshot, HighlightedChunk, ToDisplayPoint,
+        Block, BlockContext, BlockStyle, DisplaySnapshot, FoldId, HighlightedChunk, ToDisplayPoint,
     },
     editor_settings::{
         CurrentLineHighlight, DoubleClickInMultibuffer, MultiCursorModifier, ScrollBeyondLastLine,
@@ -43,12 +43,8 @@ use gpui::{
     transparent_black,
 };
 use itertools::Itertools;
-use language::{
-    ChunkRendererContext,
-    language_settings::{
-        IndentGuideBackgroundColoring, IndentGuideColoring, IndentGuideSettings,
-        ShowWhitespaceSetting,
-    },
+use language::language_settings::{
+    IndentGuideBackgroundColoring, IndentGuideColoring, IndentGuideSettings, ShowWhitespaceSetting,
 };
 use lsp::DiagnosticSeverity;
 use multi_buffer::{
@@ -5807,6 +5803,7 @@ pub(crate) struct LineWithInvisibles {
 enum LineFragment {
     Text(ShapedLine),
     Element {
+        id: FoldId,
         element: Option<AnyElement>,
         size: Size<Pixels>,
         len: usize,
@@ -5908,6 +5905,7 @@ impl LineWithInvisibles {
                         width += size.width;
                         len += highlighted_chunk.text.len();
                         fragments.push(LineFragment::Element {
+                            id: renderer.id,
                             element: Some(element),
                             size,
                             len: highlighted_chunk.text.len(),
@@ -6863,6 +6861,24 @@ impl Element for EditorElement {
                         window,
                         cx,
                     );
+                    let new_fold_widths = line_layouts
+                        .iter()
+                        .flat_map(|layout| &layout.fragments)
+                        .filter_map(|fragment| {
+                            if let LineFragment::Element { id, size, .. } = fragment {
+                                Some((*id, size.width))
+                            } else {
+                                None
+                            }
+                        });
+                    if self.editor.update(cx, |editor, cx| {
+                        editor.update_fold_widths(new_fold_widths, cx)
+                    }) {
+                        // If the fold widths have changed, we need to prepaint
+                        // the element again to account for any changes in
+                        // wrapping.
+                        return self.prepaint(None, bounds, &mut (), window, cx);
+                    }
 
                     let longest_line_blame_width = self
                         .editor

crates/gpui/src/text_system/line_wrapper.rs

@@ -32,7 +32,7 @@ impl LineWrapper {
     /// Wrap a line of text to the given width with this wrapper's font and font size.
     pub fn wrap_line<'a>(
         &'a mut self,
-        line: &'a str,
+        fragments: &'a [LineFragment],
         wrap_width: Pixels,
     ) -> impl Iterator<Item = Boundary> + 'a {
         let mut width = px(0.);
@@ -42,32 +42,61 @@ impl LineWrapper {
         let mut last_candidate_width = px(0.);
         let mut last_wrap_ix = 0;
         let mut prev_c = '\0';
-        let mut char_indices = line.char_indices();
+        let mut index = 0;
+        let mut candidates = fragments
+            .into_iter()
+            .flat_map(move |fragment| fragment.wrap_boundary_candidates())
+            .peekable();
         iter::from_fn(move || {
-            for (ix, c) in char_indices.by_ref() {
-                if c == '\n' {
-                    continue;
-                }
+            for candidate in candidates.by_ref() {
+                let ix = index;
+                index += candidate.len_utf8();
+                let mut new_prev_c = prev_c;
+                let item_width = match candidate {
+                    WrapBoundaryCandidate::Char { character: c } => {
+                        if c == '\n' {
+                            continue;
+                        }
 
-                if Self::is_word_char(c) {
-                    if prev_c == ' ' && c != ' ' && first_non_whitespace_ix.is_some() {
-                        last_candidate_ix = ix;
-                        last_candidate_width = width;
+                        if Self::is_word_char(c) {
+                            if prev_c == ' ' && c != ' ' && first_non_whitespace_ix.is_some() {
+                                last_candidate_ix = ix;
+                                last_candidate_width = width;
+                            }
+                        } else {
+                            // CJK may not be space separated, e.g.: `Hello world你好世界`
+                            if c != ' ' && first_non_whitespace_ix.is_some() {
+                                last_candidate_ix = ix;
+                                last_candidate_width = width;
+                            }
+                        }
+
+                        if c != ' ' && first_non_whitespace_ix.is_none() {
+                            first_non_whitespace_ix = Some(ix);
+                        }
+
+                        new_prev_c = c;
+
+                        self.width_for_char(c)
                     }
-                } else {
-                    // CJK may not be space separated, e.g.: `Hello world你好世界`
-                    if c != ' ' && first_non_whitespace_ix.is_some() {
-                        last_candidate_ix = ix;
-                        last_candidate_width = width;
+                    WrapBoundaryCandidate::Element {
+                        width: element_width,
+                        ..
+                    } => {
+                        if prev_c == ' ' && first_non_whitespace_ix.is_some() {
+                            last_candidate_ix = ix;
+                            last_candidate_width = width;
+                        }
+
+                        if first_non_whitespace_ix.is_none() {
+                            first_non_whitespace_ix = Some(ix);
+                        }
+
+                        element_width
                     }
-                }
+                };
 
-                if c != ' ' && first_non_whitespace_ix.is_none() {
-                    first_non_whitespace_ix = Some(ix);
-                }
-
-                let char_width = self.width_for_char(c);
-                width += char_width;
+                width += item_width;
                 if width > wrap_width && ix > last_wrap_ix {
                     if let (None, Some(first_non_whitespace_ix)) = (indent, first_non_whitespace_ix)
                     {
@@ -82,7 +111,7 @@ impl LineWrapper {
                         last_candidate_ix = 0;
                     } else {
                         last_wrap_ix = ix;
-                        width = char_width;
+                        width = item_width;
                     }
 
                     if let Some(indent) = indent {
@@ -91,7 +120,8 @@ impl LineWrapper {
 
                     return Some(Boundary::new(last_wrap_ix, indent.unwrap_or(0)));
                 }
-                prev_c = c;
+
+                prev_c = new_prev_c;
             }
 
             None
@@ -213,6 +243,65 @@ fn update_runs_after_truncation(result: &str, ellipsis: &str, runs: &mut Vec<Tex
     }
 }
 
+/// A fragment of a line that can be wrapped.
+pub enum LineFragment<'a> {
+    /// A text fragment consisting of characters.
+    Text {
+        /// The text content of the fragment.
+        text: &'a str,
+    },
+    /// A non-text element with a fixed width.
+    Element {
+        /// The width of the element in pixels.
+        width: Pixels,
+        /// The UTF-8 encoded length of the element.
+        len_utf8: usize,
+    },
+}
+
+impl<'a> LineFragment<'a> {
+    /// Creates a new text fragment from the given text.
+    pub fn text(text: &'a str) -> Self {
+        LineFragment::Text { text }
+    }
+
+    /// Creates a new non-text element with the given width and UTF-8 encoded length.
+    pub fn element(width: Pixels, len_utf8: usize) -> Self {
+        LineFragment::Element { width, len_utf8 }
+    }
+
+    fn wrap_boundary_candidates(&self) -> impl Iterator<Item = WrapBoundaryCandidate> {
+        let text = match self {
+            LineFragment::Text { text } => text,
+            LineFragment::Element { .. } => "\0",
+        };
+        text.chars().map(move |character| {
+            if let LineFragment::Element { width, len_utf8 } = self {
+                WrapBoundaryCandidate::Element {
+                    width: *width,
+                    len_utf8: *len_utf8,
+                }
+            } else {
+                WrapBoundaryCandidate::Char { character }
+            }
+        })
+    }
+}
+
+enum WrapBoundaryCandidate {
+    Char { character: char },
+    Element { width: Pixels, len_utf8: usize },
+}
+
+impl WrapBoundaryCandidate {
+    pub fn len_utf8(&self) -> usize {
+        match self {
+            WrapBoundaryCandidate::Char { character } => character.len_utf8(),
+            WrapBoundaryCandidate::Element { len_utf8: len, .. } => *len,
+        }
+    }
+}
+
 /// A boundary between two lines of text.
 #[derive(Copy, Clone, Debug, PartialEq, Eq)]
 pub struct Boundary {
@@ -278,7 +367,7 @@ mod tests {
 
         assert_eq!(
             wrapper
-                .wrap_line("aa bbb cccc ddddd eeee", px(72.))
+                .wrap_line(&[LineFragment::text("aa bbb cccc ddddd eeee")], px(72.))
                 .collect::<Vec<_>>(),
             &[
                 Boundary::new(7, 0),
@@ -288,7 +377,7 @@ mod tests {
         );
         assert_eq!(
             wrapper
-                .wrap_line("aaa aaaaaaaaaaaaaaaaaa", px(72.0))
+                .wrap_line(&[LineFragment::text("aaa aaaaaaaaaaaaaaaaaa")], px(72.0))
                 .collect::<Vec<_>>(),
             &[
                 Boundary::new(4, 0),
@@ -298,7 +387,7 @@ mod tests {
         );
         assert_eq!(
             wrapper
-                .wrap_line("     aaaaaaa", px(72.))
+                .wrap_line(&[LineFragment::text("     aaaaaaa")], px(72.))
                 .collect::<Vec<_>>(),
             &[
                 Boundary::new(7, 5),
@@ -308,7 +397,10 @@ mod tests {
         );
         assert_eq!(
             wrapper
-                .wrap_line("                            ", px(72.))
+                .wrap_line(
+                    &[LineFragment::text("                            ")],
+                    px(72.)
+                )
                 .collect::<Vec<_>>(),
             &[
                 Boundary::new(7, 0),
@@ -318,7 +410,7 @@ mod tests {
         );
         assert_eq!(
             wrapper
-                .wrap_line("          aaaaaaaaaaaaaa", px(72.))
+                .wrap_line(&[LineFragment::text("          aaaaaaaaaaaaaa")], px(72.))
                 .collect::<Vec<_>>(),
             &[
                 Boundary::new(7, 0),
@@ -327,6 +419,84 @@ mod tests {
                 Boundary::new(22, 3),
             ]
         );
+
+        // Test wrapping multiple text fragments
+        assert_eq!(
+            wrapper
+                .wrap_line(
+                    &[
+                        LineFragment::text("aa bbb "),
+                        LineFragment::text("cccc ddddd eeee")
+                    ],
+                    px(72.)
+                )
+                .collect::<Vec<_>>(),
+            &[
+                Boundary::new(7, 0),
+                Boundary::new(12, 0),
+                Boundary::new(18, 0)
+            ],
+        );
+
+        // Test wrapping with a mix of text and element fragments
+        assert_eq!(
+            wrapper
+                .wrap_line(
+                    &[
+                        LineFragment::text("aa "),
+                        LineFragment::element(px(20.), 1),
+                        LineFragment::text(" bbb "),
+                        LineFragment::element(px(30.), 1),
+                        LineFragment::text(" cccc")
+                    ],
+                    px(72.)
+                )
+                .collect::<Vec<_>>(),
+            &[
+                Boundary::new(5, 0),
+                Boundary::new(9, 0),
+                Boundary::new(11, 0)
+            ],
+        );
+
+        // Test with element at the beginning and text afterward
+        assert_eq!(
+            wrapper
+                .wrap_line(
+                    &[
+                        LineFragment::element(px(50.), 1),
+                        LineFragment::text(" aaaa bbbb cccc dddd")
+                    ],
+                    px(72.)
+                )
+                .collect::<Vec<_>>(),
+            &[
+                Boundary::new(2, 0),
+                Boundary::new(7, 0),
+                Boundary::new(12, 0),
+                Boundary::new(17, 0)
+            ],
+        );
+
+        // Test with a large element that forces wrapping by itself
+        assert_eq!(
+            wrapper
+                .wrap_line(
+                    &[
+                        LineFragment::text("short text "),
+                        LineFragment::element(px(100.), 1),
+                        LineFragment::text(" more text")
+                    ],
+                    px(72.)
+                )
+                .collect::<Vec<_>>(),
+            &[
+                Boundary::new(6, 0),
+                Boundary::new(11, 0),
+                Boundary::new(12, 0),
+                Boundary::new(18, 0)
+            ],
+        );
     }
 
     #[test]

crates/language/src/buffer.rs

@@ -25,8 +25,8 @@ use collections::HashMap;
 use fs::MTime;
 use futures::channel::oneshot;
 use gpui::{
-    AnyElement, App, AppContext as _, Context, Entity, EventEmitter, HighlightStyle, Pixels,
-    SharedString, StyledText, Task, TaskLabel, TextStyle, Window,
+    App, AppContext as _, Context, Entity, EventEmitter, HighlightStyle, SharedString, StyledText,
+    Task, TaskLabel, TextStyle,
 };
 use lsp::{LanguageServerId, NumberOrString};
 use parking_lot::Mutex;
@@ -43,14 +43,13 @@ use std::{
     cmp::{self, Ordering, Reverse},
     collections::{BTreeMap, BTreeSet},
     ffi::OsStr,
-    fmt,
     future::Future,
     iter::{self, Iterator, Peekable},
     mem,
     num::NonZeroU32,
-    ops::{Deref, DerefMut, Range},
+    ops::{Deref, Range},
     path::{Path, PathBuf},
-    rc, str,
+    rc,
     sync::{Arc, LazyLock},
     time::{Duration, Instant},
     vec,
@@ -483,45 +482,6 @@ pub struct Chunk<'a> {
     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 {
-    /// 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,
-}
-
-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
-    }
 }
 
 /// A set of edits to a given version of a buffer, computed asynchronously.