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 default fill:#f5f5f5,stroke:#bbb,text-align:left,white-space:pre
62{b:i$}linkStyle default stroke:#666666",
63            b = "",
64            i = self.base.indent
65        )?;
66        Ok(())
67    }
68
69    fn write_node_definition(
70        &mut self,
71        node_id: VizNodeKey,
72        node_label: &super::render::NodeLabel,
73        _node_type: HydroNodeType,
74        _location_id: Option<LocationKey>,
75        _location_type: Option<LocationType>,
76        _backtrace: Option<&crate::compile::ir::backtrace::Backtrace>,
77    ) -> Result<(), Self::Err> {
78        // Create the full label string using DebugExpr::Display for expressions
79        let full_label = match node_label {
80            super::render::NodeLabel::Static(s) => s.clone(),
81            super::render::NodeLabel::WithExprs { op_name, exprs } => {
82                if exprs.is_empty() {
83                    format!("{}()", op_name)
84                } else {
85                    // This is where DebugExpr::Display gets called with q! macro cleanup
86                    let expr_strs: Vec<String> = exprs.iter().map(|e| e.to_string()).collect();
87                    format!("{}({})", op_name, expr_strs.join(", "))
88                }
89            }
90        };
91
92        // Determine what label to display based on config
93        let display_label = if self.base.config.use_short_labels {
94            super::render::extract_short_label(&full_label)
95        } else {
96            full_label
97        };
98
99        writeln!(
100            self.base.write,
101            "{b:i$}n{node_id}[\"{escaped_label}\"]",
102            escaped_label = escape_mermaid(&display_label),
103            b = "",
104            i = self.base.indent
105        )?;
106        Ok(())
107    }
108
109    fn write_edge(
110        &mut self,
111        src_id: VizNodeKey,
112        dst_id: VizNodeKey,
113        edge_properties: &std::collections::HashSet<HydroEdgeProp>,
114        label: Option<&str>,
115    ) -> Result<(), Self::Err> {
116        // Use unified edge style system
117        let style = super::render::get_unified_edge_style(edge_properties, None, None);
118
119        // Determine arrow style based on edge properties
120        let arrow_style = if edge_properties.contains(&HydroEdgeProp::Network) {
121            "-.->".to_owned()
122        } else {
123            match style.line_pattern {
124                super::render::LinePattern::Dotted => "-.->".to_owned(),
125                super::render::LinePattern::Dashed => "--o".to_owned(),
126                _ => {
127                    if style.line_width > 1 {
128                        "==>".to_owned()
129                    } else {
130                        "-->".to_owned()
131                    }
132                }
133            }
134        };
135
136        // Write the edge definition on its own line
137        writeln!(
138            self.base.write,
139            "{b:i$}n{src}{arrow}{label}n{dst}",
140            src = src_id,
141            arrow = arrow_style,
142            label = if let Some(label) = label {
143                Cow::Owned(format!("|{}|", escape_mermaid(label)))
144            } else {
145                Cow::Borrowed("")
146            },
147            dst = dst_id,
148            b = "",
149            i = self.base.indent,
150        )?;
151
152        // Add styling using unified edge style color
153        writeln!(
154            self.base.write,
155            "{b:i$}linkStyle {} stroke:{}",
156            self.link_count,
157            style.color,
158            b = "",
159            i = self.base.indent,
160        )?;
161
162        self.link_count += 1;
163        Ok(())
164    }
165
166    fn write_location_start(
167        &mut self,
168        location_key: LocationKey,
169        location_type: LocationType,
170    ) -> Result<(), Self::Err> {
171        writeln!(
172            self.base.write,
173            "{b:i$}subgraph {loc} [\"{location_type:?} {loc}\"]",
174            loc = location_key,
175            b = "",
176            i = self.base.indent,
177        )?;
178        self.base.indent += 4;
179        Ok(())
180    }
181
182    fn write_node(&mut self, node_id: VizNodeKey) -> Result<(), Self::Err> {
183        writeln!(
184            self.base.write,
185            "{b:i$}n{node_id}",
186            b = "",
187            i = self.base.indent
188        )
189    }
190
191    fn write_location_end(&mut self) -> Result<(), Self::Err> {
192        self.base.indent -= 4;
193        writeln!(self.base.write, "{b:i$}end", b = "", i = self.base.indent)
194    }
195
196    fn write_epilogue(&mut self) -> Result<(), Self::Err> {
197        Ok(())
198    }
199}