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_optional_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 Optional based on OptionalBoundKind.
422fn add_optional_bound_property(
423    properties: &mut HashSet<HydroEdgeProp>,
424    bound: &crate::compile::ir::OptionalBoundKind,
425) {
426    use crate::compile::ir::OptionalBoundKind;
427
428    match bound {
429        OptionalBoundKind::Bounded => {
430            properties.insert(HydroEdgeProp::Bounded);
431        }
432        OptionalBoundKind::InitNone | OptionalBoundKind::Unbounded => {
433            properties.insert(HydroEdgeProp::Unbounded);
434        }
435    }
436}
437
438/// Helper function to add bound property for Singleton based on SingletonBoundKind.
439fn add_singleton_bound_property(
440    properties: &mut HashSet<HydroEdgeProp>,
441    bound: &crate::compile::ir::SingletonBoundKind,
442) {
443    use crate::compile::ir::SingletonBoundKind;
444
445    match bound {
446        SingletonBoundKind::Bounded => {
447            properties.insert(HydroEdgeProp::Bounded);
448        }
449        SingletonBoundKind::Monotonic | SingletonBoundKind::Unbounded => {
450            properties.insert(HydroEdgeProp::Unbounded);
451        }
452    }
453}
454
455/// Helper function to add bound property for KeyedSingleton based on KeyedSingletonBoundKind.
456fn add_keyed_singleton_bound_property(
457    properties: &mut HashSet<HydroEdgeProp>,
458    bound: &crate::compile::ir::KeyedSingletonBoundKind,
459) {
460    use crate::compile::ir::KeyedSingletonBoundKind;
461
462    match bound {
463        KeyedSingletonBoundKind::Bounded => {
464            properties.insert(HydroEdgeProp::Bounded);
465        }
466        KeyedSingletonBoundKind::BoundedValue
467        | KeyedSingletonBoundKind::MonotonicKeys
468        | KeyedSingletonBoundKind::MonotonicValue
469        | KeyedSingletonBoundKind::Unbounded => {
470            properties.insert(HydroEdgeProp::Unbounded);
471        }
472    }
473}
474
475/// Helper function to add order property based on StreamOrder.
476fn add_order_property(
477    properties: &mut HashSet<HydroEdgeProp>,
478    order: &crate::compile::ir::StreamOrder,
479) {
480    use crate::compile::ir::StreamOrder;
481
482    match order {
483        StreamOrder::TotalOrder => {
484            properties.insert(HydroEdgeProp::TotalOrder);
485        }
486        StreamOrder::NoOrder => {
487            properties.insert(HydroEdgeProp::NoOrder);
488        }
489    }
490}
491
492/// Detect if an edge crosses network boundaries by comparing source and destination locations.
493/// Returns true if the edge represents network communication between different locations.
494pub fn is_network_edge(src_location: &LocationId, dst_location: &LocationId) -> bool {
495    // Compare the root locations to determine if they differ
496    src_location.root() != dst_location.root()
497}
498
499/// Add network edge tag if source and destination locations differ.
500pub fn add_network_edge_tag(
501    properties: &mut HashSet<HydroEdgeProp>,
502    src_location: &LocationId,
503    dst_location: &LocationId,
504) {
505    if is_network_edge(src_location, dst_location) {
506        properties.insert(HydroEdgeProp::Network);
507    }
508}
509
510/// Configuration for graph writing.
511#[derive(Debug, Clone, Copy)]
512pub struct HydroWriteConfig<'a> {
513    pub show_metadata: bool,
514    pub show_location_groups: bool,
515    pub use_short_labels: bool,
516    pub location_names: &'a SecondaryMap<LocationKey, String>,
517}
518
519impl Default for HydroWriteConfig<'_> {
520    fn default() -> Self {
521        static EMPTY: OnceLock<SecondaryMap<LocationKey, String>> = OnceLock::new();
522        Self {
523            show_metadata: false,
524            show_location_groups: true,
525            use_short_labels: true, // Default to short labels for all renderers
526            location_names: EMPTY.get_or_init(SecondaryMap::new),
527        }
528    }
529}
530
531/// Node information in the Hydro graph.
532#[derive(Clone)]
533pub struct HydroGraphNode {
534    pub label: NodeLabel,
535    pub node_type: HydroNodeType,
536    pub location_key: Option<LocationKey>,
537    pub backtrace: Option<Backtrace>,
538}
539
540slotmap::new_key_type! {
541    /// Unique identifier for nodes in the visualization graph.
542    ///
543    /// This is counted/allocated separately from any other IDs within `hydro_lang`.
544    pub struct VizNodeKey;
545}
546
547impl Display for VizNodeKey {
548    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
549        write!(f, "viz{:?}", self.data()) // `"viz1v1"``
550    }
551}
552
553/// This is used by the visualizer
554/// TODO(mingwei): Make this more robust?
555impl std::str::FromStr for VizNodeKey {
556    type Err = Option<ParseIntError>;
557
558    fn from_str(s: &str) -> Result<Self, Self::Err> {
559        let nvn = s.strip_prefix("viz").ok_or(None)?;
560        let (idx, ver) = nvn.split_once("v").ok_or(None)?;
561        let idx: u64 = idx.parse()?;
562        let ver: u64 = ver.parse()?;
563        Ok(slotmap::KeyData::from_ffi((ver << 32) | idx).into())
564    }
565}
566
567impl VizNodeKey {
568    /// A key for testing with index 1.
569    #[cfg(test)]
570    pub const TEST_KEY_1: Self = Self(slotmap::KeyData::from_ffi(0x0000008F00000001)); // `1v143`
571
572    /// A key for testing with index 2.
573    #[cfg(test)]
574    pub const TEST_KEY_2: Self = Self(slotmap::KeyData::from_ffi(0x0000008F00000002)); // `2v143`
575}
576
577/// Edge information in the Hydro graph.
578#[derive(Debug, Clone)]
579pub struct HydroGraphEdge {
580    pub src: VizNodeKey,
581    pub dst: VizNodeKey,
582    pub edge_properties: HashSet<HydroEdgeProp>,
583    pub label: Option<String>,
584}
585
586/// Graph structure tracker for Hydro IR rendering.
587#[derive(Default)]
588pub struct HydroGraphStructure {
589    pub nodes: SlotMap<VizNodeKey, HydroGraphNode>,
590    pub edges: Vec<HydroGraphEdge>,
591    pub locations: SecondaryMap<LocationKey, LocationType>,
592}
593
594impl HydroGraphStructure {
595    pub fn new() -> Self {
596        Self::default()
597    }
598
599    pub fn add_node(
600        &mut self,
601        label: NodeLabel,
602        node_type: HydroNodeType,
603        location_key: Option<LocationKey>,
604    ) -> VizNodeKey {
605        self.add_node_with_backtrace(label, node_type, location_key, None)
606    }
607
608    pub fn add_node_with_backtrace(
609        &mut self,
610        label: NodeLabel,
611        node_type: HydroNodeType,
612        location_key: Option<LocationKey>,
613        backtrace: Option<Backtrace>,
614    ) -> VizNodeKey {
615        self.nodes.insert(HydroGraphNode {
616            label,
617            node_type,
618            location_key,
619            backtrace,
620        })
621    }
622
623    /// Add a node with metadata, extracting backtrace automatically
624    pub fn add_node_with_metadata(
625        &mut self,
626        label: NodeLabel,
627        node_type: HydroNodeType,
628        metadata: &HydroIrMetadata,
629    ) -> VizNodeKey {
630        let location_key = Some(setup_location(self, metadata));
631        let backtrace = Some(metadata.op.backtrace.clone());
632        self.add_node_with_backtrace(label, node_type, location_key, backtrace)
633    }
634
635    pub fn add_edge(
636        &mut self,
637        src: VizNodeKey,
638        dst: VizNodeKey,
639        edge_properties: HashSet<HydroEdgeProp>,
640        label: Option<String>,
641    ) {
642        self.edges.push(HydroGraphEdge {
643            src,
644            dst,
645            edge_properties,
646            label,
647        });
648    }
649
650    // Legacy method for backward compatibility
651    pub fn add_edge_single(
652        &mut self,
653        src: VizNodeKey,
654        dst: VizNodeKey,
655        edge_type: HydroEdgeProp,
656        label: Option<String>,
657    ) {
658        let mut properties = HashSet::new();
659        properties.insert(edge_type);
660        self.edges.push(HydroGraphEdge {
661            src,
662            dst,
663            edge_properties: properties,
664            label,
665        });
666    }
667
668    pub fn add_location(&mut self, location_key: LocationKey, location_type: LocationType) {
669        self.locations.insert(location_key, location_type);
670    }
671}
672
673/// Function to extract an op_name from a print_root() result for use in labels.
674pub fn extract_op_name(full_label: String) -> String {
675    full_label
676        .split('(')
677        .next()
678        .unwrap_or("unknown")
679        .to_lowercase()
680}
681
682/// Extract a short, readable label from the full token stream label using print_root() style naming
683pub fn extract_short_label(full_label: &str) -> String {
684    // Use the same logic as extract_op_name but handle the specific cases we need for UI display
685    if let Some(op_name) = full_label.split('(').next() {
686        let base_name = op_name.to_lowercase();
687        match base_name.as_str() {
688            // Handle special cases for UI display
689            "source" => {
690                if full_label.contains("Iter") {
691                    "source_iter".to_owned()
692                } else if full_label.contains("Stream") {
693                    "source_stream".to_owned()
694                } else if full_label.contains("ExternalNetwork") {
695                    "external_network".to_owned()
696                } else if full_label.contains("Spin") {
697                    "spin".to_owned()
698                } else {
699                    "source".to_owned()
700                }
701            }
702            "network" => {
703                if full_label.contains("deser") {
704                    "network(recv)".to_owned()
705                } else if full_label.contains("ser") {
706                    "network(send)".to_owned()
707                } else {
708                    "network".to_owned()
709                }
710            }
711            // For all other cases, just use the lowercase base name (same as extract_op_name)
712            _ => base_name,
713        }
714    } else {
715        // Fallback for labels that don't follow the pattern
716        if full_label.len() > 20 {
717            format!("{}...", &full_label[..17])
718        } else {
719            full_label.to_owned()
720        }
721    }
722}
723
724/// Helper function to set up location in structure from metadata.
725fn setup_location(structure: &mut HydroGraphStructure, metadata: &HydroIrMetadata) -> LocationKey {
726    let root = metadata.location_id.root();
727    let location_key = root.key();
728    let location_type = root.location_type().unwrap();
729    structure.add_location(location_key, location_type);
730    location_key
731}
732
733/// Helper function to add an edge with semantic tags extracted from metadata.
734/// This function combines collection kind extraction with network detection.
735fn add_edge_with_metadata(
736    structure: &mut HydroGraphStructure,
737    src_id: VizNodeKey,
738    dst_id: VizNodeKey,
739    src_metadata: Option<&HydroIrMetadata>,
740    dst_metadata: Option<&HydroIrMetadata>,
741    label: Option<String>,
742) {
743    let mut properties = HashSet::new();
744
745    // Extract semantic tags from source metadata's collection kind
746    if let Some(metadata) = src_metadata {
747        properties.extend(extract_edge_properties_from_collection_kind(
748            &metadata.collection_kind,
749        ));
750    }
751
752    // Add network edge tag if locations differ
753    if let (Some(src_meta), Some(dst_meta)) = (src_metadata, dst_metadata) {
754        add_network_edge_tag(
755            &mut properties,
756            &src_meta.location_id,
757            &dst_meta.location_id,
758        );
759    }
760
761    // If no properties were extracted, default to Stream
762    if properties.is_empty() {
763        properties.insert(HydroEdgeProp::Stream);
764    }
765
766    structure.add_edge(src_id, dst_id, properties, label);
767}
768
769/// Helper function to write a graph structure using any GraphWrite implementation
770fn write_graph_structure<W>(
771    structure: &HydroGraphStructure,
772    graph_write: W,
773    config: HydroWriteConfig<'_>,
774) -> Result<(), W::Err>
775where
776    W: HydroGraphWrite,
777{
778    let mut graph_write = graph_write;
779    // Write the graph
780    graph_write.write_prologue()?;
781
782    // Write node definitions
783    for (node_id, node) in structure.nodes.iter() {
784        let location_type = node
785            .location_key
786            .and_then(|loc_key| structure.locations.get(loc_key))
787            .copied();
788
789        graph_write.write_node_definition(
790            node_id,
791            &node.label,
792            node.node_type,
793            node.location_key,
794            location_type,
795            node.backtrace.as_ref(),
796        )?;
797    }
798
799    // Group nodes by location if requested
800    if config.show_location_groups {
801        let mut nodes_by_location = SecondaryMap::<LocationKey, Vec<VizNodeKey>>::new();
802        for (node_id, node) in structure.nodes.iter() {
803            if let Some(location_key) = node.location_key {
804                nodes_by_location
805                    .entry(location_key)
806                    .expect("location was removed")
807                    .or_default()
808                    .push(node_id);
809            }
810        }
811
812        for (location_key, node_ids) in nodes_by_location.iter() {
813            if let Some(&location_type) = structure.locations.get(location_key) {
814                graph_write.write_location_start(location_key, location_type)?;
815                for &node_id in node_ids.iter() {
816                    graph_write.write_node(node_id)?;
817                }
818                graph_write.write_location_end()?;
819            }
820        }
821    }
822
823    // Write edges
824    for edge in structure.edges.iter() {
825        graph_write.write_edge(
826            edge.src,
827            edge.dst,
828            &edge.edge_properties,
829            edge.label.as_deref(),
830        )?;
831    }
832
833    graph_write.write_epilogue()?;
834    Ok(())
835}
836
837impl HydroRoot {
838    /// Build the graph structure by traversing the IR tree.
839    pub fn build_graph_structure(
840        &self,
841        structure: &mut HydroGraphStructure,
842        seen_tees: &mut HashMap<*const std::cell::RefCell<HydroNode>, VizNodeKey>,
843        config: HydroWriteConfig<'_>,
844    ) -> VizNodeKey {
845        // Helper function for sink nodes to reduce duplication
846        fn build_sink_node(
847            structure: &mut HydroGraphStructure,
848            seen_tees: &mut HashMap<*const std::cell::RefCell<HydroNode>, VizNodeKey>,
849            config: HydroWriteConfig<'_>,
850            input: &HydroNode,
851            sink_metadata: Option<&HydroIrMetadata>,
852            label: NodeLabel,
853        ) -> VizNodeKey {
854            let input_id = input.build_graph_structure(structure, seen_tees, config);
855
856            // If no explicit metadata is provided, extract it from the input node
857            let effective_metadata = if let Some(meta) = sink_metadata {
858                Some(meta)
859            } else {
860                match input {
861                    HydroNode::Placeholder => None,
862                    // All other variants have metadata
863                    _ => Some(input.metadata()),
864                }
865            };
866
867            let location_key = effective_metadata.map(|m| setup_location(structure, m));
868            let sink_id = structure.add_node_with_backtrace(
869                label,
870                HydroNodeType::Sink,
871                location_key,
872                effective_metadata.map(|m| m.op.backtrace.clone()),
873            );
874
875            // Extract semantic tags from input metadata
876            let input_metadata = input.metadata();
877            add_edge_with_metadata(
878                structure,
879                input_id,
880                sink_id,
881                Some(input_metadata),
882                sink_metadata,
883                None,
884            );
885
886            sink_id
887        }
888
889        match self {
890            // Sink operations - semantic tags extracted from input metadata
891            HydroRoot::ForEach { f, input, .. } => build_sink_node(
892                structure,
893                seen_tees,
894                config,
895                input,
896                None,
897                NodeLabel::with_exprs("for_each".to_owned(), vec![f.expr.clone()]),
898            ),
899
900            HydroRoot::SendExternal {
901                to_external_key,
902                to_port_id,
903                input,
904                ..
905            } => build_sink_node(
906                structure,
907                seen_tees,
908                config,
909                input,
910                None,
911                NodeLabel::with_exprs(
912                    format!("send_external({}:{})", to_external_key, to_port_id),
913                    vec![],
914                ),
915            ),
916
917            HydroRoot::DestSink { sink, input, .. } => build_sink_node(
918                structure,
919                seen_tees,
920                config,
921                input,
922                None,
923                NodeLabel::with_exprs("dest_sink".to_owned(), vec![sink.clone()]),
924            ),
925
926            HydroRoot::CycleSink {
927                cycle_id, input, ..
928            } => build_sink_node(
929                structure,
930                seen_tees,
931                config,
932                input,
933                None,
934                NodeLabel::static_label(format!("cycle_sink({})", cycle_id)),
935            ),
936
937            HydroRoot::EmbeddedOutput { ident, input, .. } => build_sink_node(
938                structure,
939                seen_tees,
940                config,
941                input,
942                None,
943                NodeLabel::static_label(format!("embedded_output({})", ident)),
944            ),
945
946            HydroRoot::Null { input, .. } => build_sink_node(
947                structure,
948                seen_tees,
949                config,
950                input,
951                None,
952                NodeLabel::static_label("null".to_owned()),
953            ),
954        }
955    }
956}
957
958impl HydroNode {
959    /// Build the graph structure recursively for this node.
960    pub fn build_graph_structure(
961        &self,
962        structure: &mut HydroGraphStructure,
963        seen_tees: &mut HashMap<*const std::cell::RefCell<HydroNode>, VizNodeKey>,
964        config: HydroWriteConfig<'_>,
965    ) -> VizNodeKey {
966        // Helper functions to reduce duplication, categorized by input/expression patterns
967
968        /// Common parameters for transform builder functions to reduce argument count
969        struct TransformParams<'a> {
970            structure: &'a mut HydroGraphStructure,
971            seen_tees: &'a mut HashMap<*const std::cell::RefCell<HydroNode>, VizNodeKey>,
972            config: HydroWriteConfig<'a>,
973            input: &'a HydroNode,
974            metadata: &'a HydroIrMetadata,
975            op_name: String,
976            node_type: HydroNodeType,
977        }
978
979        // Single-input transform with no expressions
980        fn build_simple_transform(params: TransformParams<'_>) -> VizNodeKey {
981            let input_id = params.input.build_graph_structure(
982                params.structure,
983                params.seen_tees,
984                params.config,
985            );
986            let node_id = params.structure.add_node_with_metadata(
987                NodeLabel::Static(params.op_name.clone()),
988                params.node_type,
989                params.metadata,
990            );
991
992            // Extract semantic tags from input metadata
993            let input_metadata = params.input.metadata();
994            add_edge_with_metadata(
995                params.structure,
996                input_id,
997                node_id,
998                Some(input_metadata),
999                Some(params.metadata),
1000                None,
1001            );
1002
1003            node_id
1004        }
1005
1006        // Single-input transform with one expression
1007        fn build_single_expr_transform(
1008            params: TransformParams<'_>,
1009            expr: &DebugExpr,
1010        ) -> VizNodeKey {
1011            let input_id = params.input.build_graph_structure(
1012                params.structure,
1013                params.seen_tees,
1014                params.config,
1015            );
1016            let node_id = params.structure.add_node_with_metadata(
1017                NodeLabel::with_exprs(params.op_name.clone(), vec![expr.clone()]),
1018                params.node_type,
1019                params.metadata,
1020            );
1021
1022            // Extract semantic tags from input metadata
1023            let input_metadata = params.input.metadata();
1024            add_edge_with_metadata(
1025                params.structure,
1026                input_id,
1027                node_id,
1028                Some(input_metadata),
1029                Some(params.metadata),
1030                None,
1031            );
1032
1033            node_id
1034        }
1035
1036        // Single-input transform with two expressions
1037        fn build_dual_expr_transform(
1038            params: TransformParams<'_>,
1039            expr1: &DebugExpr,
1040            expr2: &DebugExpr,
1041        ) -> VizNodeKey {
1042            let input_id = params.input.build_graph_structure(
1043                params.structure,
1044                params.seen_tees,
1045                params.config,
1046            );
1047            let node_id = params.structure.add_node_with_metadata(
1048                NodeLabel::with_exprs(params.op_name.clone(), vec![expr1.clone(), expr2.clone()]),
1049                params.node_type,
1050                params.metadata,
1051            );
1052
1053            // Extract semantic tags from input metadata
1054            let input_metadata = params.input.metadata();
1055            add_edge_with_metadata(
1056                params.structure,
1057                input_id,
1058                node_id,
1059                Some(input_metadata),
1060                Some(params.metadata),
1061                None,
1062            );
1063
1064            node_id
1065        }
1066
1067        // Helper function for source nodes
1068        fn build_source_node(
1069            structure: &mut HydroGraphStructure,
1070            metadata: &HydroIrMetadata,
1071            label: String,
1072        ) -> VizNodeKey {
1073            structure.add_node_with_metadata(
1074                NodeLabel::Static(label),
1075                HydroNodeType::Source,
1076                metadata,
1077            )
1078        }
1079
1080        match self {
1081            HydroNode::Placeholder => structure.add_node(
1082                NodeLabel::Static("PLACEHOLDER".to_owned()),
1083                HydroNodeType::Transform,
1084                None,
1085            ),
1086
1087            HydroNode::Source {
1088                source, metadata, ..
1089            } => {
1090                let label = match source {
1091                    HydroSource::Stream(expr) => format!("source_stream({})", expr),
1092                    HydroSource::ExternalNetwork() => "external_network()".to_owned(),
1093                    HydroSource::Iter(expr) => format!("source_iter({})", expr),
1094                    HydroSource::Spin() => "spin()".to_owned(),
1095                    HydroSource::ClusterMembers(location_id, _) => {
1096                        format!(
1097                            "source_stream(cluster_membership_stream({:?}))",
1098                            location_id
1099                        )
1100                    }
1101                    HydroSource::Embedded(ident) => {
1102                        format!("embedded_input({})", ident)
1103                    }
1104                    HydroSource::EmbeddedSingleton(ident) => {
1105                        format!("embedded_singleton_input({})", ident)
1106                    }
1107                };
1108                build_source_node(structure, metadata, label)
1109            }
1110
1111            HydroNode::SingletonSource {
1112                value,
1113                first_tick_only,
1114                metadata,
1115            } => {
1116                let label = if *first_tick_only {
1117                    format!("singleton_first_tick({})", value)
1118                } else {
1119                    format!("singleton({})", value)
1120                };
1121                build_source_node(structure, metadata, label)
1122            }
1123
1124            HydroNode::ExternalInput {
1125                from_external_key,
1126                from_port_id,
1127                metadata,
1128                ..
1129            } => build_source_node(
1130                structure,
1131                metadata,
1132                format!("external_input({}:{})", from_external_key, from_port_id),
1133            ),
1134
1135            HydroNode::CycleSource {
1136                cycle_id, metadata, ..
1137            } => build_source_node(structure, metadata, format!("cycle_source({})", cycle_id)),
1138
1139            HydroNode::Tee { inner, metadata }
1140            | HydroNode::Reference {
1141                inner, metadata, ..
1142            } => {
1143                let ptr = inner.as_ptr();
1144                if let Some(&existing_id) = seen_tees.get(&ptr) {
1145                    return existing_id;
1146                }
1147
1148                let input_id = inner
1149                    .0
1150                    .borrow()
1151                    .build_graph_structure(structure, seen_tees, config);
1152                let node_type = if matches!(self, HydroNode::Reference { .. }) {
1153                    HydroNodeType::Aggregation
1154                } else {
1155                    HydroNodeType::Tee
1156                };
1157                let tee_id = structure.add_node_with_metadata(
1158                    NodeLabel::Static(extract_op_name(self.print_root())),
1159                    node_type,
1160                    metadata,
1161                );
1162
1163                seen_tees.insert(ptr, tee_id);
1164
1165                // Extract semantic tags from input
1166                let inner_borrow = inner.0.borrow();
1167                let input_metadata = inner_borrow.metadata();
1168                add_edge_with_metadata(
1169                    structure,
1170                    input_id,
1171                    tee_id,
1172                    Some(input_metadata),
1173                    Some(metadata),
1174                    None,
1175                );
1176                drop(inner_borrow);
1177
1178                tee_id
1179            }
1180
1181            HydroNode::PartitionSide {
1182                inner, metadata, ..
1183            } => {
1184                let ptr = inner.as_ptr();
1185                if let Some(&existing_id) = seen_tees.get(&ptr) {
1186                    return existing_id;
1187                }
1188
1189                let input_id = inner
1190                    .0
1191                    .borrow()
1192                    .build_graph_structure(structure, seen_tees, config);
1193                let partition_id = structure.add_node_with_metadata(
1194                    NodeLabel::Static(extract_op_name(self.print_root())),
1195                    HydroNodeType::Tee,
1196                    metadata,
1197                );
1198
1199                seen_tees.insert(ptr, partition_id);
1200
1201                // Extract semantic tags from input
1202                let inner_borrow = inner.0.borrow();
1203                let input_metadata = inner_borrow.metadata();
1204                add_edge_with_metadata(
1205                    structure,
1206                    input_id,
1207                    partition_id,
1208                    Some(input_metadata),
1209                    Some(metadata),
1210                    None,
1211                );
1212                drop(inner_borrow);
1213
1214                partition_id
1215            }
1216            HydroNode::PartitionShared { input, .. } => {
1217                // Transparent pass-thru.
1218                input.build_graph_structure(structure, seen_tees, config)
1219            }
1220
1221            // Non-deterministic operation
1222            HydroNode::ObserveNonDet {
1223                inner, metadata, ..
1224            } => build_simple_transform(TransformParams {
1225                structure,
1226                seen_tees,
1227                config,
1228                input: inner,
1229                metadata,
1230                op_name: extract_op_name(self.print_root()),
1231                node_type: HydroNodeType::NonDeterministic,
1232            }),
1233
1234            // Transform operations with Stream edges - grouped by node/edge type
1235            HydroNode::Cast { inner, metadata }
1236            | HydroNode::AssertIsConsistent {
1237                inner, metadata, ..
1238            }
1239            | HydroNode::DeferTick {
1240                input: inner,
1241                metadata,
1242            }
1243            | HydroNode::Enumerate {
1244                input: inner,
1245                metadata,
1246                ..
1247            }
1248            | HydroNode::Unique {
1249                input: inner,
1250                metadata,
1251            }
1252            | HydroNode::ResolveFutures {
1253                input: inner,
1254                metadata,
1255            }
1256            | HydroNode::ResolveFuturesBlocking {
1257                input: inner,
1258                metadata,
1259            }
1260            | HydroNode::ResolveFuturesOrdered {
1261                input: inner,
1262                metadata,
1263            } => build_simple_transform(TransformParams {
1264                structure,
1265                seen_tees,
1266                config,
1267                input: inner,
1268                metadata,
1269                op_name: extract_op_name(self.print_root()),
1270                node_type: HydroNodeType::Transform,
1271            }),
1272
1273            // Aggregation operation - semantic tags extracted from metadata
1274            HydroNode::Sort {
1275                input: inner,
1276                metadata,
1277            } => build_simple_transform(TransformParams {
1278                structure,
1279                seen_tees,
1280                config,
1281                input: inner,
1282                metadata,
1283                op_name: extract_op_name(self.print_root()),
1284                node_type: HydroNodeType::Aggregation,
1285            }),
1286
1287            // Single-expression Transform operations - grouped by node type
1288            HydroNode::Map {
1289                f, input, metadata, ..
1290            }
1291            | HydroNode::Filter { f, input, metadata }
1292            | HydroNode::FlatMap { f, input, metadata }
1293            | HydroNode::FlatMapStreamBlocking { f, input, metadata }
1294            | HydroNode::FilterMap { f, input, metadata }
1295            | HydroNode::Inspect { f, input, metadata } => build_single_expr_transform(
1296                TransformParams {
1297                    structure,
1298                    seen_tees,
1299                    config,
1300                    input,
1301                    metadata,
1302                    op_name: extract_op_name(self.print_root()),
1303                    node_type: HydroNodeType::Transform,
1304                },
1305                &f.expr,
1306            ),
1307
1308            // Single-expression Aggregation operations - grouped by node type
1309            HydroNode::Reduce { f, input, metadata }
1310            | HydroNode::ReduceKeyed { f, input, metadata } => build_single_expr_transform(
1311                TransformParams {
1312                    structure,
1313                    seen_tees,
1314                    config,
1315                    input,
1316                    metadata,
1317                    op_name: extract_op_name(self.print_root()),
1318                    node_type: HydroNodeType::Aggregation,
1319                },
1320                &f.expr,
1321            ),
1322
1323            // Join-like operations with left/right edge labels - grouped by edge labeling
1324            HydroNode::Join {
1325                left,
1326                right,
1327                metadata,
1328            }
1329            | HydroNode::JoinHalf {
1330                left,
1331                right,
1332                metadata,
1333            }
1334            | HydroNode::CrossProduct {
1335                left,
1336                right,
1337                metadata,
1338            }
1339            | HydroNode::CrossSingleton {
1340                left,
1341                right,
1342                metadata,
1343            } => {
1344                let left_id = left.build_graph_structure(structure, seen_tees, config);
1345                let right_id = right.build_graph_structure(structure, seen_tees, config);
1346                let node_id = structure.add_node_with_metadata(
1347                    NodeLabel::Static(extract_op_name(self.print_root())),
1348                    HydroNodeType::Join,
1349                    metadata,
1350                );
1351
1352                // Extract semantic tags for left edge
1353                let left_metadata = left.metadata();
1354                add_edge_with_metadata(
1355                    structure,
1356                    left_id,
1357                    node_id,
1358                    Some(left_metadata),
1359                    Some(metadata),
1360                    Some("left".to_owned()),
1361                );
1362
1363                // Extract semantic tags for right edge
1364                let right_metadata = right.metadata();
1365                add_edge_with_metadata(
1366                    structure,
1367                    right_id,
1368                    node_id,
1369                    Some(right_metadata),
1370                    Some(metadata),
1371                    Some("right".to_owned()),
1372                );
1373
1374                node_id
1375            }
1376
1377            // Join-like operations with pos/neg edge labels - grouped by edge labeling
1378            HydroNode::Difference {
1379                pos: left,
1380                neg: right,
1381                metadata,
1382            }
1383            | HydroNode::AntiJoin {
1384                pos: left,
1385                neg: right,
1386                metadata,
1387            } => {
1388                let left_id = left.build_graph_structure(structure, seen_tees, config);
1389                let right_id = right.build_graph_structure(structure, seen_tees, config);
1390                let node_id = structure.add_node_with_metadata(
1391                    NodeLabel::Static(extract_op_name(self.print_root())),
1392                    HydroNodeType::Join,
1393                    metadata,
1394                );
1395
1396                // Extract semantic tags for pos edge
1397                let left_metadata = left.metadata();
1398                add_edge_with_metadata(
1399                    structure,
1400                    left_id,
1401                    node_id,
1402                    Some(left_metadata),
1403                    Some(metadata),
1404                    Some("pos".to_owned()),
1405                );
1406
1407                // Extract semantic tags for neg edge
1408                let right_metadata = right.metadata();
1409                add_edge_with_metadata(
1410                    structure,
1411                    right_id,
1412                    node_id,
1413                    Some(right_metadata),
1414                    Some(metadata),
1415                    Some("neg".to_owned()),
1416                );
1417
1418                node_id
1419            }
1420
1421            // Dual expression transforms - consolidated using pattern matching
1422            HydroNode::Fold {
1423                init,
1424                acc,
1425                input,
1426                metadata,
1427                ..
1428            }
1429            | HydroNode::FoldKeyed {
1430                init,
1431                acc,
1432                input,
1433                metadata,
1434                ..
1435            }
1436            | HydroNode::Scan {
1437                init,
1438                acc,
1439                input,
1440                metadata,
1441            }
1442            | HydroNode::ScanAsyncBlocking {
1443                init,
1444                acc,
1445                input,
1446                metadata,
1447            } => {
1448                let node_type = HydroNodeType::Aggregation; // All are aggregation operations
1449
1450                build_dual_expr_transform(
1451                    TransformParams {
1452                        structure,
1453                        seen_tees,
1454                        config,
1455                        input,
1456                        metadata,
1457                        op_name: extract_op_name(self.print_root()),
1458                        node_type,
1459                    },
1460                    &init.expr,
1461                    &acc.expr,
1462                )
1463            }
1464
1465            // Combination of join and transform
1466            HydroNode::ReduceKeyedWatermark {
1467                f,
1468                input,
1469                watermark,
1470                metadata,
1471            } => {
1472                let input_id = input.build_graph_structure(structure, seen_tees, config);
1473                let watermark_id = watermark.build_graph_structure(structure, seen_tees, config);
1474                let location_key = Some(setup_location(structure, metadata));
1475                let join_node_id = structure.add_node_with_backtrace(
1476                    NodeLabel::Static(extract_op_name(self.print_root())),
1477                    HydroNodeType::Join,
1478                    location_key,
1479                    Some(metadata.op.backtrace.clone()),
1480                );
1481
1482                // Extract semantic tags for input edge
1483                let input_metadata = input.metadata();
1484                add_edge_with_metadata(
1485                    structure,
1486                    input_id,
1487                    join_node_id,
1488                    Some(input_metadata),
1489                    Some(metadata),
1490                    Some("input".to_owned()),
1491                );
1492
1493                // Extract semantic tags for watermark edge
1494                let watermark_metadata = watermark.metadata();
1495                add_edge_with_metadata(
1496                    structure,
1497                    watermark_id,
1498                    join_node_id,
1499                    Some(watermark_metadata),
1500                    Some(metadata),
1501                    Some("watermark".to_owned()),
1502                );
1503
1504                let node_id = structure.add_node_with_backtrace(
1505                    NodeLabel::with_exprs(extract_op_name(self.print_root()), vec![f.expr.clone()]),
1506                    HydroNodeType::Aggregation,
1507                    location_key,
1508                    Some(metadata.op.backtrace.clone()),
1509                );
1510
1511                // Edge from join to aggregation node
1512                let join_metadata = metadata; // Use the same metadata
1513                add_edge_with_metadata(
1514                    structure,
1515                    join_node_id,
1516                    node_id,
1517                    Some(join_metadata),
1518                    Some(metadata),
1519                    None,
1520                );
1521
1522                node_id
1523            }
1524
1525            HydroNode::Network {
1526                serialize,
1527                deserialize,
1528                input,
1529                metadata,
1530                ..
1531            } => {
1532                let input_id = input.build_graph_structure(structure, seen_tees, config);
1533                let _from_location_key = setup_location(structure, metadata);
1534
1535                let root = metadata.location_id.root();
1536                let to_location_key = root.key();
1537                let to_location_type = root.location_type().unwrap();
1538                structure.add_location(to_location_key, to_location_type);
1539
1540                let has_serialize = match serialize {
1541                    crate::compile::ir::NetworkSend::Custom { serialize_fn } => {
1542                        serialize_fn.is_some()
1543                    }
1544                    // Embedded channels still convert member-id tags on the send side.
1545                    crate::compile::ir::NetworkSend::Embedded { .. } => true,
1546                };
1547                let has_deserialize = match deserialize {
1548                    crate::compile::ir::NetworkRecv::Custom { deserialize_fn } => {
1549                        deserialize_fn.is_some()
1550                    }
1551                    crate::compile::ir::NetworkRecv::Embedded { .. } => true,
1552                };
1553
1554                let mut label = "network(".to_owned();
1555                if has_serialize {
1556                    label.push_str("send");
1557                }
1558                if has_deserialize {
1559                    if has_serialize {
1560                        label.push_str(" + ");
1561                    }
1562                    label.push_str("recv");
1563                }
1564                label.push(')');
1565
1566                let network_id = structure.add_node_with_backtrace(
1567                    NodeLabel::Static(label),
1568                    HydroNodeType::Network,
1569                    Some(to_location_key),
1570                    Some(metadata.op.backtrace.clone()),
1571                );
1572
1573                // Extract semantic tags for network edge
1574                let input_metadata = input.metadata();
1575                add_edge_with_metadata(
1576                    structure,
1577                    input_id,
1578                    network_id,
1579                    Some(input_metadata),
1580                    Some(metadata),
1581                    Some(format!("to {:?}({})", to_location_type, to_location_key)),
1582                );
1583
1584                network_id
1585            }
1586
1587            // Non-deterministic batch operation
1588            HydroNode::Batch { inner, metadata } => build_simple_transform(TransformParams {
1589                structure,
1590                seen_tees,
1591                config,
1592                input: inner,
1593                metadata,
1594                op_name: extract_op_name(self.print_root()),
1595                node_type: HydroNodeType::NonDeterministic,
1596            }),
1597
1598            HydroNode::YieldConcat { inner, .. } => {
1599                // Unpersist is typically optimized away, just pass through
1600                inner.build_graph_structure(structure, seen_tees, config)
1601            }
1602
1603            HydroNode::UnboundSingleton { inner, .. } => {
1604                inner.build_graph_structure(structure, seen_tees, config)
1605            }
1606
1607            HydroNode::BeginAtomic { inner, .. } => {
1608                inner.build_graph_structure(structure, seen_tees, config)
1609            }
1610
1611            HydroNode::EndAtomic { inner, .. } => {
1612                inner.build_graph_structure(structure, seen_tees, config)
1613            }
1614
1615            HydroNode::Chain {
1616                first,
1617                second,
1618                metadata,
1619            }
1620            | HydroNode::MergeOrdered {
1621                first,
1622                second,
1623                metadata,
1624            } => {
1625                let first_id = first.build_graph_structure(structure, seen_tees, config);
1626                let second_id = second.build_graph_structure(structure, seen_tees, config);
1627                let location_key = Some(setup_location(structure, metadata));
1628                let chain_id = structure.add_node_with_backtrace(
1629                    NodeLabel::Static(extract_op_name(self.print_root())),
1630                    HydroNodeType::Transform,
1631                    location_key,
1632                    Some(metadata.op.backtrace.clone()),
1633                );
1634
1635                // Extract semantic tags for first edge
1636                let first_metadata = first.metadata();
1637                add_edge_with_metadata(
1638                    structure,
1639                    first_id,
1640                    chain_id,
1641                    Some(first_metadata),
1642                    Some(metadata),
1643                    Some("first".to_owned()),
1644                );
1645
1646                // Extract semantic tags for second edge
1647                let second_metadata = second.metadata();
1648                add_edge_with_metadata(
1649                    structure,
1650                    second_id,
1651                    chain_id,
1652                    Some(second_metadata),
1653                    Some(metadata),
1654                    Some("second".to_owned()),
1655                );
1656
1657                chain_id
1658            }
1659
1660            HydroNode::VersionedNetworkFork {
1661                senders, metadata, ..
1662            } => {
1663                let location_key = Some(setup_location(structure, metadata));
1664                let fork_id = structure.add_node_with_backtrace(
1665                    NodeLabel::Static(extract_op_name(self.print_root())),
1666                    HydroNodeType::NonDeterministic,
1667                    location_key,
1668                    Some(metadata.op.backtrace.clone()),
1669                );
1670
1671                for (version, sender, _serialize) in senders {
1672                    let sender_id = sender.build_graph_structure(structure, seen_tees, config);
1673                    let sender_metadata = sender.metadata();
1674                    add_edge_with_metadata(
1675                        structure,
1676                        sender_id,
1677                        fork_id,
1678                        Some(sender_metadata),
1679                        Some(metadata),
1680                        Some(format!("send v{version}")),
1681                    );
1682                }
1683
1684                fork_id
1685            }
1686
1687            HydroNode::VersionedNetwork {
1688                fork,
1689                version,
1690                metadata,
1691                ..
1692            } => {
1693                let ptr = fork.as_ptr();
1694                let fork_id = if let Some(&existing_id) = seen_tees.get(&ptr) {
1695                    existing_id
1696                } else {
1697                    let built = fork
1698                        .0
1699                        .borrow()
1700                        .build_graph_structure(structure, seen_tees, config);
1701                    seen_tees.insert(ptr, built);
1702                    built
1703                };
1704
1705                let branch_location = Some(setup_location(structure, metadata));
1706                let branch_id = structure.add_node_with_backtrace(
1707                    NodeLabel::Static(extract_op_name(self.print_root())),
1708                    HydroNodeType::NonDeterministic,
1709                    branch_location,
1710                    Some(metadata.op.backtrace.clone()),
1711                );
1712
1713                add_edge_with_metadata(
1714                    structure,
1715                    fork_id,
1716                    branch_id,
1717                    Some(metadata),
1718                    Some(metadata),
1719                    Some(format!("recv v{version}")),
1720                );
1721
1722                branch_id
1723            }
1724
1725            HydroNode::ChainFirst {
1726                first,
1727                second,
1728                metadata,
1729            } => {
1730                let first_id = first.build_graph_structure(structure, seen_tees, config);
1731                let second_id = second.build_graph_structure(structure, seen_tees, config);
1732                let location_key = Some(setup_location(structure, metadata));
1733                let chain_id = structure.add_node_with_backtrace(
1734                    NodeLabel::Static(extract_op_name(self.print_root())),
1735                    HydroNodeType::Transform,
1736                    location_key,
1737                    Some(metadata.op.backtrace.clone()),
1738                );
1739
1740                // Extract semantic tags for first edge
1741                let first_metadata = first.metadata();
1742                add_edge_with_metadata(
1743                    structure,
1744                    first_id,
1745                    chain_id,
1746                    Some(first_metadata),
1747                    Some(metadata),
1748                    Some("first".to_owned()),
1749                );
1750
1751                // Extract semantic tags for second edge
1752                let second_metadata = second.metadata();
1753                add_edge_with_metadata(
1754                    structure,
1755                    second_id,
1756                    chain_id,
1757                    Some(second_metadata),
1758                    Some(metadata),
1759                    Some("second".to_owned()),
1760                );
1761
1762                chain_id
1763            }
1764
1765            HydroNode::Counter {
1766                tag: _,
1767                prefix: _,
1768                duration,
1769                input,
1770                metadata,
1771            } => build_single_expr_transform(
1772                TransformParams {
1773                    structure,
1774                    seen_tees,
1775                    config,
1776                    input,
1777                    metadata,
1778                    op_name: extract_op_name(self.print_root()),
1779                    node_type: HydroNodeType::Transform,
1780                },
1781                duration,
1782            ),
1783        }
1784    }
1785}
1786
1787/// Utility functions for rendering multiple roots as a single graph.
1788/// Macro to reduce duplication in render functions.
1789macro_rules! render_hydro_ir {
1790    ($name:ident, $write_fn:ident) => {
1791        pub fn $name(roots: &[HydroRoot], config: HydroWriteConfig<'_>) -> String {
1792            let mut output = String::new();
1793            $write_fn(&mut output, roots, config).unwrap();
1794            output
1795        }
1796    };
1797}
1798
1799/// Macro to reduce duplication in write functions.
1800macro_rules! write_hydro_ir {
1801    ($name:ident, $writer_type:ty, $constructor:expr) => {
1802        pub fn $name(
1803            output: impl std::fmt::Write,
1804            roots: &[HydroRoot],
1805            config: HydroWriteConfig<'_>,
1806        ) -> std::fmt::Result {
1807            let mut graph_write: $writer_type = $constructor(output, config);
1808            write_hydro_ir_graph(&mut graph_write, roots, config)
1809        }
1810    };
1811}
1812
1813render_hydro_ir!(render_hydro_ir_mermaid, write_hydro_ir_mermaid);
1814write_hydro_ir!(
1815    write_hydro_ir_mermaid,
1816    HydroMermaid<'_, _>,
1817    HydroMermaid::new_with_config
1818);
1819
1820render_hydro_ir!(render_hydro_ir_dot, write_hydro_ir_dot);
1821write_hydro_ir!(
1822    write_hydro_ir_dot,
1823    HydroDot<'_, _>,
1824    HydroDot::new_with_config
1825);
1826
1827// Legacy hydroscope function - now uses HydroJson for consistency
1828render_hydro_ir!(render_hydro_ir_hydroscope, write_hydro_ir_json);
1829
1830// JSON rendering
1831render_hydro_ir!(render_hydro_ir_json, write_hydro_ir_json);
1832write_hydro_ir!(write_hydro_ir_json, HydroJson<'_, _>, HydroJson::new);
1833
1834fn write_hydro_ir_graph<W>(
1835    graph_write: W,
1836    roots: &[HydroRoot],
1837    config: HydroWriteConfig<'_>,
1838) -> Result<(), W::Err>
1839where
1840    W: HydroGraphWrite,
1841{
1842    let mut structure = HydroGraphStructure::new();
1843    let mut seen_tees = HashMap::new();
1844
1845    // Build the graph structure for all roots
1846    for leaf in roots {
1847        leaf.build_graph_structure(&mut structure, &mut seen_tees, config);
1848    }
1849
1850    write_graph_structure(&structure, graph_write, config)
1851}
hydro_lang/viz/render.rs

hydro_lang/viz/
render.rs
