hydro_lang/viz/
graphviz.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 DOT graph label.
7pub fn escape_dot(string: &str, newline: &str) -> String {
8    string.replace('"', "\\\"").replace('\n', newline)
9}
10
11/// DOT/Graphviz graph writer for Hydro IR.
12pub struct HydroDot<W> {
13    base: IndentedGraphWriter<W>,
14}
15
16impl<W> HydroDot<W> {
17    pub fn new(write: W) -> Self {
18        Self {
19            base: IndentedGraphWriter::new(write),
20        }
21    }
22
23    pub fn new_with_config(write: W, config: &super::render::HydroWriteConfig) -> Self {
24        Self {
25            base: IndentedGraphWriter::new_with_config(write, config),
26        }
27    }
28}
29
30impl<W> HydroGraphWrite for HydroDot<W>
31where
32    W: Write,
33{
34    type Err = super::render::GraphWriteError;
35
36    fn write_prologue(&mut self) -> Result<(), Self::Err> {
37        writeln!(
38            self.base.write,
39            "{b:i$}digraph HydroIR {{",
40            b = "",
41            i = self.base.indent
42        )?;
43        self.base.indent += 4;
44
45        // Use dot layout for better edge routing between subgraphs
46        writeln!(
47            self.base.write,
48            "{b:i$}layout=dot;",
49            b = "",
50            i = self.base.indent
51        )?;
52        writeln!(
53            self.base.write,
54            "{b:i$}compound=true;",
55            b = "",
56            i = self.base.indent
57        )?;
58        writeln!(
59            self.base.write,
60            "{b:i$}concentrate=true;",
61            b = "",
62            i = self.base.indent
63        )?;
64
65        const FONTS: &str = "\"Monaco,Menlo,Consolas,&quot;Droid Sans Mono&quot;,Inconsolata,&quot;Courier New&quot;,monospace\"";
66        writeln!(
67            self.base.write,
68            "{b:i$}node [fontname={}, style=filled];",
69            FONTS,
70            b = "",
71            i = self.base.indent
72        )?;
73        writeln!(
74            self.base.write,
75            "{b:i$}edge [fontname={}];",
76            FONTS,
77            b = "",
78            i = self.base.indent
79        )?;
80        Ok(())
81    }
82
83    fn write_node_definition(
84        &mut self,
85        node_id: usize,
86        node_label: &super::render::NodeLabel,
87        node_type: HydroNodeType,
88        _location_id: Option<usize>,
89        _location_type: Option<&str>,
90        _backtrace: Option<&crate::compile::ir::backtrace::Backtrace>,
91    ) -> Result<(), Self::Err> {
92        // Create the full label string using DebugExpr::Display for expressions
93        let full_label = match node_label {
94            super::render::NodeLabel::Static(s) => s.clone(),
95            super::render::NodeLabel::WithExprs { op_name, exprs } => {
96                if exprs.is_empty() {
97                    format!("{}()", op_name)
98                } else {
99                    // This is where DebugExpr::Display gets called with q! macro cleanup
100                    let expr_strs: Vec<String> = exprs.iter().map(|e| e.to_string()).collect();
101                    format!("{}({})", op_name, expr_strs.join(", "))
102                }
103            }
104        };
105
106        // Determine what label to display based on config
107        let display_label = if self.base.config.use_short_labels {
108            super::render::extract_short_label(&full_label)
109        } else {
110            full_label
111        };
112
113        let escaped_label = escape_dot(&display_label, "\\l");
114        let label = format!("n{}", node_id);
115
116        let (shape_str, color_str) = match node_type {
117            // ColorBrewer Set3 palette colors (matching Mermaid and Hydroscope)
118            HydroNodeType::Source => ("ellipse", "\"#8dd3c7\""), // Light teal
119            HydroNodeType::Transform => ("box", "\"#ffffb3\""),  // Light yellow
120            HydroNodeType::Join => ("diamond", "\"#bebada\""),   // Light purple
121            HydroNodeType::Aggregation => ("house", "\"#fb8072\""), // Light red/salmon
122            HydroNodeType::Network => ("doubleoctagon", "\"#80b1d3\""), // Light blue
123            HydroNodeType::Sink => ("invhouse", "\"#fdb462\""),  // Light orange
124            HydroNodeType::Tee => ("terminator", "\"#b3de69\""), // Light green
125            HydroNodeType::NonDeterministic => ("hexagon", "\"#fccde5\""), // Light pink/magenta
126        };
127
128        write!(
129            self.base.write,
130            "{b:i$}{label} [label=\"({node_id}) {escaped_label}{}\"",
131            if escaped_label.contains("\\l") {
132                "\\l"
133            } else {
134                ""
135            },
136            b = "",
137            i = self.base.indent,
138        )?;
139        write!(
140            self.base.write,
141            ", shape={shape_str}, fillcolor={color_str}"
142        )?;
143        writeln!(self.base.write, "]")?;
144        Ok(())
145    }
146
147    fn write_edge(
148        &mut self,
149        src_id: usize,
150        dst_id: usize,
151        edge_properties: &std::collections::HashSet<HydroEdgeProp>,
152        label: Option<&str>,
153    ) -> Result<(), Self::Err> {
154        let mut properties = Vec::<Cow<'static, str>>::new();
155
156        if let Some(label) = label {
157            properties.push(format!("label=\"{}\"", escape_dot(label, "\\n")).into());
158        }
159
160        let style = super::render::get_unified_edge_style(edge_properties, None, None);
161
162        properties.push(format!("color=\"{}\"", style.color).into());
163
164        if style.line_width > 1 {
165            properties.push("style=\"bold\"".into());
166        }
167
168        match style.line_pattern {
169            super::render::LinePattern::Dotted => {
170                properties.push("style=\"dotted\"".into());
171            }
172            super::render::LinePattern::Dashed => {
173                properties.push("style=\"dashed\"".into());
174            }
175            _ => {}
176        }
177
178        write!(
179            self.base.write,
180            "{b:i$}n{} -> n{}",
181            src_id,
182            dst_id,
183            b = "",
184            i = self.base.indent,
185        )?;
186
187        if !properties.is_empty() {
188            write!(self.base.write, " [")?;
189            for prop in itertools::Itertools::intersperse(properties.into_iter(), ", ".into()) {
190                write!(self.base.write, "{}", prop)?;
191            }
192            write!(self.base.write, "]")?;
193        }
194        writeln!(self.base.write)?;
195        Ok(())
196    }
197
198    fn write_location_start(
199        &mut self,
200        location_id: usize,
201        location_type: &str,
202    ) -> Result<(), Self::Err> {
203        writeln!(
204            self.base.write,
205            "{b:i$}subgraph cluster_loc_{id} {{",
206            id = location_id,
207            b = "",
208            i = self.base.indent,
209        )?;
210        self.base.indent += 4;
211
212        // Use dot layout for interior nodes within containers
213        writeln!(
214            self.base.write,
215            "{b:i$}layout=dot;",
216            b = "",
217            i = self.base.indent
218        )?;
219        writeln!(
220            self.base.write,
221            "{b:i$}label = \"{location_type} {id}\"",
222            id = location_id,
223            b = "",
224            i = self.base.indent
225        )?;
226        writeln!(
227            self.base.write,
228            "{b:i$}style=filled",
229            b = "",
230            i = self.base.indent
231        )?;
232        writeln!(
233            self.base.write,
234            "{b:i$}fillcolor=\"#fafafa\"",
235            b = "",
236            i = self.base.indent
237        )?;
238        writeln!(
239            self.base.write,
240            "{b:i$}color=\"#e0e0e0\"",
241            b = "",
242            i = self.base.indent
243        )?;
244        Ok(())
245    }
246
247    fn write_node(&mut self, node_id: usize) -> Result<(), Self::Err> {
248        writeln!(
249            self.base.write,
250            "{b:i$}n{node_id}",
251            b = "",
252            i = self.base.indent
253        )
254    }
255
256    fn write_location_end(&mut self) -> Result<(), Self::Err> {
257        self.base.indent -= 4;
258        writeln!(self.base.write, "{b:i$}}}", b = "", i = self.base.indent)
259    }
260
261    fn write_epilogue(&mut self) -> Result<(), Self::Err> {
262        self.base.indent -= 4;
263        writeln!(self.base.write, "{b:i$}}}", b = "", i = self.base.indent)
264    }
265}
266
267/// Open DOT/Graphviz visualization in browser for a BuiltFlow
268#[cfg(feature = "build")]
269pub fn open_browser(
270    built_flow: &crate::compile::built::BuiltFlow,
271) -> Result<(), Box<dyn std::error::Error>> {
272    let config = super::render::HydroWriteConfig {
273        show_metadata: false,
274        show_location_groups: true,
275        use_short_labels: true, // Default to short labels
276        process_id_name: built_flow.process_id_name().clone(),
277        cluster_id_name: built_flow.cluster_id_name().clone(),
278        external_id_name: built_flow.external_id_name().clone(),
279    };
280
281    // Use the existing debug function
282    crate::viz::debug::open_dot(built_flow.ir(), Some(config))?;
283
284    Ok(())
285}