1use collections::HashMap;
2use gpui::{
3 Animation, AnimationExt, AnyElement, Context, ImageSource, RenderImage, StyledText, Task, img,
4 pulsating_between,
5};
6use std::collections::BTreeMap;
7use std::ops::Range;
8use std::sync::{Arc, OnceLock};
9use std::time::Duration;
10use ui::prelude::*;
11
12use crate::parser::{CodeBlockKind, MarkdownEvent, MarkdownTag};
13
14use super::{Markdown, MarkdownStyle, ParsedMarkdown};
15
16type MermaidDiagramCache = HashMap<ParsedMarkdownMermaidDiagramContents, Arc<CachedMermaidDiagram>>;
17
18#[derive(Clone, Debug)]
19pub(crate) struct ParsedMarkdownMermaidDiagram {
20 pub(crate) content_range: Range<usize>,
21 pub(crate) contents: ParsedMarkdownMermaidDiagramContents,
22}
23
24#[derive(Clone, Debug, PartialEq, Eq, Hash)]
25pub(crate) struct ParsedMarkdownMermaidDiagramContents {
26 pub(crate) contents: SharedString,
27 pub(crate) scale: u32,
28}
29
30#[derive(Default, Clone)]
31pub(crate) struct MermaidState {
32 cache: MermaidDiagramCache,
33 order: Vec<ParsedMarkdownMermaidDiagramContents>,
34}
35
36struct CachedMermaidDiagram {
37 render_image: Arc<OnceLock<anyhow::Result<Arc<RenderImage>>>>,
38 fallback_image: Option<Arc<RenderImage>>,
39 _task: Task<()>,
40}
41
42impl MermaidState {
43 pub(crate) fn clear(&mut self) {
44 self.cache.clear();
45 self.order.clear();
46 }
47
48 fn get_fallback_image(
49 idx: usize,
50 old_order: &[ParsedMarkdownMermaidDiagramContents],
51 new_order_len: usize,
52 cache: &MermaidDiagramCache,
53 ) -> Option<Arc<RenderImage>> {
54 if old_order.len() != new_order_len {
55 return None;
56 }
57
58 old_order.get(idx).and_then(|old_content| {
59 cache.get(old_content).and_then(|old_cached| {
60 old_cached
61 .render_image
62 .get()
63 .and_then(|result| result.as_ref().ok().cloned())
64 .or_else(|| old_cached.fallback_image.clone())
65 })
66 })
67 }
68
69 pub(crate) fn update(&mut self, parsed: &ParsedMarkdown, cx: &mut Context<Markdown>) {
70 let mut new_order = Vec::new();
71 for mermaid_diagram in parsed.mermaid_diagrams.values() {
72 new_order.push(mermaid_diagram.contents.clone());
73 }
74
75 for (idx, new_content) in new_order.iter().enumerate() {
76 if !self.cache.contains_key(new_content) {
77 let fallback =
78 Self::get_fallback_image(idx, &self.order, new_order.len(), &self.cache);
79 self.cache.insert(
80 new_content.clone(),
81 Arc::new(CachedMermaidDiagram::new(new_content.clone(), fallback, cx)),
82 );
83 }
84 }
85
86 let new_order_set: std::collections::HashSet<_> = new_order.iter().cloned().collect();
87 self.cache
88 .retain(|content, _| new_order_set.contains(content));
89 self.order = new_order;
90 }
91}
92
93impl CachedMermaidDiagram {
94 fn new(
95 contents: ParsedMarkdownMermaidDiagramContents,
96 fallback_image: Option<Arc<RenderImage>>,
97 cx: &mut Context<Markdown>,
98 ) -> Self {
99 let render_image = Arc::new(OnceLock::<anyhow::Result<Arc<RenderImage>>>::new());
100 let render_image_clone = render_image.clone();
101 let svg_renderer = cx.svg_renderer();
102
103 let task = cx.spawn(async move |this, cx| {
104 let value = cx
105 .background_spawn(async move {
106 let svg_string = mermaid_rs_renderer::render(&contents.contents)?;
107 let scale = contents.scale as f32 / 100.0;
108 svg_renderer
109 .render_single_frame(svg_string.as_bytes(), scale)
110 .map_err(|error| anyhow::anyhow!("{error}"))
111 })
112 .await;
113 let _ = render_image_clone.set(value);
114 this.update(cx, |_, cx| {
115 cx.notify();
116 })
117 .ok();
118 });
119
120 Self {
121 render_image,
122 fallback_image,
123 _task: task,
124 }
125 }
126
127 #[cfg(test)]
128 fn new_for_test(
129 render_image: Option<Arc<RenderImage>>,
130 fallback_image: Option<Arc<RenderImage>>,
131 ) -> Self {
132 let result = Arc::new(OnceLock::new());
133 if let Some(render_image) = render_image {
134 let _ = result.set(Ok(render_image));
135 }
136 Self {
137 render_image: result,
138 fallback_image,
139 _task: Task::ready(()),
140 }
141 }
142}
143
144fn parse_mermaid_info(info: &str) -> Option<u32> {
145 let mut parts = info.split_whitespace();
146 if parts.next()? != "mermaid" {
147 return None;
148 }
149
150 Some(
151 parts
152 .next()
153 .and_then(|scale| scale.parse().ok())
154 .unwrap_or(100)
155 .clamp(10, 500),
156 )
157}
158
159pub(crate) fn extract_mermaid_diagrams(
160 source: &str,
161 events: &[(Range<usize>, MarkdownEvent)],
162) -> BTreeMap<usize, ParsedMarkdownMermaidDiagram> {
163 let mut mermaid_diagrams = BTreeMap::default();
164
165 for (source_range, event) in events {
166 let MarkdownEvent::Start(MarkdownTag::CodeBlock { kind, metadata }) = event else {
167 continue;
168 };
169 let CodeBlockKind::FencedLang(info) = kind else {
170 continue;
171 };
172 let Some(scale) = parse_mermaid_info(info.as_ref()) else {
173 continue;
174 };
175
176 let contents = source[metadata.content_range.clone()]
177 .strip_suffix('\n')
178 .unwrap_or(&source[metadata.content_range.clone()])
179 .to_string();
180 mermaid_diagrams.insert(
181 source_range.start,
182 ParsedMarkdownMermaidDiagram {
183 content_range: metadata.content_range.clone(),
184 contents: ParsedMarkdownMermaidDiagramContents {
185 contents: contents.into(),
186 scale,
187 },
188 },
189 );
190 }
191
192 mermaid_diagrams
193}
194
195pub(crate) fn render_mermaid_diagram(
196 parsed: &ParsedMarkdownMermaidDiagram,
197 mermaid_state: &MermaidState,
198 style: &MarkdownStyle,
199) -> AnyElement {
200 let cached = mermaid_state.cache.get(&parsed.contents);
201 let mut container = div().w_full();
202 container.style().refine(&style.code_block);
203
204 if let Some(result) = cached.and_then(|cached| cached.render_image.get()) {
205 match result {
206 Ok(render_image) => container
207 .child(
208 div().w_full().child(
209 img(ImageSource::Render(render_image.clone()))
210 .max_w_full()
211 .with_fallback(|| {
212 div()
213 .child(Label::new("Failed to load mermaid diagram"))
214 .into_any_element()
215 }),
216 ),
217 )
218 .into_any_element(),
219 Err(_) => container
220 .child(StyledText::new(parsed.contents.contents.clone()))
221 .into_any_element(),
222 }
223 } else if let Some(fallback) = cached.and_then(|cached| cached.fallback_image.as_ref()) {
224 container
225 .child(
226 div()
227 .w_full()
228 .child(
229 img(ImageSource::Render(fallback.clone()))
230 .max_w_full()
231 .with_fallback(|| {
232 div()
233 .child(Label::new("Failed to load mermaid diagram"))
234 .into_any_element()
235 }),
236 )
237 .with_animation(
238 "mermaid-fallback-pulse",
239 Animation::new(Duration::from_secs(2))
240 .repeat()
241 .with_easing(pulsating_between(0.6, 1.0)),
242 |element, delta| element.opacity(delta),
243 ),
244 )
245 .into_any_element()
246 } else {
247 container
248 .child(
249 Label::new("Rendering mermaid diagram...")
250 .color(Color::Muted)
251 .with_animation(
252 "mermaid-loading-pulse",
253 Animation::new(Duration::from_secs(2))
254 .repeat()
255 .with_easing(pulsating_between(0.4, 0.8)),
256 |label, delta| label.alpha(delta),
257 ),
258 )
259 .into_any_element()
260 }
261}
262
263#[cfg(test)]
264mod tests {
265 use super::{
266 CachedMermaidDiagram, MermaidDiagramCache, MermaidState,
267 ParsedMarkdownMermaidDiagramContents, extract_mermaid_diagrams, parse_mermaid_info,
268 };
269 use crate::{CodeBlockRenderer, Markdown, MarkdownElement, MarkdownOptions, MarkdownStyle};
270 use collections::HashMap;
271 use gpui::{Context, IntoElement, Render, RenderImage, TestAppContext, Window, size};
272 use std::sync::Arc;
273 use ui::prelude::*;
274
275 fn ensure_theme_initialized(cx: &mut TestAppContext) {
276 cx.update(|cx| {
277 if !cx.has_global::<settings::SettingsStore>() {
278 settings::init(cx);
279 }
280 if !cx.has_global::<theme::GlobalTheme>() {
281 theme_settings::init(theme::LoadThemes::JustBase, cx);
282 }
283 });
284 }
285
286 fn render_markdown_with_options(
287 markdown: &str,
288 options: MarkdownOptions,
289 cx: &mut TestAppContext,
290 ) -> crate::RenderedText {
291 struct TestWindow;
292
293 impl Render for TestWindow {
294 fn render(&mut self, _: &mut Window, _: &mut Context<Self>) -> impl IntoElement {
295 div()
296 }
297 }
298
299 ensure_theme_initialized(cx);
300
301 let (_, cx) = cx.add_window_view(|_, _| TestWindow);
302 let markdown = cx.new(|cx| {
303 Markdown::new_with_options(markdown.to_string().into(), None, None, options, cx)
304 });
305 cx.run_until_parked();
306 let (rendered, _) = cx.draw(
307 Default::default(),
308 size(px(600.0), px(600.0)),
309 |_window, _cx| {
310 MarkdownElement::new(markdown, MarkdownStyle::default()).code_block_renderer(
311 CodeBlockRenderer::Default {
312 copy_button: false,
313 copy_button_on_hover: false,
314 border: false,
315 },
316 )
317 },
318 );
319 rendered.text
320 }
321
322 fn mock_render_image(cx: &mut TestAppContext) -> Arc<RenderImage> {
323 cx.update(|cx| {
324 cx.svg_renderer()
325 .render_single_frame(
326 br#"<svg xmlns="http://www.w3.org/2000/svg" width="1" height="1"></svg>"#,
327 1.0,
328 )
329 .unwrap()
330 })
331 }
332
333 fn mermaid_contents(contents: &str) -> ParsedMarkdownMermaidDiagramContents {
334 ParsedMarkdownMermaidDiagramContents {
335 contents: contents.to_string().into(),
336 scale: 100,
337 }
338 }
339
340 fn mermaid_sequence(diagrams: &[&str]) -> Vec<ParsedMarkdownMermaidDiagramContents> {
341 diagrams
342 .iter()
343 .map(|diagram| mermaid_contents(diagram))
344 .collect()
345 }
346
347 fn mermaid_fallback(
348 new_diagram: &str,
349 new_full_order: &[ParsedMarkdownMermaidDiagramContents],
350 old_full_order: &[ParsedMarkdownMermaidDiagramContents],
351 cache: &MermaidDiagramCache,
352 ) -> Option<Arc<RenderImage>> {
353 let new_content = mermaid_contents(new_diagram);
354 let idx = new_full_order
355 .iter()
356 .position(|diagram| diagram == &new_content)?;
357 MermaidState::get_fallback_image(idx, old_full_order, new_full_order.len(), cache)
358 }
359
360 #[test]
361 fn test_parse_mermaid_info() {
362 assert_eq!(parse_mermaid_info("mermaid"), Some(100));
363 assert_eq!(parse_mermaid_info("mermaid 150"), Some(150));
364 assert_eq!(parse_mermaid_info("mermaid 5"), Some(10));
365 assert_eq!(parse_mermaid_info("mermaid 999"), Some(500));
366 assert_eq!(parse_mermaid_info("rust"), None);
367 }
368
369 #[test]
370 fn test_extract_mermaid_diagrams_parses_scale() {
371 let markdown = "```mermaid 150\ngraph TD;\n```\n\n```rust\nfn main() {}\n```";
372 let events = crate::parser::parse_markdown_with_options(markdown, false).events;
373 let diagrams = extract_mermaid_diagrams(markdown, &events);
374
375 assert_eq!(diagrams.len(), 1);
376 let diagram = diagrams.values().next().unwrap();
377 assert_eq!(diagram.contents.contents, "graph TD;");
378 assert_eq!(diagram.contents.scale, 150);
379 }
380
381 #[gpui::test]
382 fn test_mermaid_fallback_on_edit(cx: &mut TestAppContext) {
383 let old_full_order = mermaid_sequence(&["graph A", "graph B", "graph C"]);
384 let new_full_order = mermaid_sequence(&["graph A", "graph B modified", "graph C"]);
385
386 let svg_b = mock_render_image(cx);
387
388 let mut cache: MermaidDiagramCache = HashMap::default();
389 cache.insert(
390 mermaid_contents("graph A"),
391 Arc::new(CachedMermaidDiagram::new_for_test(
392 Some(mock_render_image(cx)),
393 None,
394 )),
395 );
396 cache.insert(
397 mermaid_contents("graph B"),
398 Arc::new(CachedMermaidDiagram::new_for_test(
399 Some(svg_b.clone()),
400 None,
401 )),
402 );
403 cache.insert(
404 mermaid_contents("graph C"),
405 Arc::new(CachedMermaidDiagram::new_for_test(
406 Some(mock_render_image(cx)),
407 None,
408 )),
409 );
410
411 let fallback =
412 mermaid_fallback("graph B modified", &new_full_order, &old_full_order, &cache);
413
414 assert_eq!(fallback.as_ref().map(|image| image.id), Some(svg_b.id));
415 }
416
417 #[gpui::test]
418 fn test_mermaid_no_fallback_on_add_in_middle(cx: &mut TestAppContext) {
419 let old_full_order = mermaid_sequence(&["graph A", "graph C"]);
420 let new_full_order = mermaid_sequence(&["graph A", "graph NEW", "graph C"]);
421
422 let mut cache: MermaidDiagramCache = HashMap::default();
423 cache.insert(
424 mermaid_contents("graph A"),
425 Arc::new(CachedMermaidDiagram::new_for_test(
426 Some(mock_render_image(cx)),
427 None,
428 )),
429 );
430 cache.insert(
431 mermaid_contents("graph C"),
432 Arc::new(CachedMermaidDiagram::new_for_test(
433 Some(mock_render_image(cx)),
434 None,
435 )),
436 );
437
438 let fallback = mermaid_fallback("graph NEW", &new_full_order, &old_full_order, &cache);
439
440 assert!(fallback.is_none());
441 }
442
443 #[gpui::test]
444 fn test_mermaid_fallback_chains_on_rapid_edits(cx: &mut TestAppContext) {
445 let old_full_order = mermaid_sequence(&["graph A", "graph B modified", "graph C"]);
446 let new_full_order = mermaid_sequence(&["graph A", "graph B modified again", "graph C"]);
447
448 let original_svg = mock_render_image(cx);
449
450 let mut cache: MermaidDiagramCache = HashMap::default();
451 cache.insert(
452 mermaid_contents("graph A"),
453 Arc::new(CachedMermaidDiagram::new_for_test(
454 Some(mock_render_image(cx)),
455 None,
456 )),
457 );
458 cache.insert(
459 mermaid_contents("graph B modified"),
460 Arc::new(CachedMermaidDiagram::new_for_test(
461 None,
462 Some(original_svg.clone()),
463 )),
464 );
465 cache.insert(
466 mermaid_contents("graph C"),
467 Arc::new(CachedMermaidDiagram::new_for_test(
468 Some(mock_render_image(cx)),
469 None,
470 )),
471 );
472
473 let fallback = mermaid_fallback(
474 "graph B modified again",
475 &new_full_order,
476 &old_full_order,
477 &cache,
478 );
479
480 assert_eq!(
481 fallback.as_ref().map(|image| image.id),
482 Some(original_svg.id)
483 );
484 }
485
486 #[gpui::test]
487 fn test_mermaid_fallback_with_duplicate_blocks_edit_second(cx: &mut TestAppContext) {
488 let old_full_order = mermaid_sequence(&["graph A", "graph A", "graph B"]);
489 let new_full_order = mermaid_sequence(&["graph A", "graph A edited", "graph B"]);
490
491 let svg_a = mock_render_image(cx);
492
493 let mut cache: MermaidDiagramCache = HashMap::default();
494 cache.insert(
495 mermaid_contents("graph A"),
496 Arc::new(CachedMermaidDiagram::new_for_test(
497 Some(svg_a.clone()),
498 None,
499 )),
500 );
501 cache.insert(
502 mermaid_contents("graph B"),
503 Arc::new(CachedMermaidDiagram::new_for_test(
504 Some(mock_render_image(cx)),
505 None,
506 )),
507 );
508
509 let fallback = mermaid_fallback("graph A edited", &new_full_order, &old_full_order, &cache);
510
511 assert_eq!(fallback.as_ref().map(|image| image.id), Some(svg_a.id));
512 }
513
514 #[gpui::test]
515 fn test_mermaid_rendering_replaces_code_block_text(cx: &mut TestAppContext) {
516 let rendered = render_markdown_with_options(
517 "```mermaid\ngraph TD;\n```",
518 MarkdownOptions {
519 render_mermaid_diagrams: true,
520 ..Default::default()
521 },
522 cx,
523 );
524
525 let text = rendered
526 .lines
527 .iter()
528 .map(|line| line.layout.wrapped_text())
529 .collect::<Vec<_>>()
530 .join("\n");
531
532 assert!(!text.contains("graph TD;"));
533 }
534
535 #[gpui::test]
536 fn test_mermaid_source_anchor_maps_inside_block(cx: &mut TestAppContext) {
537 struct TestWindow;
538
539 impl Render for TestWindow {
540 fn render(&mut self, _: &mut Window, _: &mut Context<Self>) -> impl IntoElement {
541 div()
542 }
543 }
544
545 ensure_theme_initialized(cx);
546
547 let (_, cx) = cx.add_window_view(|_, _| TestWindow);
548 let markdown = cx.new(|cx| {
549 Markdown::new_with_options(
550 "```mermaid\ngraph TD;\n```".into(),
551 None,
552 None,
553 MarkdownOptions {
554 render_mermaid_diagrams: true,
555 ..Default::default()
556 },
557 cx,
558 )
559 });
560 cx.run_until_parked();
561 let render_image = mock_render_image(cx);
562 markdown.update(cx, |markdown, _| {
563 let contents = markdown
564 .parsed_markdown
565 .mermaid_diagrams
566 .values()
567 .next()
568 .unwrap()
569 .contents
570 .clone();
571 markdown.mermaid_state.cache.insert(
572 contents.clone(),
573 Arc::new(CachedMermaidDiagram::new_for_test(Some(render_image), None)),
574 );
575 markdown.mermaid_state.order = vec![contents];
576 });
577
578 let (rendered, _) = cx.draw(
579 Default::default(),
580 size(px(600.0), px(600.0)),
581 |_window, _cx| {
582 MarkdownElement::new(markdown.clone(), MarkdownStyle::default())
583 .code_block_renderer(CodeBlockRenderer::Default {
584 copy_button: false,
585 copy_button_on_hover: false,
586 border: false,
587 })
588 },
589 );
590
591 let mermaid_diagram = markdown.update(cx, |markdown, _| {
592 markdown
593 .parsed_markdown
594 .mermaid_diagrams
595 .values()
596 .next()
597 .unwrap()
598 .clone()
599 });
600 assert!(
601 rendered
602 .text
603 .position_for_source_index(mermaid_diagram.content_range.start)
604 .is_some()
605 );
606 assert!(
607 rendered
608 .text
609 .position_for_source_index(mermaid_diagram.content_range.end.saturating_sub(1))
610 .is_some()
611 );
612 }
613}