1use std::borrow::Cow;
2use std::fmt::Write;
3
4use super::render::{HydroEdgeProp, HydroGraphWrite, HydroNodeType, IndentedGraphWriter};
5use crate::viz::render::VizNodeKey;
6
7pub fn escape_mermaid(string: &str) -> String {
9 string
10 .replace('&', "&")
11 .replace('<', "<")
12 .replace('>', ">")
13 .replace('"', """)
14 .replace('#', "#")
15 .replace('\n', "<br>")
16 .replace("`", "`")
18 .replace('(', "(")
20 .replace(')', ")")
21 .replace('|', "|")
23}
24
25pub struct HydroMermaid<W> {
27 base: IndentedGraphWriter<W>,
28 link_count: usize,
29}
30
31impl<W> HydroMermaid<W> {
32 pub fn new(write: W) -> Self {
33 Self {
34 base: IndentedGraphWriter::new(write),
35 link_count: 0,
36 }
37 }
38
39 pub fn new_with_config(write: W, config: &super::render::HydroWriteConfig) -> Self {
40 Self {
41 base: IndentedGraphWriter::new_with_config(write, config),
42 link_count: 0,
43 }
44 }
45}
46
47impl<W> HydroGraphWrite for HydroMermaid<W>
48where
49 W: Write,
50{
51 type Err = super::render::GraphWriteError;
52
53 fn write_prologue(&mut self) -> Result<(), Self::Err> {
54 writeln!(
55 self.base.write,
56 "{b:i$}%%{{init:{{'theme':'base','themeVariables':{{'clusterBkg':'#fafafa','clusterBorder':'#e0e0e0'}},'elk':{{'algorithm':'mrtree','elk.direction':'DOWN','elk.layered.spacing.nodeNodeBetweenLayers':'30'}}}}}}%%
57{b:i$}graph TD
58{b:i$}classDef sourceClass fill:#8dd3c7,stroke:#86c8bd,text-align:left,white-space:pre
59{b:i$}classDef transformClass fill:#ffffb3,stroke:#f5f5a8,text-align:left,white-space:pre
60{b:i$}classDef joinClass fill:#bebada,stroke:#b5b1cf,text-align:left,white-space:pre
61{b:i$}classDef aggClass fill:#fb8072,stroke:#ee796b,text-align:left,white-space:pre
62{b:i$}classDef networkClass fill:#80b1d3,stroke:#79a8c8,text-align:left,white-space:pre
63{b:i$}classDef sinkClass fill:#fdb462,stroke:#f0aa5b,text-align:left,white-space:pre
64{b:i$}classDef teeClass fill:#b3de69,stroke:#aad362,text-align:left,white-space:pre
65{b:i$}classDef nondetClass fill:#fccde5,stroke:#f3c4dc,text-align:left,white-space:pre
66{b:i$}linkStyle default stroke:#666666",
67 b = "",
68 i = self.base.indent
69 )?;
70 Ok(())
71 }
72
73 fn write_node_definition(
74 &mut self,
75 node_id: VizNodeKey,
76 node_label: &super::render::NodeLabel,
77 node_type: HydroNodeType,
78 _location_id: Option<usize>,
79 _location_type: Option<&str>,
80 _backtrace: Option<&crate::compile::ir::backtrace::Backtrace>,
81 ) -> Result<(), Self::Err> {
82 let class_str = match node_type {
83 HydroNodeType::Source => "sourceClass",
84 HydroNodeType::Transform => "transformClass",
85 HydroNodeType::Join => "joinClass",
86 HydroNodeType::Aggregation => "aggClass",
87 HydroNodeType::Network => "networkClass",
88 HydroNodeType::Sink => "sinkClass",
89 HydroNodeType::Tee => "teeClass",
90 HydroNodeType::NonDeterministic => "nondetClass",
91 };
92
93 let (lbracket, rbracket) = match node_type {
94 HydroNodeType::Source => ("[[", "]]"),
95 HydroNodeType::Sink => ("[/", "/]"),
96 HydroNodeType::Network => ("[[", "]]"),
97 HydroNodeType::Tee => ("(", ")"),
98 _ => ("[", "]"),
99 };
100
101 let full_label = match node_label {
103 super::render::NodeLabel::Static(s) => s.clone(),
104 super::render::NodeLabel::WithExprs { op_name, exprs } => {
105 if exprs.is_empty() {
106 format!("{}()", op_name)
107 } else {
108 let expr_strs: Vec<String> = exprs.iter().map(|e| e.to_string()).collect();
110 format!("{}({})", op_name, expr_strs.join(", "))
111 }
112 }
113 };
114
115 let display_label = if self.base.config.use_short_labels {
117 super::render::extract_short_label(&full_label)
118 } else {
119 full_label
120 };
121
122 let label = format!(
123 r#"n{node_id}{lbracket}"{escaped_label}"{rbracket}:::{class}"#,
124 escaped_label = escape_mermaid(&display_label),
125 class = class_str,
126 );
127
128 writeln!(
129 self.base.write,
130 "{b:i$}{label}",
131 b = "",
132 i = self.base.indent
133 )?;
134 Ok(())
135 }
136
137 fn write_edge(
138 &mut self,
139 src_id: VizNodeKey,
140 dst_id: VizNodeKey,
141 edge_properties: &std::collections::HashSet<HydroEdgeProp>,
142 label: Option<&str>,
143 ) -> Result<(), Self::Err> {
144 let style = super::render::get_unified_edge_style(edge_properties, None, None);
146
147 let arrow_style = if edge_properties.contains(&HydroEdgeProp::Network) {
149 "-.->".to_string()
150 } else {
151 match style.line_pattern {
152 super::render::LinePattern::Dotted => "-.->".to_string(),
153 super::render::LinePattern::Dashed => "--o".to_string(),
154 _ => {
155 if style.line_width > 1 {
156 "==>".to_string()
157 } else {
158 "-->".to_string()
159 }
160 }
161 }
162 };
163
164 writeln!(
166 self.base.write,
167 "{b:i$}n{src}{arrow}{label}n{dst}",
168 src = src_id,
169 arrow = arrow_style,
170 label = if let Some(label) = label {
171 Cow::Owned(format!("|{}|", escape_mermaid(label)))
172 } else {
173 Cow::Borrowed("")
174 },
175 dst = dst_id,
176 b = "",
177 i = self.base.indent,
178 )?;
179
180 writeln!(
182 self.base.write,
183 "{b:i$}linkStyle {} stroke:{}",
184 self.link_count,
185 style.color,
186 b = "",
187 i = self.base.indent,
188 )?;
189
190 self.link_count += 1;
191 Ok(())
192 }
193
194 fn write_location_start(
195 &mut self,
196 location_id: usize,
197 location_type: &str,
198 ) -> Result<(), Self::Err> {
199 writeln!(
200 self.base.write,
201 "{b:i$}subgraph loc_{id} [\"{location_type} {id}\"]",
202 id = location_id,
203 b = "",
204 i = self.base.indent,
205 )?;
206 self.base.indent += 4;
207 Ok(())
208 }
209
210 fn write_node(&mut self, node_id: VizNodeKey) -> Result<(), Self::Err> {
211 writeln!(
212 self.base.write,
213 "{b:i$}n{node_id}",
214 b = "",
215 i = self.base.indent
216 )
217 }
218
219 fn write_location_end(&mut self) -> Result<(), Self::Err> {
220 self.base.indent -= 4;
221 writeln!(self.base.write, "{b:i$}end", b = "", i = self.base.indent)
222 }
223
224 fn write_epilogue(&mut self) -> Result<(), Self::Err> {
225 Ok(())
226 }
227}
228
229#[cfg(feature = "build")]
231pub fn open_browser(
232 built_flow: &crate::compile::built::BuiltFlow,
233) -> Result<(), Box<dyn std::error::Error>> {
234 let config = super::render::HydroWriteConfig {
235 show_metadata: false,
236 show_location_groups: true,
237 use_short_labels: true, process_id_name: built_flow.process_id_name().clone(),
239 cluster_id_name: built_flow.cluster_id_name().clone(),
240 external_id_name: built_flow.external_id_name().clone(),
241 };
242
243 crate::viz::debug::open_mermaid(built_flow.ir(), Some(config))?;
245
246 Ok(())
247}