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
11pub 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 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 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 #[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}