Skip to main content

hydro_lang/viz/
mermaid.rs

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
10/// Escapes a string for use in a mermaid graph label.
11pub fn escape_mermaid(string: &str) -> String {
12    string
13        .replace('&', "&")
14        .replace('<', "&lt;")
15        .replace('>', "&gt;")
16        .replace('"', "&quot;")
17        .replace('#', "&num;")
18        .replace('\n', "<br>")
19        // Handle code block markers
20        .replace("`", "&#96;")
21        // Handle parentheses that can conflict with Mermaid syntax
22        .replace('(', "&#40;")
23        .replace(')', "&#41;")
24        // Handle pipes that can conflict with Mermaid edge labels
25        .replace('|', "&#124;")
26}
27
28/// Mermaid graph writer for Hydro IR.
29pub 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        // Create the full label string using DebugExpr::Display for expressions
105        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                    // This is where DebugExpr::Display gets called with q! macro cleanup
112                    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        // Determine what label to display based on config
119        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        // Use unified edge style system
148        let style = super::render::get_unified_edge_style(edge_properties, None, None);
149
150        // Determine arrow style based on edge properties
151        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        // Write the edge definition on its own line
168        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        // Add styling using unified edge style color
184        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/// Open mermaid visualization in browser for a BuiltFlow
233#[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, // Default to short labels
241        location_names: built_flow.location_names(),
242    };
243
244    // Use the existing debug function
245    crate::viz::debug::open_mermaid(built_flow.ir(), Some(config))?;
246
247    Ok(())
248}