Skip to main content

hydro_lang/viz/
api.rs

1use std::error::Error;
2
3use slotmap::SecondaryMap;
4
5use crate::compile::ir::HydroRoot;
6use crate::location::LocationKey;
7use crate::viz::render::{
8    HydroWriteConfig, render_hydro_ir_dot, render_hydro_ir_json, render_hydro_ir_mermaid,
9};
10
11/// Graph generation API for built flows
12pub struct GraphApi<'a> {
13    ir: &'a [HydroRoot],
14    location_names: &'a SecondaryMap<LocationKey, String>,
15}
16
17/// Graph output format
18#[derive(Debug, Clone, Copy)]
19pub enum GraphFormat {
20    Mermaid,
21    Dot,
22    Hydroscope,
23}
24
25impl GraphFormat {
26    fn file_extension(self) -> &'static str {
27        match self {
28            GraphFormat::Mermaid => "mmd",
29            GraphFormat::Dot => "dot",
30            GraphFormat::Hydroscope => "json",
31        }
32    }
33
34    fn browser_message(self) -> &'static str {
35        match self {
36            GraphFormat::Mermaid => "Opening Mermaid graph in browser...",
37            GraphFormat::Dot => "Opening Graphviz/DOT graph in browser...",
38            GraphFormat::Hydroscope => "Opening Hydroscope graph in browser...",
39        }
40    }
41}
42
43impl<'a> GraphApi<'a> {
44    pub fn new(ir: &'a [HydroRoot], location_names: &'a SecondaryMap<LocationKey, String>) -> Self {
45        Self { ir, location_names }
46    }
47
48    /// Convert configuration options to HydroWriteConfig
49    fn to_hydro_config(
50        &self,
51        show_metadata: bool,
52        show_location_groups: bool,
53        use_short_labels: bool,
54    ) -> HydroWriteConfig<'a> {
55        HydroWriteConfig {
56            show_metadata,
57            show_location_groups,
58            use_short_labels,
59            location_names: self.location_names,
60        }
61    }
62
63    /// Generate graph content as string
64    fn render_graph_to_string(&self, format: GraphFormat, config: HydroWriteConfig<'_>) -> String {
65        match format {
66            GraphFormat::Mermaid => render_hydro_ir_mermaid(self.ir, config),
67            GraphFormat::Dot => render_hydro_ir_dot(self.ir, config),
68            GraphFormat::Hydroscope => render_hydro_ir_json(self.ir, config),
69        }
70    }
71
72    /// Open graph in browser
73    fn open_graph_in_browser(
74        &self,
75        format: GraphFormat,
76        config: HydroWriteConfig,
77    ) -> Result<(), Box<dyn Error>> {
78        match format {
79            GraphFormat::Mermaid => Ok(crate::viz::debug::open_mermaid(self.ir, Some(config))?),
80            GraphFormat::Dot => Ok(crate::viz::debug::open_dot(self.ir, Some(config))?),
81            GraphFormat::Hydroscope => Ok(crate::viz::debug::open_json_visualizer(
82                self.ir,
83                Some(config),
84            )?),
85        }
86    }
87
88    /// Generic method to open graph in browser
89    fn open_browser(
90        &self,
91        format: GraphFormat,
92        show_metadata: bool,
93        show_location_groups: bool,
94        use_short_labels: bool,
95        message_handler: Option<&dyn Fn(&str)>,
96    ) -> Result<(), Box<dyn Error>> {
97        let default_handler = |msg: &str| println!("{}", msg);
98        let handler = message_handler.unwrap_or(&default_handler);
99
100        let config = self.to_hydro_config(show_metadata, show_location_groups, use_short_labels);
101
102        handler(format.browser_message());
103        self.open_graph_in_browser(format, config)?;
104        Ok(())
105    }
106
107    /// Generate and save graph to file
108    fn write_graph_to_file(
109        &self,
110        format: GraphFormat,
111        filename: &str,
112        show_metadata: bool,
113        show_location_groups: bool,
114        use_short_labels: bool,
115    ) -> Result<(), Box<dyn Error>> {
116        let config = self.to_hydro_config(show_metadata, show_location_groups, use_short_labels);
117        let content = self.render_graph_to_string(format, config);
118        std::fs::write(filename, content)?;
119        println!("Generated: {}", filename);
120        Ok(())
121    }
122
123    /// Generate mermaid graph as string
124    pub fn mermaid_to_string(
125        &self,
126        show_metadata: bool,
127        show_location_groups: bool,
128        use_short_labels: bool,
129    ) -> String {
130        let config = self.to_hydro_config(show_metadata, show_location_groups, use_short_labels);
131        self.render_graph_to_string(GraphFormat::Mermaid, config)
132    }
133
134    /// Generate DOT graph as string
135    pub fn dot_to_string(
136        &self,
137        show_metadata: bool,
138        show_location_groups: bool,
139        use_short_labels: bool,
140    ) -> String {
141        let config = self.to_hydro_config(show_metadata, show_location_groups, use_short_labels);
142        self.render_graph_to_string(GraphFormat::Dot, config)
143    }
144
145    /// Generate Hydroscope graph as string
146    pub fn hydroscope_to_string(
147        &self,
148        show_metadata: bool,
149        show_location_groups: bool,
150        use_short_labels: bool,
151    ) -> String {
152        let config = self.to_hydro_config(show_metadata, show_location_groups, use_short_labels);
153        self.render_graph_to_string(GraphFormat::Hydroscope, config)
154    }
155
156    /// Write mermaid graph to file
157    pub fn mermaid_to_file(
158        &self,
159        filename: &str,
160        show_metadata: bool,
161        show_location_groups: bool,
162        use_short_labels: bool,
163    ) -> Result<(), Box<dyn Error>> {
164        self.write_graph_to_file(
165            GraphFormat::Mermaid,
166            filename,
167            show_metadata,
168            show_location_groups,
169            use_short_labels,
170        )
171    }
172
173    /// Write DOT graph to file
174    pub fn dot_to_file(
175        &self,
176        filename: &str,
177        show_metadata: bool,
178        show_location_groups: bool,
179        use_short_labels: bool,
180    ) -> Result<(), Box<dyn Error>> {
181        self.write_graph_to_file(
182            GraphFormat::Dot,
183            filename,
184            show_metadata,
185            show_location_groups,
186            use_short_labels,
187        )
188    }
189
190    /// Write Hydroscope graph to file
191    pub fn hydroscope_to_file(
192        &self,
193        filename: &str,
194        show_metadata: bool,
195        show_location_groups: bool,
196        use_short_labels: bool,
197    ) -> Result<(), Box<dyn Error>> {
198        self.write_graph_to_file(
199            GraphFormat::Hydroscope,
200            filename,
201            show_metadata,
202            show_location_groups,
203            use_short_labels,
204        )
205    }
206
207    /// Open mermaid graph in browser
208    pub fn mermaid_to_browser(
209        &self,
210        show_metadata: bool,
211        show_location_groups: bool,
212        use_short_labels: bool,
213        message_handler: Option<&dyn Fn(&str)>,
214    ) -> Result<(), Box<dyn Error>> {
215        self.open_browser(
216            GraphFormat::Mermaid,
217            show_metadata,
218            show_location_groups,
219            use_short_labels,
220            message_handler,
221        )
222    }
223
224    /// Open DOT graph in browser
225    pub fn dot_to_browser(
226        &self,
227        show_metadata: bool,
228        show_location_groups: bool,
229        use_short_labels: bool,
230        message_handler: Option<&dyn Fn(&str)>,
231    ) -> Result<(), Box<dyn Error>> {
232        self.open_browser(
233            GraphFormat::Dot,
234            show_metadata,
235            show_location_groups,
236            use_short_labels,
237            message_handler,
238        )
239    }
240
241    /// Open Hydroscope graph in browser
242    pub fn hydroscope_to_browser(
243        &self,
244        show_metadata: bool,
245        show_location_groups: bool,
246        use_short_labels: bool,
247        message_handler: Option<&dyn Fn(&str)>,
248    ) -> Result<(), Box<dyn Error>> {
249        self.open_browser(
250            GraphFormat::Hydroscope,
251            show_metadata,
252            show_location_groups,
253            use_short_labels,
254            message_handler,
255        )
256    }
257
258    /// Generate all graph types and save to files with a given prefix
259    pub fn generate_all_files(
260        &self,
261        prefix: &str,
262        show_metadata: bool,
263        show_location_groups: bool,
264        use_short_labels: bool,
265    ) -> Result<(), Box<dyn Error>> {
266        let label_suffix = if use_short_labels { "_short" } else { "_long" };
267
268        let formats = [
269            GraphFormat::Mermaid,
270            GraphFormat::Dot,
271            GraphFormat::Hydroscope,
272        ];
273
274        for format in formats {
275            let filename = format!(
276                "{}{}_labels.{}",
277                prefix,
278                label_suffix,
279                format.file_extension()
280            );
281            self.write_graph_to_file(
282                format,
283                &filename,
284                show_metadata,
285                show_location_groups,
286                use_short_labels,
287            )?;
288        }
289
290        Ok(())
291    }
292
293    /// Generate graph based on GraphConfig, delegating to the appropriate method
294    #[cfg(feature = "build")]
295    pub fn generate_graph_with_config(
296        &self,
297        config: &crate::viz::config::GraphConfig,
298        message_handler: Option<&dyn Fn(&str)>,
299    ) -> Result<(), Box<dyn Error>> {
300        if let Some(graph_type) = config.graph {
301            let format = match graph_type {
302                crate::viz::config::GraphType::Mermaid => GraphFormat::Mermaid,
303                crate::viz::config::GraphType::Dot => GraphFormat::Dot,
304                crate::viz::config::GraphType::Json => GraphFormat::Hydroscope,
305            };
306
307            if config.file {
308                let filename = format!("hydro_graph.{}", format.file_extension());
309                self.write_graph_to_file(
310                    format,
311                    &filename,
312                    !config.no_metadata,
313                    !config.no_location_groups,
314                    !config.long_labels,
315                )?;
316                println!("Graph written to {}", filename);
317            } else {
318                self.open_browser(
319                    format,
320                    !config.no_metadata,
321                    !config.no_location_groups,
322                    !config.long_labels,
323                    message_handler,
324                )?;
325            }
326        }
327        Ok(())
328    }
329
330    /// Generate all graph files based on GraphConfig
331    #[cfg(feature = "build")]
332    pub fn generate_all_files_with_config(
333        &self,
334        config: &crate::viz::config::GraphConfig,
335        prefix: &str,
336    ) -> Result<(), Box<dyn Error>> {
337        self.generate_all_files(
338            prefix,
339            !config.no_metadata,
340            !config.no_location_groups,
341            !config.long_labels,
342        )
343    }
344}
345
346#[cfg(test)]
347mod tests {
348    use super::*;
349
350    #[test]
351    fn test_graph_format() {
352        assert_eq!(GraphFormat::Mermaid.file_extension(), "mmd");
353        assert_eq!(GraphFormat::Dot.file_extension(), "dot");
354        assert_eq!(GraphFormat::Hydroscope.file_extension(), "json");
355
356        assert_eq!(
357            GraphFormat::Mermaid.browser_message(),
358            "Opening Mermaid graph in browser..."
359        );
360        assert_eq!(
361            GraphFormat::Dot.browser_message(),
362            "Opening Graphviz/DOT graph in browser..."
363        );
364        assert_eq!(
365            GraphFormat::Hydroscope.browser_message(),
366            "Opening Hydroscope graph in browser..."
367        );
368    }
369
370    #[test]
371    fn test_graph_api_creation() {
372        let ir = vec![];
373
374        let mut location_names = SecondaryMap::new();
375        let loc_key_1 = LocationKey::TEST_KEY_1;
376        location_names.insert(loc_key_1, "test_process".to_owned());
377
378        let api = GraphApi::new(&ir, &location_names);
379
380        // Test config creation
381        let config = api.to_hydro_config(true, true, false);
382        assert!(config.show_metadata);
383        assert!(config.show_location_groups);
384        assert!(!config.use_short_labels);
385        assert_eq!(config.location_names.len(), 1);
386        assert_eq!(config.location_names[loc_key_1], "test_process");
387    }
388
389    #[test]
390    fn test_string_generation() {
391        let ir = vec![];
392
393        let mut location_names = SecondaryMap::new();
394        let loc_key_1 = LocationKey::TEST_KEY_1;
395        location_names.insert(loc_key_1, "test_process".to_owned());
396
397        let api = GraphApi::new(&ir, &location_names);
398
399        // Test that string generation methods don't panic and return some content
400        let mermaid = api.mermaid_to_string(true, true, false);
401        let dot = api.dot_to_string(true, true, false);
402        let hydroscope = api.hydroscope_to_string(true, true, false);
403
404        // These should all return non-empty strings
405        assert!(!mermaid.is_empty());
406        assert!(!dot.is_empty());
407        assert!(!hydroscope.is_empty());
408    }
409}