hydro_lang/compile/ir/

mod.rs

use core::panic;
use std::cell::RefCell;
use std::collections::HashMap;
#[cfg(feature = "build")]
use std::collections::HashSet;
use std::fmt::{Debug, Display};
use std::hash::{Hash, Hasher};
use std::ops::Deref;
use std::rc::Rc;

#[cfg(feature = "build")]
use dfir_lang::graph::FlatGraphBuilder;
#[cfg(feature = "build")]
use proc_macro2::Span;
use proc_macro2::TokenStream;
use quote::ToTokens;
#[cfg(feature = "build")]
use quote::quote;
#[cfg(feature = "build")]
use slotmap::{SecondaryMap, SparseSecondaryMap};
#[cfg(feature = "build")]
use syn::parse_quote;
use syn::visit::{self, Visit};
use syn::visit_mut::VisitMut;

use crate::compile::builder::{CycleId, ExternalPortId};
#[cfg(feature = "build")]
use crate::compile::deploy_provider::{Deploy, Node, RegisterPort};
use crate::location::dynamic::LocationId;
use crate::location::{LocationKey, NetworkHint};

pub mod backtrace;
use backtrace::Backtrace;

/// Wrapper that displays only the tokens of a parsed expr.
///
/// Boxes `syn::Type` which is ~240 bytes.
#[derive(Clone, Hash)]
pub struct DebugExpr(pub Box<syn::Expr>);

impl From<syn::Expr> for DebugExpr {
    fn from(expr: syn::Expr) -> Self {
        Self(Box::new(expr))
    }
}

impl Deref for DebugExpr {
    type Target = syn::Expr;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl ToTokens for DebugExpr {
    fn to_tokens(&self, tokens: &mut TokenStream) {
        self.0.to_tokens(tokens);
    }
}

impl Debug for DebugExpr {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.0.to_token_stream())
    }
}

impl Display for DebugExpr {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let original = self.0.as_ref().clone();
        let simplified = simplify_q_macro(original);

        // For now, just use quote formatting without trying to parse as a statement
        // This avoids the syn::parse_quote! issues entirely
        write!(f, "q!({})", quote::quote!(#simplified))
    }
}

/// Simplify expanded q! macro calls back to q!(...) syntax for better readability
fn simplify_q_macro(mut expr: syn::Expr) -> syn::Expr {
    // Try to parse the token string as a syn::Expr
    // Use a visitor to simplify q! macro expansions
    let mut simplifier = QMacroSimplifier::new();
    simplifier.visit_expr_mut(&mut expr);

    // If we found and simplified a q! macro, return the simplified version
    if let Some(simplified) = simplifier.simplified_result {
        simplified
    } else {
        expr
    }
}

/// AST visitor that simplifies q! macro expansions
#[derive(Default)]
pub struct QMacroSimplifier {
    pub simplified_result: Option<syn::Expr>,
}

impl QMacroSimplifier {
    pub fn new() -> Self {
        Self::default()
    }
}

impl VisitMut for QMacroSimplifier {
    fn visit_expr_mut(&mut self, expr: &mut syn::Expr) {
        // Check if we already found a result to avoid further processing
        if self.simplified_result.is_some() {
            return;
        }

        if let syn::Expr::Call(call) = expr && let syn::Expr::Path(path_expr) = call.func.as_ref()
            // Look for calls to stageleft::runtime_support::fn*
            && self.is_stageleft_runtime_support_call(&path_expr.path)
            // Try to extract the closure from the arguments
            && let Some(closure) = self.extract_closure_from_args(&call.args)
        {
            self.simplified_result = Some(closure);
            return;
        }

        // Continue visiting child expressions using the default implementation
        // Use the default visitor to avoid infinite recursion
        syn::visit_mut::visit_expr_mut(self, expr);
    }
}

impl QMacroSimplifier {
    fn is_stageleft_runtime_support_call(&self, path: &syn::Path) -> bool {
        // Check if this is a call to stageleft::runtime_support::fn*
        if let Some(last_segment) = path.segments.last() {
            let fn_name = last_segment.ident.to_string();
            // if fn_name.starts_with("fn") && fn_name.contains("_expr") {
            fn_name.contains("_type_hint")
                && path.segments.len() > 2
                && path.segments[0].ident == "stageleft"
                && path.segments[1].ident == "runtime_support"
        } else {
            false
        }
    }

    fn extract_closure_from_args(
        &self,
        args: &syn::punctuated::Punctuated<syn::Expr, syn::Token![,]>,
    ) -> Option<syn::Expr> {
        // Look through the arguments for a closure expression
        for arg in args {
            if let syn::Expr::Closure(_) = arg {
                return Some(arg.clone());
            }
            // Also check for closures nested in other expressions (like blocks)
            if let Some(closure_expr) = self.find_closure_in_expr(arg) {
                return Some(closure_expr);
            }
        }
        None
    }

    fn find_closure_in_expr(&self, expr: &syn::Expr) -> Option<syn::Expr> {
        let mut visitor = ClosureFinder {
            found_closure: None,
            prefer_inner_blocks: true,
        };
        visitor.visit_expr(expr);
        visitor.found_closure
    }
}

/// Visitor that finds closures in expressions with special block handling
struct ClosureFinder {
    found_closure: Option<syn::Expr>,
    prefer_inner_blocks: bool,
}

impl<'ast> Visit<'ast> for ClosureFinder {
    fn visit_expr(&mut self, expr: &'ast syn::Expr) {
        // If we already found a closure, don't continue searching
        if self.found_closure.is_some() {
            return;
        }

        match expr {
            syn::Expr::Closure(_) => {
                self.found_closure = Some(expr.clone());
            }
            syn::Expr::Block(block) if self.prefer_inner_blocks => {
                // Special handling for blocks - look for inner blocks that contain closures
                for stmt in &block.block.stmts {
                    if let syn::Stmt::Expr(stmt_expr, _) = stmt
                        && let syn::Expr::Block(_) = stmt_expr
                    {
                        // Check if this nested block contains a closure
                        let mut inner_visitor = ClosureFinder {
                            found_closure: None,
                            prefer_inner_blocks: false, // Avoid infinite recursion
                        };
                        inner_visitor.visit_expr(stmt_expr);
                        if inner_visitor.found_closure.is_some() {
                            // Found a closure in an inner block, return that block
                            self.found_closure = Some(stmt_expr.clone());
                            return;
                        }
                    }
                }

                // If no inner block with closure found, continue with normal visitation
                visit::visit_expr(self, expr);

                // If we found a closure, just return the closure itself, not the whole block
                // unless we're in the special case where we want the containing block
                if self.found_closure.is_some() {
                    // The closure was found during visitation, no need to wrap in block
                }
            }
            _ => {
                // Use default visitor behavior for all other expressions
                visit::visit_expr(self, expr);
            }
        }
    }
}

/// Debug displays the type's tokens.
///
/// Boxes `syn::Type` which is ~320 bytes.
#[derive(Clone, PartialEq, Eq, Hash)]
pub struct DebugType(pub Box<syn::Type>);

impl From<syn::Type> for DebugType {
    fn from(t: syn::Type) -> Self {
        Self(Box::new(t))
    }
}

impl Deref for DebugType {
    type Target = syn::Type;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl ToTokens for DebugType {
    fn to_tokens(&self, tokens: &mut TokenStream) {
        self.0.to_tokens(tokens);
    }
}

impl Debug for DebugType {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.0.to_token_stream())
    }
}

pub enum DebugInstantiate {
    Building,
    Finalized(Box<DebugInstantiateFinalized>),
}

#[cfg_attr(
    not(feature = "build"),
    expect(
        dead_code,
        reason = "sink, source unused without `feature = \"build\"`."
    )
)]
pub struct DebugInstantiateFinalized {
    sink: syn::Expr,
    source: syn::Expr,
    connect_fn: Option<Box<dyn FnOnce()>>,
}

impl From<DebugInstantiateFinalized> for DebugInstantiate {
    fn from(f: DebugInstantiateFinalized) -> Self {
        Self::Finalized(Box::new(f))
    }
}

impl Debug for DebugInstantiate {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "<network instantiate>")
    }
}

impl Hash for DebugInstantiate {
    fn hash<H: Hasher>(&self, _state: &mut H) {
        // Do nothing
    }
}

impl Clone for DebugInstantiate {
    fn clone(&self) -> Self {
        match self {
            DebugInstantiate::Building => DebugInstantiate::Building,
            DebugInstantiate::Finalized(_) => {
                panic!("DebugInstantiate::Finalized should not be cloned")
            }
        }
    }
}

/// Tracks the instantiation state of a `ClusterMembers` source.
///
/// During `compile_network`, the first `ClusterMembers` node for a given
/// `(at_location, target_cluster)` pair is promoted to [`Self::Stream`] and
/// receives the expression returned by `Deploy::cluster_membership_stream`.
/// All subsequent nodes for the same pair are set to [`Self::Tee`] so that
/// during code-gen they simply reference the tee output of the first node
/// instead of creating a redundant `source_stream`.
#[derive(Debug, Hash, Clone)]
pub enum ClusterMembersState {
    /// Not yet instantiated.
    Uninit,
    /// The primary instance: holds the stream expression and will emit
    /// `source_stream(expr) -> tee()` during code-gen.
    Stream(DebugExpr),
    /// A secondary instance that references the tee output of the primary.
    /// Stores `(at_location_root, target_cluster_location)` so that `emit_core`
    /// can derive the deterministic tee ident without extra state.
    Tee(LocationId, LocationId),
}

/// A source in a Hydro graph, where data enters the graph.
#[derive(Debug, Hash, Clone)]
pub enum HydroSource {
    Stream(DebugExpr),
    ExternalNetwork(),
    Iter(DebugExpr),
    Spin(),
    ClusterMembers(LocationId, ClusterMembersState),
    Embedded(syn::Ident),
    EmbeddedSingleton(syn::Ident),
}

#[cfg(feature = "build")]
/// A trait that abstracts over elements of DFIR code-gen that differ between production deployment
/// and simulations.
///
/// In particular, this lets the simulator fuse together all locations into one DFIR graph, spit
/// out separate graphs for each tick, and emit hooks for controlling non-deterministic operators.
pub trait DfirBuilder {
    /// Whether the representation of singletons should include intermediate states.
    fn singleton_intermediates(&self) -> bool;

    /// Gets the DFIR builder for the given location, creating it if necessary.
    fn get_dfir_mut(&mut self, location: &LocationId) -> &mut FlatGraphBuilder;

    fn batch(
        &mut self,
        in_ident: syn::Ident,
        in_location: &LocationId,
        in_kind: &CollectionKind,
        out_ident: &syn::Ident,
        out_location: &LocationId,
        op_meta: &HydroIrOpMetadata,
    );
    fn yield_from_tick(
        &mut self,
        in_ident: syn::Ident,
        in_location: &LocationId,
        in_kind: &CollectionKind,
        out_ident: &syn::Ident,
        out_location: &LocationId,
    );

    fn begin_atomic(
        &mut self,
        in_ident: syn::Ident,
        in_location: &LocationId,
        in_kind: &CollectionKind,
        out_ident: &syn::Ident,
        out_location: &LocationId,
        op_meta: &HydroIrOpMetadata,
    );
    fn end_atomic(
        &mut self,
        in_ident: syn::Ident,
        in_location: &LocationId,
        in_kind: &CollectionKind,
        out_ident: &syn::Ident,
    );

    #[expect(clippy::too_many_arguments, reason = "TODO // internal")]
    fn observe_nondet(
        &mut self,
        trusted: bool,
        location: &LocationId,
        in_ident: syn::Ident,
        in_kind: &CollectionKind,
        out_ident: &syn::Ident,
        out_kind: &CollectionKind,
        op_meta: &HydroIrOpMetadata,
    );

    #[expect(clippy::too_many_arguments, reason = "TODO")]
    fn create_network(
        &mut self,
        from: &LocationId,
        to: &LocationId,
        input_ident: syn::Ident,
        out_ident: &syn::Ident,
        serialize: Option<&DebugExpr>,
        sink: syn::Expr,
        source: syn::Expr,
        deserialize: Option<&DebugExpr>,
        tag_id: usize,
        networking_info: &crate::networking::NetworkingInfo,
    );

    fn create_external_source(
        &mut self,
        on: &LocationId,
        source_expr: syn::Expr,
        out_ident: &syn::Ident,
        deserialize: Option<&DebugExpr>,
        tag_id: usize,
    );

    fn create_external_output(
        &mut self,
        on: &LocationId,
        sink_expr: syn::Expr,
        input_ident: &syn::Ident,
        serialize: Option<&DebugExpr>,
        tag_id: usize,
    );
}

#[cfg(feature = "build")]
impl DfirBuilder for SecondaryMap<LocationKey, FlatGraphBuilder> {
    fn singleton_intermediates(&self) -> bool {
        false
    }

    fn get_dfir_mut(&mut self, location: &LocationId) -> &mut FlatGraphBuilder {
        self.entry(location.root().key())
            .expect("location was removed")
            .or_default()
    }

    fn batch(
        &mut self,
        in_ident: syn::Ident,
        in_location: &LocationId,
        in_kind: &CollectionKind,
        out_ident: &syn::Ident,
        _out_location: &LocationId,
        _op_meta: &HydroIrOpMetadata,
    ) {
        let builder = self.get_dfir_mut(in_location.root());
        if in_kind.is_bounded()
            && matches!(
                in_kind,
                CollectionKind::Singleton { .. }
                    | CollectionKind::Optional { .. }
                    | CollectionKind::KeyedSingleton { .. }
            )
        {
            assert!(in_location.is_top_level());
            builder.add_dfir(
                parse_quote! {
                    #out_ident = #in_ident -> persist::<'static>();
                },
                None,
                None,
            );
        } else {
            builder.add_dfir(
                parse_quote! {
                    #out_ident = #in_ident;
                },
                None,
                None,
            );
        }
    }

    fn yield_from_tick(
        &mut self,
        in_ident: syn::Ident,
        in_location: &LocationId,
        _in_kind: &CollectionKind,
        out_ident: &syn::Ident,
        _out_location: &LocationId,
    ) {
        let builder = self.get_dfir_mut(in_location.root());
        builder.add_dfir(
            parse_quote! {
                #out_ident = #in_ident;
            },
            None,
            None,
        );
    }

    fn begin_atomic(
        &mut self,
        in_ident: syn::Ident,
        in_location: &LocationId,
        _in_kind: &CollectionKind,
        out_ident: &syn::Ident,
        _out_location: &LocationId,
        _op_meta: &HydroIrOpMetadata,
    ) {
        let builder = self.get_dfir_mut(in_location.root());
        builder.add_dfir(
            parse_quote! {
                #out_ident = #in_ident;
            },
            None,
            None,
        );
    }

    fn end_atomic(
        &mut self,
        in_ident: syn::Ident,
        in_location: &LocationId,
        _in_kind: &CollectionKind,
        out_ident: &syn::Ident,
    ) {
        let builder = self.get_dfir_mut(in_location.root());
        builder.add_dfir(
            parse_quote! {
                #out_ident = #in_ident;
            },
            None,
            None,
        );
    }

    fn observe_nondet(
        &mut self,
        _trusted: bool,
        location: &LocationId,
        in_ident: syn::Ident,
        _in_kind: &CollectionKind,
        out_ident: &syn::Ident,
        _out_kind: &CollectionKind,
        _op_meta: &HydroIrOpMetadata,
    ) {
        let builder = self.get_dfir_mut(location);
        builder.add_dfir(
            parse_quote! {
                #out_ident = #in_ident;
            },
            None,
            None,
        );
    }

    fn create_network(
        &mut self,
        from: &LocationId,
        to: &LocationId,
        input_ident: syn::Ident,
        out_ident: &syn::Ident,
        serialize: Option<&DebugExpr>,
        sink: syn::Expr,
        source: syn::Expr,
        deserialize: Option<&DebugExpr>,
        tag_id: usize,
        _networking_info: &crate::networking::NetworkingInfo,
    ) {
        let sender_builder = self.get_dfir_mut(from);
        if let Some(serialize_pipeline) = serialize {
            sender_builder.add_dfir(
                parse_quote! {
                    #input_ident -> map(#serialize_pipeline) -> dest_sink(#sink);
                },
                None,
                // operator tag separates send and receive, which otherwise have the same next_stmt_id
                Some(&format!("send{}", tag_id)),
            );
        } else {
            sender_builder.add_dfir(
                parse_quote! {
                    #input_ident -> dest_sink(#sink);
                },
                None,
                Some(&format!("send{}", tag_id)),
            );
        }

        let receiver_builder = self.get_dfir_mut(to);
        if let Some(deserialize_pipeline) = deserialize {
            receiver_builder.add_dfir(
                parse_quote! {
                    #out_ident = source_stream(#source) -> map(#deserialize_pipeline);
                },
                None,
                Some(&format!("recv{}", tag_id)),
            );
        } else {
            receiver_builder.add_dfir(
                parse_quote! {
                    #out_ident = source_stream(#source);
                },
                None,
                Some(&format!("recv{}", tag_id)),
            );
        }
    }

    fn create_external_source(
        &mut self,
        on: &LocationId,
        source_expr: syn::Expr,
        out_ident: &syn::Ident,
        deserialize: Option<&DebugExpr>,
        tag_id: usize,
    ) {
        let receiver_builder = self.get_dfir_mut(on);
        if let Some(deserialize_pipeline) = deserialize {
            receiver_builder.add_dfir(
                parse_quote! {
                    #out_ident = source_stream(#source_expr) -> map(#deserialize_pipeline);
                },
                None,
                Some(&format!("recv{}", tag_id)),
            );
        } else {
            receiver_builder.add_dfir(
                parse_quote! {
                    #out_ident = source_stream(#source_expr);
                },
                None,
                Some(&format!("recv{}", tag_id)),
            );
        }
    }

    fn create_external_output(
        &mut self,
        on: &LocationId,
        sink_expr: syn::Expr,
        input_ident: &syn::Ident,
        serialize: Option<&DebugExpr>,
        tag_id: usize,
    ) {
        let sender_builder = self.get_dfir_mut(on);
        if let Some(serialize_fn) = serialize {
            sender_builder.add_dfir(
                parse_quote! {
                    #input_ident -> map(#serialize_fn) -> dest_sink(#sink_expr);
                },
                None,
                // operator tag separates send and receive, which otherwise have the same next_stmt_id
                Some(&format!("send{}", tag_id)),
            );
        } else {
            sender_builder.add_dfir(
                parse_quote! {
                    #input_ident -> dest_sink(#sink_expr);
                },
                None,
                Some(&format!("send{}", tag_id)),
            );
        }
    }
}

#[cfg(feature = "build")]
pub enum BuildersOrCallback<'a, L, N>
where
    L: FnMut(&mut HydroRoot, &mut usize),
    N: FnMut(&mut HydroNode, &mut usize),
{
    Builders(&'a mut dyn DfirBuilder),
    Callback(L, N),
}

/// An root in a Hydro graph, which is an pipeline that doesn't emit
/// any downstream values. Traversals over the dataflow graph and
/// generating DFIR IR start from roots.
#[derive(Debug, Hash)]
pub enum HydroRoot {
    ForEach {
        f: DebugExpr,
        input: Box<HydroNode>,
        op_metadata: HydroIrOpMetadata,
    },
    SendExternal {
        to_external_key: LocationKey,
        to_port_id: ExternalPortId,
        to_many: bool,
        unpaired: bool,
        serialize_fn: Option<DebugExpr>,
        instantiate_fn: DebugInstantiate,
        input: Box<HydroNode>,
        op_metadata: HydroIrOpMetadata,
    },
    DestSink {
        sink: DebugExpr,
        input: Box<HydroNode>,
        op_metadata: HydroIrOpMetadata,
    },
    CycleSink {
        cycle_id: CycleId,
        input: Box<HydroNode>,
        op_metadata: HydroIrOpMetadata,
    },
    EmbeddedOutput {
        ident: syn::Ident,
        input: Box<HydroNode>,
        op_metadata: HydroIrOpMetadata,
    },
    Null {
        input: Box<HydroNode>,
        op_metadata: HydroIrOpMetadata,
    },
}

impl HydroRoot {
    #[cfg(feature = "build")]
    #[expect(clippy::too_many_arguments, reason = "TODO(internal)")]
    pub fn compile_network<'a, D>(
        &mut self,
        extra_stmts: &mut SparseSecondaryMap<LocationKey, Vec<syn::Stmt>>,
        seen_tees: &mut SeenSharedNodes,
        seen_cluster_members: &mut HashSet<(LocationId, LocationId)>,
        processes: &SparseSecondaryMap<LocationKey, D::Process>,
        clusters: &SparseSecondaryMap<LocationKey, D::Cluster>,
        externals: &SparseSecondaryMap<LocationKey, D::External>,
        env: &mut D::InstantiateEnv,
    ) where
        D: Deploy<'a>,
    {
        let refcell_extra_stmts = RefCell::new(extra_stmts);
        let refcell_env = RefCell::new(env);
        let refcell_seen_cluster_members = RefCell::new(seen_cluster_members);
        self.transform_bottom_up(
            &mut |l| {
                if let HydroRoot::SendExternal {
                    input,
                    to_external_key,
                    to_port_id,
                    to_many,
                    unpaired,
                    instantiate_fn,
                    ..
                } = l
                {
                    let ((sink_expr, source_expr), connect_fn) = match instantiate_fn {
                        DebugInstantiate::Building => {
                            let to_node = externals
                                .get(*to_external_key)
                                .unwrap_or_else(|| {
                                    panic!("A external used in the graph was not instantiated: {}", to_external_key)
                                })
                                .clone();

                            match input.metadata().location_id.root() {
                                &LocationId::Process(process_key) => {
                                    if *to_many {
                                        (
                                            (
                                                D::e2o_many_sink(format!("{}_{}", *to_external_key, *to_port_id)),
                                                parse_quote!(DUMMY),
                                            ),
                                            Box::new(|| {}) as Box<dyn FnOnce()>,
                                        )
                                    } else {
                                        let from_node = processes
                                            .get(process_key)
                                            .unwrap_or_else(|| {
                                                panic!("A process used in the graph was not instantiated: {}", process_key)
                                            })
                                            .clone();

                                        let sink_port = from_node.next_port();
                                        let source_port = to_node.next_port();

                                        if *unpaired {
                                            use stageleft::quote_type;
                                            use tokio_util::codec::LengthDelimitedCodec;

                                            to_node.register(*to_port_id, source_port.clone());

                                            let _ = D::e2o_source(
                                                refcell_extra_stmts.borrow_mut().entry(process_key).expect("location was removed").or_default(),
                                                &to_node, &source_port,
                                                &from_node, &sink_port,
                                                &quote_type::<LengthDelimitedCodec>(),
                                                format!("{}_{}", *to_external_key, *to_port_id)
                                            );
                                        }

                                        (
                                            (
                                                D::o2e_sink(
                                                    &from_node,
                                                    &sink_port,
                                                    &to_node,
                                                    &source_port,
                                                    format!("{}_{}", *to_external_key, *to_port_id)
                                                ),
                                                parse_quote!(DUMMY),
                                            ),
                                            if *unpaired {
                                                D::e2o_connect(
                                                    &to_node,
                                                    &source_port,
                                                    &from_node,
                                                    &sink_port,
                                                    *to_many,
                                                    NetworkHint::Auto,
                                                )
                                            } else {
                                                Box::new(|| {}) as Box<dyn FnOnce()>
                                            },
                                        )
                                    }
                                }
                                LocationId::Cluster(_) => todo!("SendExternal from a cluster location is not yet supported"),
                                _ => panic!()
                            }
                        },

                        DebugInstantiate::Finalized(_) => panic!("network already finalized"),
                    };

                    *instantiate_fn = DebugInstantiateFinalized {
                        sink: sink_expr,
                        source: source_expr,
                        connect_fn: Some(connect_fn),
                    }
                    .into();
                } else if let HydroRoot::EmbeddedOutput { ident, input, .. } = l {
                    let element_type = match &input.metadata().collection_kind {
                        CollectionKind::Stream { element_type, .. } => element_type.0.as_ref().clone(),
                        _ => panic!("Embedded output must have Stream collection kind"),
                    };
                    let location_key = match input.metadata().location_id.root() {
                        LocationId::Process(key) | LocationId::Cluster(key) => *key,
                        _ => panic!("Embedded output must be on a process or cluster"),
                    };
                    D::register_embedded_output(
                        &mut refcell_env.borrow_mut(),
                        location_key,
                        ident,
                        &element_type,
                    );
                }
            },
            &mut |n| {
                if let HydroNode::Network {
                    name,
                    networking_info,
                    input,
                    instantiate_fn,
                    metadata,
                    ..
                } = n
                {
                    let (sink_expr, source_expr, connect_fn) = match instantiate_fn {
                        DebugInstantiate::Building => instantiate_network::<D>(
                            &mut refcell_env.borrow_mut(),
                            input.metadata().location_id.root(),
                            metadata.location_id.root(),
                            processes,
                            clusters,
                            name.as_deref(),
                            networking_info,
                        ),

                        DebugInstantiate::Finalized(_) => panic!("network already finalized"),
                    };

                    *instantiate_fn = DebugInstantiateFinalized {
                        sink: sink_expr,
                        source: source_expr,
                        connect_fn: Some(connect_fn),
                    }
                    .into();
                } else if let HydroNode::ExternalInput {
                    from_external_key,
                    from_port_id,
                    from_many,
                    codec_type,
                    port_hint,
                    instantiate_fn,
                    metadata,
                    ..
                } = n
                {
                    let ((sink_expr, source_expr), connect_fn) = match instantiate_fn {
                        DebugInstantiate::Building => {
                            let from_node = externals
                                .get(*from_external_key)
                                .unwrap_or_else(|| {
                                    panic!(
                                        "A external used in the graph was not instantiated: {}",
                                        from_external_key,
                                    )
                                })
                                .clone();

                            match metadata.location_id.root() {
                                &LocationId::Process(process_key) => {
                                    let to_node = processes
                                        .get(process_key)
                                        .unwrap_or_else(|| {
                                            panic!("A process used in the graph was not instantiated: {}", process_key)
                                        })
                                        .clone();

                                    let sink_port = from_node.next_port();
                                    let source_port = to_node.next_port();

                                    from_node.register(*from_port_id, sink_port.clone());

                                    (
                                        (
                                            parse_quote!(DUMMY),
                                            if *from_many {
                                                D::e2o_many_source(
                                                    refcell_extra_stmts.borrow_mut().entry(process_key).expect("location was removed").or_default(),
                                                    &to_node, &source_port,
                                                    codec_type.0.as_ref(),
                                                    format!("{}_{}", *from_external_key, *from_port_id)
                                                )
                                            } else {
                                                D::e2o_source(
                                                    refcell_extra_stmts.borrow_mut().entry(process_key).expect("location was removed").or_default(),
                                                    &from_node, &sink_port,
                                                    &to_node, &source_port,
                                                    codec_type.0.as_ref(),
                                                    format!("{}_{}", *from_external_key, *from_port_id)
                                                )
                                            },
                                        ),
                                        D::e2o_connect(&from_node, &sink_port, &to_node, &source_port, *from_many, *port_hint),
                                    )
                                }
                                LocationId::Cluster(_) => todo!("ExternalInput to a cluster location is not yet supported"),
                                _ => panic!()
                            }
                        },

                        DebugInstantiate::Finalized(_) => panic!("network already finalized"),
                    };

                    *instantiate_fn = DebugInstantiateFinalized {
                        sink: sink_expr,
                        source: source_expr,
                        connect_fn: Some(connect_fn),
                    }
                    .into();
                } else if let HydroNode::Source { source: HydroSource::Embedded(ident), metadata } = n {
                    let element_type = match &metadata.collection_kind {
                        CollectionKind::Stream { element_type, .. } => element_type.0.as_ref().clone(),
                        _ => panic!("Embedded source must have Stream collection kind"),
                    };
                    let location_key = match metadata.location_id.root() {
                        LocationId::Process(key) | LocationId::Cluster(key) => *key,
                        _ => panic!("Embedded source must be on a process or cluster"),
                    };
                    D::register_embedded_stream_input(
                        &mut refcell_env.borrow_mut(),
                        location_key,
                        ident,
                        &element_type,
                    );
                } else if let HydroNode::Source { source: HydroSource::EmbeddedSingleton(ident), metadata } = n {
                    let element_type = match &metadata.collection_kind {
                        CollectionKind::Singleton { element_type, .. } => element_type.0.as_ref().clone(),
                        _ => panic!("EmbeddedSingleton source must have Singleton collection kind"),
                    };
                    let location_key = match metadata.location_id.root() {
                        LocationId::Process(key) | LocationId::Cluster(key) => *key,
                        _ => panic!("EmbeddedSingleton source must be on a process or cluster"),
                    };
                    D::register_embedded_singleton_input(
                        &mut refcell_env.borrow_mut(),
                        location_key,
                        ident,
                        &element_type,
                    );
                } else if let HydroNode::Source { source: HydroSource::ClusterMembers(location_id, state), metadata } = n {
                    match state {
                        ClusterMembersState::Uninit => {
                            let at_location = metadata.location_id.root().clone();
                            let key = (at_location.clone(), LocationId::Cluster(location_id.key()));
                            if refcell_seen_cluster_members.borrow_mut().insert(key) {
                                // First occurrence: call cluster_membership_stream and mark as Stream.
                                let expr = stageleft::QuotedWithContext::splice_untyped_ctx(
                                    D::cluster_membership_stream(&mut refcell_env.borrow_mut(), &at_location, location_id),
                                    &(),
                                );
                                *state = ClusterMembersState::Stream(expr.into());
                            } else {
                                // Already instantiated for this (at, target) pair: just tee.
                                *state = ClusterMembersState::Tee(at_location, location_id.clone());
                            }
                        }
                        ClusterMembersState::Stream(_) | ClusterMembersState::Tee(..) => {
                            panic!("cluster members already finalized");
                        }
                    }
                }
            },
            seen_tees,
            false,
        );
    }

    pub fn connect_network(&mut self, seen_tees: &mut SeenSharedNodes) {
        self.transform_bottom_up(
            &mut |l| {
                if let HydroRoot::SendExternal { instantiate_fn, .. } = l {
                    match instantiate_fn {
                        DebugInstantiate::Building => panic!("network not built"),

                        DebugInstantiate::Finalized(finalized) => {
                            (finalized.connect_fn.take().unwrap())();
                        }
                    }
                }
            },
            &mut |n| {
                if let HydroNode::Network { instantiate_fn, .. }
                | HydroNode::ExternalInput { instantiate_fn, .. } = n
                {
                    match instantiate_fn {
                        DebugInstantiate::Building => panic!("network not built"),

                        DebugInstantiate::Finalized(finalized) => {
                            (finalized.connect_fn.take().unwrap())();
                        }
                    }
                }
            },
            seen_tees,
            false,
        );
    }

    pub fn transform_bottom_up(
        &mut self,
        transform_root: &mut impl FnMut(&mut HydroRoot),
        transform_node: &mut impl FnMut(&mut HydroNode),
        seen_tees: &mut SeenSharedNodes,
        check_well_formed: bool,
    ) {
        self.transform_children(
            |n, s| n.transform_bottom_up(transform_node, s, check_well_formed),
            seen_tees,
        );

        transform_root(self);
    }

    pub fn transform_children(
        &mut self,
        mut transform: impl FnMut(&mut HydroNode, &mut SeenSharedNodes),
        seen_tees: &mut SeenSharedNodes,
    ) {
        match self {
            HydroRoot::ForEach { input, .. }
            | HydroRoot::SendExternal { input, .. }
            | HydroRoot::DestSink { input, .. }
            | HydroRoot::CycleSink { input, .. }
            | HydroRoot::EmbeddedOutput { input, .. }
            | HydroRoot::Null { input, .. } => {
                transform(input, seen_tees);
            }
        }
    }

    pub fn deep_clone(&self, seen_tees: &mut SeenSharedNodes) -> HydroRoot {
        match self {
            HydroRoot::ForEach {
                f,
                input,
                op_metadata,
            } => HydroRoot::ForEach {
                f: f.clone(),
                input: Box::new(input.deep_clone(seen_tees)),
                op_metadata: op_metadata.clone(),
            },
            HydroRoot::SendExternal {
                to_external_key,
                to_port_id,
                to_many,
                unpaired,
                serialize_fn,
                instantiate_fn,
                input,
                op_metadata,
            } => HydroRoot::SendExternal {
                to_external_key: *to_external_key,
                to_port_id: *to_port_id,
                to_many: *to_many,
                unpaired: *unpaired,
                serialize_fn: serialize_fn.clone(),
                instantiate_fn: instantiate_fn.clone(),
                input: Box::new(input.deep_clone(seen_tees)),
                op_metadata: op_metadata.clone(),
            },
            HydroRoot::DestSink {
                sink,
                input,
                op_metadata,
            } => HydroRoot::DestSink {
                sink: sink.clone(),
                input: Box::new(input.deep_clone(seen_tees)),
                op_metadata: op_metadata.clone(),
            },
            HydroRoot::CycleSink {
                cycle_id,
                input,
                op_metadata,
            } => HydroRoot::CycleSink {
                cycle_id: *cycle_id,
                input: Box::new(input.deep_clone(seen_tees)),
                op_metadata: op_metadata.clone(),
            },
            HydroRoot::EmbeddedOutput {
                ident,
                input,
                op_metadata,
            } => HydroRoot::EmbeddedOutput {
                ident: ident.clone(),
                input: Box::new(input.deep_clone(seen_tees)),
                op_metadata: op_metadata.clone(),
            },
            HydroRoot::Null { input, op_metadata } => HydroRoot::Null {
                input: Box::new(input.deep_clone(seen_tees)),
                op_metadata: op_metadata.clone(),
            },
        }
    }

    #[cfg(feature = "build")]
    pub fn emit(
        &mut self,
        graph_builders: &mut dyn DfirBuilder,
        seen_tees: &mut SeenSharedNodes,
        built_tees: &mut HashMap<*const RefCell<HydroNode>, Vec<syn::Ident>>,
        next_stmt_id: &mut usize,
    ) {
        self.emit_core(
            &mut BuildersOrCallback::<
                fn(&mut HydroRoot, &mut usize),
                fn(&mut HydroNode, &mut usize),
            >::Builders(graph_builders),
            seen_tees,
            built_tees,
            next_stmt_id,
        );
    }

    #[cfg(feature = "build")]
    pub fn emit_core(
        &mut self,
        builders_or_callback: &mut BuildersOrCallback<
            impl FnMut(&mut HydroRoot, &mut usize),
            impl FnMut(&mut HydroNode, &mut usize),
        >,
        seen_tees: &mut SeenSharedNodes,
        built_tees: &mut HashMap<*const RefCell<HydroNode>, Vec<syn::Ident>>,
        next_stmt_id: &mut usize,
    ) {
        match self {
            HydroRoot::ForEach { f, input, .. } => {
                let input_ident =
                    input.emit_core(builders_or_callback, seen_tees, built_tees, next_stmt_id);

                match builders_or_callback {
                    BuildersOrCallback::Builders(graph_builders) => {
                        graph_builders
                            .get_dfir_mut(&input.metadata().location_id)
                            .add_dfir(
                                parse_quote! {
                                    #input_ident -> for_each(#f);
                                },
                                None,
                                Some(&next_stmt_id.to_string()),
                            );
                    }
                    BuildersOrCallback::Callback(leaf_callback, _) => {
                        leaf_callback(self, next_stmt_id);
                    }
                }

                *next_stmt_id += 1;
            }

            HydroRoot::SendExternal {
                serialize_fn,
                instantiate_fn,
                input,
                ..
            } => {
                let input_ident =
                    input.emit_core(builders_or_callback, seen_tees, built_tees, next_stmt_id);

                match builders_or_callback {
                    BuildersOrCallback::Builders(graph_builders) => {
                        let (sink_expr, _) = match instantiate_fn {
                            DebugInstantiate::Building => (
                                syn::parse_quote!(DUMMY_SINK),
                                syn::parse_quote!(DUMMY_SOURCE),
                            ),

                            DebugInstantiate::Finalized(finalized) => {
                                (finalized.sink.clone(), finalized.source.clone())
                            }
                        };

                        graph_builders.create_external_output(
                            &input.metadata().location_id,
                            sink_expr,
                            &input_ident,
                            serialize_fn.as_ref(),
                            *next_stmt_id,
                        );
                    }
                    BuildersOrCallback::Callback(leaf_callback, _) => {
                        leaf_callback(self, next_stmt_id);
                    }
                }

                *next_stmt_id += 1;
            }

            HydroRoot::DestSink { sink, input, .. } => {
                let input_ident =
                    input.emit_core(builders_or_callback, seen_tees, built_tees, next_stmt_id);

                match builders_or_callback {
                    BuildersOrCallback::Builders(graph_builders) => {
                        graph_builders
                            .get_dfir_mut(&input.metadata().location_id)
                            .add_dfir(
                                parse_quote! {
                                    #input_ident -> dest_sink(#sink);
                                },
                                None,
                                Some(&next_stmt_id.to_string()),
                            );
                    }
                    BuildersOrCallback::Callback(leaf_callback, _) => {
                        leaf_callback(self, next_stmt_id);
                    }
                }

                *next_stmt_id += 1;
            }

            HydroRoot::CycleSink {
                cycle_id, input, ..
            } => {
                let input_ident =
                    input.emit_core(builders_or_callback, seen_tees, built_tees, next_stmt_id);

                match builders_or_callback {
                    BuildersOrCallback::Builders(graph_builders) => {
                        let elem_type: syn::Type = match &input.metadata().collection_kind {
                            CollectionKind::KeyedSingleton {
                                key_type,
                                value_type,
                                ..
                            }
                            | CollectionKind::KeyedStream {
                                key_type,
                                value_type,
                                ..
                            } => {
                                parse_quote!((#key_type, #value_type))
                            }
                            CollectionKind::Stream { element_type, .. }
                            | CollectionKind::Singleton { element_type, .. }
                            | CollectionKind::Optional { element_type, .. } => {
                                parse_quote!(#element_type)
                            }
                        };

                        let cycle_id_ident = cycle_id.as_ident();
                        graph_builders
                            .get_dfir_mut(&input.metadata().location_id)
                            .add_dfir(
                                parse_quote! {
                                    #cycle_id_ident = #input_ident -> identity::<#elem_type>();
                                },
                                None,
                                None,
                            );
                    }
                    // No ID, no callback
                    BuildersOrCallback::Callback(_, _) => {}
                }
            }

            HydroRoot::EmbeddedOutput { ident, input, .. } => {
                let input_ident =
                    input.emit_core(builders_or_callback, seen_tees, built_tees, next_stmt_id);

                match builders_or_callback {
                    BuildersOrCallback::Builders(graph_builders) => {
                        graph_builders
                            .get_dfir_mut(&input.metadata().location_id)
                            .add_dfir(
                                parse_quote! {
                                    #input_ident -> for_each(&mut #ident);
                                },
                                None,
                                Some(&next_stmt_id.to_string()),
                            );
                    }
                    BuildersOrCallback::Callback(leaf_callback, _) => {
                        leaf_callback(self, next_stmt_id);
                    }
                }

                *next_stmt_id += 1;
            }

            HydroRoot::Null { input, .. } => {
                let input_ident =
                    input.emit_core(builders_or_callback, seen_tees, built_tees, next_stmt_id);

                match builders_or_callback {
                    BuildersOrCallback::Builders(graph_builders) => {
                        graph_builders
                            .get_dfir_mut(&input.metadata().location_id)
                            .add_dfir(
                                parse_quote! {
                                    #input_ident -> for_each(|_| {});
                                },
                                None,
                                Some(&next_stmt_id.to_string()),
                            );
                    }
                    BuildersOrCallback::Callback(leaf_callback, _) => {
                        leaf_callback(self, next_stmt_id);
                    }
                }

                *next_stmt_id += 1;
            }
        }
    }

    pub fn op_metadata(&self) -> &HydroIrOpMetadata {
        match self {
            HydroRoot::ForEach { op_metadata, .. }
            | HydroRoot::SendExternal { op_metadata, .. }
            | HydroRoot::DestSink { op_metadata, .. }
            | HydroRoot::CycleSink { op_metadata, .. }
            | HydroRoot::EmbeddedOutput { op_metadata, .. }
            | HydroRoot::Null { op_metadata, .. } => op_metadata,
        }
    }

    pub fn op_metadata_mut(&mut self) -> &mut HydroIrOpMetadata {
        match self {
            HydroRoot::ForEach { op_metadata, .. }
            | HydroRoot::SendExternal { op_metadata, .. }
            | HydroRoot::DestSink { op_metadata, .. }
            | HydroRoot::CycleSink { op_metadata, .. }
            | HydroRoot::EmbeddedOutput { op_metadata, .. }
            | HydroRoot::Null { op_metadata, .. } => op_metadata,
        }
    }

    pub fn input(&self) -> &HydroNode {
        match self {
            HydroRoot::ForEach { input, .. }
            | HydroRoot::SendExternal { input, .. }
            | HydroRoot::DestSink { input, .. }
            | HydroRoot::CycleSink { input, .. }
            | HydroRoot::EmbeddedOutput { input, .. }
            | HydroRoot::Null { input, .. } => input,
        }
    }

    pub fn input_metadata(&self) -> &HydroIrMetadata {
        self.input().metadata()
    }

    pub fn print_root(&self) -> String {
        match self {
            HydroRoot::ForEach { f, .. } => format!("ForEach({:?})", f),
            HydroRoot::SendExternal { .. } => "SendExternal".to_owned(),
            HydroRoot::DestSink { sink, .. } => format!("DestSink({:?})", sink),
            HydroRoot::CycleSink { cycle_id, .. } => format!("CycleSink({})", cycle_id),
            HydroRoot::EmbeddedOutput { ident, .. } => {
                format!("EmbeddedOutput({})", ident)
            }
            HydroRoot::Null { .. } => "Null".to_owned(),
        }
    }

    pub fn visit_debug_expr(&mut self, mut transform: impl FnMut(&mut DebugExpr)) {
        match self {
            HydroRoot::ForEach { f, .. } | HydroRoot::DestSink { sink: f, .. } => {
                transform(f);
            }
            HydroRoot::SendExternal { .. }
            | HydroRoot::CycleSink { .. }
            | HydroRoot::EmbeddedOutput { .. }
            | HydroRoot::Null { .. } => {}
        }
    }
}

#[cfg(feature = "build")]
pub fn emit(ir: &mut Vec<HydroRoot>) -> SecondaryMap<LocationKey, FlatGraphBuilder> {
    let mut builders = SecondaryMap::new();
    let mut seen_tees = HashMap::new();
    let mut built_tees = HashMap::new();
    let mut next_stmt_id = 0;
    for leaf in ir {
        leaf.emit(
            &mut builders,
            &mut seen_tees,
            &mut built_tees,
            &mut next_stmt_id,
        );
    }
    builders
}

#[cfg(feature = "build")]
pub fn traverse_dfir(
    ir: &mut [HydroRoot],
    transform_root: impl FnMut(&mut HydroRoot, &mut usize),
    transform_node: impl FnMut(&mut HydroNode, &mut usize),
) {
    let mut seen_tees = HashMap::new();
    let mut built_tees = HashMap::new();
    let mut next_stmt_id = 0;
    let mut callback = BuildersOrCallback::Callback(transform_root, transform_node);
    ir.iter_mut().for_each(|leaf| {
        leaf.emit_core(
            &mut callback,
            &mut seen_tees,
            &mut built_tees,
            &mut next_stmt_id,
        );
    });
}

pub fn transform_bottom_up(
    ir: &mut [HydroRoot],
    transform_root: &mut impl FnMut(&mut HydroRoot),
    transform_node: &mut impl FnMut(&mut HydroNode),
    check_well_formed: bool,
) {
    let mut seen_tees = HashMap::new();
    ir.iter_mut().for_each(|leaf| {
        leaf.transform_bottom_up(
            transform_root,
            transform_node,
            &mut seen_tees,
            check_well_formed,
        );
    });
}

pub fn deep_clone(ir: &[HydroRoot]) -> Vec<HydroRoot> {
    let mut seen_tees = HashMap::new();
    ir.iter()
        .map(|leaf| leaf.deep_clone(&mut seen_tees))
        .collect()
}

type PrintedTees = RefCell<Option<(usize, HashMap<*const RefCell<HydroNode>, usize>)>>;
thread_local! {
    static PRINTED_TEES: PrintedTees = const { RefCell::new(None) };
}

pub fn dbg_dedup_tee<T>(f: impl FnOnce() -> T) -> T {
    PRINTED_TEES.with(|printed_tees| {
        let mut printed_tees_mut = printed_tees.borrow_mut();
        *printed_tees_mut = Some((0, HashMap::new()));
        drop(printed_tees_mut);

        let ret = f();

        let mut printed_tees_mut = printed_tees.borrow_mut();
        *printed_tees_mut = None;

        ret
    })
}

pub struct SharedNode(pub Rc<RefCell<HydroNode>>);

impl SharedNode {
    pub fn as_ptr(&self) -> *const RefCell<HydroNode> {
        Rc::as_ptr(&self.0)
    }
}

impl Debug for SharedNode {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        PRINTED_TEES.with(|printed_tees| {
            let mut printed_tees_mut_borrow = printed_tees.borrow_mut();
            let printed_tees_mut = printed_tees_mut_borrow.as_mut();

            if let Some(printed_tees_mut) = printed_tees_mut {
                if let Some(existing) = printed_tees_mut
                    .1
                    .get(&(self.0.as_ref() as *const RefCell<HydroNode>))
                {
                    write!(f, "<shared {}>", existing)
                } else {
                    let next_id = printed_tees_mut.0;
                    printed_tees_mut.0 += 1;
                    printed_tees_mut
                        .1
                        .insert(self.0.as_ref() as *const RefCell<HydroNode>, next_id);
                    drop(printed_tees_mut_borrow);
                    write!(f, "<shared {}>: ", next_id)?;
                    Debug::fmt(&self.0.borrow(), f)
                }
            } else {
                drop(printed_tees_mut_borrow);
                write!(f, "<shared>: ")?;
                Debug::fmt(&self.0.borrow(), f)
            }
        })
    }
}

impl Hash for SharedNode {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.0.borrow_mut().hash(state);
    }
}

#[derive(Clone, PartialEq, Eq, Debug)]
pub enum BoundKind {
    Unbounded,
    Bounded,
}

#[derive(Clone, PartialEq, Eq, Debug)]
pub enum StreamOrder {
    NoOrder,
    TotalOrder,
}

#[derive(Clone, PartialEq, Eq, Debug)]
pub enum StreamRetry {
    AtLeastOnce,
    ExactlyOnce,
}

#[derive(Clone, PartialEq, Eq, Debug)]
pub enum KeyedSingletonBoundKind {
    Unbounded,
    BoundedValue,
    Bounded,
}

#[derive(Clone, PartialEq, Eq, Debug)]
pub enum CollectionKind {
    Stream {
        bound: BoundKind,
        order: StreamOrder,
        retry: StreamRetry,
        element_type: DebugType,
    },
    Singleton {
        bound: BoundKind,
        element_type: DebugType,
    },
    Optional {
        bound: BoundKind,
        element_type: DebugType,
    },
    KeyedStream {
        bound: BoundKind,
        value_order: StreamOrder,
        value_retry: StreamRetry,
        key_type: DebugType,
        value_type: DebugType,
    },
    KeyedSingleton {
        bound: KeyedSingletonBoundKind,
        key_type: DebugType,
        value_type: DebugType,
    },
}

impl CollectionKind {
    pub fn is_bounded(&self) -> bool {
        matches!(
            self,
            CollectionKind::Stream {
                bound: BoundKind::Bounded,
                ..
            } | CollectionKind::Singleton {
                bound: BoundKind::Bounded,
                ..
            } | CollectionKind::Optional {
                bound: BoundKind::Bounded,
                ..
            } | CollectionKind::KeyedStream {
                bound: BoundKind::Bounded,
                ..
            } | CollectionKind::KeyedSingleton {
                bound: KeyedSingletonBoundKind::Bounded,
                ..
            }
        )
    }
}

#[derive(Clone)]
pub struct HydroIrMetadata {
    pub location_id: LocationId,
    pub collection_kind: CollectionKind,
    pub cardinality: Option<usize>,
    pub tag: Option<String>,
    pub op: HydroIrOpMetadata,
}

// HydroIrMetadata shouldn't be used to hash or compare
impl Hash for HydroIrMetadata {
    fn hash<H: Hasher>(&self, _: &mut H) {}
}

impl PartialEq for HydroIrMetadata {
    fn eq(&self, _: &Self) -> bool {
        true
    }
}

impl Eq for HydroIrMetadata {}

impl Debug for HydroIrMetadata {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("HydroIrMetadata")
            .field("location_id", &self.location_id)
            .field("collection_kind", &self.collection_kind)
            .finish()
    }
}

/// Metadata that is specific to the operator itself, rather than its outputs.
/// This is available on _both_ inner nodes and roots.
#[derive(Clone)]
pub struct HydroIrOpMetadata {
    pub backtrace: Backtrace,
    pub cpu_usage: Option<f64>,
    pub network_recv_cpu_usage: Option<f64>,
    pub id: Option<usize>,
}

impl HydroIrOpMetadata {
    #[expect(
        clippy::new_without_default,
        reason = "explicit calls to new ensure correct backtrace bounds"
    )]
    pub fn new() -> HydroIrOpMetadata {
        Self::new_with_skip(1)
    }

    fn new_with_skip(skip_count: usize) -> HydroIrOpMetadata {
        HydroIrOpMetadata {
            backtrace: Backtrace::get_backtrace(2 + skip_count),
            cpu_usage: None,
            network_recv_cpu_usage: None,
            id: None,
        }
    }
}

impl Debug for HydroIrOpMetadata {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("HydroIrOpMetadata").finish()
    }
}

impl Hash for HydroIrOpMetadata {
    fn hash<H: Hasher>(&self, _: &mut H) {}
}

/// An intermediate node in a Hydro graph, which consumes data
/// from upstream nodes and emits data to downstream nodes.
#[derive(Debug, Hash)]
pub enum HydroNode {
    Placeholder,

    /// Manually "casts" between two different collection kinds.
    ///
    /// Using this IR node requires special care, since it bypasses many of Hydro's core
    /// correctness checks. In particular, the user must ensure that every possible
    /// "interpretation" of the input corresponds to a distinct "interpretation" of the output,
    /// where an "interpretation" is a possible output of `ObserveNonDet` applied to the
    /// collection. This ensures that the simulator does not miss any possible outputs.
    Cast {
        inner: Box<HydroNode>,
        metadata: HydroIrMetadata,
    },

    /// Strengthens the guarantees of a stream by non-deterministically selecting a possible
    /// interpretation of the input stream.
    ///
    /// In production, this simply passes through the input, but in simulation, this operator
    /// explicitly selects a randomized interpretation.
    ObserveNonDet {
        inner: Box<HydroNode>,
        trusted: bool, // if true, we do not need to simulate non-determinism
        metadata: HydroIrMetadata,
    },

    Source {
        source: HydroSource,
        metadata: HydroIrMetadata,
    },

    SingletonSource {
        value: DebugExpr,
        first_tick_only: bool,
        metadata: HydroIrMetadata,
    },

    CycleSource {
        cycle_id: CycleId,
        metadata: HydroIrMetadata,
    },

    Tee {
        inner: SharedNode,
        metadata: HydroIrMetadata,
    },

    Partition {
        inner: SharedNode,
        f: DebugExpr,
        is_true: bool,
        metadata: HydroIrMetadata,
    },

    BeginAtomic {
        inner: Box<HydroNode>,
        metadata: HydroIrMetadata,
    },

    EndAtomic {
        inner: Box<HydroNode>,
        metadata: HydroIrMetadata,
    },

    Batch {
        inner: Box<HydroNode>,
        metadata: HydroIrMetadata,
    },

    YieldConcat {
        inner: Box<HydroNode>,
        metadata: HydroIrMetadata,
    },

    Chain {
        first: Box<HydroNode>,
        second: Box<HydroNode>,
        metadata: HydroIrMetadata,
    },

    ChainFirst {
        first: Box<HydroNode>,
        second: Box<HydroNode>,
        metadata: HydroIrMetadata,
    },

    CrossProduct {
        left: Box<HydroNode>,
        right: Box<HydroNode>,
        metadata: HydroIrMetadata,
    },

    CrossSingleton {
        left: Box<HydroNode>,
        right: Box<HydroNode>,
        metadata: HydroIrMetadata,
    },

    Join {
        left: Box<HydroNode>,
        right: Box<HydroNode>,
        metadata: HydroIrMetadata,
    },

    Difference {
        pos: Box<HydroNode>,
        neg: Box<HydroNode>,
        metadata: HydroIrMetadata,
    },

    AntiJoin {
        pos: Box<HydroNode>,
        neg: Box<HydroNode>,
        metadata: HydroIrMetadata,
    },

    ResolveFutures {
        input: Box<HydroNode>,
        metadata: HydroIrMetadata,
    },
    ResolveFuturesOrdered {
        input: Box<HydroNode>,
        metadata: HydroIrMetadata,
    },

    Map {
        f: DebugExpr,
        input: Box<HydroNode>,
        metadata: HydroIrMetadata,
    },
    FlatMap {
        f: DebugExpr,
        input: Box<HydroNode>,
        metadata: HydroIrMetadata,
    },
    Filter {
        f: DebugExpr,
        input: Box<HydroNode>,
        metadata: HydroIrMetadata,
    },
    FilterMap {
        f: DebugExpr,
        input: Box<HydroNode>,
        metadata: HydroIrMetadata,
    },

    DeferTick {
        input: Box<HydroNode>,
        metadata: HydroIrMetadata,
    },
    Enumerate {
        input: Box<HydroNode>,
        metadata: HydroIrMetadata,
    },
    Inspect {
        f: DebugExpr,
        input: Box<HydroNode>,
        metadata: HydroIrMetadata,
    },

    Unique {
        input: Box<HydroNode>,
        metadata: HydroIrMetadata,
    },

    Sort {
        input: Box<HydroNode>,
        metadata: HydroIrMetadata,
    },
    Fold {
        init: DebugExpr,
        acc: DebugExpr,
        input: Box<HydroNode>,
        metadata: HydroIrMetadata,
    },

    Scan {
        init: DebugExpr,
        acc: DebugExpr,
        input: Box<HydroNode>,
        metadata: HydroIrMetadata,
    },
    FoldKeyed {
        init: DebugExpr,
        acc: DebugExpr,
        input: Box<HydroNode>,
        metadata: HydroIrMetadata,
    },

    Reduce {
        f: DebugExpr,
        input: Box<HydroNode>,
        metadata: HydroIrMetadata,
    },
    ReduceKeyed {
        f: DebugExpr,
        input: Box<HydroNode>,
        metadata: HydroIrMetadata,
    },
    ReduceKeyedWatermark {
        f: DebugExpr,
        input: Box<HydroNode>,
        watermark: Box<HydroNode>,
        metadata: HydroIrMetadata,
    },

    Network {
        name: Option<String>,
        networking_info: crate::networking::NetworkingInfo,
        serialize_fn: Option<DebugExpr>,
        instantiate_fn: DebugInstantiate,
        deserialize_fn: Option<DebugExpr>,
        input: Box<HydroNode>,
        metadata: HydroIrMetadata,
    },

    ExternalInput {
        from_external_key: LocationKey,
        from_port_id: ExternalPortId,
        from_many: bool,
        codec_type: DebugType,
        port_hint: NetworkHint,
        instantiate_fn: DebugInstantiate,
        deserialize_fn: Option<DebugExpr>,
        metadata: HydroIrMetadata,
    },

    Counter {
        tag: String,
        duration: DebugExpr,
        prefix: String,
        input: Box<HydroNode>,
        metadata: HydroIrMetadata,
    },
}

pub type SeenSharedNodes = HashMap<*const RefCell<HydroNode>, Rc<RefCell<HydroNode>>>;
pub type SeenSharedNodeLocations = HashMap<*const RefCell<HydroNode>, LocationId>;

impl HydroNode {
    pub fn transform_bottom_up(
        &mut self,
        transform: &mut impl FnMut(&mut HydroNode),
        seen_tees: &mut SeenSharedNodes,
        check_well_formed: bool,
    ) {
        self.transform_children(
            |n, s| n.transform_bottom_up(transform, s, check_well_formed),
            seen_tees,
        );

        transform(self);

        let self_location = self.metadata().location_id.root();

        if check_well_formed {
            match &*self {
                HydroNode::Network { .. } => {}
                _ => {
                    self.input_metadata().iter().for_each(|i| {
                        if i.location_id.root() != self_location {
                            panic!(
                                "Mismatching IR locations, child: {:?} ({:?}) of: {:?} ({:?})",
                                i,
                                i.location_id.root(),
                                self,
                                self_location
                            )
                        }
                    });
                }
            }
        }
    }

    #[inline(always)]
    pub fn transform_children(
        &mut self,
        mut transform: impl FnMut(&mut HydroNode, &mut SeenSharedNodes),
        seen_tees: &mut SeenSharedNodes,
    ) {
        match self {
            HydroNode::Placeholder => {
                panic!();
            }

            HydroNode::Source { .. }
            | HydroNode::SingletonSource { .. }
            | HydroNode::CycleSource { .. }
            | HydroNode::ExternalInput { .. } => {}

            HydroNode::Tee { inner, .. } => {
                if let Some(transformed) = seen_tees.get(&inner.as_ptr()) {
                    *inner = SharedNode(transformed.clone());
                } else {
                    let transformed_cell = Rc::new(RefCell::new(HydroNode::Placeholder));
                    seen_tees.insert(inner.as_ptr(), transformed_cell.clone());
                    let mut orig = inner.0.replace(HydroNode::Placeholder);
                    transform(&mut orig, seen_tees);
                    *transformed_cell.borrow_mut() = orig;
                    *inner = SharedNode(transformed_cell);
                }
            }

            HydroNode::Partition { inner, .. } => {
                if let Some(transformed) = seen_tees.get(&inner.as_ptr()) {
                    *inner = SharedNode(transformed.clone());
                } else {
                    let transformed_cell = Rc::new(RefCell::new(HydroNode::Placeholder));
                    seen_tees.insert(inner.as_ptr(), transformed_cell.clone());
                    let mut orig = inner.0.replace(HydroNode::Placeholder);
                    transform(&mut orig, seen_tees);
                    *transformed_cell.borrow_mut() = orig;
                    *inner = SharedNode(transformed_cell);
                }
            }

            HydroNode::Cast { inner, .. }
            | HydroNode::ObserveNonDet { inner, .. }
            | HydroNode::BeginAtomic { inner, .. }
            | HydroNode::EndAtomic { inner, .. }
            | HydroNode::Batch { inner, .. }
            | HydroNode::YieldConcat { inner, .. } => {
                transform(inner.as_mut(), seen_tees);
            }

            HydroNode::Chain { first, second, .. } => {
                transform(first.as_mut(), seen_tees);
                transform(second.as_mut(), seen_tees);
            }

            HydroNode::ChainFirst { first, second, .. } => {
                transform(first.as_mut(), seen_tees);
                transform(second.as_mut(), seen_tees);
            }

            HydroNode::CrossSingleton { left, right, .. }
            | HydroNode::CrossProduct { left, right, .. }
            | HydroNode::Join { left, right, .. } => {
                transform(left.as_mut(), seen_tees);
                transform(right.as_mut(), seen_tees);
            }

            HydroNode::Difference { pos, neg, .. } | HydroNode::AntiJoin { pos, neg, .. } => {
                transform(pos.as_mut(), seen_tees);
                transform(neg.as_mut(), seen_tees);
            }

            HydroNode::ReduceKeyedWatermark {
                input, watermark, ..
            } => {
                transform(input.as_mut(), seen_tees);
                transform(watermark.as_mut(), seen_tees);
            }

            HydroNode::Map { input, .. }
            | HydroNode::ResolveFutures { input, .. }
            | HydroNode::ResolveFuturesOrdered { input, .. }
            | HydroNode::FlatMap { input, .. }
            | HydroNode::Filter { input, .. }
            | HydroNode::FilterMap { input, .. }
            | HydroNode::Sort { input, .. }
            | HydroNode::DeferTick { input, .. }
            | HydroNode::Enumerate { input, .. }
            | HydroNode::Inspect { input, .. }
            | HydroNode::Unique { input, .. }
            | HydroNode::Network { input, .. }
            | HydroNode::Fold { input, .. }
            | HydroNode::Scan { input, .. }
            | HydroNode::FoldKeyed { input, .. }
            | HydroNode::Reduce { input, .. }
            | HydroNode::ReduceKeyed { input, .. }
            | HydroNode::Counter { input, .. } => {
                transform(input.as_mut(), seen_tees);
            }
        }
    }

    pub fn deep_clone(&self, seen_tees: &mut SeenSharedNodes) -> HydroNode {
        match self {
            HydroNode::Placeholder => HydroNode::Placeholder,
            HydroNode::Cast { inner, metadata } => HydroNode::Cast {
                inner: Box::new(inner.deep_clone(seen_tees)),
                metadata: metadata.clone(),
            },
            HydroNode::ObserveNonDet {
                inner,
                trusted,
                metadata,
            } => HydroNode::ObserveNonDet {
                inner: Box::new(inner.deep_clone(seen_tees)),
                trusted: *trusted,
                metadata: metadata.clone(),
            },
            HydroNode::Source { source, metadata } => HydroNode::Source {
                source: source.clone(),
                metadata: metadata.clone(),
            },
            HydroNode::SingletonSource {
                value,
                first_tick_only,
                metadata,
            } => HydroNode::SingletonSource {
                value: value.clone(),
                first_tick_only: *first_tick_only,
                metadata: metadata.clone(),
            },
            HydroNode::CycleSource { cycle_id, metadata } => HydroNode::CycleSource {
                cycle_id: *cycle_id,
                metadata: metadata.clone(),
            },
            HydroNode::Tee { inner, metadata } => {
                if let Some(transformed) = seen_tees.get(&inner.as_ptr()) {
                    HydroNode::Tee {
                        inner: SharedNode(transformed.clone()),
                        metadata: metadata.clone(),
                    }
                } else {
                    let new_rc = Rc::new(RefCell::new(HydroNode::Placeholder));
                    seen_tees.insert(inner.as_ptr(), new_rc.clone());
                    let cloned = inner.0.borrow().deep_clone(seen_tees);
                    *new_rc.borrow_mut() = cloned;
                    HydroNode::Tee {
                        inner: SharedNode(new_rc),
                        metadata: metadata.clone(),
                    }
                }
            }
            HydroNode::Partition {
                inner,
                f,
                is_true,
                metadata,
            } => {
                if let Some(transformed) = seen_tees.get(&inner.as_ptr()) {
                    HydroNode::Partition {
                        inner: SharedNode(transformed.clone()),
                        f: f.clone(),
                        is_true: *is_true,
                        metadata: metadata.clone(),
                    }
                } else {
                    let new_rc = Rc::new(RefCell::new(HydroNode::Placeholder));
                    seen_tees.insert(inner.as_ptr(), new_rc.clone());
                    let cloned = inner.0.borrow().deep_clone(seen_tees);
                    *new_rc.borrow_mut() = cloned;
                    HydroNode::Partition {
                        inner: SharedNode(new_rc),
                        f: f.clone(),
                        is_true: *is_true,
                        metadata: metadata.clone(),
                    }
                }
            }
            HydroNode::YieldConcat { inner, metadata } => HydroNode::YieldConcat {
                inner: Box::new(inner.deep_clone(seen_tees)),
                metadata: metadata.clone(),
            },
            HydroNode::BeginAtomic { inner, metadata } => HydroNode::BeginAtomic {
                inner: Box::new(inner.deep_clone(seen_tees)),
                metadata: metadata.clone(),
            },
            HydroNode::EndAtomic { inner, metadata } => HydroNode::EndAtomic {
                inner: Box::new(inner.deep_clone(seen_tees)),
                metadata: metadata.clone(),
            },
            HydroNode::Batch { inner, metadata } => HydroNode::Batch {
                inner: Box::new(inner.deep_clone(seen_tees)),
                metadata: metadata.clone(),
            },
            HydroNode::Chain {
                first,
                second,
                metadata,
            } => HydroNode::Chain {
                first: Box::new(first.deep_clone(seen_tees)),
                second: Box::new(second.deep_clone(seen_tees)),
                metadata: metadata.clone(),
            },
            HydroNode::ChainFirst {
                first,
                second,
                metadata,
            } => HydroNode::ChainFirst {
                first: Box::new(first.deep_clone(seen_tees)),
                second: Box::new(second.deep_clone(seen_tees)),
                metadata: metadata.clone(),
            },
            HydroNode::CrossProduct {
                left,
                right,
                metadata,
            } => HydroNode::CrossProduct {
                left: Box::new(left.deep_clone(seen_tees)),
                right: Box::new(right.deep_clone(seen_tees)),
                metadata: metadata.clone(),
            },
            HydroNode::CrossSingleton {
                left,
                right,
                metadata,
            } => HydroNode::CrossSingleton {
                left: Box::new(left.deep_clone(seen_tees)),
                right: Box::new(right.deep_clone(seen_tees)),
                metadata: metadata.clone(),
            },
            HydroNode::Join {
                left,
                right,
                metadata,
            } => HydroNode::Join {
                left: Box::new(left.deep_clone(seen_tees)),
                right: Box::new(right.deep_clone(seen_tees)),
                metadata: metadata.clone(),
            },
            HydroNode::Difference { pos, neg, metadata } => HydroNode::Difference {
                pos: Box::new(pos.deep_clone(seen_tees)),
                neg: Box::new(neg.deep_clone(seen_tees)),
                metadata: metadata.clone(),
            },
            HydroNode::AntiJoin { pos, neg, metadata } => HydroNode::AntiJoin {
                pos: Box::new(pos.deep_clone(seen_tees)),
                neg: Box::new(neg.deep_clone(seen_tees)),
                metadata: metadata.clone(),
            },
            HydroNode::ResolveFutures { input, metadata } => HydroNode::ResolveFutures {
                input: Box::new(input.deep_clone(seen_tees)),
                metadata: metadata.clone(),
            },
            HydroNode::ResolveFuturesOrdered { input, metadata } => {
                HydroNode::ResolveFuturesOrdered {
                    input: Box::new(input.deep_clone(seen_tees)),
                    metadata: metadata.clone(),
                }
            }
            HydroNode::Map { f, input, metadata } => HydroNode::Map {
                f: f.clone(),
                input: Box::new(input.deep_clone(seen_tees)),
                metadata: metadata.clone(),
            },
            HydroNode::FlatMap { f, input, metadata } => HydroNode::FlatMap {
                f: f.clone(),
                input: Box::new(input.deep_clone(seen_tees)),
                metadata: metadata.clone(),
            },
            HydroNode::Filter { f, input, metadata } => HydroNode::Filter {
                f: f.clone(),
                input: Box::new(input.deep_clone(seen_tees)),
                metadata: metadata.clone(),
            },
            HydroNode::FilterMap { f, input, metadata } => HydroNode::FilterMap {
                f: f.clone(),
                input: Box::new(input.deep_clone(seen_tees)),
                metadata: metadata.clone(),
            },
            HydroNode::DeferTick { input, metadata } => HydroNode::DeferTick {
                input: Box::new(input.deep_clone(seen_tees)),
                metadata: metadata.clone(),
            },
            HydroNode::Enumerate { input, metadata } => HydroNode::Enumerate {
                input: Box::new(input.deep_clone(seen_tees)),
                metadata: metadata.clone(),
            },
            HydroNode::Inspect { f, input, metadata } => HydroNode::Inspect {
                f: f.clone(),
                input: Box::new(input.deep_clone(seen_tees)),
                metadata: metadata.clone(),
            },
            HydroNode::Unique { input, metadata } => HydroNode::Unique {
                input: Box::new(input.deep_clone(seen_tees)),
                metadata: metadata.clone(),
            },
            HydroNode::Sort { input, metadata } => HydroNode::Sort {
                input: Box::new(input.deep_clone(seen_tees)),
                metadata: metadata.clone(),
            },
            HydroNode::Fold {
                init,
                acc,
                input,
                metadata,
            } => HydroNode::Fold {
                init: init.clone(),
                acc: acc.clone(),
                input: Box::new(input.deep_clone(seen_tees)),
                metadata: metadata.clone(),
            },
            HydroNode::Scan {
                init,
                acc,
                input,
                metadata,
            } => HydroNode::Scan {
                init: init.clone(),
                acc: acc.clone(),
                input: Box::new(input.deep_clone(seen_tees)),
                metadata: metadata.clone(),
            },
            HydroNode::FoldKeyed {
                init,
                acc,
                input,
                metadata,
            } => HydroNode::FoldKeyed {
                init: init.clone(),
                acc: acc.clone(),
                input: Box::new(input.deep_clone(seen_tees)),
                metadata: metadata.clone(),
            },
            HydroNode::ReduceKeyedWatermark {
                f,
                input,
                watermark,
                metadata,
            } => HydroNode::ReduceKeyedWatermark {
                f: f.clone(),
                input: Box::new(input.deep_clone(seen_tees)),
                watermark: Box::new(watermark.deep_clone(seen_tees)),
                metadata: metadata.clone(),
            },
            HydroNode::Reduce { f, input, metadata } => HydroNode::Reduce {
                f: f.clone(),
                input: Box::new(input.deep_clone(seen_tees)),
                metadata: metadata.clone(),
            },
            HydroNode::ReduceKeyed { f, input, metadata } => HydroNode::ReduceKeyed {
                f: f.clone(),
                input: Box::new(input.deep_clone(seen_tees)),
                metadata: metadata.clone(),
            },
            HydroNode::Network {
                name,
                networking_info,
                serialize_fn,
                instantiate_fn,
                deserialize_fn,
                input,
                metadata,
            } => HydroNode::Network {
                name: name.clone(),
                networking_info: networking_info.clone(),
                serialize_fn: serialize_fn.clone(),
                instantiate_fn: instantiate_fn.clone(),
                deserialize_fn: deserialize_fn.clone(),
                input: Box::new(input.deep_clone(seen_tees)),
                metadata: metadata.clone(),
            },
            HydroNode::ExternalInput {
                from_external_key,
                from_port_id,
                from_many,
                codec_type,
                port_hint,
                instantiate_fn,
                deserialize_fn,
                metadata,
            } => HydroNode::ExternalInput {
                from_external_key: *from_external_key,
                from_port_id: *from_port_id,
                from_many: *from_many,
                codec_type: codec_type.clone(),
                port_hint: *port_hint,
                instantiate_fn: instantiate_fn.clone(),
                deserialize_fn: deserialize_fn.clone(),
                metadata: metadata.clone(),
            },
            HydroNode::Counter {
                tag,
                duration,
                prefix,
                input,
                metadata,
            } => HydroNode::Counter {
                tag: tag.clone(),
                duration: duration.clone(),
                prefix: prefix.clone(),
                input: Box::new(input.deep_clone(seen_tees)),
                metadata: metadata.clone(),
            },
        }
    }

    #[cfg(feature = "build")]
    pub fn emit_core(
        &mut self,
        builders_or_callback: &mut BuildersOrCallback<
            impl FnMut(&mut HydroRoot, &mut usize),
            impl FnMut(&mut HydroNode, &mut usize),
        >,
        seen_tees: &mut SeenSharedNodes,
        built_tees: &mut HashMap<*const RefCell<HydroNode>, Vec<syn::Ident>>,
        next_stmt_id: &mut usize,
    ) -> syn::Ident {
        let mut ident_stack: Vec<syn::Ident> = Vec::new();

        self.transform_bottom_up(
            &mut |node: &mut HydroNode| {
                let out_location = node.metadata().location_id.clone();
                match node {
                    HydroNode::Placeholder => {
                        panic!()
                    }

                    HydroNode::Cast { .. } => {
                        // Cast passes through the input ident unchanged
                        // The input ident is already on the stack from processing the child
                        match builders_or_callback {
                            BuildersOrCallback::Builders(_) => {}
                            BuildersOrCallback::Callback(_, node_callback) => {
                                node_callback(node, next_stmt_id);
                            }
                        }

                        *next_stmt_id += 1;
                        // input_ident stays on stack as output
                    }

                    HydroNode::ObserveNonDet {
                        inner,
                        trusted,
                        metadata,
                        ..
                    } => {
                        let inner_ident = ident_stack.pop().unwrap();

                        let observe_ident =
                            syn::Ident::new(&format!("stream_{}", *next_stmt_id), Span::call_site());

                        match builders_or_callback {
                            BuildersOrCallback::Builders(graph_builders) => {
                                graph_builders.observe_nondet(
                                    *trusted,
                                    &inner.metadata().location_id,
                                    inner_ident,
                                    &inner.metadata().collection_kind,
                                    &observe_ident,
                                    &metadata.collection_kind,
                                    &metadata.op,
                                );
                            }
                            BuildersOrCallback::Callback(_, node_callback) => {
                                node_callback(node, next_stmt_id);
                            }
                        }

                        *next_stmt_id += 1;

                        ident_stack.push(observe_ident);
                    }

                    HydroNode::Batch {
                        inner, metadata, ..
                    } => {
                        let inner_ident = ident_stack.pop().unwrap();

                        let batch_ident =
                            syn::Ident::new(&format!("stream_{}", *next_stmt_id), Span::call_site());

                        match builders_or_callback {
                            BuildersOrCallback::Builders(graph_builders) => {
                                graph_builders.batch(
                                    inner_ident,
                                    &inner.metadata().location_id,
                                    &inner.metadata().collection_kind,
                                    &batch_ident,
                                    &out_location,
                                    &metadata.op,
                                );
                            }
                            BuildersOrCallback::Callback(_, node_callback) => {
                                node_callback(node, next_stmt_id);
                            }
                        }

                        *next_stmt_id += 1;

                        ident_stack.push(batch_ident);
                    }

                    HydroNode::YieldConcat { inner, .. } => {
                        let inner_ident = ident_stack.pop().unwrap();

                        let yield_ident =
                            syn::Ident::new(&format!("stream_{}", *next_stmt_id), Span::call_site());

                        match builders_or_callback {
                            BuildersOrCallback::Builders(graph_builders) => {
                                graph_builders.yield_from_tick(
                                    inner_ident,
                                    &inner.metadata().location_id,
                                    &inner.metadata().collection_kind,
                                    &yield_ident,
                                    &out_location,
                                );
                            }
                            BuildersOrCallback::Callback(_, node_callback) => {
                                node_callback(node, next_stmt_id);
                            }
                        }

                        *next_stmt_id += 1;

                        ident_stack.push(yield_ident);
                    }

                    HydroNode::BeginAtomic { inner, metadata } => {
                        let inner_ident = ident_stack.pop().unwrap();

                        let begin_ident =
                            syn::Ident::new(&format!("stream_{}", *next_stmt_id), Span::call_site());

                        match builders_or_callback {
                            BuildersOrCallback::Builders(graph_builders) => {
                                graph_builders.begin_atomic(
                                    inner_ident,
                                    &inner.metadata().location_id,
                                    &inner.metadata().collection_kind,
                                    &begin_ident,
                                    &out_location,
                                    &metadata.op,
                                );
                            }
                            BuildersOrCallback::Callback(_, node_callback) => {
                                node_callback(node, next_stmt_id);
                            }
                        }

                        *next_stmt_id += 1;

                        ident_stack.push(begin_ident);
                    }

                    HydroNode::EndAtomic { inner, .. } => {
                        let inner_ident = ident_stack.pop().unwrap();

                        let end_ident =
                            syn::Ident::new(&format!("stream_{}", *next_stmt_id), Span::call_site());

                        match builders_or_callback {
                            BuildersOrCallback::Builders(graph_builders) => {
                                graph_builders.end_atomic(
                                    inner_ident,
                                    &inner.metadata().location_id,
                                    &inner.metadata().collection_kind,
                                    &end_ident,
                                );
                            }
                            BuildersOrCallback::Callback(_, node_callback) => {
                                node_callback(node, next_stmt_id);
                            }
                        }

                        *next_stmt_id += 1;

                        ident_stack.push(end_ident);
                    }

                    HydroNode::Source {
                        source, metadata, ..
                    } => {
                        if let HydroSource::ExternalNetwork() = source {
                            ident_stack.push(syn::Ident::new("DUMMY", Span::call_site()));
                        } else {
                            let source_ident =
                                syn::Ident::new(&format!("stream_{}", *next_stmt_id), Span::call_site());

                            let source_stmt = match source {
                                HydroSource::Stream(expr) => {
                                    debug_assert!(metadata.location_id.is_top_level());
                                    parse_quote! {
                                        #source_ident = source_stream(#expr);
                                    }
                                }

                                HydroSource::ExternalNetwork() => {
                                    unreachable!()
                                }

                                HydroSource::Iter(expr) => {
                                    if metadata.location_id.is_top_level() {
                                        parse_quote! {
                                            #source_ident = source_iter(#expr);
                                        }
                                    } else {
                                        // TODO(shadaj): a more natural semantics would be to to re-evaluate the expression on each tick
                                        parse_quote! {
                                            #source_ident = source_iter(#expr) -> persist::<'static>();
                                        }
                                    }
                                }

                                HydroSource::Spin() => {
                                    debug_assert!(metadata.location_id.is_top_level());
                                    parse_quote! {
                                        #source_ident = spin();
                                    }
                                }

                                HydroSource::ClusterMembers(target_loc, state) => {
                                    debug_assert!(metadata.location_id.is_top_level());

                                    let members_tee_ident = syn::Ident::new(
                                        &format!(
                                            "__cluster_members_tee_{}_{}",
                                            metadata.location_id.root().key(),
                                            target_loc.key(),
                                        ),
                                        Span::call_site(),
                                    );

                                    match state {
                                        ClusterMembersState::Stream(d) => {
                                            parse_quote! {
                                                #members_tee_ident = source_stream(#d) -> tee();
                                                #source_ident = #members_tee_ident;
                                            }
                                        },
                                        ClusterMembersState::Uninit => syn::parse_quote! {
                                            #source_ident = source_stream(DUMMY);
                                        },
                                        ClusterMembersState::Tee(..) => parse_quote! {
                                            #source_ident = #members_tee_ident;
                                        },
                                    }
                                }

                                HydroSource::Embedded(ident) => {
                                    parse_quote! {
                                        #source_ident = source_stream(#ident);
                                    }
                                }

                                HydroSource::EmbeddedSingleton(ident) => {
                                    parse_quote! {
                                        #source_ident = source_iter([#ident]);
                                    }
                                }
                            };

                            match builders_or_callback {
                                BuildersOrCallback::Builders(graph_builders) => {
                                    let builder = graph_builders.get_dfir_mut(&out_location);
                                    builder.add_dfir(source_stmt, None, Some(&next_stmt_id.to_string()));
                                }
                                BuildersOrCallback::Callback(_, node_callback) => {
                                    node_callback(node, next_stmt_id);
                                }
                            }

                            *next_stmt_id += 1;

                            ident_stack.push(source_ident);
                        }
                    }

                    HydroNode::SingletonSource { value, first_tick_only, metadata } => {
                        let source_ident =
                            syn::Ident::new(&format!("stream_{}", *next_stmt_id), Span::call_site());

                        match builders_or_callback {
                            BuildersOrCallback::Builders(graph_builders) => {
                                let builder = graph_builders.get_dfir_mut(&out_location);

                                if *first_tick_only {
                                    assert!(
                                        !metadata.location_id.is_top_level(),
                                        "first_tick_only SingletonSource must be inside a tick"
                                    );
                                }

                                if *first_tick_only
                                    || (metadata.location_id.is_top_level()
                                        && metadata.collection_kind.is_bounded())
                                {
                                    builder.add_dfir(
                                        parse_quote! {
                                            #source_ident = source_iter([#value]);
                                        },
                                        None,
                                        Some(&next_stmt_id.to_string()),
                                    );
                                } else {
                                    builder.add_dfir(
                                        parse_quote! {
                                            #source_ident = source_iter([#value]) -> persist::<'static>();
                                        },
                                        None,
                                        Some(&next_stmt_id.to_string()),
                                    );
                                }
                            }
                            BuildersOrCallback::Callback(_, node_callback) => {
                                node_callback(node, next_stmt_id);
                            }
                        }

                        *next_stmt_id += 1;

                        ident_stack.push(source_ident);
                    }

                    HydroNode::CycleSource { cycle_id, .. } => {
                        let ident = cycle_id.as_ident();

                        match builders_or_callback {
                            BuildersOrCallback::Builders(_) => {}
                            BuildersOrCallback::Callback(_, node_callback) => {
                                node_callback(node, next_stmt_id);
                            }
                        }

                        // consume a stmt id even though we did not emit anything so that we can instrument this
                        *next_stmt_id += 1;

                        ident_stack.push(ident);
                    }

                    HydroNode::Tee { inner, .. } => {
                        let ret_ident = if let Some(built_idents) =
                            built_tees.get(&(inner.0.as_ref() as *const RefCell<HydroNode>))
                        {
                            match builders_or_callback {
                                BuildersOrCallback::Builders(_) => {}
                                BuildersOrCallback::Callback(_, node_callback) => {
                                    node_callback(node, next_stmt_id);
                                }
                            }

                            built_idents[0].clone()
                        } else {
                            // The inner node was already processed by transform_bottom_up,
                            // so its ident is on the stack
                            let inner_ident = ident_stack.pop().unwrap();

                            let tee_ident =
                                syn::Ident::new(&format!("stream_{}", *next_stmt_id), Span::call_site());

                            built_tees.insert(
                                inner.0.as_ref() as *const RefCell<HydroNode>,
                                vec![tee_ident.clone()],
                            );

                            match builders_or_callback {
                                BuildersOrCallback::Builders(graph_builders) => {
                                    let builder = graph_builders.get_dfir_mut(&out_location);
                                    builder.add_dfir(
                                        parse_quote! {
                                            #tee_ident = #inner_ident -> tee();
                                        },
                                        None,
                                        Some(&next_stmt_id.to_string()),
                                    );
                                }
                                BuildersOrCallback::Callback(_, node_callback) => {
                                    node_callback(node, next_stmt_id);
                                }
                            }

                            tee_ident
                        };

                        // we consume a stmt id regardless of if we emit the tee() operator,
                        // so that during rewrites we touch all recipients of the tee()

                        *next_stmt_id += 1;
                        ident_stack.push(ret_ident);
                    }

                    HydroNode::Partition {
                        inner, f, is_true, ..
                    } => {
                        let is_true = *is_true; // need to copy early to avoid borrow checking issues with node
                        let ptr = inner.0.as_ref() as *const RefCell<HydroNode>;
                        let ret_ident = if let Some(built_idents) = built_tees.get(&ptr) {
                            match builders_or_callback {
                                BuildersOrCallback::Builders(_) => {}
                                BuildersOrCallback::Callback(_, node_callback) => {
                                    node_callback(node, next_stmt_id);
                                }
                            }

                            let idx = if is_true { 0 } else { 1 };
                            built_idents[idx].clone()
                        } else {
                            // The inner node was already processed by transform_bottom_up,
                            // so its ident is on the stack
                            let inner_ident = ident_stack.pop().unwrap();

                            let partition_ident = syn::Ident::new(
                                &format!("stream_{}_partition", *next_stmt_id),
                                Span::call_site(),
                            );
                            let true_ident = syn::Ident::new(
                                &format!("stream_{}_true", *next_stmt_id),
                                Span::call_site(),
                            );
                            let false_ident = syn::Ident::new(
                                &format!("stream_{}_false", *next_stmt_id),
                                Span::call_site(),
                            );

                            built_tees.insert(
                                ptr,
                                vec![true_ident.clone(), false_ident.clone()],
                            );

                            match builders_or_callback {
                                BuildersOrCallback::Builders(graph_builders) => {
                                    let builder = graph_builders.get_dfir_mut(&out_location);
                                    builder.add_dfir(
                                        parse_quote! {
                                            #partition_ident = #inner_ident -> partition(|__item, __num_outputs| if (#f)(__item) { 0_usize } else { 1_usize });
                                            #true_ident = #partition_ident[0];
                                            #false_ident = #partition_ident[1];
                                        },
                                        None,
                                        Some(&next_stmt_id.to_string()),
                                    );
                                }
                                BuildersOrCallback::Callback(_, node_callback) => {
                                    node_callback(node, next_stmt_id);
                                }
                            }

                            if is_true { true_ident } else { false_ident }
                        };

                        *next_stmt_id += 1;
                        ident_stack.push(ret_ident);
                    }

                    HydroNode::Chain { .. } => {
                        // Children are processed left-to-right, so second is on top
                        let second_ident = ident_stack.pop().unwrap();
                        let first_ident = ident_stack.pop().unwrap();

                        let chain_ident =
                            syn::Ident::new(&format!("stream_{}", *next_stmt_id), Span::call_site());

                        match builders_or_callback {
                            BuildersOrCallback::Builders(graph_builders) => {
                                let builder = graph_builders.get_dfir_mut(&out_location);
                                builder.add_dfir(
                                    parse_quote! {
                                        #chain_ident = chain();
                                        #first_ident -> [0]#chain_ident;
                                        #second_ident -> [1]#chain_ident;
                                    },
                                    None,
                                    Some(&next_stmt_id.to_string()),
                                );
                            }
                            BuildersOrCallback::Callback(_, node_callback) => {
                                node_callback(node, next_stmt_id);
                            }
                        }

                        *next_stmt_id += 1;

                        ident_stack.push(chain_ident);
                    }

                    HydroNode::ChainFirst { .. } => {
                        let second_ident = ident_stack.pop().unwrap();
                        let first_ident = ident_stack.pop().unwrap();

                        let chain_ident =
                            syn::Ident::new(&format!("stream_{}", *next_stmt_id), Span::call_site());

                        match builders_or_callback {
                            BuildersOrCallback::Builders(graph_builders) => {
                                let builder = graph_builders.get_dfir_mut(&out_location);
                                builder.add_dfir(
                                    parse_quote! {
                                        #chain_ident = chain_first_n(1);
                                        #first_ident -> [0]#chain_ident;
                                        #second_ident -> [1]#chain_ident;
                                    },
                                    None,
                                    Some(&next_stmt_id.to_string()),
                                );
                            }
                            BuildersOrCallback::Callback(_, node_callback) => {
                                node_callback(node, next_stmt_id);
                            }
                        }

                        *next_stmt_id += 1;

                        ident_stack.push(chain_ident);
                    }

                    HydroNode::CrossSingleton { right, .. } => {
                        let right_ident = ident_stack.pop().unwrap();
                        let left_ident = ident_stack.pop().unwrap();

                        let cross_ident =
                            syn::Ident::new(&format!("stream_{}", *next_stmt_id), Span::call_site());

                        match builders_or_callback {
                            BuildersOrCallback::Builders(graph_builders) => {
                                let builder = graph_builders.get_dfir_mut(&out_location);

                                if right.metadata().location_id.is_top_level()
                                    && right.metadata().collection_kind.is_bounded()
                                {
                                    builder.add_dfir(
                                        parse_quote! {
                                            #cross_ident = cross_singleton();
                                            #left_ident -> [input]#cross_ident;
                                            #right_ident -> persist::<'static>() -> [single]#cross_ident;
                                        },
                                        None,
                                        Some(&next_stmt_id.to_string()),
                                    );
                                } else {
                                    builder.add_dfir(
                                        parse_quote! {
                                            #cross_ident = cross_singleton();
                                            #left_ident -> [input]#cross_ident;
                                            #right_ident -> [single]#cross_ident;
                                        },
                                        None,
                                        Some(&next_stmt_id.to_string()),
                                    );
                                }
                            }
                            BuildersOrCallback::Callback(_, node_callback) => {
                                node_callback(node, next_stmt_id);
                            }
                        }

                        *next_stmt_id += 1;

                        ident_stack.push(cross_ident);
                    }

                    HydroNode::CrossProduct { .. } | HydroNode::Join { .. } => {
                        let operator: syn::Ident = if matches!(node, HydroNode::CrossProduct { .. }) {
                            parse_quote!(cross_join_multiset)
                        } else {
                            parse_quote!(join_multiset)
                        };

                        let (HydroNode::CrossProduct { left, right, .. }
                        | HydroNode::Join { left, right, .. }) = node
                        else {
                            unreachable!()
                        };

                        let is_top_level = left.metadata().location_id.is_top_level()
                            && right.metadata().location_id.is_top_level();
                        let left_lifetime = if left.metadata().location_id.is_top_level() {
                            quote!('static)
                        } else {
                            quote!('tick)
                        };

                        let right_lifetime = if right.metadata().location_id.is_top_level() {
                            quote!('static)
                        } else {
                            quote!('tick)
                        };

                        let right_ident = ident_stack.pop().unwrap();
                        let left_ident = ident_stack.pop().unwrap();

                        let stream_ident =
                            syn::Ident::new(&format!("stream_{}", *next_stmt_id), Span::call_site());

                        match builders_or_callback {
                            BuildersOrCallback::Builders(graph_builders) => {
                                let builder = graph_builders.get_dfir_mut(&out_location);
                                builder.add_dfir(
                                    if is_top_level {
                                        // if both inputs are root, the output is expected to have streamy semantics, so we need
                                        // a multiset_delta() to negate the replay behavior
                                        parse_quote! {
                                            #stream_ident = #operator::<#left_lifetime, #right_lifetime>() -> multiset_delta();
                                            #left_ident -> [0]#stream_ident;
                                            #right_ident -> [1]#stream_ident;
                                        }
                                    } else {
                                        parse_quote! {
                                            #stream_ident = #operator::<#left_lifetime, #right_lifetime>();
                                            #left_ident -> [0]#stream_ident;
                                            #right_ident -> [1]#stream_ident;
                                        }
                                    }
                                    ,
                                    None,
                                    Some(&next_stmt_id.to_string()),
                                );
                            }
                            BuildersOrCallback::Callback(_, node_callback) => {
                                node_callback(node, next_stmt_id);
                            }
                        }

                        *next_stmt_id += 1;

                        ident_stack.push(stream_ident);
                    }

                    HydroNode::Difference { .. } | HydroNode::AntiJoin { .. } => {
                        let operator: syn::Ident = if matches!(node, HydroNode::Difference { .. }) {
                            parse_quote!(difference)
                        } else {
                            parse_quote!(anti_join)
                        };

                        let (HydroNode::Difference { neg, .. } | HydroNode::AntiJoin { neg, .. }) =
                            node
                        else {
                            unreachable!()
                        };

                        let neg_lifetime = if neg.metadata().location_id.is_top_level() {
                            quote!('static)
                        } else {
                            quote!('tick)
                        };

                        let neg_ident = ident_stack.pop().unwrap();
                        let pos_ident = ident_stack.pop().unwrap();

                        let stream_ident =
                            syn::Ident::new(&format!("stream_{}", *next_stmt_id), Span::call_site());

                        match builders_or_callback {
                            BuildersOrCallback::Builders(graph_builders) => {
                                let builder = graph_builders.get_dfir_mut(&out_location);
                                builder.add_dfir(
                                    parse_quote! {
                                        #stream_ident = #operator::<'tick, #neg_lifetime>();
                                        #pos_ident -> [pos]#stream_ident;
                                        #neg_ident -> [neg]#stream_ident;
                                    },
                                    None,
                                    Some(&next_stmt_id.to_string()),
                                );
                            }
                            BuildersOrCallback::Callback(_, node_callback) => {
                                node_callback(node, next_stmt_id);
                            }
                        }

                        *next_stmt_id += 1;

                        ident_stack.push(stream_ident);
                    }

                    HydroNode::ResolveFutures { .. } => {
                        let input_ident = ident_stack.pop().unwrap();

                        let futures_ident =
                            syn::Ident::new(&format!("stream_{}", *next_stmt_id), Span::call_site());

                        match builders_or_callback {
                            BuildersOrCallback::Builders(graph_builders) => {
                                let builder = graph_builders.get_dfir_mut(&out_location);
                                builder.add_dfir(
                                    parse_quote! {
                                        #futures_ident = #input_ident -> resolve_futures();
                                    },
                                    None,
                                    Some(&next_stmt_id.to_string()),
                                );
                            }
                            BuildersOrCallback::Callback(_, node_callback) => {
                                node_callback(node, next_stmt_id);
                            }
                        }

                        *next_stmt_id += 1;

                        ident_stack.push(futures_ident);
                    }

                    HydroNode::ResolveFuturesOrdered { .. } => {
                        let input_ident = ident_stack.pop().unwrap();

                        let futures_ident =
                            syn::Ident::new(&format!("stream_{}", *next_stmt_id), Span::call_site());

                        match builders_or_callback {
                            BuildersOrCallback::Builders(graph_builders) => {
                                let builder = graph_builders.get_dfir_mut(&out_location);
                                builder.add_dfir(
                                    parse_quote! {
                                        #futures_ident = #input_ident -> resolve_futures_ordered();
                                    },
                                    None,
                                    Some(&next_stmt_id.to_string()),
                                );
                            }
                            BuildersOrCallback::Callback(_, node_callback) => {
                                node_callback(node, next_stmt_id);
                            }
                        }

                        *next_stmt_id += 1;

                        ident_stack.push(futures_ident);
                    }

                    HydroNode::Map { f, .. } => {
                        let input_ident = ident_stack.pop().unwrap();

                        let map_ident =
                            syn::Ident::new(&format!("stream_{}", *next_stmt_id), Span::call_site());

                        match builders_or_callback {
                            BuildersOrCallback::Builders(graph_builders) => {
                                let builder = graph_builders.get_dfir_mut(&out_location);
                                builder.add_dfir(
                                    parse_quote! {
                                        #map_ident = #input_ident -> map(#f);
                                    },
                                    None,
                                    Some(&next_stmt_id.to_string()),
                                );
                            }
                            BuildersOrCallback::Callback(_, node_callback) => {
                                node_callback(node, next_stmt_id);
                            }
                        }

                        *next_stmt_id += 1;

                        ident_stack.push(map_ident);
                    }

                    HydroNode::FlatMap { f, .. } => {
                        let input_ident = ident_stack.pop().unwrap();

                        let flat_map_ident =
                            syn::Ident::new(&format!("stream_{}", *next_stmt_id), Span::call_site());

                        match builders_or_callback {
                            BuildersOrCallback::Builders(graph_builders) => {
                                let builder = graph_builders.get_dfir_mut(&out_location);
                                builder.add_dfir(
                                    parse_quote! {
                                        #flat_map_ident = #input_ident -> flat_map(#f);
                                    },
                                    None,
                                    Some(&next_stmt_id.to_string()),
                                );
                            }
                            BuildersOrCallback::Callback(_, node_callback) => {
                                node_callback(node, next_stmt_id);
                            }
                        }

                        *next_stmt_id += 1;

                        ident_stack.push(flat_map_ident);
                    }

                    HydroNode::Filter { f, .. } => {
                        let input_ident = ident_stack.pop().unwrap();

                        let filter_ident =
                            syn::Ident::new(&format!("stream_{}", *next_stmt_id), Span::call_site());

                        match builders_or_callback {
                            BuildersOrCallback::Builders(graph_builders) => {
                                let builder = graph_builders.get_dfir_mut(&out_location);
                                builder.add_dfir(
                                    parse_quote! {
                                        #filter_ident = #input_ident -> filter(#f);
                                    },
                                    None,
                                    Some(&next_stmt_id.to_string()),
                                );
                            }
                            BuildersOrCallback::Callback(_, node_callback) => {
                                node_callback(node, next_stmt_id);
                            }
                        }

                        *next_stmt_id += 1;

                        ident_stack.push(filter_ident);
                    }

                    HydroNode::FilterMap { f, .. } => {
                        let input_ident = ident_stack.pop().unwrap();

                        let filter_map_ident =
                            syn::Ident::new(&format!("stream_{}", *next_stmt_id), Span::call_site());

                        match builders_or_callback {
                            BuildersOrCallback::Builders(graph_builders) => {
                                let builder = graph_builders.get_dfir_mut(&out_location);
                                builder.add_dfir(
                                    parse_quote! {
                                        #filter_map_ident = #input_ident -> filter_map(#f);
                                    },
                                    None,
                                    Some(&next_stmt_id.to_string()),
                                );
                            }
                            BuildersOrCallback::Callback(_, node_callback) => {
                                node_callback(node, next_stmt_id);
                            }
                        }

                        *next_stmt_id += 1;

                        ident_stack.push(filter_map_ident);
                    }

                    HydroNode::Sort { .. } => {
                        let input_ident = ident_stack.pop().unwrap();

                        let sort_ident =
                            syn::Ident::new(&format!("stream_{}", *next_stmt_id), Span::call_site());

                        match builders_or_callback {
                            BuildersOrCallback::Builders(graph_builders) => {
                                let builder = graph_builders.get_dfir_mut(&out_location);
                                builder.add_dfir(
                                    parse_quote! {
                                        #sort_ident = #input_ident -> sort();
                                    },
                                    None,
                                    Some(&next_stmt_id.to_string()),
                                );
                            }
                            BuildersOrCallback::Callback(_, node_callback) => {
                                node_callback(node, next_stmt_id);
                            }
                        }

                        *next_stmt_id += 1;

                        ident_stack.push(sort_ident);
                    }

                    HydroNode::DeferTick { .. } => {
                        let input_ident = ident_stack.pop().unwrap();

                        let defer_tick_ident =
                            syn::Ident::new(&format!("stream_{}", *next_stmt_id), Span::call_site());

                        match builders_or_callback {
                            BuildersOrCallback::Builders(graph_builders) => {
                                let builder = graph_builders.get_dfir_mut(&out_location);
                                builder.add_dfir(
                                    parse_quote! {
                                        #defer_tick_ident = #input_ident -> defer_tick_lazy();
                                    },
                                    None,
                                    Some(&next_stmt_id.to_string()),
                                );
                            }
                            BuildersOrCallback::Callback(_, node_callback) => {
                                node_callback(node, next_stmt_id);
                            }
                        }

                        *next_stmt_id += 1;

                        ident_stack.push(defer_tick_ident);
                    }

                    HydroNode::Enumerate { input, .. } => {
                        let input_ident = ident_stack.pop().unwrap();

                        let enumerate_ident =
                            syn::Ident::new(&format!("stream_{}", *next_stmt_id), Span::call_site());

                        match builders_or_callback {
                            BuildersOrCallback::Builders(graph_builders) => {
                                let builder = graph_builders.get_dfir_mut(&out_location);
                                let lifetime = if input.metadata().location_id.is_top_level() {
                                    quote!('static)
                                } else {
                                    quote!('tick)
                                };
                                builder.add_dfir(
                                    parse_quote! {
                                        #enumerate_ident = #input_ident -> enumerate::<#lifetime>();
                                    },
                                    None,
                                    Some(&next_stmt_id.to_string()),
                                );
                            }
                            BuildersOrCallback::Callback(_, node_callback) => {
                                node_callback(node, next_stmt_id);
                            }
                        }

                        *next_stmt_id += 1;

                        ident_stack.push(enumerate_ident);
                    }

                    HydroNode::Inspect { f, .. } => {
                        let input_ident = ident_stack.pop().unwrap();

                        let inspect_ident =
                            syn::Ident::new(&format!("stream_{}", *next_stmt_id), Span::call_site());

                        match builders_or_callback {
                            BuildersOrCallback::Builders(graph_builders) => {
                                let builder = graph_builders.get_dfir_mut(&out_location);
                                builder.add_dfir(
                                    parse_quote! {
                                        #inspect_ident = #input_ident -> inspect(#f);
                                    },
                                    None,
                                    Some(&next_stmt_id.to_string()),
                                );
                            }
                            BuildersOrCallback::Callback(_, node_callback) => {
                                node_callback(node, next_stmt_id);
                            }
                        }

                        *next_stmt_id += 1;

                        ident_stack.push(inspect_ident);
                    }

                    HydroNode::Unique { input, .. } => {
                        let input_ident = ident_stack.pop().unwrap();

                        let unique_ident =
                            syn::Ident::new(&format!("stream_{}", *next_stmt_id), Span::call_site());

                        match builders_or_callback {
                            BuildersOrCallback::Builders(graph_builders) => {
                                let builder = graph_builders.get_dfir_mut(&out_location);
                                let lifetime = if input.metadata().location_id.is_top_level() {
                                    quote!('static)
                                } else {
                                    quote!('tick)
                                };

                                builder.add_dfir(
                                    parse_quote! {
                                        #unique_ident = #input_ident -> unique::<#lifetime>();
                                    },
                                    None,
                                    Some(&next_stmt_id.to_string()),
                                );
                            }
                            BuildersOrCallback::Callback(_, node_callback) => {
                                node_callback(node, next_stmt_id);
                            }
                        }

                        *next_stmt_id += 1;

                        ident_stack.push(unique_ident);
                    }

                    HydroNode::Fold { .. } | HydroNode::FoldKeyed { .. } | HydroNode::Scan { .. } => {
                        let operator: syn::Ident = if let HydroNode::Fold { input, .. } = node {
                            if input.metadata().location_id.is_top_level()
                                && input.metadata().collection_kind.is_bounded()
                            {
                                parse_quote!(fold_no_replay)
                            } else {
                                parse_quote!(fold)
                            }
                        } else if matches!(node, HydroNode::Scan { .. }) {
                            parse_quote!(scan)
                        } else if let HydroNode::FoldKeyed { input, .. } = node {
                            if input.metadata().location_id.is_top_level()
                                && input.metadata().collection_kind.is_bounded()
                            {
                                todo!("Fold keyed on a top-level bounded collection is not yet supported")
                            } else {
                                parse_quote!(fold_keyed)
                            }
                        } else {
                            unreachable!()
                        };

                        let (HydroNode::Fold { input, .. }
                        | HydroNode::FoldKeyed { input, .. }
                        | HydroNode::Scan { input, .. }) = node
                        else {
                            unreachable!()
                        };

                        let lifetime = if input.metadata().location_id.is_top_level() {
                            quote!('static)
                        } else {
                            quote!('tick)
                        };

                        let input_ident = ident_stack.pop().unwrap();

                        let (HydroNode::Fold { init, acc, .. }
                        | HydroNode::FoldKeyed { init, acc, .. }
                        | HydroNode::Scan { init, acc, .. }) = &*node
                        else {
                            unreachable!()
                        };

                        let fold_ident =
                            syn::Ident::new(&format!("stream_{}", *next_stmt_id), Span::call_site());

                        match builders_or_callback {
                            BuildersOrCallback::Builders(graph_builders) => {
                                if matches!(node, HydroNode::Fold { .. })
                                    && node.metadata().location_id.is_top_level()
                                    && !(matches!(node.metadata().location_id, LocationId::Atomic(_)))
                                    && graph_builders.singleton_intermediates()
                                    && !node.metadata().collection_kind.is_bounded()
                                {
                                    let builder = graph_builders.get_dfir_mut(&out_location);

                                    let acc: syn::Expr = parse_quote!({
                                        let mut __inner = #acc;
                                        move |__state, __value| {
                                            __inner(__state, __value);
                                            Some(__state.clone())
                                        }
                                    });

                                    builder.add_dfir(
                                        parse_quote! {
                                            source_iter([(#init)()]) -> [0]#fold_ident;
                                            #input_ident -> scan::<#lifetime>(#init, #acc) -> [1]#fold_ident;
                                            #fold_ident = chain();
                                        },
                                        None,
                                        Some(&next_stmt_id.to_string()),
                                    );
                                } else if matches!(node, HydroNode::FoldKeyed { .. })
                                    && node.metadata().location_id.is_top_level()
                                    && !(matches!(node.metadata().location_id, LocationId::Atomic(_)))
                                    && graph_builders.singleton_intermediates()
                                    && !node.metadata().collection_kind.is_bounded()
                                {
                                    let builder = graph_builders.get_dfir_mut(&out_location);

                                    let acc: syn::Expr = parse_quote!({
                                        let mut __init = #init;
                                        let mut __inner = #acc;
                                        move |__state, __kv: (_, _)| {
                                            // TODO(shadaj): we can avoid the clone when the entry exists
                                            let __state = __state
                                                .entry(::std::clone::Clone::clone(&__kv.0))
                                                .or_insert_with(|| (__init)());
                                            __inner(__state, __kv.1);
                                            Some((__kv.0, ::std::clone::Clone::clone(&*__state)))
                                        }
                                    });

                                    builder.add_dfir(
                                        parse_quote! {
                                            #fold_ident = #input_ident -> scan::<#lifetime>(|| ::std::collections::HashMap::new(), #acc);
                                        },
                                        None,
                                        Some(&next_stmt_id.to_string()),
                                    );
                                } else {
                                    let builder = graph_builders.get_dfir_mut(&out_location);
                                    builder.add_dfir(
                                        parse_quote! {
                                            #fold_ident = #input_ident -> #operator::<#lifetime>(#init, #acc);
                                        },
                                        None,
                                        Some(&next_stmt_id.to_string()),
                                    );
                                }
                            }
                            BuildersOrCallback::Callback(_, node_callback) => {
                                node_callback(node, next_stmt_id);
                            }
                        }

                        *next_stmt_id += 1;

                        ident_stack.push(fold_ident);
                    }

                    HydroNode::Reduce { .. } | HydroNode::ReduceKeyed { .. } => {
                        let operator: syn::Ident = if let HydroNode::Reduce { input, .. } = node {
                            if input.metadata().location_id.is_top_level()
                                && input.metadata().collection_kind.is_bounded()
                            {
                                parse_quote!(reduce_no_replay)
                            } else {
                                parse_quote!(reduce)
                            }
                        } else if let HydroNode::ReduceKeyed { input, .. } = node {
                            if input.metadata().location_id.is_top_level()
                                && input.metadata().collection_kind.is_bounded()
                            {
                                todo!(
                                    "Calling keyed reduce on a top-level bounded collection is not supported"
                                )
                            } else {
                                parse_quote!(reduce_keyed)
                            }
                        } else {
                            unreachable!()
                        };

                        let (HydroNode::Reduce { input, .. } | HydroNode::ReduceKeyed { input, .. }) = node
                        else {
                            unreachable!()
                        };

                        let lifetime = if input.metadata().location_id.is_top_level() {
                            quote!('static)
                        } else {
                            quote!('tick)
                        };

                        let input_ident = ident_stack.pop().unwrap();

                        let (HydroNode::Reduce { f, .. } | HydroNode::ReduceKeyed { f, .. }) = &*node
                        else {
                            unreachable!()
                        };

                        let reduce_ident =
                            syn::Ident::new(&format!("stream_{}", *next_stmt_id), Span::call_site());

                        match builders_or_callback {
                            BuildersOrCallback::Builders(graph_builders) => {
                                if matches!(node, HydroNode::Reduce { .. })
                                    && node.metadata().location_id.is_top_level()
                                    && !(matches!(node.metadata().location_id, LocationId::Atomic(_)))
                                    && graph_builders.singleton_intermediates()
                                    && !node.metadata().collection_kind.is_bounded()
                                {
                                    todo!(
                                        "Reduce with optional intermediates is not yet supported in simulator"
                                    );
                                } else if matches!(node, HydroNode::ReduceKeyed { .. })
                                    && node.metadata().location_id.is_top_level()
                                    && !(matches!(node.metadata().location_id, LocationId::Atomic(_)))
                                    && graph_builders.singleton_intermediates()
                                    && !node.metadata().collection_kind.is_bounded()
                                {
                                    todo!(
                                        "Reduce keyed with optional intermediates is not yet supported in simulator"
                                    );
                                } else {
                                    let builder = graph_builders.get_dfir_mut(&out_location);
                                    builder.add_dfir(
                                        parse_quote! {
                                            #reduce_ident = #input_ident -> #operator::<#lifetime>(#f);
                                        },
                                        None,
                                        Some(&next_stmt_id.to_string()),
                                    );
                                }
                            }
                            BuildersOrCallback::Callback(_, node_callback) => {
                                node_callback(node, next_stmt_id);
                            }
                        }

                        *next_stmt_id += 1;

                        ident_stack.push(reduce_ident);
                    }

                    HydroNode::ReduceKeyedWatermark {
                        f,
                        input,
                        metadata,
                        ..
                    } => {
                        let lifetime = if input.metadata().location_id.is_top_level() {
                            quote!('static)
                        } else {
                            quote!('tick)
                        };

                        // watermark is processed second, so it's on top
                        let watermark_ident = ident_stack.pop().unwrap();
                        let input_ident = ident_stack.pop().unwrap();

                        let chain_ident = syn::Ident::new(
                            &format!("reduce_keyed_watermark_chain_{}", *next_stmt_id),
                            Span::call_site(),
                        );

                        let fold_ident =
                            syn::Ident::new(&format!("stream_{}", *next_stmt_id), Span::call_site());

                        let agg_operator: syn::Ident = if input.metadata().location_id.is_top_level()
                            && input.metadata().collection_kind.is_bounded()
                        {
                            parse_quote!(fold_no_replay)
                        } else {
                            parse_quote!(fold)
                        };

                        match builders_or_callback {
                            BuildersOrCallback::Builders(graph_builders) => {
                                if metadata.location_id.is_top_level()
                                    && !(matches!(metadata.location_id, LocationId::Atomic(_)))
                                    && graph_builders.singleton_intermediates()
                                    && !metadata.collection_kind.is_bounded()
                                {
                                    todo!(
                                        "Reduce keyed watermarked on a top-level bounded collection is not yet supported"
                                    )
                                } else {
                                    let builder = graph_builders.get_dfir_mut(&out_location);
                                    builder.add_dfir(
                                        parse_quote! {
                                            #chain_ident = chain();
                                            #input_ident
                                                -> map(|x| (Some(x), None))
                                                -> [0]#chain_ident;
                                            #watermark_ident
                                                -> map(|watermark| (None, Some(watermark)))
                                                -> [1]#chain_ident;

                                            #fold_ident = #chain_ident
                                                -> #agg_operator::<#lifetime>(|| (::std::collections::HashMap::new(), None), {
                                                    let __reduce_keyed_fn = #f;
                                                    move |(map, opt_curr_watermark), (opt_payload, opt_watermark)| {
                                                        if let Some((k, v)) = opt_payload {
                                                            if let Some(curr_watermark) = *opt_curr_watermark {
                                                                if k < curr_watermark {
                                                                    return;
                                                                }
                                                            }
                                                            match map.entry(k) {
                                                                ::std::collections::hash_map::Entry::Vacant(e) => {
                                                                    e.insert(v);
                                                                }
                                                                ::std::collections::hash_map::Entry::Occupied(mut e) => {
                                                                    __reduce_keyed_fn(e.get_mut(), v);
                                                                }
                                                            }
                                                        } else {
                                                            let watermark = opt_watermark.unwrap();
                                                            if let Some(curr_watermark) = *opt_curr_watermark {
                                                                if watermark <= curr_watermark {
                                                                    return;
                                                                }
                                                            }
                                                            *opt_curr_watermark = opt_watermark;
                                                            map.retain(|k, _| *k >= watermark);
                                                        }
                                                    }
                                                })
                                                -> flat_map(|(map, _curr_watermark)| map);
                                        },
                                        None,
                                        Some(&next_stmt_id.to_string()),
                                    );
                                }
                            }
                            BuildersOrCallback::Callback(_, node_callback) => {
                                node_callback(node, next_stmt_id);
                            }
                        }

                        *next_stmt_id += 1;

                        ident_stack.push(fold_ident);
                    }

                    HydroNode::Network {
                        networking_info,
                        serialize_fn: serialize_pipeline,
                        instantiate_fn,
                        deserialize_fn: deserialize_pipeline,
                        input,
                        ..
                    } => {
                        let input_ident = ident_stack.pop().unwrap();

                        let receiver_stream_ident =
                            syn::Ident::new(&format!("stream_{}", *next_stmt_id), Span::call_site());

                        match builders_or_callback {
                            BuildersOrCallback::Builders(graph_builders) => {
                                let (sink_expr, source_expr) = match instantiate_fn {
                                    DebugInstantiate::Building => (
                                        syn::parse_quote!(DUMMY_SINK),
                                        syn::parse_quote!(DUMMY_SOURCE),
                                    ),

                                    DebugInstantiate::Finalized(finalized) => {
                                        (finalized.sink.clone(), finalized.source.clone())
                                    }
                                };

                                graph_builders.create_network(
                                    &input.metadata().location_id,
                                    &out_location,
                                    input_ident,
                                    &receiver_stream_ident,
                                    serialize_pipeline.as_ref(),
                                    sink_expr,
                                    source_expr,
                                    deserialize_pipeline.as_ref(),
                                    *next_stmt_id,
                                    networking_info,
                                );
                            }
                            BuildersOrCallback::Callback(_, node_callback) => {
                                node_callback(node, next_stmt_id);
                            }
                        }

                        *next_stmt_id += 1;

                        ident_stack.push(receiver_stream_ident);
                    }

                    HydroNode::ExternalInput {
                        instantiate_fn,
                        deserialize_fn: deserialize_pipeline,
                        ..
                    } => {
                        let receiver_stream_ident =
                            syn::Ident::new(&format!("stream_{}", *next_stmt_id), Span::call_site());

                        match builders_or_callback {
                            BuildersOrCallback::Builders(graph_builders) => {
                                let (_, source_expr) = match instantiate_fn {
                                    DebugInstantiate::Building => (
                                        syn::parse_quote!(DUMMY_SINK),
                                        syn::parse_quote!(DUMMY_SOURCE),
                                    ),

                                    DebugInstantiate::Finalized(finalized) => {
                                        (finalized.sink.clone(), finalized.source.clone())
                                    }
                                };

                                graph_builders.create_external_source(
                                    &out_location,
                                    source_expr,
                                    &receiver_stream_ident,
                                    deserialize_pipeline.as_ref(),
                                    *next_stmt_id,
                                );
                            }
                            BuildersOrCallback::Callback(_, node_callback) => {
                                node_callback(node, next_stmt_id);
                            }
                        }

                        *next_stmt_id += 1;

                        ident_stack.push(receiver_stream_ident);
                    }

                    HydroNode::Counter {
                        tag,
                        duration,
                        prefix,
                        ..
                    } => {
                        let input_ident = ident_stack.pop().unwrap();

                        let counter_ident =
                            syn::Ident::new(&format!("stream_{}", *next_stmt_id), Span::call_site());

                        match builders_or_callback {
                            BuildersOrCallback::Builders(graph_builders) => {
                                let arg = format!("{}({})", prefix, tag);
                                let builder = graph_builders.get_dfir_mut(&out_location);
                                builder.add_dfir(
                                    parse_quote! {
                                        #counter_ident = #input_ident -> _counter(#arg, #duration);
                                    },
                                    None,
                                    Some(&next_stmt_id.to_string()),
                                );
                            }
                            BuildersOrCallback::Callback(_, node_callback) => {
                                node_callback(node, next_stmt_id);
                            }
                        }

                        *next_stmt_id += 1;

                        ident_stack.push(counter_ident);
                    }
                }
            },
            seen_tees,
            false,
        );

        ident_stack
            .pop()
            .expect("ident_stack should have exactly one element after traversal")
    }

    pub fn visit_debug_expr(&mut self, mut transform: impl FnMut(&mut DebugExpr)) {
        match self {
            HydroNode::Placeholder => {
                panic!()
            }
            HydroNode::Cast { .. } | HydroNode::ObserveNonDet { .. } => {}
            HydroNode::Source { source, .. } => match source {
                HydroSource::Stream(expr) | HydroSource::Iter(expr) => transform(expr),
                HydroSource::ExternalNetwork()
                | HydroSource::Spin()
                | HydroSource::ClusterMembers(_, _)
                | HydroSource::Embedded(_)
                | HydroSource::EmbeddedSingleton(_) => {} // TODO: what goes here?
            },
            HydroNode::SingletonSource { value, .. } => {
                transform(value);
            }
            HydroNode::CycleSource { .. }
            | HydroNode::Tee { .. }
            | HydroNode::YieldConcat { .. }
            | HydroNode::BeginAtomic { .. }
            | HydroNode::EndAtomic { .. }
            | HydroNode::Batch { .. }
            | HydroNode::Chain { .. }
            | HydroNode::ChainFirst { .. }
            | HydroNode::CrossProduct { .. }
            | HydroNode::CrossSingleton { .. }
            | HydroNode::ResolveFutures { .. }
            | HydroNode::ResolveFuturesOrdered { .. }
            | HydroNode::Join { .. }
            | HydroNode::Difference { .. }
            | HydroNode::AntiJoin { .. }
            | HydroNode::DeferTick { .. }
            | HydroNode::Enumerate { .. }
            | HydroNode::Unique { .. }
            | HydroNode::Sort { .. } => {}
            HydroNode::Map { f, .. }
            | HydroNode::FlatMap { f, .. }
            | HydroNode::Filter { f, .. }
            | HydroNode::FilterMap { f, .. }
            | HydroNode::Inspect { f, .. }
            | HydroNode::Partition { f, .. }
            | HydroNode::Reduce { f, .. }
            | HydroNode::ReduceKeyed { f, .. }
            | HydroNode::ReduceKeyedWatermark { f, .. } => {
                transform(f);
            }
            HydroNode::Fold { init, acc, .. }
            | HydroNode::Scan { init, acc, .. }
            | HydroNode::FoldKeyed { init, acc, .. } => {
                transform(init);
                transform(acc);
            }
            HydroNode::Network {
                serialize_fn,
                deserialize_fn,
                ..
            } => {
                if let Some(serialize_fn) = serialize_fn {
                    transform(serialize_fn);
                }
                if let Some(deserialize_fn) = deserialize_fn {
                    transform(deserialize_fn);
                }
            }
            HydroNode::ExternalInput { deserialize_fn, .. } => {
                if let Some(deserialize_fn) = deserialize_fn {
                    transform(deserialize_fn);
                }
            }
            HydroNode::Counter { duration, .. } => {
                transform(duration);
            }
        }
    }

    pub fn op_metadata(&self) -> &HydroIrOpMetadata {
        &self.metadata().op
    }

    pub fn metadata(&self) -> &HydroIrMetadata {
        match self {
            HydroNode::Placeholder => {
                panic!()
            }
            HydroNode::Cast { metadata, .. } => metadata,
            HydroNode::ObserveNonDet { metadata, .. } => metadata,
            HydroNode::Source { metadata, .. } => metadata,
            HydroNode::SingletonSource { metadata, .. } => metadata,
            HydroNode::CycleSource { metadata, .. } => metadata,
            HydroNode::Tee { metadata, .. } => metadata,
            HydroNode::Partition { metadata, .. } => metadata,
            HydroNode::YieldConcat { metadata, .. } => metadata,
            HydroNode::BeginAtomic { metadata, .. } => metadata,
            HydroNode::EndAtomic { metadata, .. } => metadata,
            HydroNode::Batch { metadata, .. } => metadata,
            HydroNode::Chain { metadata, .. } => metadata,
            HydroNode::ChainFirst { metadata, .. } => metadata,
            HydroNode::CrossProduct { metadata, .. } => metadata,
            HydroNode::CrossSingleton { metadata, .. } => metadata,
            HydroNode::Join { metadata, .. } => metadata,
            HydroNode::Difference { metadata, .. } => metadata,
            HydroNode::AntiJoin { metadata, .. } => metadata,
            HydroNode::ResolveFutures { metadata, .. } => metadata,
            HydroNode::ResolveFuturesOrdered { metadata, .. } => metadata,
            HydroNode::Map { metadata, .. } => metadata,
            HydroNode::FlatMap { metadata, .. } => metadata,
            HydroNode::Filter { metadata, .. } => metadata,
            HydroNode::FilterMap { metadata, .. } => metadata,
            HydroNode::DeferTick { metadata, .. } => metadata,
            HydroNode::Enumerate { metadata, .. } => metadata,
            HydroNode::Inspect { metadata, .. } => metadata,
            HydroNode::Unique { metadata, .. } => metadata,
            HydroNode::Sort { metadata, .. } => metadata,
            HydroNode::Scan { metadata, .. } => metadata,
            HydroNode::Fold { metadata, .. } => metadata,
            HydroNode::FoldKeyed { metadata, .. } => metadata,
            HydroNode::Reduce { metadata, .. } => metadata,
            HydroNode::ReduceKeyed { metadata, .. } => metadata,
            HydroNode::ReduceKeyedWatermark { metadata, .. } => metadata,
            HydroNode::ExternalInput { metadata, .. } => metadata,
            HydroNode::Network { metadata, .. } => metadata,
            HydroNode::Counter { metadata, .. } => metadata,
        }
    }

    pub fn op_metadata_mut(&mut self) -> &mut HydroIrOpMetadata {
        &mut self.metadata_mut().op
    }

    pub fn metadata_mut(&mut self) -> &mut HydroIrMetadata {
        match self {
            HydroNode::Placeholder => {
                panic!()
            }
            HydroNode::Cast { metadata, .. } => metadata,
            HydroNode::ObserveNonDet { metadata, .. } => metadata,
            HydroNode::Source { metadata, .. } => metadata,
            HydroNode::SingletonSource { metadata, .. } => metadata,
            HydroNode::CycleSource { metadata, .. } => metadata,
            HydroNode::Tee { metadata, .. } => metadata,
            HydroNode::Partition { metadata, .. } => metadata,
            HydroNode::YieldConcat { metadata, .. } => metadata,
            HydroNode::BeginAtomic { metadata, .. } => metadata,
            HydroNode::EndAtomic { metadata, .. } => metadata,
            HydroNode::Batch { metadata, .. } => metadata,
            HydroNode::Chain { metadata, .. } => metadata,
            HydroNode::ChainFirst { metadata, .. } => metadata,
            HydroNode::CrossProduct { metadata, .. } => metadata,
            HydroNode::CrossSingleton { metadata, .. } => metadata,
            HydroNode::Join { metadata, .. } => metadata,
            HydroNode::Difference { metadata, .. } => metadata,
            HydroNode::AntiJoin { metadata, .. } => metadata,
            HydroNode::ResolveFutures { metadata, .. } => metadata,
            HydroNode::ResolveFuturesOrdered { metadata, .. } => metadata,
            HydroNode::Map { metadata, .. } => metadata,
            HydroNode::FlatMap { metadata, .. } => metadata,
            HydroNode::Filter { metadata, .. } => metadata,
            HydroNode::FilterMap { metadata, .. } => metadata,
            HydroNode::DeferTick { metadata, .. } => metadata,
            HydroNode::Enumerate { metadata, .. } => metadata,
            HydroNode::Inspect { metadata, .. } => metadata,
            HydroNode::Unique { metadata, .. } => metadata,
            HydroNode::Sort { metadata, .. } => metadata,
            HydroNode::Scan { metadata, .. } => metadata,
            HydroNode::Fold { metadata, .. } => metadata,
            HydroNode::FoldKeyed { metadata, .. } => metadata,
            HydroNode::Reduce { metadata, .. } => metadata,
            HydroNode::ReduceKeyed { metadata, .. } => metadata,
            HydroNode::ReduceKeyedWatermark { metadata, .. } => metadata,
            HydroNode::ExternalInput { metadata, .. } => metadata,
            HydroNode::Network { metadata, .. } => metadata,
            HydroNode::Counter { metadata, .. } => metadata,
        }
    }

    pub fn input(&self) -> Vec<&HydroNode> {
        match self {
            HydroNode::Placeholder => {
                panic!()
            }
            HydroNode::Source { .. }
            | HydroNode::SingletonSource { .. }
            | HydroNode::ExternalInput { .. }
            | HydroNode::CycleSource { .. }
            | HydroNode::Tee { .. }
            | HydroNode::Partition { .. } => {
                // Tee/Partition should find their input in separate special ways
                vec![]
            }
            HydroNode::Cast { inner, .. }
            | HydroNode::ObserveNonDet { inner, .. }
            | HydroNode::YieldConcat { inner, .. }
            | HydroNode::BeginAtomic { inner, .. }
            | HydroNode::EndAtomic { inner, .. }
            | HydroNode::Batch { inner, .. } => {
                vec![inner]
            }
            HydroNode::Chain { first, second, .. } => {
                vec![first, second]
            }
            HydroNode::ChainFirst { first, second, .. } => {
                vec![first, second]
            }
            HydroNode::CrossProduct { left, right, .. }
            | HydroNode::CrossSingleton { left, right, .. }
            | HydroNode::Join { left, right, .. } => {
                vec![left, right]
            }
            HydroNode::Difference { pos, neg, .. } | HydroNode::AntiJoin { pos, neg, .. } => {
                vec![pos, neg]
            }
            HydroNode::Map { input, .. }
            | HydroNode::FlatMap { input, .. }
            | HydroNode::Filter { input, .. }
            | HydroNode::FilterMap { input, .. }
            | HydroNode::Sort { input, .. }
            | HydroNode::DeferTick { input, .. }
            | HydroNode::Enumerate { input, .. }
            | HydroNode::Inspect { input, .. }
            | HydroNode::Unique { input, .. }
            | HydroNode::Network { input, .. }
            | HydroNode::Counter { input, .. }
            | HydroNode::ResolveFutures { input, .. }
            | HydroNode::ResolveFuturesOrdered { input, .. }
            | HydroNode::Fold { input, .. }
            | HydroNode::FoldKeyed { input, .. }
            | HydroNode::Reduce { input, .. }
            | HydroNode::ReduceKeyed { input, .. }
            | HydroNode::Scan { input, .. } => {
                vec![input]
            }
            HydroNode::ReduceKeyedWatermark {
                input, watermark, ..
            } => {
                vec![input, watermark]
            }
        }
    }

    pub fn input_metadata(&self) -> Vec<&HydroIrMetadata> {
        self.input()
            .iter()
            .map(|input_node| input_node.metadata())
            .collect()
    }

    /// Returns `true` if this node is a Tee or Partition whose inner Rc
    /// has other live references, meaning the upstream is already driven
    /// by another consumer and does not need a Null sink.
    pub fn is_shared_with_others(&self) -> bool {
        match self {
            HydroNode::Tee { inner, .. } | HydroNode::Partition { inner, .. } => {
                Rc::strong_count(&inner.0) > 1
            }
            _ => false,
        }
    }

    pub fn print_root(&self) -> String {
        match self {
            HydroNode::Placeholder => {
                panic!()
            }
            HydroNode::Cast { .. } => "Cast()".to_owned(),
            HydroNode::ObserveNonDet { .. } => "ObserveNonDet()".to_owned(),
            HydroNode::Source { source, .. } => format!("Source({:?})", source),
            HydroNode::SingletonSource {
                value,
                first_tick_only,
                ..
            } => format!(
                "SingletonSource({:?}, first_tick_only={})",
                value, first_tick_only
            ),
            HydroNode::CycleSource { cycle_id, .. } => format!("CycleSource({})", cycle_id),
            HydroNode::Tee { inner, .. } => format!("Tee({})", inner.0.borrow().print_root()),
            HydroNode::Partition { f, is_true, .. } => {
                format!("Partition({:?}, is_true={})", f, is_true)
            }
            HydroNode::YieldConcat { .. } => "YieldConcat()".to_owned(),
            HydroNode::BeginAtomic { .. } => "BeginAtomic()".to_owned(),
            HydroNode::EndAtomic { .. } => "EndAtomic()".to_owned(),
            HydroNode::Batch { .. } => "Batch()".to_owned(),
            HydroNode::Chain { first, second, .. } => {
                format!("Chain({}, {})", first.print_root(), second.print_root())
            }
            HydroNode::ChainFirst { first, second, .. } => {
                format!(
                    "ChainFirst({}, {})",
                    first.print_root(),
                    second.print_root()
                )
            }
            HydroNode::CrossProduct { left, right, .. } => {
                format!(
                    "CrossProduct({}, {})",
                    left.print_root(),
                    right.print_root()
                )
            }
            HydroNode::CrossSingleton { left, right, .. } => {
                format!(
                    "CrossSingleton({}, {})",
                    left.print_root(),
                    right.print_root()
                )
            }
            HydroNode::Join { left, right, .. } => {
                format!("Join({}, {})", left.print_root(), right.print_root())
            }
            HydroNode::Difference { pos, neg, .. } => {
                format!("Difference({}, {})", pos.print_root(), neg.print_root())
            }
            HydroNode::AntiJoin { pos, neg, .. } => {
                format!("AntiJoin({}, {})", pos.print_root(), neg.print_root())
            }
            HydroNode::ResolveFutures { .. } => "ResolveFutures()".to_owned(),
            HydroNode::ResolveFuturesOrdered { .. } => "ResolveFuturesOrdered()".to_owned(),
            HydroNode::Map { f, .. } => format!("Map({:?})", f),
            HydroNode::FlatMap { f, .. } => format!("FlatMap({:?})", f),
            HydroNode::Filter { f, .. } => format!("Filter({:?})", f),
            HydroNode::FilterMap { f, .. } => format!("FilterMap({:?})", f),
            HydroNode::DeferTick { .. } => "DeferTick()".to_owned(),
            HydroNode::Enumerate { .. } => "Enumerate()".to_owned(),
            HydroNode::Inspect { f, .. } => format!("Inspect({:?})", f),
            HydroNode::Unique { .. } => "Unique()".to_owned(),
            HydroNode::Sort { .. } => "Sort()".to_owned(),
            HydroNode::Fold { init, acc, .. } => format!("Fold({:?}, {:?})", init, acc),
            HydroNode::Scan { init, acc, .. } => format!("Scan({:?}, {:?})", init, acc),
            HydroNode::FoldKeyed { init, acc, .. } => format!("FoldKeyed({:?}, {:?})", init, acc),
            HydroNode::Reduce { f, .. } => format!("Reduce({:?})", f),
            HydroNode::ReduceKeyed { f, .. } => format!("ReduceKeyed({:?})", f),
            HydroNode::ReduceKeyedWatermark { f, .. } => format!("ReduceKeyedWatermark({:?})", f),
            HydroNode::Network { .. } => "Network()".to_owned(),
            HydroNode::ExternalInput { .. } => "ExternalInput()".to_owned(),
            HydroNode::Counter { tag, duration, .. } => {
                format!("Counter({:?}, {:?})", tag, duration)
            }
        }
    }
}

#[cfg(feature = "build")]
fn instantiate_network<'a, D>(
    env: &mut D::InstantiateEnv,
    from_location: &LocationId,
    to_location: &LocationId,
    processes: &SparseSecondaryMap<LocationKey, D::Process>,
    clusters: &SparseSecondaryMap<LocationKey, D::Cluster>,
    name: Option<&str>,
    networking_info: &crate::networking::NetworkingInfo,
) -> (syn::Expr, syn::Expr, Box<dyn FnOnce()>)
where
    D: Deploy<'a>,
{
    let ((sink, source), connect_fn) = match (from_location, to_location) {
        (&LocationId::Process(from), &LocationId::Process(to)) => {
            let from_node = processes
                .get(from)
                .unwrap_or_else(|| {
                    panic!("A process used in the graph was not instantiated: {}", from)
                })
                .clone();
            let to_node = processes
                .get(to)
                .unwrap_or_else(|| {
                    panic!("A process used in the graph was not instantiated: {}", to)
                })
                .clone();

            let sink_port = from_node.next_port();
            let source_port = to_node.next_port();

            (
                D::o2o_sink_source(
                    env,
                    &from_node,
                    &sink_port,
                    &to_node,
                    &source_port,
                    name,
                    networking_info,
                ),
                D::o2o_connect(&from_node, &sink_port, &to_node, &source_port),
            )
        }
        (&LocationId::Process(from), &LocationId::Cluster(to)) => {
            let from_node = processes
                .get(from)
                .unwrap_or_else(|| {
                    panic!("A process used in the graph was not instantiated: {}", from)
                })
                .clone();
            let to_node = clusters
                .get(to)
                .unwrap_or_else(|| {
                    panic!("A cluster used in the graph was not instantiated: {}", to)
                })
                .clone();

            let sink_port = from_node.next_port();
            let source_port = to_node.next_port();

            (
                D::o2m_sink_source(
                    env,
                    &from_node,
                    &sink_port,
                    &to_node,
                    &source_port,
                    name,
                    networking_info,
                ),
                D::o2m_connect(&from_node, &sink_port, &to_node, &source_port),
            )
        }
        (&LocationId::Cluster(from), &LocationId::Process(to)) => {
            let from_node = clusters
                .get(from)
                .unwrap_or_else(|| {
                    panic!("A cluster used in the graph was not instantiated: {}", from)
                })
                .clone();
            let to_node = processes
                .get(to)
                .unwrap_or_else(|| {
                    panic!("A process used in the graph was not instantiated: {}", to)
                })
                .clone();

            let sink_port = from_node.next_port();
            let source_port = to_node.next_port();

            (
                D::m2o_sink_source(
                    env,
                    &from_node,
                    &sink_port,
                    &to_node,
                    &source_port,
                    name,
                    networking_info,
                ),
                D::m2o_connect(&from_node, &sink_port, &to_node, &source_port),
            )
        }
        (&LocationId::Cluster(from), &LocationId::Cluster(to)) => {
            let from_node = clusters
                .get(from)
                .unwrap_or_else(|| {
                    panic!("A cluster used in the graph was not instantiated: {}", from)
                })
                .clone();
            let to_node = clusters
                .get(to)
                .unwrap_or_else(|| {
                    panic!("A cluster used in the graph was not instantiated: {}", to)
                })
                .clone();

            let sink_port = from_node.next_port();
            let source_port = to_node.next_port();

            (
                D::m2m_sink_source(
                    env,
                    &from_node,
                    &sink_port,
                    &to_node,
                    &source_port,
                    name,
                    networking_info,
                ),
                D::m2m_connect(&from_node, &sink_port, &to_node, &source_port),
            )
        }
        (LocationId::Tick(_, _), _) => panic!(),
        (_, LocationId::Tick(_, _)) => panic!(),
        (LocationId::Atomic(_), _) => panic!(),
        (_, LocationId::Atomic(_)) => panic!(),
    };
    (sink, source, connect_fn)
}

#[cfg(test)]
mod test {
    use std::mem::size_of;

    use stageleft::{QuotedWithContext, q};

    use super::*;

    #[test]
    #[cfg_attr(
        not(feature = "build"),
        ignore = "expects inclusion of feature-gated fields"
    )]
    fn hydro_node_size() {
        assert_eq!(size_of::<HydroNode>(), 248);
    }

    #[test]
    #[cfg_attr(
        not(feature = "build"),
        ignore = "expects inclusion of feature-gated fields"
    )]
    fn hydro_root_size() {
        assert_eq!(size_of::<HydroRoot>(), 136);
    }

    #[test]
    fn test_simplify_q_macro_basic() {
        // Test basic non-q! expression
        let simple_expr: syn::Expr = syn::parse_str("x + y").unwrap();
        let result = simplify_q_macro(simple_expr.clone());
        assert_eq!(result, simple_expr);
    }

    #[test]
    fn test_simplify_q_macro_actual_stageleft_call() {
        // Test a simplified version of what a real stageleft call might look like
        let stageleft_call = q!(|x: usize| x + 1).splice_fn1_ctx(&());
        let result = simplify_q_macro(stageleft_call);
        // This should be processed by our visitor and simplified to q!(...)
        // since we detect the stageleft::runtime_support::fn_* pattern
        hydro_build_utils::assert_snapshot!(result.to_token_stream().to_string());
    }

    #[test]
    fn test_closure_no_pipe_at_start() {
        // Test a closure that does not start with a pipe
        let stageleft_call = q!({
            let foo = 123;
            move |b: usize| b + foo
        })
        .splice_fn1_ctx(&());
        let result = simplify_q_macro(stageleft_call);
        hydro_build_utils::assert_snapshot!(result.to_token_stream().to_string());
    }
}