Skip to main content

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