1use std::borrow::Cow;
2use std::fmt::Write;
3
4use super::render::{
5 HydroEdgeProp, HydroGraphWrite, HydroNodeType, HydroWriteConfig, IndentedGraphWriter,
6};
7use crate::location::{LocationKey, LocationType};
8use crate::viz::render::VizNodeKey;
9
10pub fn escape_mermaid(string: &str) -> String {
12 string
13 .replace('&', "&")
14 .replace('<', "<")
15 .replace('>', ">")
16 .replace('"', """)
17 .replace('#', "#")
18 .replace('\n', "<br>")
19 .replace("`", "`")
21 .replace('(', "(")
23 .replace(')', ")")
24 .replace('|', "|")
26}
27
28pub struct HydroMermaid<'a, W> {
30 base: IndentedGraphWriter<'a, W>,
31 link_count: usize,
32}
33
34impl<'a, W> HydroMermaid<'a, W> {
35 pub fn new(write: W) -> Self {
36 Self {
37 base: IndentedGraphWriter::new(write),
38 link_count: 0,
39 }
40 }
41
42 pub fn new_with_config(write: W, config: HydroWriteConfig<'a>) -> Self {
43 Self {
44 base: IndentedGraphWriter::new_with_config(write, config),
45 link_count: 0,
46 }
47 }
48}
49
50impl<W> HydroGraphWrite for HydroMermaid<'_, W>
51where
52 W: Write,
53{
54 type Err = super::render::GraphWriteError;
55
56 fn write_prologue(&mut self) -> Result<(), Self::Err> {
57 writeln!(
58 self.base.write,
59 "{b:i$}%%{{init:{{'theme':'base','themeVariables':{{'clusterBkg':'#fafafa','clusterBorder':'#e0e0e0'}},'elk':{{'algorithm':'mrtree','elk.direction':'DOWN','elk.layered.spacing.nodeNodeBetweenLayers':'30'}}}}}}%%
60{b:i$}graph TD
61{b:i$}classDef default fill:#f5f5f5,stroke:#bbb,text-align:left,white-space:pre
62{b:i$}linkStyle default stroke:#666666",
63 b = "",
64 i = self.base.indent
65 )?;
66 Ok(())
67 }
68
69 fn write_node_definition(
70 &mut self,
71 node_id: VizNodeKey,
72 node_label: &super::render::NodeLabel,
73 _node_type: HydroNodeType,
74 _location_id: Option<LocationKey>,
75 _location_type: Option<LocationType>,
76 _backtrace: Option<&crate::compile::ir::backtrace::Backtrace>,
77 ) -> Result<(), Self::Err> {
78 let full_label = match node_label {
80 super::render::NodeLabel::Static(s) => s.clone(),
81 super::render::NodeLabel::WithExprs { op_name, exprs } => {
82 if exprs.is_empty() {
83 format!("{}()", op_name)
84 } else {
85 let expr_strs: Vec<String> = exprs.iter().map(|e| e.to_string()).collect();
87 format!("{}({})", op_name, expr_strs.join(", "))
88 }
89 }
90 };
91
92 let display_label = if self.base.config.use_short_labels {
94 super::render::extract_short_label(&full_label)
95 } else {
96 full_label
97 };
98
99 writeln!(
100 self.base.write,
101 "{b:i$}n{node_id}[\"{escaped_label}\"]",
102 escaped_label = escape_mermaid(&display_label),
103 b = "",
104 i = self.base.indent
105 )?;
106 Ok(())
107 }
108
109 fn write_edge(
110 &mut self,
111 src_id: VizNodeKey,
112 dst_id: VizNodeKey,
113 edge_properties: &std::collections::HashSet<HydroEdgeProp>,
114 label: Option<&str>,
115 ) -> Result<(), Self::Err> {
116 let style = super::render::get_unified_edge_style(edge_properties, None, None);
118
119 let arrow_style = if edge_properties.contains(&HydroEdgeProp::Network) {
121 "-.->".to_owned()
122 } else {
123 match style.line_pattern {
124 super::render::LinePattern::Dotted => "-.->".to_owned(),
125 super::render::LinePattern::Dashed => "--o".to_owned(),
126 _ => {
127 if style.line_width > 1 {
128 "==>".to_owned()
129 } else {
130 "-->".to_owned()
131 }
132 }
133 }
134 };
135
136 writeln!(
138 self.base.write,
139 "{b:i$}n{src}{arrow}{label}n{dst}",
140 src = src_id,
141 arrow = arrow_style,
142 label = if let Some(label) = label {
143 Cow::Owned(format!("|{}|", escape_mermaid(label)))
144 } else {
145 Cow::Borrowed("")
146 },
147 dst = dst_id,
148 b = "",
149 i = self.base.indent,
150 )?;
151
152 writeln!(
154 self.base.write,
155 "{b:i$}linkStyle {} stroke:{}",
156 self.link_count,
157 style.color,
158 b = "",
159 i = self.base.indent,
160 )?;
161
162 self.link_count += 1;
163 Ok(())
164 }
165
166 fn write_location_start(
167 &mut self,
168 location_key: LocationKey,
169 location_type: LocationType,
170 ) -> Result<(), Self::Err> {
171 writeln!(
172 self.base.write,
173 "{b:i$}subgraph {loc} [\"{location_type:?} {loc}\"]",
174 loc = location_key,
175 b = "",
176 i = self.base.indent,
177 )?;
178 self.base.indent += 4;
179 Ok(())
180 }
181
182 fn write_node(&mut self, node_id: VizNodeKey) -> Result<(), Self::Err> {
183 writeln!(
184 self.base.write,
185 "{b:i$}n{node_id}",
186 b = "",
187 i = self.base.indent
188 )
189 }
190
191 fn write_location_end(&mut self) -> Result<(), Self::Err> {
192 self.base.indent -= 4;
193 writeln!(self.base.write, "{b:i$}end", b = "", i = self.base.indent)
194 }
195
196 fn write_epilogue(&mut self) -> Result<(), Self::Err> {
197 Ok(())
198 }
199}