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$}classDef nondetClass fill:#fccde5,stroke:#f3c4dc,text-align:left,white-space:pre
65{b:i$}linkStyle default stroke:#666666",
66            b = "",
67            i = self.base.indent
68        )?;
69        Ok(())
70    }
71
72    fn write_node_definition(
73        &mut self,
74        node_id: usize,
75        node_label: &super::render::NodeLabel,
76        node_type: HydroNodeType,
77        _location_id: Option<usize>,
78        _location_type: Option<&str>,
79        _backtrace: Option<&crate::compile::ir::backtrace::Backtrace>,
80    ) -> Result<(), Self::Err> {
81        let class_str = match node_type {
82            HydroNodeType::Source => "sourceClass",
83            HydroNodeType::Transform => "transformClass",
84            HydroNodeType::Join => "joinClass",
85            HydroNodeType::Aggregation => "aggClass",
86            HydroNodeType::Network => "networkClass",
87            HydroNodeType::Sink => "sinkClass",
88            HydroNodeType::Tee => "teeClass",
89            HydroNodeType::NonDeterministic => "nondetClass",
90        };
91
92        let (lbracket, rbracket) = match node_type {
93            HydroNodeType::Source => ("[[", "]]"),
94            HydroNodeType::Sink => ("[/", "/]"),
95            HydroNodeType::Network => ("[[", "]]"),
96            HydroNodeType::Tee => ("(", ")"),
97            _ => ("[", "]"),
98        };
99
100        // Create the full label string using DebugExpr::Display for expressions
101        let full_label = match node_label {
102            super::render::NodeLabel::Static(s) => s.clone(),
103            super::render::NodeLabel::WithExprs { op_name, exprs } => {
104                if exprs.is_empty() {
105                    format!("{}()", op_name)
106                } else {
107                    // This is where DebugExpr::Display gets called with q! macro cleanup
108                    let expr_strs: Vec<String> = exprs.iter().map(|e| e.to_string()).collect();
109                    format!("{}({})", op_name, expr_strs.join(", "))
110                }
111            }
112        };
113
114        // Determine what label to display based on config
115        let display_label = if self.base.config.use_short_labels {
116            super::render::extract_short_label(&full_label)
117        } else {
118            full_label
119        };
120
121        let label = format!(
122            r#"n{node_id}{lbracket}"{escaped_label}"{rbracket}:::{class}"#,
123            escaped_label = escape_mermaid(&display_label),
124            class = class_str,
125        );
126
127        writeln!(
128            self.base.write,
129            "{b:i$}{label}",
130            b = "",
131            i = self.base.indent
132        )?;
133        Ok(())
134    }
135
136    fn write_edge(
137        &mut self,
138        src_id: usize,
139        dst_id: usize,
140        edge_properties: &std::collections::HashSet<HydroEdgeProp>,
141        label: Option<&str>,
142    ) -> Result<(), Self::Err> {
143        // Use unified edge style system
144        let style = super::render::get_unified_edge_style(edge_properties, None, None);
145
146        // Determine arrow style based on edge properties
147        let arrow_style = if edge_properties.contains(&HydroEdgeProp::Network) {
148            "-.->".to_string()
149        } else {
150            match style.line_pattern {
151                super::render::LinePattern::Dotted => "-.->".to_string(),
152                super::render::LinePattern::Dashed => "--o".to_string(),
153                _ => {
154                    if style.line_width > 1 {
155                        "==>".to_string()
156                    } else {
157                        "-->".to_string()
158                    }
159                }
160            }
161        };
162
163        // Write the edge definition on its own line
164        writeln!(
165            self.base.write,
166            "{b:i$}n{src}{arrow}{label}n{dst}",
167            src = src_id,
168            arrow = arrow_style,
169            label = if let Some(label) = label {
170                Cow::Owned(format!("|{}|", escape_mermaid(label)))
171            } else {
172                Cow::Borrowed("")
173            },
174            dst = dst_id,
175            b = "",
176            i = self.base.indent,
177        )?;
178
179        // Add styling using unified edge style color
180        writeln!(
181            self.base.write,
182            "{b:i$}linkStyle {} stroke:{}",
183            self.link_count,
184            style.color,
185            b = "",
186            i = self.base.indent,
187        )?;
188
189        self.link_count += 1;
190        Ok(())
191    }
192
193    fn write_location_start(
194        &mut self,
195        location_id: usize,
196        location_type: &str,
197    ) -> Result<(), Self::Err> {
198        writeln!(
199            self.base.write,
200            "{b:i$}subgraph loc_{id} [\"{location_type} {id}\"]",
201            id = location_id,
202            b = "",
203            i = self.base.indent,
204        )?;
205        self.base.indent += 4;
206        Ok(())
207    }
208
209    fn write_node(&mut self, node_id: usize) -> Result<(), Self::Err> {
210        writeln!(
211            self.base.write,
212            "{b:i$}n{node_id}",
213            b = "",
214            i = self.base.indent
215        )
216    }
217
218    fn write_location_end(&mut self) -> Result<(), Self::Err> {
219        self.base.indent -= 4;
220        writeln!(self.base.write, "{b:i$}end", b = "", i = self.base.indent)
221    }
222
223    fn write_epilogue(&mut self) -> Result<(), Self::Err> {
224        Ok(())
225    }
226}
227
228/// Open mermaid visualization in browser for a BuiltFlow
229#[cfg(feature = "build")]
230pub fn open_browser(
231    built_flow: &crate::compile::built::BuiltFlow,
232) -> Result<(), Box<dyn std::error::Error>> {
233    let config = super::render::HydroWriteConfig {
234        show_metadata: false,
235        show_location_groups: true,
236        use_short_labels: true, // Default to short labels
237        process_id_name: built_flow.process_id_name().clone(),
238        cluster_id_name: built_flow.cluster_id_name().clone(),
239        external_id_name: built_flow.external_id_name().clone(),
240    };
241
242    // Use the existing debug function
243    crate::viz::debug::open_mermaid(built_flow.ir(), Some(config))?;
244
245    Ok(())
246}