hydro_lang/graph/
mermaid.rs

1use std::borrow::Cow;
2use std::fmt::Write;
3
4use super::render::{HydroEdgeType, 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('&', "&amp;")
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    ) -> Result<(), Self::Err> {
79        let class_str = match node_type {
80            HydroNodeType::Source => "sourceClass",
81            HydroNodeType::Transform => "transformClass",
82            HydroNodeType::Join => "joinClass",
83            HydroNodeType::Aggregation => "aggClass",
84            HydroNodeType::Network => "networkClass",
85            HydroNodeType::Sink => "sinkClass",
86            HydroNodeType::Tee => "teeClass",
87        };
88
89        let (lbracket, rbracket) = match node_type {
90            HydroNodeType::Source => ("[[", "]]"),
91            HydroNodeType::Sink => ("[/", "/]"),
92            HydroNodeType::Network => ("[[", "]]"),
93            HydroNodeType::Tee => ("(", ")"),
94            _ => ("[", "]"),
95        };
96
97        // Create the full label string using DebugExpr::Display for expressions
98        let full_label = match node_label {
99            super::render::NodeLabel::Static(s) => s.clone(),
100            super::render::NodeLabel::WithExprs { op_name, exprs } => {
101                if exprs.is_empty() {
102                    format!("{}()", op_name)
103                } else {
104                    // This is where DebugExpr::Display gets called with q! macro cleanup
105                    let expr_strs: Vec<String> = exprs.iter().map(|e| e.to_string()).collect();
106                    format!("{}({})", op_name, expr_strs.join(", "))
107                }
108            }
109        };
110
111        // Determine what label to display based on config
112        let display_label = if self.base.config.use_short_labels {
113            super::render::extract_short_label(&full_label)
114        } else {
115            full_label
116        };
117
118        let label = format!(
119            r#"n{node_id}{lbracket}"{escaped_label}"{rbracket}:::{class}"#,
120            escaped_label = escape_mermaid(&display_label),
121            class = class_str,
122        );
123
124        writeln!(
125            self.base.write,
126            "{b:i$}{label}",
127            b = "",
128            i = self.base.indent
129        )?;
130        Ok(())
131    }
132
133    fn write_edge(
134        &mut self,
135        src_id: usize,
136        dst_id: usize,
137        edge_type: HydroEdgeType,
138        label: Option<&str>,
139    ) -> Result<(), Self::Err> {
140        let arrow_style = match edge_type {
141            HydroEdgeType::Stream => "-->",
142            HydroEdgeType::Persistent => "==>",
143            HydroEdgeType::Network => "-.->",
144            HydroEdgeType::Cycle => "--o",
145        };
146
147        // Write the edge definition on its own line
148        writeln!(
149            self.base.write,
150            "{b:i$}n{src}{arrow}{label}n{dst}",
151            src = src_id,
152            arrow = arrow_style,
153            label = if let Some(label) = label {
154                Cow::Owned(format!("|{}|", escape_mermaid(label)))
155            } else {
156                Cow::Borrowed("")
157            },
158            dst = dst_id,
159            b = "",
160            i = self.base.indent,
161        )?;
162
163        // Add styling for different edge types on a separate line
164        if !matches!(edge_type, HydroEdgeType::Stream) {
165            writeln!(
166                self.base.write,
167                "{b:i$}linkStyle {} stroke:{}",
168                self.link_count,
169                match edge_type {
170                    HydroEdgeType::Persistent => "#008800",
171                    HydroEdgeType::Network => "#880088",
172                    HydroEdgeType::Cycle => "#ff0000",
173                    HydroEdgeType::Stream => "#666666", /* Should not be used here, but for completeness. */
174                },
175                b = "",
176                i = self.base.indent,
177            )?;
178        }
179
180        self.link_count += 1;
181        Ok(())
182    }
183
184    fn write_location_start(
185        &mut self,
186        location_id: usize,
187        location_type: &str,
188    ) -> Result<(), Self::Err> {
189        writeln!(
190            self.base.write,
191            "{b:i$}subgraph loc_{id} [\"{location_type} {id}\"]",
192            id = location_id,
193            b = "",
194            i = self.base.indent,
195        )?;
196        self.base.indent += 4;
197        Ok(())
198    }
199
200    fn write_node(&mut self, node_id: usize) -> Result<(), Self::Err> {
201        writeln!(
202            self.base.write,
203            "{b:i$}n{node_id}",
204            b = "",
205            i = self.base.indent
206        )
207    }
208
209    fn write_location_end(&mut self) -> Result<(), Self::Err> {
210        self.base.indent -= 4;
211        writeln!(self.base.write, "{b:i$}end", b = "", i = self.base.indent)
212    }
213
214    fn write_epilogue(&mut self) -> Result<(), Self::Err> {
215        Ok(())
216    }
217}
218
219/// Open mermaid visualization in browser for a BuiltFlow
220#[cfg(feature = "build")]
221pub fn open_browser(
222    built_flow: &crate::builder::built::BuiltFlow,
223) -> Result<(), Box<dyn std::error::Error>> {
224    let config = super::render::HydroWriteConfig {
225        show_metadata: false,
226        show_location_groups: true,
227        use_short_labels: true, // Default to short labels
228        process_id_name: built_flow.process_id_name().clone(),
229        cluster_id_name: built_flow.cluster_id_name().clone(),
230        external_id_name: built_flow.external_id_name().clone(),
231    };
232
233    // Use the existing debug function
234    crate::graph::debug::open_mermaid(built_flow.ir(), Some(config))?;
235
236    Ok(())
237}
hydro_lang/graph/mermaid.rs

hydro_lang/graph/
mermaid.rs
