hydro_lang/viz/
mermaid.rs

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