Fix panic in WrapMap::highlighted_chunks_for_rows

Also, add a unit test for `DisplayMap` with syntax highlighting when
soft wrap is on.

zed/src/editor/display_map.rs

@@ -518,28 +518,84 @@ mod tests {
                 ("() {}\n}".to_string(), Some("mod.body")),
             ]
         );
+    }
 
-        fn highlighted_chunks<'a>(
-            rows: Range<u32>,
-            map: &DisplayMap,
-            theme: &'a Theme,
-            cx: &AppContext,
-        ) -> Vec<(String, Option<&'a str>)> {
-            let mut chunks: Vec<(String, Option<&str>)> = Vec::new();
-            for (chunk, style_id) in map.snapshot(cx).highlighted_chunks_for_rows(rows) {
-                let style_name = theme.syntax_style_name(style_id);
-                if let Some((last_chunk, last_style_name)) = chunks.last_mut() {
-                    if style_name == *last_style_name {
-                        last_chunk.push_str(chunk);
-                    } else {
-                        chunks.push((chunk.to_string(), style_name));
-                    }
-                } else {
-                    chunks.push((chunk.to_string(), style_name));
-                }
-            }
-            chunks
-        }
+    #[gpui::test]
+    async fn test_highlighted_chunks_with_soft_wrapping(mut cx: gpui::TestAppContext) {
+        use unindent::Unindent as _;
+
+        cx.foreground().set_block_on_ticks(usize::MAX..=usize::MAX);
+
+        let grammar = tree_sitter_rust::language();
+        let text = r#"
+            fn outer() {}
+
+            mod module {
+                fn inner() {}
+            }"#
+        .unindent();
+        let highlight_query = tree_sitter::Query::new(
+            grammar,
+            r#"
+            (mod_item name: (identifier) body: _ @mod.body)
+            (function_item name: (identifier) @fn.name)"#,
+        )
+        .unwrap();
+        let theme = Theme::parse(
+            r#"
+            [syntax]
+            "mod.body" = 0xff0000
+            "fn.name" = 0x00ff00"#,
+        )
+        .unwrap();
+        let lang = Arc::new(Language {
+            config: LanguageConfig {
+                name: "Test".to_string(),
+                path_suffixes: vec![".test".to_string()],
+                ..Default::default()
+            },
+            grammar: grammar.clone(),
+            highlight_query,
+            brackets_query: tree_sitter::Query::new(grammar, "").unwrap(),
+            theme_mapping: Default::default(),
+        });
+        lang.set_theme(&theme);
+
+        let buffer = cx.add_model(|cx| {
+            Buffer::from_history(0, History::new(text.into()), None, Some(lang), cx)
+        });
+        buffer.condition(&cx, |buf, _| !buf.is_parsing()).await;
+
+        let font_cache = cx.font_cache();
+        let settings = Settings {
+            tab_size: 4,
+            buffer_font_family: font_cache.load_family(&["Courier"]).unwrap(),
+            buffer_font_size: 16.0,
+            ..Settings::new(&font_cache).unwrap()
+        };
+        let mut map = cx.read(|cx| DisplayMap::new(buffer, settings, Some(40.0), cx));
+        assert_eq!(
+            cx.read(|cx| highlighted_chunks(0..5, &map, &theme, cx)),
+            [
+                ("fn \n".to_string(), None),
+                ("oute\nr".to_string(), Some("fn.name")),
+                ("() \n{}\n\n".to_string(), None),
+            ]
+        );
+        assert_eq!(
+            cx.read(|cx| highlighted_chunks(3..5, &map, &theme, cx)),
+            [("{}\n\n".to_string(), None)]
+        );
+
+        cx.read(|cx| map.fold(vec![Point::new(0, 6)..Point::new(3, 2)], cx));
+        assert_eq!(
+            cx.read(|cx| highlighted_chunks(1..4, &map, &theme, cx)),
+            [
+                ("out".to_string(), Some("fn.name")),
+                ("…\n".to_string(), None),
+                ("  fn\n \n".to_string(), Some("mod.body"))
+            ]
+        );
     }
 
     #[gpui::test]
@@ -667,4 +723,26 @@ mod tests {
             DisplayPoint::new(1, 11)
         )
     }
+
+    fn highlighted_chunks<'a>(
+        rows: Range<u32>,
+        map: &DisplayMap,
+        theme: &'a Theme,
+        cx: &AppContext,
+    ) -> Vec<(String, Option<&'a str>)> {
+        let mut chunks: Vec<(String, Option<&str>)> = Vec::new();
+        for (chunk, style_id) in map.snapshot(cx).highlighted_chunks_for_rows(rows) {
+            let style_name = theme.syntax_style_name(style_id);
+            if let Some((last_chunk, last_style_name)) = chunks.last_mut() {
+                if style_name == *last_style_name {
+                    last_chunk.push_str(chunk);
+                } else {
+                    chunks.push((chunk.to_string(), style_name));
+                }
+            } else {
+                chunks.push((chunk.to_string(), style_name));
+            }
+        }
+        chunks
+    }
 }

zed/src/editor/display_map/wrap_map.rs

@@ -158,9 +158,10 @@ impl Snapshot {
         let input_end = self.to_input_point(output_end).min(self.input.max_point());
         HighlightedChunks {
             input_chunks: self.input.highlighted_chunks(input_start..input_end),
-            input_position: input_start,
-            style_id: StyleId::default(),
             input_chunk: "",
+            style_id: StyleId::default(),
+            output_position: output_start,
+            max_output_row: rows.end,
             transforms,
         }
     }
@@ -238,7 +239,8 @@ pub struct HighlightedChunks<'a> {
     input_chunks: tab_map::HighlightedChunks<'a>,
     input_chunk: &'a str,
     style_id: StyleId,
-    input_position: InputPoint,
+    output_position: OutputPoint,
+    max_output_row: u32,
     transforms: Cursor<'a, Transform, OutputPoint, InputPoint>,
 }
 
@@ -292,8 +294,13 @@ impl<'a> Iterator for HighlightedChunks<'a> {
     type Item = (&'a str, StyleId);
 
     fn next(&mut self) -> Option<Self::Item> {
+        if self.output_position.row() >= self.max_output_row {
+            return None;
+        }
+
         let transform = self.transforms.item()?;
         if let Some(display_text) = transform.display_text {
+            self.output_position.0 += transform.summary.output.lines;
             self.transforms.next(&());
             return Some((display_text, self.style_id));
         }
@@ -305,18 +312,18 @@ impl<'a> Iterator for HighlightedChunks<'a> {
         }
 
         let mut input_len = 0;
-        let transform_end = self.transforms.sum_end(&());
+        let transform_end = self.transforms.seek_end(&());
         for c in self.input_chunk.chars() {
             let char_len = c.len_utf8();
             input_len += char_len;
             if c == '\n' {
-                *self.input_position.row_mut() += 1;
-                *self.input_position.column_mut() = 0;
+                *self.output_position.row_mut() += 1;
+                *self.output_position.column_mut() = 0;
             } else {
-                *self.input_position.column_mut() += char_len as u32;
+                *self.output_position.column_mut() += char_len as u32;
             }
 
-            if self.input_position >= transform_end {
+            if self.output_position >= transform_end {
                 self.transforms.next(&());
                 break;
             }