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        };
126
127        write!(
128            self.base.write,
129            "{b:i$}{label} [label=\"({node_id}) {escaped_label}{}\"",
130            if escaped_label.contains("\\l") {
131                "\\l"
132            } else {
133                ""
134            },
135            b = "",
136            i = self.base.indent,
137        )?;
138        write!(
139            self.base.write,
140            ", shape={shape_str}, fillcolor={color_str}"
141        )?;
142        writeln!(self.base.write, "]")?;
143        Ok(())
144    }
145
146    fn write_edge(
147        &mut self,
148        src_id: usize,
149        dst_id: usize,
150        edge_properties: &std::collections::HashSet<HydroEdgeProp>,
151        label: Option<&str>,
152    ) -> Result<(), Self::Err> {
153        let mut properties = Vec::<Cow<'static, str>>::new();
154
155        if let Some(label) = label {
156            properties.push(format!("label=\"{}\"", escape_dot(label, "\\n")).into());
157        }
158
159        let style = super::render::get_unified_edge_style(edge_properties, None, None);
160
161        properties.push(format!("color=\"{}\"", style.color).into());
162
163        if style.line_width > 1 {
164            properties.push("style=\"bold\"".into());
165        }
166
167        match style.line_pattern {
168            super::render::LinePattern::Dotted => {
169                properties.push("style=\"dotted\"".into());
170            }
171            super::render::LinePattern::Dashed => {
172                properties.push("style=\"dashed\"".into());
173            }
174            _ => {}
175        }
176
177        write!(
178            self.base.write,
179            "{b:i$}n{} -> n{}",
180            src_id,
181            dst_id,
182            b = "",
183            i = self.base.indent,
184        )?;
185
186        if !properties.is_empty() {
187            write!(self.base.write, " [")?;
188            for prop in itertools::Itertools::intersperse(properties.into_iter(), ", ".into()) {
189                write!(self.base.write, "{}", prop)?;
190            }
191            write!(self.base.write, "]")?;
192        }
193        writeln!(self.base.write)?;
194        Ok(())
195    }
196
197    fn write_location_start(
198        &mut self,
199        location_id: usize,
200        location_type: &str,
201    ) -> Result<(), Self::Err> {
202        writeln!(
203            self.base.write,
204            "{b:i$}subgraph cluster_loc_{id} {{",
205            id = location_id,
206            b = "",
207            i = self.base.indent,
208        )?;
209        self.base.indent += 4;
210
211        // Use dot layout for interior nodes within containers
212        writeln!(
213            self.base.write,
214            "{b:i$}layout=dot;",
215            b = "",
216            i = self.base.indent
217        )?;
218        writeln!(
219            self.base.write,
220            "{b:i$}label = \"{location_type} {id}\"",
221            id = location_id,
222            b = "",
223            i = self.base.indent
224        )?;
225        writeln!(
226            self.base.write,
227            "{b:i$}style=filled",
228            b = "",
229            i = self.base.indent
230        )?;
231        writeln!(
232            self.base.write,
233            "{b:i$}fillcolor=\"#fafafa\"",
234            b = "",
235            i = self.base.indent
236        )?;
237        writeln!(
238            self.base.write,
239            "{b:i$}color=\"#e0e0e0\"",
240            b = "",
241            i = self.base.indent
242        )?;
243        Ok(())
244    }
245
246    fn write_node(&mut self, node_id: usize) -> Result<(), Self::Err> {
247        writeln!(
248            self.base.write,
249            "{b:i$}n{node_id}",
250            b = "",
251            i = self.base.indent
252        )
253    }
254
255    fn write_location_end(&mut self) -> Result<(), Self::Err> {
256        self.base.indent -= 4;
257        writeln!(self.base.write, "{b:i$}}}", b = "", i = self.base.indent)
258    }
259
260    fn write_epilogue(&mut self) -> Result<(), Self::Err> {
261        self.base.indent -= 4;
262        writeln!(self.base.write, "{b:i$}}}", b = "", i = self.base.indent)
263    }
264}
265
266/// Open DOT/Graphviz visualization in browser for a BuiltFlow
267#[cfg(feature = "build")]
268pub fn open_browser(
269    built_flow: &crate::compile::built::BuiltFlow,
270) -> Result<(), Box<dyn std::error::Error>> {
271    let config = super::render::HydroWriteConfig {
272        show_metadata: false,
273        show_location_groups: true,
274        use_short_labels: true, // Default to short labels
275        process_id_name: built_flow.process_id_name().clone(),
276        cluster_id_name: built_flow.cluster_id_name().clone(),
277        external_id_name: built_flow.external_id_name().clone(),
278    };
279
280    // Use the existing debug function
281    crate::viz::debug::open_dot(built_flow.ir(), Some(config))?;
282
283    Ok(())
284}