hydro_lang/graph/
graphviz.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 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    ) -> Result<(), Self::Err> {
91        // Create the full label string using DebugExpr::Display for expressions
92        let full_label = match node_label {
93            super::render::NodeLabel::Static(s) => s.clone(),
94            super::render::NodeLabel::WithExprs { op_name, exprs } => {
95                if exprs.is_empty() {
96                    format!("{}()", op_name)
97                } else {
98                    // This is where DebugExpr::Display gets called with q! macro cleanup
99                    let expr_strs: Vec<String> = exprs.iter().map(|e| e.to_string()).collect();
100                    format!("{}({})", op_name, expr_strs.join(", "))
101                }
102            }
103        };
104
105        // Determine what label to display based on config
106        let display_label = if self.base.config.use_short_labels {
107            super::render::extract_short_label(&full_label)
108        } else {
109            full_label
110        };
111
112        let escaped_label = escape_dot(&display_label, "\\l");
113        let label = format!("n{}", node_id);
114
115        let (shape_str, color_str) = match node_type {
116            // ColorBrewer Set3 palette colors (matching Mermaid and ReactFlow)
117            HydroNodeType::Source => ("ellipse", "\"#8dd3c7\""), // Light teal
118            HydroNodeType::Transform => ("box", "\"#ffffb3\""),  // Light yellow
119            HydroNodeType::Join => ("diamond", "\"#bebada\""),   // Light purple
120            HydroNodeType::Aggregation => ("house", "\"#fb8072\""), // Light red/salmon
121            HydroNodeType::Network => ("doubleoctagon", "\"#80b1d3\""), // Light blue
122            HydroNodeType::Sink => ("invhouse", "\"#fdb462\""),  // Light orange
123            HydroNodeType::Tee => ("terminator", "\"#b3de69\""), // Light green
124        };
125
126        write!(
127            self.base.write,
128            "{b:i$}{label} [label=\"({node_id}) {escaped_label}{}\"",
129            if escaped_label.contains("\\l") {
130                "\\l"
131            } else {
132                ""
133            },
134            b = "",
135            i = self.base.indent,
136        )?;
137        write!(
138            self.base.write,
139            ", shape={shape_str}, fillcolor={color_str}"
140        )?;
141        writeln!(self.base.write, "]")?;
142        Ok(())
143    }
144
145    fn write_edge(
146        &mut self,
147        src_id: usize,
148        dst_id: usize,
149        edge_type: HydroEdgeType,
150        label: Option<&str>,
151    ) -> Result<(), Self::Err> {
152        let mut properties = Vec::<Cow<'static, str>>::new();
153
154        if let Some(label) = label {
155            properties.push(format!("label=\"{}\"", escape_dot(label, "\\n")).into());
156        }
157
158        // Styling based on edge type
159        match edge_type {
160            HydroEdgeType::Persistent => {
161                properties.push("color=\"#008800\"".into());
162                properties.push("style=\"bold\"".into());
163            }
164            HydroEdgeType::Network => {
165                properties.push("color=\"#880088\"".into());
166                properties.push("style=\"dashed\"".into());
167            }
168            HydroEdgeType::Cycle => {
169                properties.push("color=\"#ff8800\"".into());
170                properties.push("style=\"dotted\"".into());
171            }
172            HydroEdgeType::Stream => {}
173        }
174
175        write!(
176            self.base.write,
177            "{b:i$}n{} -> n{}",
178            src_id,
179            dst_id,
180            b = "",
181            i = self.base.indent,
182        )?;
183
184        if !properties.is_empty() {
185            write!(self.base.write, " [")?;
186            for prop in itertools::Itertools::intersperse(properties.into_iter(), ", ".into()) {
187                write!(self.base.write, "{}", prop)?;
188            }
189            write!(self.base.write, "]")?;
190        }
191        writeln!(self.base.write)?;
192        Ok(())
193    }
194
195    fn write_location_start(
196        &mut self,
197        location_id: usize,
198        location_type: &str,
199    ) -> Result<(), Self::Err> {
200        writeln!(
201            self.base.write,
202            "{b:i$}subgraph cluster_loc_{id} {{",
203            id = location_id,
204            b = "",
205            i = self.base.indent,
206        )?;
207        self.base.indent += 4;
208
209        // Use dot layout for interior nodes within containers
210        writeln!(
211            self.base.write,
212            "{b:i$}layout=dot;",
213            b = "",
214            i = self.base.indent
215        )?;
216        writeln!(
217            self.base.write,
218            "{b:i$}label = \"{location_type} {id}\"",
219            id = location_id,
220            b = "",
221            i = self.base.indent
222        )?;
223        writeln!(
224            self.base.write,
225            "{b:i$}style=filled",
226            b = "",
227            i = self.base.indent
228        )?;
229        writeln!(
230            self.base.write,
231            "{b:i$}fillcolor=\"#fafafa\"",
232            b = "",
233            i = self.base.indent
234        )?;
235        writeln!(
236            self.base.write,
237            "{b:i$}color=\"#e0e0e0\"",
238            b = "",
239            i = self.base.indent
240        )?;
241        Ok(())
242    }
243
244    fn write_node(&mut self, node_id: usize) -> Result<(), Self::Err> {
245        writeln!(
246            self.base.write,
247            "{b:i$}n{node_id}",
248            b = "",
249            i = self.base.indent
250        )
251    }
252
253    fn write_location_end(&mut self) -> Result<(), Self::Err> {
254        self.base.indent -= 4;
255        writeln!(self.base.write, "{b:i$}}}", b = "", i = self.base.indent)
256    }
257
258    fn write_epilogue(&mut self) -> Result<(), Self::Err> {
259        self.base.indent -= 4;
260        writeln!(self.base.write, "{b:i$}}}", b = "", i = self.base.indent)
261    }
262}
263
264/// Open DOT/Graphviz visualization in browser for a BuiltFlow
265#[cfg(feature = "build")]
266pub fn open_browser(
267    built_flow: &crate::builder::built::BuiltFlow,
268) -> Result<(), Box<dyn std::error::Error>> {
269    let config = super::render::HydroWriteConfig {
270        show_metadata: false,
271        show_location_groups: true,
272        use_short_labels: true, // Default to short labels
273        process_id_name: built_flow.process_id_name().clone(),
274        cluster_id_name: built_flow.cluster_id_name().clone(),
275        external_id_name: built_flow.external_id_name().clone(),
276    };
277
278    // Use the existing debug function
279    crate::graph::debug::open_dot(built_flow.ir(), Some(config))?;
280
281    Ok(())
282}