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
17impl<'a> GraphApi<'a> {
18    pub fn new(ir: &'a [HydroRoot], location_names: &'a SecondaryMap<LocationKey, String>) -> Self {
19        Self { ir, location_names }
20    }
21
22    fn config(&self, use_short_labels: bool, show_metadata: bool) -> HydroWriteConfig<'a> {
23        HydroWriteConfig {
24            show_metadata,
25            show_location_groups: true,
26            use_short_labels,
27            location_names: self.location_names,
28        }
29    }
30
31    /// Render graph to string in the given format.
32    pub fn render(
33        &self,
34        format: crate::viz::config::GraphType,
35        use_short_labels: bool,
36        show_metadata: bool,
37    ) -> String {
38        let config = self.config(use_short_labels, show_metadata);
39        match format {
40            crate::viz::config::GraphType::Mermaid => render_hydro_ir_mermaid(self.ir, config),
41            crate::viz::config::GraphType::Dot => render_hydro_ir_dot(self.ir, config),
42            crate::viz::config::GraphType::Json => render_hydro_ir_json(self.ir, config),
43        }
44    }
45
46    /// Write graph to file.
47    pub fn write_to_file(
48        &self,
49        format: crate::viz::config::GraphType,
50        filename: &str,
51        use_short_labels: bool,
52        show_metadata: bool,
53    ) -> Result<(), Box<dyn Error>> {
54        let content = self.render(format, use_short_labels, show_metadata);
55        std::fs::write(filename, content)?;
56        Ok(())
57    }
58
59    /// Generate graph based on CLI GraphConfig. Returns Some(path) if a file was written.
60    #[cfg(feature = "build")]
61    pub fn generate_graph(
62        &self,
63        config: &crate::viz::config::GraphConfig,
64    ) -> Result<Option<String>, Box<dyn Error>> {
65        let Some(graph_type) = config.graph else {
66            return Ok(None);
67        };
68        let filename = config
69            .output
70            .clone()
71            .unwrap_or_else(|| format!("hydro_graph.{}", graph_type.file_extension()));
72        self.write_to_file(
73            graph_type,
74            &filename,
75            !config.long_labels,
76            !config.no_metadata,
77        )?;
78        Ok(Some(filename))
79    }
80}
81
82#[cfg(test)]
83mod tests {
84    use super::*;
85
86    #[test]
87    fn test_string_generation() {
88        let ir = vec![];
89        let mut location_names = SecondaryMap::new();
90        let loc_key_1 = LocationKey::TEST_KEY_1;
91        location_names.insert(loc_key_1, "test_process".to_owned());
92
93        let api = GraphApi::new(&ir, &location_names);
94
95        let mermaid = api.render(crate::viz::config::GraphType::Mermaid, true, true);
96        let dot = api.render(crate::viz::config::GraphType::Dot, true, true);
97        let json = api.render(crate::viz::config::GraphType::Json, true, true);
98
99        assert!(!mermaid.is_empty());
100        assert!(!dot.is_empty());
101        assert!(!json.is_empty());
102    }
103}