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

hydro_lang/viz/
api.rs
