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 sourceClass fill:#8dd3c7,stroke:#86c8bd,text-align:left,white-space:pre
62{b:i$}classDef transformClass fill:#ffffb3,stroke:#f5f5a8,text-align:left,white-space:pre
63{b:i$}classDef joinClass fill:#bebada,stroke:#b5b1cf,text-align:left,white-space:pre
64{b:i$}classDef aggClass fill:#fb8072,stroke:#ee796b,text-align:left,white-space:pre
65{b:i$}classDef networkClass fill:#80b1d3,stroke:#79a8c8,text-align:left,white-space:pre
66{b:i$}classDef sinkClass fill:#fdb462,stroke:#f0aa5b,text-align:left,white-space:pre
67{b:i$}classDef teeClass fill:#b3de69,stroke:#aad362,text-align:left,white-space:pre
68{b:i$}classDef nondetClass fill:#fccde5,stroke:#f3c4dc,text-align:left,white-space:pre
69{b:i$}linkStyle default stroke:#666666",
70 b = "",
71 i = self.base.indent
72 )?;
73 Ok(())
74 }
75
76 fn write_node_definition(
77 &mut self,
78 node_id: VizNodeKey,
79 node_label: &super::render::NodeLabel,
80 node_type: HydroNodeType,
81 _location_id: Option<LocationKey>,
82 _location_type: Option<LocationType>,
83 _backtrace: Option<&crate::compile::ir::backtrace::Backtrace>,
84 ) -> Result<(), Self::Err> {
85 let class_str = match node_type {
86 HydroNodeType::Source => "sourceClass",
87 HydroNodeType::Transform => "transformClass",
88 HydroNodeType::Join => "joinClass",
89 HydroNodeType::Aggregation => "aggClass",
90 HydroNodeType::Network => "networkClass",
91 HydroNodeType::Sink => "sinkClass",
92 HydroNodeType::Tee => "teeClass",
93 HydroNodeType::NonDeterministic => "nondetClass",
94 };
95
96 let (lbracket, rbracket) = match node_type {
97 HydroNodeType::Source => ("[[", "]]"),
98 HydroNodeType::Sink => ("[/", "/]"),
99 HydroNodeType::Network => ("[[", "]]"),
100 HydroNodeType::Tee => ("(", ")"),
101 _ => ("[", "]"),
102 };
103
104 let full_label = match node_label {
106 super::render::NodeLabel::Static(s) => s.clone(),
107 super::render::NodeLabel::WithExprs { op_name, exprs } => {
108 if exprs.is_empty() {
109 format!("{}()", op_name)
110 } else {
111 let expr_strs: Vec<String> = exprs.iter().map(|e| e.to_string()).collect();
113 format!("{}({})", op_name, expr_strs.join(", "))
114 }
115 }
116 };
117
118 let display_label = if self.base.config.use_short_labels {
120 super::render::extract_short_label(&full_label)
121 } else {
122 full_label
123 };
124
125 let label = format!(
126 r#"n{node_id}{lbracket}"{escaped_label}"{rbracket}:::{class}"#,
127 escaped_label = escape_mermaid(&display_label),
128 class = class_str,
129 );
130
131 writeln!(
132 self.base.write,
133 "{b:i$}{label}",
134 b = "",
135 i = self.base.indent
136 )?;
137 Ok(())
138 }
139
140 fn write_edge(
141 &mut self,
142 src_id: VizNodeKey,
143 dst_id: VizNodeKey,
144 edge_properties: &std::collections::HashSet<HydroEdgeProp>,
145 label: Option<&str>,
146 ) -> Result<(), Self::Err> {
147 let style = super::render::get_unified_edge_style(edge_properties, None, None);
149
150 let arrow_style = if edge_properties.contains(&HydroEdgeProp::Network) {
152 "-.->".to_owned()
153 } else {
154 match style.line_pattern {
155 super::render::LinePattern::Dotted => "-.->".to_owned(),
156 super::render::LinePattern::Dashed => "--o".to_owned(),
157 _ => {
158 if style.line_width > 1 {
159 "==>".to_owned()
160 } else {
161 "-->".to_owned()
162 }
163 }
164 }
165 };
166
167 writeln!(
169 self.base.write,
170 "{b:i$}n{src}{arrow}{label}n{dst}",
171 src = src_id,
172 arrow = arrow_style,
173 label = if let Some(label) = label {
174 Cow::Owned(format!("|{}|", escape_mermaid(label)))
175 } else {
176 Cow::Borrowed("")
177 },
178 dst = dst_id,
179 b = "",
180 i = self.base.indent,
181 )?;
182
183 writeln!(
185 self.base.write,
186 "{b:i$}linkStyle {} stroke:{}",
187 self.link_count,
188 style.color,
189 b = "",
190 i = self.base.indent,
191 )?;
192
193 self.link_count += 1;
194 Ok(())
195 }
196
197 fn write_location_start(
198 &mut self,
199 location_key: LocationKey,
200 location_type: LocationType,
201 ) -> Result<(), Self::Err> {
202 writeln!(
203 self.base.write,
204 "{b:i$}subgraph {loc} [\"{location_type:?} {loc}\"]",
205 loc = location_key,
206 b = "",
207 i = self.base.indent,
208 )?;
209 self.base.indent += 4;
210 Ok(())
211 }
212
213 fn write_node(&mut self, node_id: VizNodeKey) -> Result<(), Self::Err> {
214 writeln!(
215 self.base.write,
216 "{b:i$}n{node_id}",
217 b = "",
218 i = self.base.indent
219 )
220 }
221
222 fn write_location_end(&mut self) -> Result<(), Self::Err> {
223 self.base.indent -= 4;
224 writeln!(self.base.write, "{b:i$}end", b = "", i = self.base.indent)
225 }
226
227 fn write_epilogue(&mut self) -> Result<(), Self::Err> {
228 Ok(())
229 }
230}
231
232#[cfg(feature = "build")]
234pub fn open_browser(
235 built_flow: &crate::compile::built::BuiltFlow,
236) -> Result<(), Box<dyn std::error::Error>> {
237 let config = HydroWriteConfig {
238 show_metadata: false,
239 show_location_groups: true,
240 use_short_labels: true, location_names: built_flow.location_names(),
242 };
243
244 crate::viz::debug::open_mermaid(built_flow.ir(), Some(config))?;
246
247 Ok(())
248}