hydro_lang/viz/
mermaid.rs

1use std::borrow::Cow;
2use std::fmt::Write;
3
4use super::render::{HydroEdgeProp, HydroGraphWrite, HydroNodeType, IndentedGraphWriter};
5use crate::viz::render::VizNodeKey;
6
7/// Escapes a string for use in a mermaid graph label.
8pub fn escape_mermaid(string: &str) -> String {
9    string
10        .replace('&', "&")
11        .replace('<', "&lt;")
12        .replace('>', "&gt;")
13        .replace('"', "&quot;")
14        .replace('#', "&num;")
15        .replace('\n', "<br>")
16        // Handle code block markers
17        .replace("`", "&#96;")
18        // Handle parentheses that can conflict with Mermaid syntax
19        .replace('(', "&#40;")
20        .replace(')', "&#41;")
21        // Handle pipes that can conflict with Mermaid edge labels
22        .replace('|', "&#124;")
23}
24
25/// Mermaid graph writer for Hydro IR.
26pub 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        // Create the full label string using DebugExpr::Display for expressions
102        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                    // This is where DebugExpr::Display gets called with q! macro cleanup
109                    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        // Determine what label to display based on config
116        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        // Use unified edge style system
145        let style = super::render::get_unified_edge_style(edge_properties, None, None);
146
147        // Determine arrow style based on edge properties
148        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        // Write the edge definition on its own line
165        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        // Add styling using unified edge style color
181        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/// Open mermaid visualization in browser for a BuiltFlow
230#[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, // Default to short labels
238        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    // Use the existing debug function
244    crate::viz::debug::open_mermaid(built_flow.ir(), Some(config))?;
245
246    Ok(())
247}