hydro_lang/viz/
render.rs

1use std::collections::{HashMap, HashSet};
2use std::error::Error;
3use std::fmt::{Display, Write};
4use std::num::ParseIntError;
5use std::sync::OnceLock;
6
7use auto_impl::auto_impl;
8use slotmap::{Key, SecondaryMap, SlotMap};
9
10pub use super::graphviz::{HydroDot, escape_dot};
11pub use super::json::HydroJson;
12// Re-export specific implementations
13pub use super::mermaid::{HydroMermaid, escape_mermaid};
14use crate::compile::ir::backtrace::Backtrace;
15use crate::compile::ir::{DebugExpr, HydroIrMetadata, HydroNode, HydroRoot, HydroSource};
16use crate::location::dynamic::LocationId;
17use crate::location::{LocationKey, LocationType};
18
19/// Label for a graph node - can be either a static string or contain expressions.
20#[derive(Debug, Clone)]
21pub enum NodeLabel {
22    /// A static string label
23    Static(String),
24    /// A label with an operation name and expression arguments
25    WithExprs {
26        op_name: String,
27        exprs: Vec<DebugExpr>,
28    },
29}
30
31impl NodeLabel {
32    /// Create a static label
33    pub fn static_label(s: String) -> Self {
34        Self::Static(s)
35    }
36
37    /// Create a label for an operation with multiple expression
38    pub fn with_exprs(op_name: String, exprs: Vec<DebugExpr>) -> Self {
39        Self::WithExprs { op_name, exprs }
40    }
41}
42
43impl Display for NodeLabel {
44    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
45        match self {
46            Self::Static(s) => write!(f, "{}", s),
47            Self::WithExprs { op_name, exprs } => {
48                if exprs.is_empty() {
49                    write!(f, "{}()", op_name)
50                } else {
51                    let expr_strs: Vec<_> = exprs.iter().map(|e| e.to_string()).collect();
52                    write!(f, "{}({})", op_name, expr_strs.join(", "))
53                }
54            }
55        }
56    }
57}
58
59/// Base struct for text-based graph writers that use indentation.
60/// Contains common fields shared by DOT and Mermaid writers.
61pub struct IndentedGraphWriter<'a, W> {
62    pub write: W,
63    pub indent: usize,
64    pub config: HydroWriteConfig<'a>,
65}
66
67impl<'a, W> IndentedGraphWriter<'a, W> {
68    /// Create a new writer with default configuration.
69    pub fn new(write: W) -> Self {
70        Self {
71            write,
72            indent: 0,
73            config: HydroWriteConfig::default(),
74        }
75    }
76
77    /// Create a new writer with the given configuration.
78    pub fn new_with_config(write: W, config: HydroWriteConfig<'a>) -> Self {
79        Self {
80            write,
81            indent: 0,
82            config,
83        }
84    }
85}
86
87impl<W: Write> IndentedGraphWriter<'_, W> {
88    /// Write an indented line using the current indentation level.
89    pub fn writeln_indented(&mut self, content: &str) -> Result<(), std::fmt::Error> {
90        writeln!(self.write, "{b:i$}{content}", b = "", i = self.indent)
91    }
92}
93
94/// Common error type used by all graph writers.
95pub type GraphWriteError = std::fmt::Error;
96
97/// Trait for writing textual representations of Hydro IR graphs, i.e. mermaid or dot graphs.
98#[auto_impl(&mut, Box)]
99pub trait HydroGraphWrite {
100    /// Error type emitted by writing.
101    type Err: Error;
102
103    /// Begin the graph. First method called.
104    fn write_prologue(&mut self) -> Result<(), Self::Err>;
105
106    /// Write a node definition with styling.
107    fn write_node_definition(
108        &mut self,
109        node_id: VizNodeKey,
110        node_label: &NodeLabel,
111        node_type: HydroNodeType,
112        location_key: Option<LocationKey>,
113        location_type: Option<LocationType>,
114        backtrace: Option<&Backtrace>,
115    ) -> Result<(), Self::Err>;
116
117    /// Write an edge between nodes with optional labeling.
118    fn write_edge(
119        &mut self,
120        src_id: VizNodeKey,
121        dst_id: VizNodeKey,
122        edge_properties: &HashSet<HydroEdgeProp>,
123        label: Option<&str>,
124    ) -> Result<(), Self::Err>;
125
126    /// Begin writing a location grouping (process/cluster).
127    fn write_location_start(
128        &mut self,
129        location_key: LocationKey,
130        location_type: LocationType,
131    ) -> Result<(), Self::Err>;
132
133    /// Write a node within a location.
134    fn write_node(&mut self, node_id: VizNodeKey) -> Result<(), Self::Err>;
135
136    /// End writing a location grouping.
137    fn write_location_end(&mut self) -> Result<(), Self::Err>;
138
139    /// End the graph. Last method called.
140    fn write_epilogue(&mut self) -> Result<(), Self::Err>;
141}
142
143/// Node type utilities - centralized handling of HydroNodeType operations
144pub mod node_type_utils {
145    use super::HydroNodeType;
146
147    /// All node types with their string names
148    const NODE_TYPE_DATA: &[(HydroNodeType, &str)] = &[
149        (HydroNodeType::Source, "Source"),
150        (HydroNodeType::Transform, "Transform"),
151        (HydroNodeType::Join, "Join"),
152        (HydroNodeType::Aggregation, "Aggregation"),
153        (HydroNodeType::Network, "Network"),
154        (HydroNodeType::Sink, "Sink"),
155        (HydroNodeType::Tee, "Tee"),
156        (HydroNodeType::NonDeterministic, "NonDeterministic"),
157    ];
158
159    /// Convert HydroNodeType to string representation (used by JSON format)
160    pub fn to_string(node_type: HydroNodeType) -> &'static str {
161        NODE_TYPE_DATA
162            .iter()
163            .find(|(nt, _)| *nt == node_type)
164            .map(|(_, name)| *name)
165            .unwrap_or("Unknown")
166    }
167
168    /// Get all node types with their string representations (used by JSON format)
169    pub fn all_types_with_strings() -> Vec<(HydroNodeType, &'static str)> {
170        NODE_TYPE_DATA.to_vec()
171    }
172}
173
174/// Types of nodes in Hydro IR for styling purposes.
175#[derive(Debug, Clone, Copy, PartialEq, Eq)]
176pub enum HydroNodeType {
177    Source,
178    Transform,
179    Join,
180    Aggregation,
181    Network,
182    Sink,
183    Tee,
184    NonDeterministic,
185}
186
187/// Types of edges in Hydro IR representing stream properties.
188#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
189pub enum HydroEdgeProp {
190    Bounded,
191    Unbounded,
192    TotalOrder,
193    NoOrder,
194    Keyed,
195    // Collection type tags for styling
196    Stream,
197    KeyedSingleton,
198    KeyedStream,
199    Singleton,
200    Optional,
201    Network,
202    Cycle,
203}
204
205/// Unified edge style representation for all graph formats.
206/// This intermediate format allows consistent styling across JSON, DOT, and Mermaid.
207#[derive(Debug, Clone, PartialEq, Eq)]
208pub struct UnifiedEdgeStyle {
209    /// Line pattern (solid, dashed)
210    pub line_pattern: LinePattern,
211    /// Line width (1 = thin, 3 = thick)
212    pub line_width: u8,
213    /// Arrowhead style
214    pub arrowhead: ArrowheadStyle,
215    /// Line style (single plain line, or line with hash marks/dots for keyed streams)
216    pub line_style: LineStyle,
217    /// Halo/background effect for boundedness
218    pub halo: HaloStyle,
219    /// Line waviness for ordering information
220    pub waviness: WavinessStyle,
221    /// Whether animation is enabled (JSON only)
222    pub animation: AnimationStyle,
223    /// Color for the edge
224    pub color: &'static str,
225}
226
227#[derive(Debug, Clone, Copy, PartialEq, Eq)]
228pub enum LinePattern {
229    Solid,
230    Dotted,
231    Dashed,
232}
233
234#[derive(Debug, Clone, Copy, PartialEq, Eq)]
235pub enum ArrowheadStyle {
236    TriangleFilled,
237    CircleFilled,
238    DiamondOpen,
239    Default,
240}
241
242#[derive(Debug, Clone, Copy, PartialEq, Eq)]
243pub enum LineStyle {
244    /// Plain single line
245    Single,
246    /// Single line with hash marks/dots (for keyed streams)
247    HashMarks,
248}
249
250#[derive(Debug, Clone, Copy, PartialEq, Eq)]
251pub enum HaloStyle {
252    None,
253    LightBlue,
254}
255
256#[derive(Debug, Clone, Copy, PartialEq, Eq)]
257pub enum WavinessStyle {
258    None,
259    Wavy,
260}
261
262#[derive(Debug, Clone, Copy, PartialEq, Eq)]
263pub enum AnimationStyle {
264    Static,
265    Animated,
266}
267
268impl Default for UnifiedEdgeStyle {
269    fn default() -> Self {
270        Self {
271            line_pattern: LinePattern::Solid,
272            line_width: 1,
273            arrowhead: ArrowheadStyle::Default,
274            line_style: LineStyle::Single,
275            halo: HaloStyle::None,
276            waviness: WavinessStyle::None,
277            animation: AnimationStyle::Static,
278            color: "#666666",
279        }
280    }
281}
282
283/// Convert HydroEdgeType properties to unified edge style.
284/// This is the core logic for determining edge visual properties.
285///
286/// # Visual Encoding Mapping
287///
288/// | Semantic Property | Visual Channel | Values |
289/// |------------------|----------------|---------|
290/// | Network | Line Pattern + Animation | Local (solid, static), Network (dashed, animated) |
291/// | Ordering | Waviness | TotalOrder (straight), NoOrder (wavy) |
292/// | Boundedness | Halo | Bounded (none), Unbounded (light-blue transparent) |
293/// | Keyedness | Line Style | NotKeyed (plain line), Keyed (line with hash marks/dots) |
294/// | Collection Type | Color + Arrowhead | Stream (blue #2563eb, triangle), Singleton (black, circle), Optional (gray, diamond) |
295pub fn get_unified_edge_style(
296    edge_properties: &HashSet<HydroEdgeProp>,
297    src_location: Option<usize>,
298    dst_location: Option<usize>,
299) -> UnifiedEdgeStyle {
300    let mut style = UnifiedEdgeStyle::default();
301
302    // Network communication group - controls line pattern AND animation
303    let is_network = edge_properties.contains(&HydroEdgeProp::Network)
304        || (src_location.is_some() && dst_location.is_some() && src_location != dst_location);
305
306    if is_network {
307        style.line_pattern = LinePattern::Dashed;
308        style.animation = AnimationStyle::Animated;
309    } else {
310        style.line_pattern = LinePattern::Solid;
311        style.animation = AnimationStyle::Static;
312    }
313
314    // Boundedness group - controls halo
315    if edge_properties.contains(&HydroEdgeProp::Unbounded) {
316        style.halo = HaloStyle::LightBlue;
317    } else {
318        style.halo = HaloStyle::None;
319    }
320
321    // Collection type group - controls arrowhead and color
322    if edge_properties.contains(&HydroEdgeProp::Stream) {
323        style.arrowhead = ArrowheadStyle::TriangleFilled;
324        style.color = "#2563eb"; // Bright blue for Stream
325    } else if edge_properties.contains(&HydroEdgeProp::KeyedStream) {
326        style.arrowhead = ArrowheadStyle::TriangleFilled;
327        style.color = "#2563eb"; // Bright blue for Stream (keyed variant)
328    } else if edge_properties.contains(&HydroEdgeProp::KeyedSingleton) {
329        style.arrowhead = ArrowheadStyle::TriangleFilled;
330        style.color = "#000000"; // Black for Singleton (keyed variant)
331    } else if edge_properties.contains(&HydroEdgeProp::Singleton) {
332        style.arrowhead = ArrowheadStyle::CircleFilled;
333        style.color = "#000000"; // Black for Singleton
334    } else if edge_properties.contains(&HydroEdgeProp::Optional) {
335        style.arrowhead = ArrowheadStyle::DiamondOpen;
336        style.color = "#6b7280"; // Gray for Optional
337    }
338
339    // Keyedness group - controls hash marks on the line
340    if edge_properties.contains(&HydroEdgeProp::Keyed) {
341        style.line_style = LineStyle::HashMarks; // Renders as hash marks/dots on the line in hydroscope
342    } else {
343        style.line_style = LineStyle::Single;
344    }
345
346    // Ordering group - waviness channel
347    if edge_properties.contains(&HydroEdgeProp::NoOrder) {
348        style.waviness = WavinessStyle::Wavy;
349    } else if edge_properties.contains(&HydroEdgeProp::TotalOrder) {
350        style.waviness = WavinessStyle::None;
351    }
352
353    style
354}
355
356/// Extract semantic edge properties from CollectionKind metadata.
357/// This function analyzes the collection type and extracts relevant semantic tags
358/// for visualization purposes.
359pub fn extract_edge_properties_from_collection_kind(
360    collection_kind: &crate::compile::ir::CollectionKind,
361) -> HashSet<HydroEdgeProp> {
362    use crate::compile::ir::CollectionKind;
363
364    let mut properties = HashSet::new();
365
366    match collection_kind {
367        CollectionKind::Stream { bound, order, .. } => {
368            properties.insert(HydroEdgeProp::Stream);
369            add_bound_property(&mut properties, bound);
370            add_order_property(&mut properties, order);
371        }
372        CollectionKind::KeyedStream {
373            bound, value_order, ..
374        } => {
375            properties.insert(HydroEdgeProp::KeyedStream);
376            properties.insert(HydroEdgeProp::Keyed);
377            add_bound_property(&mut properties, bound);
378            add_order_property(&mut properties, value_order);
379        }
380        CollectionKind::Singleton { bound, .. } => {
381            properties.insert(HydroEdgeProp::Singleton);
382            add_singleton_bound_property(&mut properties, bound);
383            // Singletons have implicit TotalOrder
384            properties.insert(HydroEdgeProp::TotalOrder);
385        }
386        CollectionKind::Optional { bound, .. } => {
387            properties.insert(HydroEdgeProp::Optional);
388            add_bound_property(&mut properties, bound);
389            // Optionals have implicit TotalOrder
390            properties.insert(HydroEdgeProp::TotalOrder);
391        }
392        CollectionKind::KeyedSingleton { bound, .. } => {
393            properties.insert(HydroEdgeProp::Singleton);
394            properties.insert(HydroEdgeProp::Keyed);
395            // KeyedSingletons boundedness depends on the bound kind
396            add_keyed_singleton_bound_property(&mut properties, bound);
397            properties.insert(HydroEdgeProp::TotalOrder);
398        }
399    }
400
401    properties
402}
403
404/// Helper function to add bound property based on BoundKind.
405fn add_bound_property(
406    properties: &mut HashSet<HydroEdgeProp>,
407    bound: &crate::compile::ir::BoundKind,
408) {
409    use crate::compile::ir::BoundKind;
410
411    match bound {
412        BoundKind::Bounded => {
413            properties.insert(HydroEdgeProp::Bounded);
414        }
415        BoundKind::Unbounded => {
416            properties.insert(HydroEdgeProp::Unbounded);
417        }
418    }
419}
420
421/// Helper function to add bound property for Singleton based on SingletonBoundKind.
422fn add_singleton_bound_property(
423    properties: &mut HashSet<HydroEdgeProp>,
424    bound: &crate::compile::ir::SingletonBoundKind,
425) {
426    use crate::compile::ir::SingletonBoundKind;
427
428    match bound {
429        SingletonBoundKind::Bounded => {
430            properties.insert(HydroEdgeProp::Bounded);
431        }
432        SingletonBoundKind::Monotonic | SingletonBoundKind::Unbounded => {
433            properties.insert(HydroEdgeProp::Unbounded);
434        }
435    }
436}
437
438/// Helper function to add bound property for KeyedSingleton based on KeyedSingletonBoundKind.
439fn add_keyed_singleton_bound_property(
440    properties: &mut HashSet<HydroEdgeProp>,
441    bound: &crate::compile::ir::KeyedSingletonBoundKind,
442) {
443    use crate::compile::ir::KeyedSingletonBoundKind;
444
445    match bound {
446        KeyedSingletonBoundKind::Bounded => {
447            properties.insert(HydroEdgeProp::Bounded);
448        }
449        KeyedSingletonBoundKind::BoundedValue
450        | KeyedSingletonBoundKind::MonotonicKeys
451        | KeyedSingletonBoundKind::MonotonicValue
452        | KeyedSingletonBoundKind::Unbounded => {
453            properties.insert(HydroEdgeProp::Unbounded);
454        }
455    }
456}
457
458/// Helper function to add order property based on StreamOrder.
459fn add_order_property(
460    properties: &mut HashSet<HydroEdgeProp>,
461    order: &crate::compile::ir::StreamOrder,
462) {
463    use crate::compile::ir::StreamOrder;
464
465    match order {
466        StreamOrder::TotalOrder => {
467            properties.insert(HydroEdgeProp::TotalOrder);
468        }
469        StreamOrder::NoOrder => {
470            properties.insert(HydroEdgeProp::NoOrder);
471        }
472    }
473}
474
475/// Detect if an edge crosses network boundaries by comparing source and destination locations.
476/// Returns true if the edge represents network communication between different locations.
477pub fn is_network_edge(src_location: &LocationId, dst_location: &LocationId) -> bool {
478    // Compare the root locations to determine if they differ
479    src_location.root() != dst_location.root()
480}
481
482/// Add network edge tag if source and destination locations differ.
483pub fn add_network_edge_tag(
484    properties: &mut HashSet<HydroEdgeProp>,
485    src_location: &LocationId,
486    dst_location: &LocationId,
487) {
488    if is_network_edge(src_location, dst_location) {
489        properties.insert(HydroEdgeProp::Network);
490    }
491}
492
493/// Configuration for graph writing.
494#[derive(Debug, Clone, Copy)]
495pub struct HydroWriteConfig<'a> {
496    pub show_metadata: bool,
497    pub show_location_groups: bool,
498    pub use_short_labels: bool,
499    pub location_names: &'a SecondaryMap<LocationKey, String>,
500}
501
502impl Default for HydroWriteConfig<'_> {
503    fn default() -> Self {
504        static EMPTY: OnceLock<SecondaryMap<LocationKey, String>> = OnceLock::new();
505        Self {
506            show_metadata: false,
507            show_location_groups: true,
508            use_short_labels: true, // Default to short labels for all renderers
509            location_names: EMPTY.get_or_init(SecondaryMap::new),
510        }
511    }
512}
513
514/// Node information in the Hydro graph.
515#[derive(Clone)]
516pub struct HydroGraphNode {
517    pub label: NodeLabel,
518    pub node_type: HydroNodeType,
519    pub location_key: Option<LocationKey>,
520    pub backtrace: Option<Backtrace>,
521}
522
523slotmap::new_key_type! {
524    /// Unique identifier for nodes in the visualization graph.
525    ///
526    /// This is counted/allocated separately from any other IDs within `hydro_lang`.
527    pub struct VizNodeKey;
528}
529
530impl Display for VizNodeKey {
531    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
532        write!(f, "viz{:?}", self.data()) // `"viz1v1"``
533    }
534}
535
536/// This is used by the visualizer
537/// TODO(mingwei): Make this more robust?
538impl std::str::FromStr for VizNodeKey {
539    type Err = Option<ParseIntError>;
540
541    fn from_str(s: &str) -> Result<Self, Self::Err> {
542        let nvn = s.strip_prefix("viz").ok_or(None)?;
543        let (idx, ver) = nvn.split_once("v").ok_or(None)?;
544        let idx: u64 = idx.parse()?;
545        let ver: u64 = ver.parse()?;
546        Ok(slotmap::KeyData::from_ffi((ver << 32) | idx).into())
547    }
548}
549
550impl VizNodeKey {
551    /// A key for testing with index 1.
552    #[cfg(test)]
553    pub const TEST_KEY_1: Self = Self(slotmap::KeyData::from_ffi(0x0000008F00000001)); // `1v143`
554
555    /// A key for testing with index 2.
556    #[cfg(test)]
557    pub const TEST_KEY_2: Self = Self(slotmap::KeyData::from_ffi(0x0000008F00000002)); // `2v143`
558}
559
560/// Edge information in the Hydro graph.
561#[derive(Debug, Clone)]
562pub struct HydroGraphEdge {
563    pub src: VizNodeKey,
564    pub dst: VizNodeKey,
565    pub edge_properties: HashSet<HydroEdgeProp>,
566    pub label: Option<String>,
567}
568
569/// Graph structure tracker for Hydro IR rendering.
570#[derive(Default)]
571pub struct HydroGraphStructure {
572    pub nodes: SlotMap<VizNodeKey, HydroGraphNode>,
573    pub edges: Vec<HydroGraphEdge>,
574    pub locations: SecondaryMap<LocationKey, LocationType>,
575}
576
577impl HydroGraphStructure {
578    pub fn new() -> Self {
579        Self::default()
580    }
581
582    pub fn add_node(
583        &mut self,
584        label: NodeLabel,
585        node_type: HydroNodeType,
586        location_key: Option<LocationKey>,
587    ) -> VizNodeKey {
588        self.add_node_with_backtrace(label, node_type, location_key, None)
589    }
590
591    pub fn add_node_with_backtrace(
592        &mut self,
593        label: NodeLabel,
594        node_type: HydroNodeType,
595        location_key: Option<LocationKey>,
596        backtrace: Option<Backtrace>,
597    ) -> VizNodeKey {
598        self.nodes.insert(HydroGraphNode {
599            label,
600            node_type,
601            location_key,
602            backtrace,
603        })
604    }
605
606    /// Add a node with metadata, extracting backtrace automatically
607    pub fn add_node_with_metadata(
608        &mut self,
609        label: NodeLabel,
610        node_type: HydroNodeType,
611        metadata: &HydroIrMetadata,
612    ) -> VizNodeKey {
613        let location_key = Some(setup_location(self, metadata));
614        let backtrace = Some(metadata.op.backtrace.clone());
615        self.add_node_with_backtrace(label, node_type, location_key, backtrace)
616    }
617
618    pub fn add_edge(
619        &mut self,
620        src: VizNodeKey,
621        dst: VizNodeKey,
622        edge_properties: HashSet<HydroEdgeProp>,
623        label: Option<String>,
624    ) {
625        self.edges.push(HydroGraphEdge {
626            src,
627            dst,
628            edge_properties,
629            label,
630        });
631    }
632
633    // Legacy method for backward compatibility
634    pub fn add_edge_single(
635        &mut self,
636        src: VizNodeKey,
637        dst: VizNodeKey,
638        edge_type: HydroEdgeProp,
639        label: Option<String>,
640    ) {
641        let mut properties = HashSet::new();
642        properties.insert(edge_type);
643        self.edges.push(HydroGraphEdge {
644            src,
645            dst,
646            edge_properties: properties,
647            label,
648        });
649    }
650
651    pub fn add_location(&mut self, location_key: LocationKey, location_type: LocationType) {
652        self.locations.insert(location_key, location_type);
653    }
654}
655
656/// Function to extract an op_name from a print_root() result for use in labels.
657pub fn extract_op_name(full_label: String) -> String {
658    full_label
659        .split('(')
660        .next()
661        .unwrap_or("unknown")
662        .to_lowercase()
663}
664
665/// Extract a short, readable label from the full token stream label using print_root() style naming
666pub fn extract_short_label(full_label: &str) -> String {
667    // Use the same logic as extract_op_name but handle the specific cases we need for UI display
668    if let Some(op_name) = full_label.split('(').next() {
669        let base_name = op_name.to_lowercase();
670        match base_name.as_str() {
671            // Handle special cases for UI display
672            "source" => {
673                if full_label.contains("Iter") {
674                    "source_iter".to_owned()
675                } else if full_label.contains("Stream") {
676                    "source_stream".to_owned()
677                } else if full_label.contains("ExternalNetwork") {
678                    "external_network".to_owned()
679                } else if full_label.contains("Spin") {
680                    "spin".to_owned()
681                } else {
682                    "source".to_owned()
683                }
684            }
685            "network" => {
686                if full_label.contains("deser") {
687                    "network(recv)".to_owned()
688                } else if full_label.contains("ser") {
689                    "network(send)".to_owned()
690                } else {
691                    "network".to_owned()
692                }
693            }
694            // For all other cases, just use the lowercase base name (same as extract_op_name)
695            _ => base_name,
696        }
697    } else {
698        // Fallback for labels that don't follow the pattern
699        if full_label.len() > 20 {
700            format!("{}...", &full_label[..17])
701        } else {
702            full_label.to_owned()
703        }
704    }
705}
706
707/// Helper function to set up location in structure from metadata.
708fn setup_location(structure: &mut HydroGraphStructure, metadata: &HydroIrMetadata) -> LocationKey {
709    let root = metadata.location_id.root();
710    let location_key = root.key();
711    let location_type = root.location_type().unwrap();
712    structure.add_location(location_key, location_type);
713    location_key
714}
715
716/// Helper function to add an edge with semantic tags extracted from metadata.
717/// This function combines collection kind extraction with network detection.
718fn add_edge_with_metadata(
719    structure: &mut HydroGraphStructure,
720    src_id: VizNodeKey,
721    dst_id: VizNodeKey,
722    src_metadata: Option<&HydroIrMetadata>,
723    dst_metadata: Option<&HydroIrMetadata>,
724    label: Option<String>,
725) {
726    let mut properties = HashSet::new();
727
728    // Extract semantic tags from source metadata's collection kind
729    if let Some(metadata) = src_metadata {
730        properties.extend(extract_edge_properties_from_collection_kind(
731            &metadata.collection_kind,
732        ));
733    }
734
735    // Add network edge tag if locations differ
736    if let (Some(src_meta), Some(dst_meta)) = (src_metadata, dst_metadata) {
737        add_network_edge_tag(
738            &mut properties,
739            &src_meta.location_id,
740            &dst_meta.location_id,
741        );
742    }
743
744    // If no properties were extracted, default to Stream
745    if properties.is_empty() {
746        properties.insert(HydroEdgeProp::Stream);
747    }
748
749    structure.add_edge(src_id, dst_id, properties, label);
750}
751
752/// Helper function to write a graph structure using any GraphWrite implementation
753fn write_graph_structure<W>(
754    structure: &HydroGraphStructure,
755    graph_write: W,
756    config: HydroWriteConfig<'_>,
757) -> Result<(), W::Err>
758where
759    W: HydroGraphWrite,
760{
761    let mut graph_write = graph_write;
762    // Write the graph
763    graph_write.write_prologue()?;
764
765    // Write node definitions
766    for (node_id, node) in structure.nodes.iter() {
767        let location_type = node
768            .location_key
769            .and_then(|loc_key| structure.locations.get(loc_key))
770            .copied();
771
772        graph_write.write_node_definition(
773            node_id,
774            &node.label,
775            node.node_type,
776            node.location_key,
777            location_type,
778            node.backtrace.as_ref(),
779        )?;
780    }
781
782    // Group nodes by location if requested
783    if config.show_location_groups {
784        let mut nodes_by_location = SecondaryMap::<LocationKey, Vec<VizNodeKey>>::new();
785        for (node_id, node) in structure.nodes.iter() {
786            if let Some(location_key) = node.location_key {
787                nodes_by_location
788                    .entry(location_key)
789                    .expect("location was removed")
790                    .or_default()
791                    .push(node_id);
792            }
793        }
794
795        for (location_key, node_ids) in nodes_by_location.iter() {
796            if let Some(&location_type) = structure.locations.get(location_key) {
797                graph_write.write_location_start(location_key, location_type)?;
798                for &node_id in node_ids.iter() {
799                    graph_write.write_node(node_id)?;
800                }
801                graph_write.write_location_end()?;
802            }
803        }
804    }
805
806    // Write edges
807    for edge in structure.edges.iter() {
808        graph_write.write_edge(
809            edge.src,
810            edge.dst,
811            &edge.edge_properties,
812            edge.label.as_deref(),
813        )?;
814    }
815
816    graph_write.write_epilogue()?;
817    Ok(())
818}
819
820impl HydroRoot {
821    /// Build the graph structure by traversing the IR tree.
822    pub fn build_graph_structure(
823        &self,
824        structure: &mut HydroGraphStructure,
825        seen_tees: &mut HashMap<*const std::cell::RefCell<HydroNode>, VizNodeKey>,
826        config: HydroWriteConfig<'_>,
827    ) -> VizNodeKey {
828        // Helper function for sink nodes to reduce duplication
829        fn build_sink_node(
830            structure: &mut HydroGraphStructure,
831            seen_tees: &mut HashMap<*const std::cell::RefCell<HydroNode>, VizNodeKey>,
832            config: HydroWriteConfig<'_>,
833            input: &HydroNode,
834            sink_metadata: Option<&HydroIrMetadata>,
835            label: NodeLabel,
836        ) -> VizNodeKey {
837            let input_id = input.build_graph_structure(structure, seen_tees, config);
838
839            // If no explicit metadata is provided, extract it from the input node
840            let effective_metadata = if let Some(meta) = sink_metadata {
841                Some(meta)
842            } else {
843                match input {
844                    HydroNode::Placeholder => None,
845                    // All other variants have metadata
846                    _ => Some(input.metadata()),
847                }
848            };
849
850            let location_key = effective_metadata.map(|m| setup_location(structure, m));
851            let sink_id = structure.add_node_with_backtrace(
852                label,
853                HydroNodeType::Sink,
854                location_key,
855                effective_metadata.map(|m| m.op.backtrace.clone()),
856            );
857
858            // Extract semantic tags from input metadata
859            let input_metadata = input.metadata();
860            add_edge_with_metadata(
861                structure,
862                input_id,
863                sink_id,
864                Some(input_metadata),
865                sink_metadata,
866                None,
867            );
868
869            sink_id
870        }
871
872        match self {
873            // Sink operations - semantic tags extracted from input metadata
874            HydroRoot::ForEach { f, input, .. } => build_sink_node(
875                structure,
876                seen_tees,
877                config,
878                input,
879                None,
880                NodeLabel::with_exprs("for_each".to_owned(), vec![f.expr.clone()]),
881            ),
882
883            HydroRoot::SendExternal {
884                to_external_key,
885                to_port_id,
886                input,
887                ..
888            } => build_sink_node(
889                structure,
890                seen_tees,
891                config,
892                input,
893                None,
894                NodeLabel::with_exprs(
895                    format!("send_external({}:{})", to_external_key, to_port_id),
896                    vec![],
897                ),
898            ),
899
900            HydroRoot::DestSink { sink, input, .. } => build_sink_node(
901                structure,
902                seen_tees,
903                config,
904                input,
905                None,
906                NodeLabel::with_exprs("dest_sink".to_owned(), vec![sink.clone()]),
907            ),
908
909            HydroRoot::CycleSink {
910                cycle_id, input, ..
911            } => build_sink_node(
912                structure,
913                seen_tees,
914                config,
915                input,
916                None,
917                NodeLabel::static_label(format!("cycle_sink({})", cycle_id)),
918            ),
919
920            HydroRoot::EmbeddedOutput { ident, input, .. } => build_sink_node(
921                structure,
922                seen_tees,
923                config,
924                input,
925                None,
926                NodeLabel::static_label(format!("embedded_output({})", ident)),
927            ),
928
929            HydroRoot::Null { input, .. } => build_sink_node(
930                structure,
931                seen_tees,
932                config,
933                input,
934                None,
935                NodeLabel::static_label("null".to_owned()),
936            ),
937        }
938    }
939}
940
941impl HydroNode {
942    /// Build the graph structure recursively for this node.
943    pub fn build_graph_structure(
944        &self,
945        structure: &mut HydroGraphStructure,
946        seen_tees: &mut HashMap<*const std::cell::RefCell<HydroNode>, VizNodeKey>,
947        config: HydroWriteConfig<'_>,
948    ) -> VizNodeKey {
949        // Helper functions to reduce duplication, categorized by input/expression patterns
950
951        /// Common parameters for transform builder functions to reduce argument count
952        struct TransformParams<'a> {
953            structure: &'a mut HydroGraphStructure,
954            seen_tees: &'a mut HashMap<*const std::cell::RefCell<HydroNode>, VizNodeKey>,
955            config: HydroWriteConfig<'a>,
956            input: &'a HydroNode,
957            metadata: &'a HydroIrMetadata,
958            op_name: String,
959            node_type: HydroNodeType,
960        }
961
962        // Single-input transform with no expressions
963        fn build_simple_transform(params: TransformParams<'_>) -> VizNodeKey {
964            let input_id = params.input.build_graph_structure(
965                params.structure,
966                params.seen_tees,
967                params.config,
968            );
969            let node_id = params.structure.add_node_with_metadata(
970                NodeLabel::Static(params.op_name.clone()),
971                params.node_type,
972                params.metadata,
973            );
974
975            // Extract semantic tags from input metadata
976            let input_metadata = params.input.metadata();
977            add_edge_with_metadata(
978                params.structure,
979                input_id,
980                node_id,
981                Some(input_metadata),
982                Some(params.metadata),
983                None,
984            );
985
986            node_id
987        }
988
989        // Single-input transform with one expression
990        fn build_single_expr_transform(
991            params: TransformParams<'_>,
992            expr: &DebugExpr,
993        ) -> VizNodeKey {
994            let input_id = params.input.build_graph_structure(
995                params.structure,
996                params.seen_tees,
997                params.config,
998            );
999            let node_id = params.structure.add_node_with_metadata(
1000                NodeLabel::with_exprs(params.op_name.clone(), vec![expr.clone()]),
1001                params.node_type,
1002                params.metadata,
1003            );
1004
1005            // Extract semantic tags from input metadata
1006            let input_metadata = params.input.metadata();
1007            add_edge_with_metadata(
1008                params.structure,
1009                input_id,
1010                node_id,
1011                Some(input_metadata),
1012                Some(params.metadata),
1013                None,
1014            );
1015
1016            node_id
1017        }
1018
1019        // Single-input transform with two expressions
1020        fn build_dual_expr_transform(
1021            params: TransformParams<'_>,
1022            expr1: &DebugExpr,
1023            expr2: &DebugExpr,
1024        ) -> VizNodeKey {
1025            let input_id = params.input.build_graph_structure(
1026                params.structure,
1027                params.seen_tees,
1028                params.config,
1029            );
1030            let node_id = params.structure.add_node_with_metadata(
1031                NodeLabel::with_exprs(params.op_name.clone(), vec![expr1.clone(), expr2.clone()]),
1032                params.node_type,
1033                params.metadata,
1034            );
1035
1036            // Extract semantic tags from input metadata
1037            let input_metadata = params.input.metadata();
1038            add_edge_with_metadata(
1039                params.structure,
1040                input_id,
1041                node_id,
1042                Some(input_metadata),
1043                Some(params.metadata),
1044                None,
1045            );
1046
1047            node_id
1048        }
1049
1050        // Helper function for source nodes
1051        fn build_source_node(
1052            structure: &mut HydroGraphStructure,
1053            metadata: &HydroIrMetadata,
1054            label: String,
1055        ) -> VizNodeKey {
1056            structure.add_node_with_metadata(
1057                NodeLabel::Static(label),
1058                HydroNodeType::Source,
1059                metadata,
1060            )
1061        }
1062
1063        match self {
1064            HydroNode::Placeholder => structure.add_node(
1065                NodeLabel::Static("PLACEHOLDER".to_owned()),
1066                HydroNodeType::Transform,
1067                None,
1068            ),
1069
1070            HydroNode::Source {
1071                source, metadata, ..
1072            } => {
1073                let label = match source {
1074                    HydroSource::Stream(expr) => format!("source_stream({})", expr),
1075                    HydroSource::ExternalNetwork() => "external_network()".to_owned(),
1076                    HydroSource::Iter(expr) => format!("source_iter({})", expr),
1077                    HydroSource::Spin() => "spin()".to_owned(),
1078                    HydroSource::ClusterMembers(location_id, _) => {
1079                        format!(
1080                            "source_stream(cluster_membership_stream({:?}))",
1081                            location_id
1082                        )
1083                    }
1084                    HydroSource::Embedded(ident) => {
1085                        format!("embedded_input({})", ident)
1086                    }
1087                    HydroSource::EmbeddedSingleton(ident) => {
1088                        format!("embedded_singleton_input({})", ident)
1089                    }
1090                };
1091                build_source_node(structure, metadata, label)
1092            }
1093
1094            HydroNode::SingletonSource {
1095                value,
1096                first_tick_only,
1097                metadata,
1098            } => {
1099                let label = if *first_tick_only {
1100                    format!("singleton_first_tick({})", value)
1101                } else {
1102                    format!("singleton({})", value)
1103                };
1104                build_source_node(structure, metadata, label)
1105            }
1106
1107            HydroNode::ExternalInput {
1108                from_external_key,
1109                from_port_id,
1110                metadata,
1111                ..
1112            } => build_source_node(
1113                structure,
1114                metadata,
1115                format!("external_input({}:{})", from_external_key, from_port_id),
1116            ),
1117
1118            HydroNode::CycleSource {
1119                cycle_id, metadata, ..
1120            } => build_source_node(structure, metadata, format!("cycle_source({})", cycle_id)),
1121
1122            HydroNode::Tee { inner, metadata }
1123            | HydroNode::Reference {
1124                inner, metadata, ..
1125            } => {
1126                let ptr = inner.as_ptr();
1127                if let Some(&existing_id) = seen_tees.get(&ptr) {
1128                    return existing_id;
1129                }
1130
1131                let input_id = inner
1132                    .0
1133                    .borrow()
1134                    .build_graph_structure(structure, seen_tees, config);
1135                let node_type = if matches!(self, HydroNode::Reference { .. }) {
1136                    HydroNodeType::Aggregation
1137                } else {
1138                    HydroNodeType::Tee
1139                };
1140                let tee_id = structure.add_node_with_metadata(
1141                    NodeLabel::Static(extract_op_name(self.print_root())),
1142                    node_type,
1143                    metadata,
1144                );
1145
1146                seen_tees.insert(ptr, tee_id);
1147
1148                // Extract semantic tags from input
1149                let inner_borrow = inner.0.borrow();
1150                let input_metadata = inner_borrow.metadata();
1151                add_edge_with_metadata(
1152                    structure,
1153                    input_id,
1154                    tee_id,
1155                    Some(input_metadata),
1156                    Some(metadata),
1157                    None,
1158                );
1159                drop(inner_borrow);
1160
1161                tee_id
1162            }
1163
1164            HydroNode::PartitionSide {
1165                inner, metadata, ..
1166            } => {
1167                let ptr = inner.as_ptr();
1168                if let Some(&existing_id) = seen_tees.get(&ptr) {
1169                    return existing_id;
1170                }
1171
1172                let input_id = inner
1173                    .0
1174                    .borrow()
1175                    .build_graph_structure(structure, seen_tees, config);
1176                let partition_id = structure.add_node_with_metadata(
1177                    NodeLabel::Static(extract_op_name(self.print_root())),
1178                    HydroNodeType::Tee,
1179                    metadata,
1180                );
1181
1182                seen_tees.insert(ptr, partition_id);
1183
1184                // Extract semantic tags from input
1185                let inner_borrow = inner.0.borrow();
1186                let input_metadata = inner_borrow.metadata();
1187                add_edge_with_metadata(
1188                    structure,
1189                    input_id,
1190                    partition_id,
1191                    Some(input_metadata),
1192                    Some(metadata),
1193                    None,
1194                );
1195                drop(inner_borrow);
1196
1197                partition_id
1198            }
1199            HydroNode::PartitionShared { input, .. } => {
1200                // Transparent pass-thru.
1201                input.build_graph_structure(structure, seen_tees, config)
1202            }
1203
1204            // Non-deterministic operation
1205            HydroNode::ObserveNonDet {
1206                inner, metadata, ..
1207            } => build_simple_transform(TransformParams {
1208                structure,
1209                seen_tees,
1210                config,
1211                input: inner,
1212                metadata,
1213                op_name: extract_op_name(self.print_root()),
1214                node_type: HydroNodeType::NonDeterministic,
1215            }),
1216
1217            // Transform operations with Stream edges - grouped by node/edge type
1218            HydroNode::Cast { inner, metadata }
1219            | HydroNode::AssertIsConsistent {
1220                inner, metadata, ..
1221            }
1222            | HydroNode::DeferTick {
1223                input: inner,
1224                metadata,
1225            }
1226            | HydroNode::Enumerate {
1227                input: inner,
1228                metadata,
1229                ..
1230            }
1231            | HydroNode::Unique {
1232                input: inner,
1233                metadata,
1234            }
1235            | HydroNode::ResolveFutures {
1236                input: inner,
1237                metadata,
1238            }
1239            | HydroNode::ResolveFuturesBlocking {
1240                input: inner,
1241                metadata,
1242            }
1243            | HydroNode::ResolveFuturesOrdered {
1244                input: inner,
1245                metadata,
1246            } => build_simple_transform(TransformParams {
1247                structure,
1248                seen_tees,
1249                config,
1250                input: inner,
1251                metadata,
1252                op_name: extract_op_name(self.print_root()),
1253                node_type: HydroNodeType::Transform,
1254            }),
1255
1256            // Aggregation operation - semantic tags extracted from metadata
1257            HydroNode::Sort {
1258                input: inner,
1259                metadata,
1260            } => build_simple_transform(TransformParams {
1261                structure,
1262                seen_tees,
1263                config,
1264                input: inner,
1265                metadata,
1266                op_name: extract_op_name(self.print_root()),
1267                node_type: HydroNodeType::Aggregation,
1268            }),
1269
1270            // Single-expression Transform operations - grouped by node type
1271            HydroNode::Map {
1272                f, input, metadata, ..
1273            }
1274            | HydroNode::Filter { f, input, metadata }
1275            | HydroNode::FlatMap { f, input, metadata }
1276            | HydroNode::FlatMapStreamBlocking { f, input, metadata }
1277            | HydroNode::FilterMap { f, input, metadata }
1278            | HydroNode::Inspect { f, input, metadata } => build_single_expr_transform(
1279                TransformParams {
1280                    structure,
1281                    seen_tees,
1282                    config,
1283                    input,
1284                    metadata,
1285                    op_name: extract_op_name(self.print_root()),
1286                    node_type: HydroNodeType::Transform,
1287                },
1288                &f.expr,
1289            ),
1290
1291            // Single-expression Aggregation operations - grouped by node type
1292            HydroNode::Reduce { f, input, metadata }
1293            | HydroNode::ReduceKeyed { f, input, metadata } => build_single_expr_transform(
1294                TransformParams {
1295                    structure,
1296                    seen_tees,
1297                    config,
1298                    input,
1299                    metadata,
1300                    op_name: extract_op_name(self.print_root()),
1301                    node_type: HydroNodeType::Aggregation,
1302                },
1303                &f.expr,
1304            ),
1305
1306            // Join-like operations with left/right edge labels - grouped by edge labeling
1307            HydroNode::Join {
1308                left,
1309                right,
1310                metadata,
1311            }
1312            | HydroNode::JoinHalf {
1313                left,
1314                right,
1315                metadata,
1316            }
1317            | HydroNode::CrossProduct {
1318                left,
1319                right,
1320                metadata,
1321            }
1322            | HydroNode::CrossSingleton {
1323                left,
1324                right,
1325                metadata,
1326            } => {
1327                let left_id = left.build_graph_structure(structure, seen_tees, config);
1328                let right_id = right.build_graph_structure(structure, seen_tees, config);
1329                let node_id = structure.add_node_with_metadata(
1330                    NodeLabel::Static(extract_op_name(self.print_root())),
1331                    HydroNodeType::Join,
1332                    metadata,
1333                );
1334
1335                // Extract semantic tags for left edge
1336                let left_metadata = left.metadata();
1337                add_edge_with_metadata(
1338                    structure,
1339                    left_id,
1340                    node_id,
1341                    Some(left_metadata),
1342                    Some(metadata),
1343                    Some("left".to_owned()),
1344                );
1345
1346                // Extract semantic tags for right edge
1347                let right_metadata = right.metadata();
1348                add_edge_with_metadata(
1349                    structure,
1350                    right_id,
1351                    node_id,
1352                    Some(right_metadata),
1353                    Some(metadata),
1354                    Some("right".to_owned()),
1355                );
1356
1357                node_id
1358            }
1359
1360            // Join-like operations with pos/neg edge labels - grouped by edge labeling
1361            HydroNode::Difference {
1362                pos: left,
1363                neg: right,
1364                metadata,
1365            }
1366            | HydroNode::AntiJoin {
1367                pos: left,
1368                neg: right,
1369                metadata,
1370            } => {
1371                let left_id = left.build_graph_structure(structure, seen_tees, config);
1372                let right_id = right.build_graph_structure(structure, seen_tees, config);
1373                let node_id = structure.add_node_with_metadata(
1374                    NodeLabel::Static(extract_op_name(self.print_root())),
1375                    HydroNodeType::Join,
1376                    metadata,
1377                );
1378
1379                // Extract semantic tags for pos edge
1380                let left_metadata = left.metadata();
1381                add_edge_with_metadata(
1382                    structure,
1383                    left_id,
1384                    node_id,
1385                    Some(left_metadata),
1386                    Some(metadata),
1387                    Some("pos".to_owned()),
1388                );
1389
1390                // Extract semantic tags for neg edge
1391                let right_metadata = right.metadata();
1392                add_edge_with_metadata(
1393                    structure,
1394                    right_id,
1395                    node_id,
1396                    Some(right_metadata),
1397                    Some(metadata),
1398                    Some("neg".to_owned()),
1399                );
1400
1401                node_id
1402            }
1403
1404            // Dual expression transforms - consolidated using pattern matching
1405            HydroNode::Fold {
1406                init,
1407                acc,
1408                input,
1409                metadata,
1410                ..
1411            }
1412            | HydroNode::FoldKeyed {
1413                init,
1414                acc,
1415                input,
1416                metadata,
1417                ..
1418            }
1419            | HydroNode::Scan {
1420                init,
1421                acc,
1422                input,
1423                metadata,
1424            }
1425            | HydroNode::ScanAsyncBlocking {
1426                init,
1427                acc,
1428                input,
1429                metadata,
1430            } => {
1431                let node_type = HydroNodeType::Aggregation; // All are aggregation operations
1432
1433                build_dual_expr_transform(
1434                    TransformParams {
1435                        structure,
1436                        seen_tees,
1437                        config,
1438                        input,
1439                        metadata,
1440                        op_name: extract_op_name(self.print_root()),
1441                        node_type,
1442                    },
1443                    &init.expr,
1444                    &acc.expr,
1445                )
1446            }
1447
1448            // Combination of join and transform
1449            HydroNode::ReduceKeyedWatermark {
1450                f,
1451                input,
1452                watermark,
1453                metadata,
1454            } => {
1455                let input_id = input.build_graph_structure(structure, seen_tees, config);
1456                let watermark_id = watermark.build_graph_structure(structure, seen_tees, config);
1457                let location_key = Some(setup_location(structure, metadata));
1458                let join_node_id = structure.add_node_with_backtrace(
1459                    NodeLabel::Static(extract_op_name(self.print_root())),
1460                    HydroNodeType::Join,
1461                    location_key,
1462                    Some(metadata.op.backtrace.clone()),
1463                );
1464
1465                // Extract semantic tags for input edge
1466                let input_metadata = input.metadata();
1467                add_edge_with_metadata(
1468                    structure,
1469                    input_id,
1470                    join_node_id,
1471                    Some(input_metadata),
1472                    Some(metadata),
1473                    Some("input".to_owned()),
1474                );
1475
1476                // Extract semantic tags for watermark edge
1477                let watermark_metadata = watermark.metadata();
1478                add_edge_with_metadata(
1479                    structure,
1480                    watermark_id,
1481                    join_node_id,
1482                    Some(watermark_metadata),
1483                    Some(metadata),
1484                    Some("watermark".to_owned()),
1485                );
1486
1487                let node_id = structure.add_node_with_backtrace(
1488                    NodeLabel::with_exprs(extract_op_name(self.print_root()), vec![f.expr.clone()]),
1489                    HydroNodeType::Aggregation,
1490                    location_key,
1491                    Some(metadata.op.backtrace.clone()),
1492                );
1493
1494                // Edge from join to aggregation node
1495                let join_metadata = metadata; // Use the same metadata
1496                add_edge_with_metadata(
1497                    structure,
1498                    join_node_id,
1499                    node_id,
1500                    Some(join_metadata),
1501                    Some(metadata),
1502                    None,
1503                );
1504
1505                node_id
1506            }
1507
1508            HydroNode::Network {
1509                serialize,
1510                deserialize,
1511                input,
1512                metadata,
1513                ..
1514            } => {
1515                let input_id = input.build_graph_structure(structure, seen_tees, config);
1516                let _from_location_key = setup_location(structure, metadata);
1517
1518                let root = metadata.location_id.root();
1519                let to_location_key = root.key();
1520                let to_location_type = root.location_type().unwrap();
1521                structure.add_location(to_location_key, to_location_type);
1522
1523                let has_serialize = match serialize {
1524                    crate::compile::ir::NetworkSend::Custom { serialize_fn } => {
1525                        serialize_fn.is_some()
1526                    }
1527                    // Embedded channels still convert member-id tags on the send side.
1528                    crate::compile::ir::NetworkSend::Embedded { .. } => true,
1529                };
1530                let has_deserialize = match deserialize {
1531                    crate::compile::ir::NetworkRecv::Custom { deserialize_fn } => {
1532                        deserialize_fn.is_some()
1533                    }
1534                    crate::compile::ir::NetworkRecv::Embedded { .. } => true,
1535                };
1536
1537                let mut label = "network(".to_owned();
1538                if has_serialize {
1539                    label.push_str("send");
1540                }
1541                if has_deserialize {
1542                    if has_serialize {
1543                        label.push_str(" + ");
1544                    }
1545                    label.push_str("recv");
1546                }
1547                label.push(')');
1548
1549                let network_id = structure.add_node_with_backtrace(
1550                    NodeLabel::Static(label),
1551                    HydroNodeType::Network,
1552                    Some(to_location_key),
1553                    Some(metadata.op.backtrace.clone()),
1554                );
1555
1556                // Extract semantic tags for network edge
1557                let input_metadata = input.metadata();
1558                add_edge_with_metadata(
1559                    structure,
1560                    input_id,
1561                    network_id,
1562                    Some(input_metadata),
1563                    Some(metadata),
1564                    Some(format!("to {:?}({})", to_location_type, to_location_key)),
1565                );
1566
1567                network_id
1568            }
1569
1570            // Non-deterministic batch operation
1571            HydroNode::Batch { inner, metadata } => build_simple_transform(TransformParams {
1572                structure,
1573                seen_tees,
1574                config,
1575                input: inner,
1576                metadata,
1577                op_name: extract_op_name(self.print_root()),
1578                node_type: HydroNodeType::NonDeterministic,
1579            }),
1580
1581            HydroNode::YieldConcat { inner, .. } => {
1582                // Unpersist is typically optimized away, just pass through
1583                inner.build_graph_structure(structure, seen_tees, config)
1584            }
1585
1586            HydroNode::UnboundSingleton { inner, .. } => {
1587                inner.build_graph_structure(structure, seen_tees, config)
1588            }
1589
1590            HydroNode::BeginAtomic { inner, .. } => {
1591                inner.build_graph_structure(structure, seen_tees, config)
1592            }
1593
1594            HydroNode::EndAtomic { inner, .. } => {
1595                inner.build_graph_structure(structure, seen_tees, config)
1596            }
1597
1598            HydroNode::Chain {
1599                first,
1600                second,
1601                metadata,
1602            }
1603            | HydroNode::MergeOrdered {
1604                first,
1605                second,
1606                metadata,
1607            } => {
1608                let first_id = first.build_graph_structure(structure, seen_tees, config);
1609                let second_id = second.build_graph_structure(structure, seen_tees, config);
1610                let location_key = Some(setup_location(structure, metadata));
1611                let chain_id = structure.add_node_with_backtrace(
1612                    NodeLabel::Static(extract_op_name(self.print_root())),
1613                    HydroNodeType::Transform,
1614                    location_key,
1615                    Some(metadata.op.backtrace.clone()),
1616                );
1617
1618                // Extract semantic tags for first edge
1619                let first_metadata = first.metadata();
1620                add_edge_with_metadata(
1621                    structure,
1622                    first_id,
1623                    chain_id,
1624                    Some(first_metadata),
1625                    Some(metadata),
1626                    Some("first".to_owned()),
1627                );
1628
1629                // Extract semantic tags for second edge
1630                let second_metadata = second.metadata();
1631                add_edge_with_metadata(
1632                    structure,
1633                    second_id,
1634                    chain_id,
1635                    Some(second_metadata),
1636                    Some(metadata),
1637                    Some("second".to_owned()),
1638                );
1639
1640                chain_id
1641            }
1642
1643            HydroNode::VersionedNetworkFork {
1644                senders, metadata, ..
1645            } => {
1646                let location_key = Some(setup_location(structure, metadata));
1647                let fork_id = structure.add_node_with_backtrace(
1648                    NodeLabel::Static(extract_op_name(self.print_root())),
1649                    HydroNodeType::NonDeterministic,
1650                    location_key,
1651                    Some(metadata.op.backtrace.clone()),
1652                );
1653
1654                for (version, sender, _serialize) in senders {
1655                    let sender_id = sender.build_graph_structure(structure, seen_tees, config);
1656                    let sender_metadata = sender.metadata();
1657                    add_edge_with_metadata(
1658                        structure,
1659                        sender_id,
1660                        fork_id,
1661                        Some(sender_metadata),
1662                        Some(metadata),
1663                        Some(format!("send v{version}")),
1664                    );
1665                }
1666
1667                fork_id
1668            }
1669
1670            HydroNode::VersionedNetwork {
1671                fork,
1672                version,
1673                metadata,
1674                ..
1675            } => {
1676                let ptr = fork.as_ptr();
1677                let fork_id = if let Some(&existing_id) = seen_tees.get(&ptr) {
1678                    existing_id
1679                } else {
1680                    let built = fork
1681                        .0
1682                        .borrow()
1683                        .build_graph_structure(structure, seen_tees, config);
1684                    seen_tees.insert(ptr, built);
1685                    built
1686                };
1687
1688                let branch_location = Some(setup_location(structure, metadata));
1689                let branch_id = structure.add_node_with_backtrace(
1690                    NodeLabel::Static(extract_op_name(self.print_root())),
1691                    HydroNodeType::NonDeterministic,
1692                    branch_location,
1693                    Some(metadata.op.backtrace.clone()),
1694                );
1695
1696                add_edge_with_metadata(
1697                    structure,
1698                    fork_id,
1699                    branch_id,
1700                    Some(metadata),
1701                    Some(metadata),
1702                    Some(format!("recv v{version}")),
1703                );
1704
1705                branch_id
1706            }
1707
1708            HydroNode::ChainFirst {
1709                first,
1710                second,
1711                metadata,
1712            } => {
1713                let first_id = first.build_graph_structure(structure, seen_tees, config);
1714                let second_id = second.build_graph_structure(structure, seen_tees, config);
1715                let location_key = Some(setup_location(structure, metadata));
1716                let chain_id = structure.add_node_with_backtrace(
1717                    NodeLabel::Static(extract_op_name(self.print_root())),
1718                    HydroNodeType::Transform,
1719                    location_key,
1720                    Some(metadata.op.backtrace.clone()),
1721                );
1722
1723                // Extract semantic tags for first edge
1724                let first_metadata = first.metadata();
1725                add_edge_with_metadata(
1726                    structure,
1727                    first_id,
1728                    chain_id,
1729                    Some(first_metadata),
1730                    Some(metadata),
1731                    Some("first".to_owned()),
1732                );
1733
1734                // Extract semantic tags for second edge
1735                let second_metadata = second.metadata();
1736                add_edge_with_metadata(
1737                    structure,
1738                    second_id,
1739                    chain_id,
1740                    Some(second_metadata),
1741                    Some(metadata),
1742                    Some("second".to_owned()),
1743                );
1744
1745                chain_id
1746            }
1747
1748            HydroNode::Counter {
1749                tag: _,
1750                prefix: _,
1751                duration,
1752                input,
1753                metadata,
1754            } => build_single_expr_transform(
1755                TransformParams {
1756                    structure,
1757                    seen_tees,
1758                    config,
1759                    input,
1760                    metadata,
1761                    op_name: extract_op_name(self.print_root()),
1762                    node_type: HydroNodeType::Transform,
1763                },
1764                duration,
1765            ),
1766        }
1767    }
1768}
1769
1770/// Utility functions for rendering multiple roots as a single graph.
1771/// Macro to reduce duplication in render functions.
1772macro_rules! render_hydro_ir {
1773    ($name:ident, $write_fn:ident) => {
1774        pub fn $name(roots: &[HydroRoot], config: HydroWriteConfig<'_>) -> String {
1775            let mut output = String::new();
1776            $write_fn(&mut output, roots, config).unwrap();
1777            output
1778        }
1779    };
1780}
1781
1782/// Macro to reduce duplication in write functions.
1783macro_rules! write_hydro_ir {
1784    ($name:ident, $writer_type:ty, $constructor:expr) => {
1785        pub fn $name(
1786            output: impl std::fmt::Write,
1787            roots: &[HydroRoot],
1788            config: HydroWriteConfig<'_>,
1789        ) -> std::fmt::Result {
1790            let mut graph_write: $writer_type = $constructor(output, config);
1791            write_hydro_ir_graph(&mut graph_write, roots, config)
1792        }
1793    };
1794}
1795
1796render_hydro_ir!(render_hydro_ir_mermaid, write_hydro_ir_mermaid);
1797write_hydro_ir!(
1798    write_hydro_ir_mermaid,
1799    HydroMermaid<'_, _>,
1800    HydroMermaid::new_with_config
1801);
1802
1803render_hydro_ir!(render_hydro_ir_dot, write_hydro_ir_dot);
1804write_hydro_ir!(
1805    write_hydro_ir_dot,
1806    HydroDot<'_, _>,
1807    HydroDot::new_with_config
1808);
1809
1810// Legacy hydroscope function - now uses HydroJson for consistency
1811render_hydro_ir!(render_hydro_ir_hydroscope, write_hydro_ir_json);
1812
1813// JSON rendering
1814render_hydro_ir!(render_hydro_ir_json, write_hydro_ir_json);
1815write_hydro_ir!(write_hydro_ir_json, HydroJson<'_, _>, HydroJson::new);
1816
1817fn write_hydro_ir_graph<W>(
1818    graph_write: W,
1819    roots: &[HydroRoot],
1820    config: HydroWriteConfig<'_>,
1821) -> Result<(), W::Err>
1822where
1823    W: HydroGraphWrite,
1824{
1825    let mut structure = HydroGraphStructure::new();
1826    let mut seen_tees = HashMap::new();
1827
1828    // Build the graph structure for all roots
1829    for leaf in roots {
1830        leaf.build_graph_structure(&mut structure, &mut seen_tees, config);
1831    }
1832
1833    write_graph_structure(&structure, graph_write, config)
1834}
hydro_lang/viz/render.rs

hydro_lang/viz/
render.rs
