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            crate::viz::config::GraphType::IrJson => {
44                crate::compile::ir::ir_to_json(self.ir).expect("failed to serialize IR to JSON")
45            }
46        }
47    }
48
49    /// Write graph to file.
50    pub fn write_to_file(
51        &self,
52        format: crate::viz::config::GraphType,
53        filename: &str,
54        use_short_labels: bool,
55        show_metadata: bool,
56    ) -> Result<(), Box<dyn Error>> {
57        let content = self.render(format, use_short_labels, show_metadata);
58        std::fs::write(filename, content)?;
59        Ok(())
60    }
61
62    /// Generate graph based on CLI GraphConfig. Returns Some(path) if a file was written.
63    #[cfg(feature = "build")]
64    pub fn generate_graph(
65        &self,
66        config: &crate::viz::config::GraphConfig,
67    ) -> Result<Option<String>, Box<dyn Error>> {
68        let Some(graph_type) = config.graph else {
69            return Ok(None);
70        };
71        let filename = config
72            .output
73            .clone()
74            .unwrap_or_else(|| format!("hydro_graph.{}", graph_type.file_extension()));
75        self.write_to_file(
76            graph_type,
77            &filename,
78            !config.long_labels,
79            !config.no_metadata,
80        )?;
81        Ok(Some(filename))
82    }
83}
84
85#[cfg(test)]
86mod tests {
87    use super::*;
88
89    #[test]
90    fn test_string_generation() {
91        let ir = vec![];
92        let mut location_names = SecondaryMap::new();
93        let loc_key_1 = LocationKey::TEST_KEY_1;
94        location_names.insert(loc_key_1, "test_process".to_owned());
95
96        let api = GraphApi::new(&ir, &location_names);
97
98        let mermaid = api.render(crate::viz::config::GraphType::Mermaid, true, true);
99        let dot = api.render(crate::viz::config::GraphType::Dot, true, true);
100        let json = api.render(crate::viz::config::GraphType::Json, true, true);
101
102        assert!(!mermaid.is_empty());
103        assert!(!dot.is_empty());
104        assert!(!json.is_empty());
105    }
106}
hydro_lang/viz/api.rs

hydro_lang/viz/
api.rs
