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
17#[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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 #[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 #[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 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 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 assert!(!mermaid.is_empty());
406 assert!(!dot.is_empty());
407 assert!(!hydroscope.is_empty());
408 }
409}