hydro_lang/compile/ir/

mod.rs

use core::panic;
use std::cell::{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;

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

pub mod backtrace;
use backtrace::Backtrace;

/// A closure expression bundled with any singleton references it captures.
///
/// When a `q!()` closure captures a `SingletonRef`, the reference is recorded here
/// alongside the closure's expression. This allows per-closure tracking of singleton
/// captures, which is important for nodes with multiple closures (e.g. Fold has `init` and `acc`).
pub struct ClosureExpr {
    pub(crate) expr: DebugExpr,
    /// Each entry is `(HydroNode::Reference, is_mut: bool)`.
    /// The index in the Vec determines the ident name via [`handoff_ref_ident`].
    /// The `access_counter` was assigned at staging time in code order.
    pub(crate) singleton_refs: Vec<(HydroNode, bool)>,
}

impl Clone for ClosureExpr {
    fn clone(&self) -> Self {
        Self {
            expr: self.expr.clone(),
            singleton_refs: self
                .singleton_refs
                .iter()
                .map(|(node, is_mut)| {
                    let HydroNode::Reference {
                        inner,
                        kind,
                        access_counter,
                        metadata,
                    } = node
                    else {
                        panic!("singleton_refs should only contain HydroNode::Reference");
                    };
                    (
                        HydroNode::Reference {
                            inner: SharedNode(Rc::clone(&inner.0)),
                            kind: *kind,
                            access_counter: access_counter.freeze(),
                            metadata: metadata.clone(),
                        },
                        *is_mut,
                    )
                })
                .collect(),
        }
    }
}

impl Hash for ClosureExpr {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.expr.hash(state);
        // singleton_refs are structural children (like HydroIrMetadata), not
        // identity-defining. Two closures with the same expr but different
        // captured refs are the same closure text — the refs only affect codegen.
    }
}

impl serde::Serialize for ClosureExpr {
    fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
        use serde::ser::SerializeStruct;
        let mut s = serializer.serialize_struct("ClosureExpr", 2)?;
        s.serialize_field("expr", &self.expr)?;
        s.serialize_field(
            "singleton_refs",
            &SerializableSingletonRefs(&self.singleton_refs),
        )?;
        s.end()
    }
}

struct SerializableSingletonRefs<'a>(&'a [(HydroNode, bool)]);

impl serde::Serialize for SerializableSingletonRefs<'_> {
    fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
        use serde::ser::SerializeSeq;
        let mut seq = serializer.serialize_seq(Some(self.0.len()))?;
        for (node, is_mut) in self.0.iter() {
            seq.serialize_element(&(node, is_mut))?;
        }
        seq.end()
    }
}

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

impl Display for ClosureExpr {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        Display::fmt(&self.expr, f)
    }
}

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

impl From<DebugExpr> for ClosureExpr {
    fn from(expr: DebugExpr) -> Self {
        Self {
            expr,
            singleton_refs: Vec::new(),
        }
    }
}

impl ClosureExpr {
    pub fn new(expr: DebugExpr, singleton_refs: Vec<(HydroNode, bool)>) -> Self {
        Self {
            expr,
            singleton_refs,
        }
    }

    pub fn has_mut_ref(&self) -> bool {
        self.singleton_refs.iter().any(|(_, is_mut)| *is_mut)
    }

    pub fn deep_clone(&self, seen_tees: &mut SeenSharedNodes) -> Self {
        Self {
            expr: self.expr.clone(),
            singleton_refs: self
                .singleton_refs
                .iter()
                .map(|(node, is_mut)| (node.deep_clone(seen_tees), *is_mut))
                .collect(),
        }
    }

    pub fn transform_children(
        &mut self,
        transform: &mut impl FnMut(&mut HydroNode, &mut SeenSharedNodes),
        seen_tees: &mut SeenSharedNodes,
    ) {
        for (ref_node, _is_mut) in self.singleton_refs.iter_mut() {
            transform(ref_node, seen_tees);
        }
    }

    /// Pop singleton ref idents from the stack and rewrite the closure's token stream,
    /// replacing local singleton ref idents with `#{N} dfir_ident` or `#{N} mut dfir_ident` references.
    #[cfg(feature = "build")]
    pub fn emit_tokens(&self, ident_stack: &mut Vec<syn::Ident>) -> TokenStream {
        if self.singleton_refs.is_empty() {
            self.expr.0.to_token_stream()
        } else {
            assert!(
                ident_stack.len() >= self.singleton_refs.len(),
                "ident_stack has {} entries but expected at least {} for singleton_refs",
                ident_stack.len(),
                self.singleton_refs.len()
            );
            let ref_idents = ident_stack.drain(ident_stack.len() - self.singleton_refs.len()..);

            let mut let_bindings = Vec::new();
            for ((i, (ref_node, is_mut)), ref_ident) in
                self.singleton_refs.iter().enumerate().zip(ref_idents)
            {
                let HydroNode::Reference { access_counter, .. } = ref_node else {
                    panic!("ClosureExpression expected references to `HydroNode::Reference`");
                };
                let group = access_counter.frozen_group();
                // TODO(mingwei): proper spanning?
                let local_ident = handoff_ref_ident(i);
                let hash = proc_macro2::Punct::new('#', proc_macro2::Spacing::Alone);
                let group_lit = proc_macro2::Literal::u32_unsuffixed(group);
                let mut_token = is_mut.then(|| quote!(mut));
                let binding = quote! {
                    let #local_ident = #hash {#group_lit} #mut_token #ref_ident;
                };
                let_bindings.push(binding);
            }

            let expr = &self.expr.0;
            quote! {
                {
                    #( #let_bindings )*
                    #expr
                }
            }
        }
    }
}

/// 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 serde::Serialize for DebugExpr {
    fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
        serializer.serialize_str(&self.to_string())
    }
}

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(expr: syn::Expr) -> syn::Expr {
    if let syn::Expr::Call(ref call) = expr && let syn::Expr::Path(path_expr) = call.func.as_ref()
        // Look for calls to stageleft::runtime_support::fn*
        && is_stageleft_runtime_support_call(&path_expr.path)
        && let syn::Expr::Block(b) = &call.args[0]
        && b.block.stmts.len() == 3
        && let Some(syn::Stmt::Expr(e, _)) = b.block.stmts.get(2)
    // skip the first two, which are imports
    {
        let mut e = e.clone();
        while let syn::Expr::Block(ref mut block) = e
            && block.block.stmts.len() == 1
            && let syn::Stmt::Expr(inner_e, _) = block.block.stmts.remove(0)
        {
            e = inner_e;
        }

        e
    } else {
        expr
    }
}

fn is_stageleft_runtime_support_call(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();
        path.segments.len() > 2
            && path.segments[0].ident == "stageleft"
            && path.segments[1].ident == "runtime_support"
            && fn_name.contains("_type_hint")
    } else {
        false
    }
}

/// 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())
    }
}

impl serde::Serialize for DebugType {
    fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
        serializer.serialize_str(&format!("{}", self.0.to_token_stream()))
    }
}

fn serialize_backtrace_as_span<S: serde::Serializer>(
    backtrace: &Backtrace,
    serializer: S,
) -> Result<S::Ok, S::Error> {
    match backtrace.format_span() {
        Some(span) => serializer.serialize_some(&span),
        None => serializer.serialize_none(),
    }
}

fn serialize_ident<S: serde::Serializer>(
    ident: &syn::Ident,
    serializer: S,
) -> Result<S::Ok, S::Error> {
    serializer.serialize_str(&ident.to_string())
}

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

impl serde::Serialize for DebugInstantiate {
    fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
        match self {
            DebugInstantiate::Building => {
                serializer.serialize_unit_variant("DebugInstantiate", 0, "Building")
            }
            DebugInstantiate::Finalized(_) => {
                panic!(
                    "cannot serialize DebugInstantiate::Finalized: contains non-serializable runtime state (closures)"
                )
            }
        }
    }
}

#[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, serde::Serialize)]
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, serde::Serialize)]
pub enum HydroSource {
    Stream(DebugExpr),
    ExternalNetwork(),
    Iter(DebugExpr),
    Spin(),
    ClusterMembers(LocationId, ClusterMembersState),
    Embedded(#[serde(serialize_with = "serialize_ident")] syn::Ident),
    EmbeddedSingleton(#[serde(serialize_with = "serialize_ident")] 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;

    /// Adds the DFIR statements to the graph for the given location.
    ///
    /// The location determines which DFIR graph the statements are placed in (for production,
    /// the graph of the location's root; for simulation, either the fused async graph or the
    /// tick's separate graph). In the future (#2902), production codegen will also use the
    /// location to place tick-located statements inside the tick's `loop { ... }` context.
    fn add_dfir_at(
        &mut self,
        location: &LocationId,
        dfir: dfir_lang::parse::DfirCode,
        operator_tag: Option<&str>,
    );

    /// Emits a source operator (e.g. `source_iter(...)`) that logically lives inside a tick.
    ///
    /// * `source_rhs` is the source pipeline (the right-hand side of the `=` assignment, e.g.
    ///   `source_iter([123])`).
    /// * `replay_each_tick` selects whether the value is re-emitted on every tick/firing
    ///   (persisted) or delivered only on the first tick/firing.
    ///
    /// The default (simulation) emission places the source directly in the tick's own graph.
    /// Production overrides this since DFIR sources must be at the root level, outside of the
    /// tick's `loop { ... }` context: it emits the source at the root and windows it into the
    /// loop.
    fn add_tick_source(
        &mut self,
        location: &LocationId,
        source_rhs: TokenStream,
        out_ident: &syn::Ident,
        replay_each_tick: bool,
        operator_tag: Option<&str>,
    );

    /// The DFIR persistence lifetime for operator state scoped to a single tick, for an operator
    /// at `op_location`.
    ///
    /// Returns `'tick`. In the future (#2902), production codegen will emit tick regions as DFIR
    /// `loop { ... }` blocks, where this must instead be `'none` when `op_location` is a tick.
    fn tick_state_lifetime(&self, _op_location: &LocationId) -> TokenStream {
        quote!('tick)
    }

    /// The DFIR persistence lifetime for operator state that accumulates across ticks, for an
    /// operator at `op_location`.
    ///
    /// Returns `'static`. In the future (#2902), production codegen will emit tick regions as
    /// DFIR `loop { ... }` blocks, where this must instead be `'loop` when `op_location` is a
    /// tick.
    fn cross_tick_state_lifetime(&self, _op_location: &LocationId) -> TokenStream {
        quote!('static)
    }

    #[expect(clippy::too_many_arguments, reason = "TODO")]
    fn batch(
        &mut self,
        in_ident: syn::Ident,
        in_location: &LocationId,
        in_kind: &CollectionKind,
        out_ident: &syn::Ident,
        out_location: &LocationId,
        op_meta: &HydroIrOpMetadata,
        fold_hooked_idents: &HashSet<String>,
    );

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

    /// Un-windows an ident when it is produced inside a tick's `loop { ... }` context
    /// (`in_location`) but is about to be consumed by an operator emitted at a location
    /// (`out_location`) outside that loop. Returns the ident to use downstream.
    ///
    /// This is needed for operators (e.g. [`HydroNode::ReduceKeyedWatermark`])
    /// where the generated code must sit inside a logical tick, while the output is not in a
    /// tick, and so we need this method.
    fn unwindow_for_consume(
        &mut self,
        in_ident: syn::Ident,
        in_location: &LocationId,
        out_location: &LocationId,
    ) -> syn::Ident;

    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 merge_ordered(
        &mut self,
        location: &LocationId,
        first_ident: syn::Ident,
        second_ident: syn::Ident,
        out_ident: &syn::Ident,
        in_kind: &CollectionKind,
        op_meta: &HydroIrOpMetadata,
        operator_tag: Option<&str>,
    );

    #[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>,
        external_element_type: Option<&syn::Type>,
        tag_id: StmtId,
        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: StmtId,
    );

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

    /// Optionally emit a fold hook that buffers and permutes inputs before the fold.
    /// Returns the new input ident to use for the fold if a hook was emitted.
    fn emit_fold_hook(
        &mut self,
        location: &LocationId,
        in_ident: &syn::Ident,
        in_kind: &CollectionKind,
        op_meta: &HydroIrOpMetadata,
    ) -> Option<syn::Ident>;

    /// Inserts necessary code to validate a manual assertion that at this point the
    /// input live collection is consistent. In production, this is a no-op, but in simulation
    /// this will (not yet implemented) inject assertions that validate consistency.
    fn assert_is_consistent(
        &mut self,
        trusted: bool,
        location: &LocationId,
        in_ident: syn::Ident,
        out_ident: &syn::Ident,
    );

    /// Observes non-determinism introduced by a mut closure operating on a non-strict
    /// (unordered / at-least-once) input. In production this is identity; in simulation
    /// it delegates to `observe_nondet` with the strict output kind.
    fn observe_for_mut(
        &mut self,
        location: &LocationId,
        in_ident: syn::Ident,
        in_kind: &CollectionKind,
        out_ident: &syn::Ident,
        op_meta: &HydroIrOpMetadata,
    );

    fn create_versioned_network_fork(
        &mut self,
        channel_id: u32,
        dest: &LocationId,
        senders: Vec<(LocationId, syn::Ident, Option<DebugExpr>)>,
        external_element_type: Option<&syn::Type>,
        tag_id: StmtId,
    );

    #[expect(clippy::too_many_arguments, reason = "networking codegen")]
    fn create_versioned_network(
        &mut self,
        channel_id: u32,
        source: &LocationId,
        dest: &LocationId,
        out_ident: &syn::Ident,
        deserialize: Option<&DebugExpr>,
        external_element_type: Option<&syn::Type>,
        tag_id: StmtId,
    );
}

/// The production (deployment) DFIR builder: emits one DFIR graph per root location
/// (process/cluster).
///
/// Tick and atomic locations are collapsed onto their root location's graph. In the future
/// (#2902), this builder will additionally emit each (unified) tick as a root-level
/// `loop {{ ... }}` context within its root location's graph.
#[cfg(feature = "build")]
#[derive(Default)]
pub struct ProdDfirBuilder {
    /// The DFIR graph builder for each root location.
    pub graphs: SecondaryMap<LocationKey, FlatGraphBuilder>,
    /// The `loop { ... }` context emitted for each (unified) tick location. Keyed by the tick's
    /// [`LocationId`] (which carries the `ClockId`), so multiple ticks on the same root location
    /// each get their own sibling root-level loop within that root's graph.
    tick_loops: HashMap<LocationId, dfir_lang::graph::GraphLoopId>,
    /// Counter for generating unique intermediate idents at loop boundaries.
    next_intermediate_id: usize,
}

#[cfg(feature = "build")]
impl ProdDfirBuilder {
    /// Gets the DFIR builder for the given location's root, creating it if necessary.
    fn graph_mut(&mut self, location: &LocationId) -> &mut FlatGraphBuilder {
        self.graphs
            .entry(location.root().key())
            .expect("location was removed")
            .or_default()
    }

    /// Returns the (unified) tick location that `location` belongs to: the location itself for
    /// tick locations, the wrapped tick for atomic locations, and `None` for root (top-level)
    /// locations.
    ///
    /// Note this is distinct from the module-level `tick_of` (which extracts a raw [`ClockId`]);
    /// here we need the full [`LocationId`] so it can be used as the loop-map key.
    fn tick_of(location: &LocationId) -> Option<&LocationId> {
        match location {
            LocationId::Tick {
                tick: Some(_),
                parent_location: _,
            } => Some(location),
            // Tick around atomic: the clock lives in the parent location.
            LocationId::Tick {
                tick: None,
                parent_location,
            } => Self::tick_of(parent_location),
            LocationId::Atomic(tick) => Self::tick_of(tick),
            LocationId::Process(_) | LocationId::Cluster(_) => None,
        }
    }

    /// Returns the `loop { ... }` context for the given location, creating it (as a root-level
    /// loop in the location's root graph) if necessary. Returns `None` for top-level locations.
    fn loop_context(&mut self, location: &LocationId) -> Option<dfir_lang::graph::GraphLoopId> {
        let tick_location = Self::tick_of(location)?.clone();
        if let Some(&loop_id) = self.tick_loops.get(&tick_location) {
            return Some(loop_id);
        }
        let loop_id = self.graph_mut(location).insert_loop(None);
        self.tick_loops.insert(tick_location, loop_id);
        Some(loop_id)
    }

    /// Adds the DFIR statements to the graph for the given location, inside the tick's loop
    /// context if the location is a tick or atomic location (otherwise at the root level).
    fn add_dfir_in(
        &mut self,
        location: &LocationId,
        dfir: dfir_lang::parse::DfirCode,
        operator_tag: Option<&str>,
    ) {
        let loop_context = self.loop_context(location);
        self.graph_mut(location)
            .add_dfir(dfir, loop_context, operator_tag);
    }

    /// Generates a unique intermediate identifier for loop-boundary plumbing.
    fn intermediate_ident(&mut self) -> syn::Ident {
        let id = self.next_intermediate_id;
        self.next_intermediate_id += 1;
        syn::Ident::new(&format!("__loop_boundary_{}", id), Span::call_site())
    }
}

#[cfg(feature = "build")]
impl DfirBuilder for ProdDfirBuilder {
    fn singleton_intermediates(&self) -> bool {
        false
    }

    fn add_dfir_at(
        &mut self,
        location: &LocationId,
        dfir: dfir_lang::parse::DfirCode,
        operator_tag: Option<&str>,
    ) {
        // Place tick/atomic-located statements inside the tick's loop context; everything else
        // at the root level.
        self.add_dfir_in(location, dfir, operator_tag);
    }

    /// In production, DFIR sources must be at the root level, outside of the tick's
    /// `loop { ... }` context: it emits the source at the root and windows it into the loop.
    fn add_tick_source(
        &mut self,
        location: &LocationId,
        source_rhs: TokenStream,
        out_ident: &syn::Ident,
        replay_each_tick: bool,
        operator_tag: Option<&str>,
    ) {
        // DFIR sources must be at the root level, not inside a `loop { ... }` context. Emit the
        // source at the root (persisting it there when it should replay every tick, since
        // `persist` is not permitted inside a loop), then window it into the tick's loop.
        let source_ident = self.intermediate_ident();
        let source_stmt = if replay_each_tick {
            parse_quote! {
                #source_ident = #source_rhs -> persist::<'static>();
            }
        } else {
            parse_quote! {
                #source_ident = #source_rhs;
            }
        };
        self.graph_mut(location)
            .add_dfir(source_stmt, None, operator_tag);
        self.add_dfir_in(
            location,
            parse_quote! {
                #out_ident = #source_ident -> batch_eager();
            },
            operator_tag,
        );
    }

    fn batch(
        &mut self,
        in_ident: syn::Ident,
        in_location: &LocationId,
        in_kind: &CollectionKind,
        out_ident: &syn::Ident,
        out_location: &LocationId,
        _op_meta: &HydroIrOpMetadata,
        _fold_hooked_idents: &HashSet<String>,
    ) {
        let is_singleton_like = matches!(
            in_kind,
            CollectionKind::Singleton { .. }
                | CollectionKind::Optional { .. }
                | CollectionKind::KeyedSingleton { .. }
        );

        match (Self::tick_of(in_location), Self::tick_of(out_location)) {
            (Some(in_tick), Some(out_tick)) => {
                // Within the same (unified) tick, e.g. entering the tick from its associated
                // atomic region: both sides live in the same loop.
                assert_eq!(
                    in_tick, out_tick,
                    "batch between distinct ticks should have been unified"
                );
                // NOTE(#2902): a bounded singleton-like value produced once inside the
                // loop (e.g. by `fold_no_replay` in an atomic region) needs to be held across
                // firings. Operators in atomic regions have `'static` lifetimes so they should
                // be persisted and replayed properly.
                self.add_dfir_in(
                    out_location,
                    parse_quote! {
                        #out_ident = #in_ident;
                    },
                    None,
                );
            }
            (None, Some(_)) => {
                // Entering the tick's loop from the top level: emit a windowing operator.
                // `batch_eager()` preserves the pre-loop tick semantics: the loop fires on every
                // tick even when the windowed input is empty.
                let in_ident = if is_singleton_like && in_kind.is_bounded() {
                    // The bounded value is produced exactly once. Persist it at the root (a
                    // `loop { ... }` context cannot contain `persist`) so it remains available,
                    // then window it into the loop on each firing.
                    let persisted_ident = self.intermediate_ident();
                    self.add_dfir_in(
                        in_location,
                        parse_quote! {
                            #persisted_ident = #in_ident -> persist::<'static>();
                        },
                        None,
                    );
                    persisted_ident
                } else {
                    in_ident
                };
                self.add_dfir_in(
                    out_location,
                    parse_quote! {
                        #out_ident = #in_ident -> batch_eager();
                    },
                    None,
                );
            }
            (Some(_), None) | (None, None) => {
                unreachable!("batch must target a tick location");
            }
        }
    }

    fn yield_from_tick(
        &mut self,
        in_ident: syn::Ident,
        in_location: &LocationId,
        _in_kind: &CollectionKind,
        out_ident: &syn::Ident,
        out_location: &LocationId,
    ) {
        // A `YieldConcat` may target either the top level (`latest`/`all_ticks`) or an atomic
        // region wrapping the *same* tick (`latest_atomic`/`all_ticks_atomic`). The atomic region
        // is fused with (runs synchronously inside) its tick's loop, so a yield into it stays in
        // the loop and is emitted as identity — emitting `all_iterations()` there would illegally
        // exit the loop even though the consumer lives inside it. This mirrors the simulation
        // builder, which also emits identity for a same-tick atomic yield.
        //
        // TODO(#2902 phase 2): singleton/optional yields to the top level need held-state
        // semantics (observe the latest value *between* firings), not plain event semantics.
        match (Self::tick_of(in_location), Self::tick_of(out_location)) {
            (Some(in_tick), Some(out_tick)) => {
                assert_eq!(
                    in_tick, out_tick,
                    "atomic yield to a different tick should have been unified"
                );
                self.add_dfir_in(
                    out_location,
                    parse_quote! {
                        #out_ident = #in_ident;
                    },
                    None,
                );
            }
            _ => {
                // Exit the tick's loop back to the top level. `in_ident` is produced inside the
                // loop; `all_iterations()` is the un-windowing operator, emitted at the root level.
                self.graph_mut(in_location).add_dfir(
                    parse_quote! {
                        #out_ident = #in_ident -> all_iterations();
                    },
                    None,
                    None,
                );
            }
        }
    }

    fn unwindow_for_consume(
        &mut self,
        in_ident: syn::Ident,
        in_location: &LocationId,
        out_location: &LocationId,
    ) -> syn::Ident {
        // If the input lives inside a tick's loop context but the consumer is emitted outside
        // that loop, the raw edge would illegally exit the loop. Insert an `all_iterations()`
        // un-windowing operator (emitted at the root level of the input's graph, just like
        // `yield_from_tick`) so the boundary crossing is explicit and legal.
        let in_loop = self.loop_context(in_location);
        let out_loop = self.loop_context(out_location);
        if in_loop.is_some() && in_loop != out_loop {
            let out_ident = self.intermediate_ident();
            self.graph_mut(in_location).add_dfir(
                parse_quote! {
                    #out_ident = #in_ident -> all_iterations();
                },
                None,
                None,
            );
            out_ident
        } else {
            in_ident
        }
    }

    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,
    ) {
        // An atomic region is fused with (runs synchronously inside) its tick's loop. Entering it
        // from the top level windows data in; entering from within the same tick is identity.
        match (Self::tick_of(in_location), Self::tick_of(out_location)) {
            (None, Some(_)) => {
                self.add_dfir_in(
                    out_location,
                    parse_quote! {
                        #out_ident = #in_ident -> batch_eager();
                    },
                    None,
                );
            }
            _ => {
                self.add_dfir_in(
                    out_location,
                    parse_quote! {
                        #out_ident = #in_ident;
                    },
                    None,
                );
            }
        }
    }

    fn end_atomic(
        &mut self,
        in_ident: syn::Ident,
        in_location: &LocationId,
        _in_kind: &CollectionKind,
        out_ident: &syn::Ident,
    ) {
        // Exit the atomic region (and thus the tick's loop) back to the top level.
        self.graph_mut(in_location).add_dfir(
            parse_quote! {
                #out_ident = #in_ident -> all_iterations();
            },
            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.graph_mut(location);
        builder.add_dfir(
            parse_quote! {
                #out_ident = #in_ident;
            },
            None,
            None,
        );
    }

    fn merge_ordered(
        &mut self,
        location: &LocationId,
        first_ident: syn::Ident,
        second_ident: syn::Ident,
        out_ident: &syn::Ident,
        _in_kind: &CollectionKind,
        _op_meta: &HydroIrOpMetadata,
        operator_tag: Option<&str>,
    ) {
        let builder = self.graph_mut(location);
        builder.add_dfir(
            parse_quote! {
                #out_ident = union();
                #first_ident -> [0]#out_ident;
                #second_ident -> [1]#out_ident;
            },
            None,
            operator_tag,
        );
    }

    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>,
        _external_element_type: Option<&syn::Type>,
        tag_id: StmtId,
        _networking_info: &crate::networking::NetworkingInfo,
    ) {
        let sender_builder = self.graph_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.graph_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: StmtId,
    ) {
        let receiver_builder = self.graph_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: StmtId,
    ) {
        let sender_builder = self.graph_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)),
            );
        }
    }

    fn emit_fold_hook(
        &mut self,
        _location: &LocationId,
        _in_ident: &syn::Ident,
        _in_kind: &CollectionKind,
        _op_meta: &HydroIrOpMetadata,
    ) -> Option<syn::Ident> {
        None
    }

    fn assert_is_consistent(
        &mut self,
        _trusted: bool,
        location: &LocationId,
        in_ident: syn::Ident,
        out_ident: &syn::Ident,
    ) {
        let builder = self.graph_mut(location);
        builder.add_dfir(
            parse_quote! {
                #out_ident = #in_ident;
            },
            None,
            None,
        );
    }

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

    fn create_versioned_network_fork(
        &mut self,
        _channel_id: u32,
        _dest: &LocationId,
        _senders: Vec<(LocationId, syn::Ident, Option<DebugExpr>)>,
        _external_element_type: Option<&syn::Type>,
        _tag_id: StmtId,
    ) {
        unreachable!(
            "HydroNode::VersionedNetworkFork is only produced by the multi-version simulator merge \
             pass and cannot be emitted by the non-simulation builder"
        );
    }

    fn create_versioned_network(
        &mut self,
        _channel_id: u32,
        _source: &LocationId,
        _dest: &LocationId,
        _out_ident: &syn::Ident,
        _deserialize: Option<&DebugExpr>,
        _external_element_type: Option<&syn::Type>,
        _tag_id: StmtId,
    ) {
        unreachable!(
            "HydroNode::VersionedNetwork is only produced by the multi-version simulator merge \
             pass and cannot be emitted by the non-simulation builder"
        );
    }
}

#[cfg(feature = "build")]
pub enum BuildersOrCallback<'a, L, N>
where
    L: FnMut(&mut HydroRoot, &mut crate::Counter<StmtId>),
    N: FnMut(&mut HydroNode, &mut crate::Counter<StmtId>),
{
    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, serde::Serialize)]
pub enum HydroRoot {
    ForEach {
        f: ClosureExpr,
        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 {
        #[serde(serialize_with = "serialize_ident")]
        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, LocationKey)>,
        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 {
                    #[cfg(feature = "tokio")]
                    input,
                    #[cfg(feature = "tokio")]
                    to_external_key,
                    #[cfg(feature = "tokio")]
                    to_port_id,
                    #[cfg(feature = "tokio")]
                    to_many,
                    #[cfg(feature = "tokio")]
                    unpaired,
                    #[cfg(feature = "tokio")]
                    instantiate_fn,
                    ..
                } = l
                {
                    #[cfg(feature = "tokio")]
                    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(cluster_key) => {
                                    let from_node = clusters
                                        .get(*cluster_key)
                                        .unwrap_or_else(|| {
                                            panic!("A cluster used in the graph was not instantiated: {}", cluster_key)
                                        })
                                        .clone();

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

                                    if *unpaired {
                                        to_node.register(*to_port_id, source_port.clone());
                                    }

                                    (
                                        (
                                            D::m2e_sink(
                                                &from_node,
                                                &sink_port,
                                                &to_node,
                                                &source_port,
                                                format!("{}_{}", *to_external_key, *to_port_id)
                                            ),
                                            parse_quote!(DUMMY),
                                        ),
                                        Box::new(|| {}) as Box<dyn FnOnce()>,
                                    )
                                }
                                _ => panic!()
                            }
                        },

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

                    #[cfg(not(feature = "tokio"))]
                    {
                        panic!("Cannot instantiate external inputs without tokio");
                    };

                    #[cfg(feature = "tokio")]
                    {
                        *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,
                    serialize,
                    deserialize,
                    metadata,
                    ..
                } = n
                {
                    let external_types = match (
                        serialize.external_element_type(),
                        deserialize.external_element_type(),
                    ) {
                        (Some(input_type), Some(output_type)) => Some((input_type, output_type)),
                        _ => None,
                    };
                    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,
                            external_types,
                        ),

                        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(cluster_key) => {
                                    let to_node = clusters
                                        .get(*cluster_key)
                                        .unwrap_or_else(|| {
                                            panic!("A cluster used in the graph was not instantiated: {}", cluster_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),
                                            D::e2m_source(
                                                refcell_extra_stmts.borrow_mut().entry(*cluster_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::e2m_connect(&from_node, &sink_port, &to_node, &source_port, *port_hint),
                                    )
                                }
                                _ => 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(), 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 { f, input, .. } => {
                f.transform_children(&mut transform, seen_tees);
                transform(input, seen_tees);
            }
            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.deep_clone(seen_tees),
                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 crate::Counter<StmtId>,
        fold_hooked_idents: &mut HashSet<String>,
    ) {
        self.emit_core(
            &mut BuildersOrCallback::<
                fn(&mut HydroRoot, &mut crate::Counter<StmtId>),
                fn(&mut HydroNode, &mut crate::Counter<StmtId>),
            >::Builders(graph_builders),
            seen_tees,
            built_tees,
            next_stmt_id,
            fold_hooked_idents,
        );
    }

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

                // for_each is always side-effecting, so we observe non-determinism
                // even when the closure does not capture a mut ref (unlike map/filter
                // which only observe when they have a mut ref).
                let input_ident = if !input.metadata().collection_kind.is_strict() {
                    let observe_stmt_id = next_stmt_id.get_and_increment();
                    let observe_ident =
                        syn::Ident::new(&format!("stream_{}", observe_stmt_id), Span::call_site());
                    if let BuildersOrCallback::Builders(graph_builders) = builders_or_callback {
                        graph_builders.observe_for_mut(
                            &input.metadata().location_id,
                            input_ident,
                            &input.metadata().collection_kind,
                            &observe_ident,
                            &input.metadata().op,
                        );
                    }
                    observe_ident
                } else {
                    input_ident
                };

                // Emit each captured handoff reference (deduplicated via `built_tees` in the
                // `HydroNode::Reference` arm), so that references captured *only* by this
                // `for_each` closure are still materialized. This mirrors how node-level
                // operators (e.g. `map`) emit their closures' captured references as part of
                // their bottom-up traversal. This is done in both the Builders and Callback
                // paths so that statement IDs stay consistent between them.
                let mut ref_idents = Vec::new();
                for (ref_node, _is_mut) in f.singleton_refs.iter_mut() {
                    assert!(
                        matches!(ref_node, HydroNode::Reference { .. }),
                        "singleton_refs should only contain HydroNode::Reference"
                    );
                    ref_idents.push(ref_node.emit_core(
                        builders_or_callback,
                        seen_tees,
                        built_tees,
                        next_stmt_id,
                        fold_hooked_idents,
                    ));
                }

                // Mint the root's statement ID only after emitting the captured refs, so that
                // statement IDs follow emission order and (in the Callback path) the callback
                // observes this root's ID as the most recently allocated one, consistent with
                // the other `HydroRoot` variants.
                let stmt_id = next_stmt_id.get_and_increment();

                match builders_or_callback {
                    BuildersOrCallback::Builders(graph_builders) => {
                        // The refs' idents are in `singleton_refs` order, matching what
                        // `emit_tokens` expects on the ident stack.
                        let mut ident_stack: Vec<syn::Ident> = ref_idents;

                        let f_tokens = f.emit_tokens(&mut ident_stack);

                        graph_builders.add_dfir_at(
                            &input.metadata().location_id,
                            parse_quote! {
                                #input_ident -> for_each(#f_tokens);
                            },
                            Some(&stmt_id.to_string()),
                        );
                    }
                    BuildersOrCallback::Callback(leaf_callback, _) => {
                        leaf_callback(self, next_stmt_id);
                    }
                }
            }

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

                let stmt_id = next_stmt_id.get_and_increment();

                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(),
                            stmt_id,
                        );
                    }
                    BuildersOrCallback::Callback(leaf_callback, _) => {
                        leaf_callback(self, next_stmt_id);
                    }
                }
            }

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

                let stmt_id = next_stmt_id.get_and_increment();

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

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

                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.add_dfir_at(
                            &input.metadata().location_id,
                            parse_quote! {
                                #cycle_id_ident = #input_ident -> identity::<#elem_type>();
                            },
                            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,
                    fold_hooked_idents,
                );

                let stmt_id = next_stmt_id.get_and_increment();

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

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

                let stmt_id = next_stmt_id.get_and_increment();

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

    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, .. } => {
                transform(&mut f.expr);
            }
            HydroRoot::DestSink { sink, .. } => {
                transform(sink);
            }
            HydroRoot::SendExternal { .. }
            | HydroRoot::CycleSink { .. }
            | HydroRoot::EmbeddedOutput { .. }
            | HydroRoot::Null { .. } => {}
        }
    }
}

#[cfg(feature = "build")]
fn tick_of(loc: &LocationId) -> Option<ClockId> {
    match loc {
        // Regular tick.
        &LocationId::Tick {
            tick: Some(tick),
            parent_location: _,
        } => Some(tick),
        // Tick around atomic.
        LocationId::Tick {
            tick: None,
            parent_location,
        } => Some(
            tick_of(parent_location)
                .expect("Tick should have either own clock ID or clock ID within parent_location."),
        ),
        LocationId::Atomic(inner) => tick_of(inner),
        _ => None,
    }
}

#[cfg(feature = "build")]
fn remap_location(loc: &mut LocationId, uf: &mut HashMap<ClockId, ClockId>) {
    match loc {
        LocationId::Tick {
            tick,
            parent_location,
        } => {
            if let Some(tick) = tick {
                *tick = uf_find(uf, *tick);
            }
            remap_location(parent_location, uf);
        }
        LocationId::Atomic(inner) => {
            remap_location(inner, uf);
        }
        LocationId::Process(_) | LocationId::Cluster(_) => {}
    }
}

#[cfg(feature = "build")]
fn uf_find(parent: &mut HashMap<ClockId, ClockId>, x: ClockId) -> ClockId {
    let p = *parent.get(&x).unwrap_or(&x);
    if p == x {
        return x;
    }
    let root = uf_find(parent, p);
    parent.insert(x, root);
    root
}

#[cfg(feature = "build")]
fn uf_union(parent: &mut HashMap<ClockId, ClockId>, a: ClockId, b: ClockId) {
    let ra = uf_find(parent, a);
    let rb = uf_find(parent, b);
    if ra != rb {
        parent.insert(ra, rb);
    }
}

/// Traverse the IR to build a union-find that unifies tick IDs connected
/// through `Batch` and `YieldConcat` nodes at atomic boundaries, then
/// rewrite all `LocationId`s to use the representative tick ID.
#[cfg(feature = "build")]
pub fn unify_atomic_ticks(ir: &mut [HydroRoot]) {
    let mut uf: HashMap<ClockId, ClockId> = HashMap::new();

    // Pass 1: collect unifications.
    transform_bottom_up(
        ir,
        &mut |_| {},
        &mut |node: &mut HydroNode| match node {
            HydroNode::Batch { inner, metadata } | HydroNode::YieldConcat { inner, metadata } => {
                if let (Some(a), Some(b)) = (
                    tick_of(&inner.metadata().location_id),
                    tick_of(&metadata.location_id),
                ) {
                    uf_union(&mut uf, a, b);
                }
            }
            HydroNode::Chain {
                first,
                second,
                metadata,
            }
            | HydroNode::ChainFirst {
                first,
                second,
                metadata,
            }
            | HydroNode::MergeOrdered {
                first,
                second,
                metadata,
            } => {
                if let (Some(a), Some(b)) = (
                    tick_of(&first.metadata().location_id),
                    tick_of(&metadata.location_id),
                ) {
                    uf_union(&mut uf, a, b);
                }
                if let (Some(a), Some(b)) = (
                    tick_of(&second.metadata().location_id),
                    tick_of(&metadata.location_id),
                ) {
                    uf_union(&mut uf, a, b);
                }
            }
            _ => {}
        },
        false,
    );

    // Pass 2: rewrite all LocationIds.
    transform_bottom_up(
        ir,
        &mut |_| {},
        &mut |node: &mut HydroNode| {
            remap_location(&mut node.metadata_mut().location_id, &mut uf);
        },
        false,
    );
}

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

#[cfg(feature = "build")]
pub fn traverse_dfir(
    ir: &mut [HydroRoot],
    transform_root: impl FnMut(&mut HydroRoot, &mut crate::Counter<StmtId>),
    transform_node: impl FnMut(&mut HydroNode, &mut crate::Counter<StmtId>),
) {
    let mut seen_tees = HashMap::new();
    let mut built_tees = HashMap::new();
    let mut next_stmt_id = crate::Counter::<StmtId>::default();
    let mut fold_hooked_idents = HashSet::new();
    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,
            &mut fold_hooked_idents,
        );
    });
}

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) };
    /// Tracks shared nodes already serialized so that `SharedNode::serialize`
    /// emits the full subtree only once and uses a `"<shared N>"` back-reference
    /// on subsequent encounters, preventing infinite loops.
    static SERIALIZED_SHARED: 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
    })
}

/// Runs `f` with a fresh shared-node deduplication scope for serialization.
/// Any `SharedNode` serialized inside `f` will be tracked; the first occurrence
/// emits the full subtree while later occurrences emit a `{"$shared_ref": id}`
/// back-reference.  The tracking state is restored when `f` returns or panics.
pub fn serialize_dedup_shared<T>(f: impl FnOnce() -> T) -> T {
    let _guard = SerializedSharedGuard::enter();
    f()
}

/// RAII guard that saves/restores the `SERIALIZED_SHARED` thread-local,
/// making `serialize_dedup_shared` re-entrant and panic-safe.
struct SerializedSharedGuard {
    previous: Option<(usize, HashMap<*const RefCell<HydroNode>, usize>)>,
}

impl SerializedSharedGuard {
    fn enter() -> Self {
        let previous = SERIALIZED_SHARED.with(|cell| {
            let mut guard = cell.borrow_mut();
            guard.replace((0, HashMap::new()))
        });
        Self { previous }
    }
}

impl Drop for SerializedSharedGuard {
    fn drop(&mut self) {
        SERIALIZED_SHARED.with(|cell| {
            *cell.borrow_mut() = self.previous.take();
        });
    }
}

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

impl serde::Serialize for SharedNode {
    /// Multiple `SharedNode`s can point to the same underlying `HydroNode` (via
    /// `Tee` / `Partition`).  A naïve recursive serialization would revisit the
    /// same subtree every time and, if the graph ever contains a cycle, loop
    /// forever.
    ///
    /// We keep a thread-local map (`SERIALIZED_SHARED`) from raw `Rc` pointer →
    /// integer id.  The first time we see a pointer we assign it the next id and
    /// emit the full subtree as `{"$shared": <id>, "node": …}`.  Every later
    /// encounter of the same pointer emits `{"$shared_ref": <id>}`, cutting the
    /// recursion.  Requires an active `serialize_dedup_shared` scope.
    fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
        SERIALIZED_SHARED.with(|cell| {
            let mut guard = cell.borrow_mut();
            // (next_id, pointer → assigned_id)
            let state = guard.as_mut().ok_or_else(|| {
                serde::ser::Error::custom(
                    "SharedNode serialization requires an active serialize_dedup_shared scope",
                )
            })?;
            let ptr = self.0.as_ptr() as *const RefCell<HydroNode>;

            if let Some(&id) = state.1.get(&ptr) {
                drop(guard);
                use serde::ser::SerializeMap;
                let mut map = serializer.serialize_map(Some(1))?;
                map.serialize_entry("$shared_ref", &id)?;
                map.end()
            } else {
                let id = state.0;
                state.0 += 1;
                state.1.insert(ptr, id);
                drop(guard);

                use serde::ser::SerializeMap;
                let mut map = serializer.serialize_map(Some(2))?;
                map.serialize_entry("$shared", &id)?;
                map.serialize_entry("node", &*self.0.borrow())?;
                map.end()
            }
        })
    }
}

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(&(std::ptr::from_ref(self.0.as_ref())))
                {
                    write!(f, "<shared {}>", existing)
                } else {
                    let next_id = printed_tees_mut.0;
                    printed_tees_mut.0 += 1;
                    printed_tees_mut
                        .1
                        .insert(std::ptr::from_ref(self.0.as_ref()), 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);
    }
}

/// A counter for tracking singleton access groups on a `HydroNode::Reference`.
///
/// Each mutable access increments the counter (before and after) to isolate itself in its own group;
/// immutable accesses share the current group.
#[derive(Debug)]
pub enum AccessCounter {
    Counting(Cell<u32>),
    Frozen(u32),
}

impl AccessCounter {
    pub fn new() -> Self {
        Self::Counting(Cell::new(0))
    }

    /// Assign the next access group for this reference.
    /// Mutable accesses get an isolated group (counter increments before and after).
    /// Immutable accesses share the current group.
    pub fn next_group(&self, is_mut: bool) -> Self {
        let AccessCounter::Counting(count) = self else {
            panic!("Cannot count on `AccessCounter::Frozen`");
        };
        let c = if is_mut {
            let c = count.get() + 1;
            count.set(c + 1);
            c
        } else {
            count.get()
        };
        Self::Frozen(c)
    }

    /// Creates a frozen counter to prevent further counting.
    pub fn freeze(&self) -> Self {
        Self::Frozen(match self {
            Self::Counting(count) => count.get(),
            Self::Frozen(count) => *count,
        })
    }

    pub fn frozen_group(&self) -> u32 {
        let Self::Frozen(count) = self else {
            panic!("`AccessCounter` not frozen");
        };
        *count
    }
}

impl Default for AccessCounter {
    fn default() -> Self {
        Self::new()
    }
}

impl Hash for AccessCounter {
    fn hash<H: Hasher>(&self, _state: &mut H) {
        // Access counter does not participate in hashing — it is runtime bookkeeping.
    }
}

impl serde::Serialize for AccessCounter {
    fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
        let count = match self {
            AccessCounter::Counting(count) => count.get(),
            AccessCounter::Frozen(count) => *count,
        };
        count.serialize(serializer)
    }
}

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

#[derive(serde::Serialize, Clone, PartialEq, Eq, Debug)]
pub enum OptionalBoundKind {
    Unbounded,
    /// The optional starts out null, but once it becomes non-null it will remain non-null
    /// forever (though the non-null value may change arbitrarily). Erases to [`BoundKind::Unbounded`].
    InitNone,
    Bounded,
}

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

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

#[derive(serde::Serialize, Clone, PartialEq, Eq, Debug)]
pub enum KeyedSingletonBoundKind {
    Unbounded,
    MonotonicKeys,
    MonotonicValue,
    BoundedValue,
    Bounded,
}

#[derive(serde::Serialize, Clone, PartialEq, Eq, Debug)]
pub enum SingletonBoundKind {
    Unbounded,
    Monotonic,
    Bounded,
}

#[derive(Clone, PartialEq, Eq, Debug, serde::Serialize)]
pub enum CollectionKind {
    Stream {
        bound: BoundKind,
        order: StreamOrder,
        retry: StreamRetry,
        element_type: DebugType,
    },
    Singleton {
        bound: SingletonBoundKind,
        element_type: DebugType,
    },
    Optional {
        bound: OptionalBoundKind,
        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: SingletonBoundKind::Bounded,
                ..
            } | CollectionKind::Optional {
                bound: OptionalBoundKind::Bounded,
                ..
            } | CollectionKind::KeyedStream {
                bound: BoundKind::Bounded,
                ..
            } | CollectionKind::KeyedSingleton {
                bound: KeyedSingletonBoundKind::Bounded,
                ..
            }
        )
    }

    /// Returns whether this collection kind is already "strict" (TotalOrder + ExactlyOnce),
    /// meaning no non-determinism needs to be observed for mut closures.
    pub fn is_strict(&self) -> bool {
        match self {
            CollectionKind::Stream { order, retry, .. } => {
                *order == StreamOrder::TotalOrder && *retry == StreamRetry::ExactlyOnce
            }
            CollectionKind::KeyedStream {
                value_order,
                value_retry,
                ..
            } => {
                *value_order == StreamOrder::TotalOrder && *value_retry == StreamRetry::ExactlyOnce
            }
            // Singletons/Optionals/KeyedSingletons do not have observable
            // non-determinism other than snapshots / batching
            CollectionKind::Singleton { .. }
            | CollectionKind::Optional { .. }
            | CollectionKind::KeyedSingleton { .. } => true,
        }
    }

    /// Creates a "strict" version of this kind with TotalOrder and ExactlyOnce.
    pub fn strict_kind(&self) -> CollectionKind {
        match self {
            CollectionKind::Stream {
                bound,
                element_type,
                ..
            } => CollectionKind::Stream {
                bound: bound.clone(),
                order: StreamOrder::TotalOrder,
                retry: StreamRetry::ExactlyOnce,
                element_type: element_type.clone(),
            },
            CollectionKind::KeyedStream {
                bound,
                key_type,
                value_type,
                ..
            } => CollectionKind::KeyedStream {
                bound: bound.clone(),
                value_order: StreamOrder::TotalOrder,
                value_retry: StreamRetry::ExactlyOnce,
                key_type: key_type.clone(),
                value_type: value_type.clone(),
            },
            other => other.clone(),
        }
    }
}

#[derive(Clone, serde::Serialize)]
pub struct HydroIrMetadata {
    pub location_id: LocationId,
    pub collection_kind: CollectionKind,
    pub consistency: Option<ClusterConsistency>,
    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, serde::Serialize)]
pub struct HydroIrOpMetadata {
    #[serde(rename = "span", serialize_with = "serialize_backtrace_as_span")]
    pub backtrace: Backtrace,
    pub cpu_usage: Option<f64>,
    pub network_recv_cpu_usage: Option<f64>,
    pub id: Option<usize>,
    /// When set, this unsafe operator (e.g. `batch` / `snapshot`) is bound to a simulator
    /// hook handle with this ID, letting simulation tests script its decisions. Ignored by
    /// non-simulator backends.
    #[serde(skip)]
    pub sim_hook_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,
            sim_hook_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) {}
}

/// How a network channel's *sender* prepares each message before it is handed to the transport.
///
/// A channel's serialization is split into a send half ([`NetworkSend`]) and a receive half
/// ([`NetworkRecv`]) so that the multi-version simulation merge can reason about each side
/// independently (the sender fork and the receiver are separate IR nodes).
#[derive(Debug, Clone, Hash, serde::Serialize)]
pub enum NetworkSend {
    /// Serialization is performed within the Hydro dataflow using the provided serialize
    /// expression. This is how channels using [`crate::networking::Bincode`] are lowered.
    Custom { serialize_fn: Option<DebugExpr> },
    /// Serialization is left to code outside of Hydro (see [`crate::networking::Embedded`]). The
    /// raw `element_type` is passed through unserialized; the only transformation is converting a
    /// routing [`crate::location::MemberId`] (the destination cluster `tag`, when demuxing) into
    /// the raw `TaglessMemberId` used by the transport. Only supported by the embedded backend.
    ///
    /// Stored as structured info (rather than a pre-baked expression) so that the code can be
    /// synthesized in a post-IR codegen pass.
    Embedded {
        tag: Option<DebugType>,
        element_type: DebugType,
    },
}

/// How a network channel's *receiver* recovers each message from the transport. See
/// [`NetworkSend`] for the sender half.
#[derive(Debug, Clone, Hash, serde::Serialize)]
pub enum NetworkRecv {
    /// Deserialization is performed within the Hydro dataflow using the provided deserialize
    /// expression. This is how channels using [`crate::networking::Bincode`] are lowered.
    Custom { deserialize_fn: Option<DebugExpr> },
    /// Deserialization is left to code outside of Hydro (see [`crate::networking::Embedded`]). The
    /// raw `element_type` is delivered to the receiver directly, with no transport `Result` to
    /// unwrap (the external code that produces the stream decides how to handle faults). The only
    /// transformation is converting a `TaglessMemberId` back into a typed
    /// [`crate::location::MemberId`] (the sender cluster `tag`, when the receiver is keyed by
    /// sender). Only supported by the embedded backend.
    Embedded {
        tag: Option<DebugType>,
        element_type: DebugType,
    },
}

#[cfg(feature = "build")]
impl NetworkSend {
    /// The raw payload type flowing across the channel when serialization is left to external code,
    /// or [`None`] when the channel serializes internally.
    pub(crate) fn external_element_type(&self) -> Option<&syn::Type> {
        match self {
            NetworkSend::Custom { .. } => None,
            NetworkSend::Embedded { element_type, .. } => Some(&element_type.0),
        }
    }
}

#[cfg(feature = "build")]
impl NetworkRecv {
    /// See [`NetworkSend::external_element_type`].
    pub(crate) fn external_element_type(&self) -> Option<&syn::Type> {
        match self {
            NetworkRecv::Custom { .. } => None,
            NetworkRecv::Embedded { element_type, .. } => Some(&element_type.0),
        }
    }
}

#[cfg(feature = "build")]
impl NetworkSend {
    /// The expression applied on the sender to prepare each message for the transport, if any.
    pub(crate) fn pipeline(&self) -> Option<DebugExpr> {
        match self {
            NetworkSend::Custom { serialize_fn } => serialize_fn.clone(),
            NetworkSend::Embedded { tag, element_type } => {
                let root = crate::staging_util::get_this_crate();
                let element_type = &element_type.0;
                let expr: syn::Expr = if let Some(tag) = tag {
                    let tag = &tag.0;
                    parse_quote! {
                        #root::runtime_support::stageleft::runtime_support::fn1_type_hint::<(#root::__staged::location::MemberId<#tag>, #element_type), _>(
                            |(id, data)| (id.into_tagless(), data)
                        )
                    }
                } else {
                    parse_quote! {
                        #root::runtime_support::stageleft::runtime_support::fn1_type_hint::<#element_type, _>(
                            |data| data
                        )
                    }
                };
                Some(expr.into())
            }
        }
    }
}

#[cfg(feature = "build")]
impl NetworkRecv {
    /// The expression applied on the receiver to recover each message from the transport, if any.
    pub(crate) fn pipeline(&self) -> Option<DebugExpr> {
        match self {
            NetworkRecv::Custom { deserialize_fn } => deserialize_fn.clone(),
            // Embedded channels hand the raw payload to the receiver directly (no transport
            // `Result`), so the developer's external code decides how to handle serialization
            // faults. The only transformation is restoring the typed `MemberId` when the receiver
            // is keyed by the sender.
            NetworkRecv::Embedded { tag, .. } => {
                let tag = tag.as_ref()?;
                let root = crate::staging_util::get_this_crate();
                let tag = &tag.0;
                let expr: syn::Expr = parse_quote! {
                    |(id, b)| (#root::__staged::location::MemberId::<#tag>::from_tagless(id as #root::__staged::location::TaglessMemberId), b)
                };
                Some(expr.into())
            }
        }
    }
}

/// An intermediate node in a Hydro graph, which consumes data
/// from upstream nodes and emits data to downstream nodes.
#[derive(Debug, Hash, serde::Serialize)]
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,
    },

    /// A reference materialization point. Wraps a SharedNode so that:
    /// - The pipe output delivers data to one consumer
    /// - `#var` references can borrow the value from the slot
    ///
    /// In DFIR codegen, emits `ident = inner_ident -> singleton()` or `-> optional()` or
    /// `-> handoff()` depending on `kind`.
    ///
    /// Uses the same `built_tees` dedup pattern as `Tee`.
    Reference {
        inner: SharedNode,
        kind: crate::handoff_ref::HandoffRefKind,
        access_counter: AccessCounter,
        metadata: HydroIrMetadata,
    },

    /// An output side of the partition operator.
    PartitionSide {
        inner: SharedNode,
        is_true: bool,
        metadata: HydroIrMetadata,
    },

    /// The inner input of partitioning, shared between two `PartitionSide`.
    PartitionShared {
        input: Box<HydroNode>,
        f: ClosureExpr,
        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,
    },

    MergeOrdered {
        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,
    },

    /// Asymmetric join where the right (build) side is bounded.
    /// The build side is accumulated (stratum-delayed) into a hash table,
    /// then the left (probe) side streams through preserving its ordering.
    JoinHalf {
        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,
    },
    ResolveFuturesBlocking {
        input: Box<HydroNode>,
        metadata: HydroIrMetadata,
    },
    ResolveFuturesOrdered {
        input: Box<HydroNode>,
        metadata: HydroIrMetadata,
    },

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

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

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

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

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

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

    Network {
        name: Option<String>,
        networking_info: crate::networking::NetworkingInfo,
        serialize: NetworkSend,
        deserialize: NetworkRecv,
        instantiate_fn: DebugInstantiate,
        input: Box<HydroNode>,
        metadata: HydroIrMetadata,
    },

    VersionedNetworkFork {
        channel_id: u32,
        channel_name: String,
        senders: Vec<(u32, Box<HydroNode>, NetworkSend)>,
        metadata: HydroIrMetadata,
    },

    VersionedNetwork {
        fork: SharedNode,
        version: u32,
        deserialize: NetworkRecv,
        metadata: HydroIrMetadata,
    },

    ExternalInput {
        from_external_key: LocationKey,
        from_port_id: ExternalPortId,
        from_many: bool,
        codec_type: DebugType,
        #[serde(skip)]
        port_hint: NetworkHint,
        instantiate_fn: DebugInstantiate,
        deserialize_fn: Option<DebugExpr>,
        metadata: HydroIrMetadata,
    },

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

    AssertIsConsistent {
        inner: Box<HydroNode>,
        trusted: bool,
        metadata: HydroIrMetadata,
    },

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

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

/// If `f` has a mut singleton ref and `in_kind` is non-strict, emits an
/// `observe_for_mut` node and returns the new ident. Otherwise returns
/// `in_ident` unchanged. Always consumes a stmt_id when applicable.
#[cfg(feature = "build")]
fn maybe_observe_for_mut(
    f: &ClosureExpr,
    in_ident: syn::Ident,
    in_location: &LocationId,
    in_kind: &CollectionKind,
    op_meta: &HydroIrOpMetadata,
    builders_or_callback: &mut BuildersOrCallback<
        '_,
        impl FnMut(&mut HydroRoot, &mut crate::Counter<StmtId>),
        impl FnMut(&mut HydroNode, &mut crate::Counter<StmtId>),
    >,
    next_stmt_id: &mut crate::Counter<StmtId>,
) -> syn::Ident {
    if f.has_mut_ref() && !in_kind.is_strict() {
        let observe_stmt_id = next_stmt_id.get_and_increment();
        let observe_ident =
            syn::Ident::new(&format!("stream_{}", observe_stmt_id), Span::call_site());
        if let BuildersOrCallback::Builders(graph_builders) = builders_or_callback {
            graph_builders.observe_for_mut(in_location, in_ident, in_kind, &observe_ident, op_meta);
        }
        observe_ident
    } else {
        in_ident
    }
}

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, .. } | HydroNode::Reference { 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::PartitionSide { 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: HydroNode = inner.0.replace(HydroNode::Placeholder);
                    transform(&mut orig, seen_tees);
                    *transformed_cell.borrow_mut() = orig;
                    *inner = SharedNode(transformed_cell);
                }
            }
            HydroNode::PartitionShared { input, f, .. } => {
                f.transform_children(&mut transform, seen_tees);
                transform(input.as_mut(), seen_tees);
            }

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

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

            HydroNode::MergeOrdered { 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, .. }
            | HydroNode::JoinHalf { 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::Map { f, input, .. } => {
                f.transform_children(&mut transform, seen_tees);
                transform(input.as_mut(), seen_tees);
            }
            HydroNode::FlatMap { f, input, .. }
            | HydroNode::FlatMapStreamBlocking { f, input, .. }
            | HydroNode::Filter { f, input, .. }
            | HydroNode::FilterMap { f, input, .. }
            | HydroNode::Inspect { f, input, .. }
            | HydroNode::Reduce { f, input, .. }
            | HydroNode::ReduceKeyed { f, input, .. } => {
                f.transform_children(&mut transform, seen_tees);
                transform(input.as_mut(), seen_tees);
            }
            HydroNode::ReduceKeyedWatermark {
                f,
                input,
                watermark,
                ..
            } => {
                f.transform_children(&mut transform, seen_tees);
                transform(input.as_mut(), seen_tees);
                transform(watermark.as_mut(), seen_tees);
            }
            HydroNode::Fold {
                init, acc, input, ..
            }
            | HydroNode::Scan {
                init, acc, input, ..
            }
            | HydroNode::ScanAsyncBlocking {
                init, acc, input, ..
            }
            | HydroNode::FoldKeyed {
                init, acc, input, ..
            } => {
                init.transform_children(&mut transform, seen_tees);
                acc.transform_children(&mut transform, seen_tees);
                transform(input.as_mut(), seen_tees);
            }
            HydroNode::ResolveFutures { input, .. }
            | HydroNode::ResolveFuturesBlocking { input, .. }
            | HydroNode::ResolveFuturesOrdered { input, .. }
            | HydroNode::Sort { input, .. }
            | HydroNode::DeferTick { input, .. }
            | HydroNode::Enumerate { input, .. }
            | HydroNode::Unique { input, .. }
            | HydroNode::Network { input, .. }
            | HydroNode::Counter { input, .. } => {
                transform(input.as_mut(), seen_tees);
            }

            HydroNode::VersionedNetworkFork { senders, .. } => {
                for (_version, sender, _serialize) in senders.iter_mut() {
                    transform(sender.as_mut(), seen_tees);
                }
            }

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

    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::UnboundSingleton { inner, metadata } => HydroNode::UnboundSingleton {
                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::AssertIsConsistent {
                inner,
                trusted,
                metadata,
            } => HydroNode::AssertIsConsistent {
                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 }
            | HydroNode::Reference {
                inner, metadata, ..
            } => {
                let cloned_inner = if let Some(transformed) = seen_tees.get(&inner.as_ptr()) {
                    SharedNode(transformed.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;
                    SharedNode(new_rc)
                };
                if let HydroNode::Reference {
                    kind,
                    access_counter,
                    ..
                } = self
                {
                    HydroNode::Reference {
                        inner: cloned_inner,
                        kind: *kind,
                        access_counter: access_counter.freeze(),
                        metadata: metadata.clone(),
                    }
                } else {
                    HydroNode::Tee {
                        inner: cloned_inner,
                        metadata: metadata.clone(),
                    }
                }
            }
            HydroNode::PartitionSide {
                inner,
                is_true,
                metadata,
            } => {
                if let Some(transformed) = seen_tees.get(&inner.as_ptr()) {
                    HydroNode::PartitionSide {
                        inner: SharedNode(transformed.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::PartitionSide {
                        inner: SharedNode(new_rc),
                        is_true: *is_true,
                        metadata: metadata.clone(),
                    }
                }
            }
            HydroNode::PartitionShared { input, f, metadata } => HydroNode::PartitionShared {
                input: Box::new(input.deep_clone(seen_tees)),
                f: f.deep_clone(seen_tees),
                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::MergeOrdered {
                first,
                second,
                metadata,
            } => HydroNode::MergeOrdered {
                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::JoinHalf {
                left,
                right,
                metadata,
            } => HydroNode::JoinHalf {
                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::ResolveFuturesBlocking { input, metadata } => {
                HydroNode::ResolveFuturesBlocking {
                    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.deep_clone(seen_tees),
                input: Box::new(input.deep_clone(seen_tees)),
                metadata: metadata.clone(),
            },
            HydroNode::FlatMap { f, input, metadata } => HydroNode::FlatMap {
                f: f.deep_clone(seen_tees),
                input: Box::new(input.deep_clone(seen_tees)),
                metadata: metadata.clone(),
            },
            HydroNode::FlatMapStreamBlocking { f, input, metadata } => {
                HydroNode::FlatMapStreamBlocking {
                    f: f.deep_clone(seen_tees),
                    input: Box::new(input.deep_clone(seen_tees)),
                    metadata: metadata.clone(),
                }
            }
            HydroNode::Filter { f, input, metadata } => HydroNode::Filter {
                f: f.deep_clone(seen_tees),
                input: Box::new(input.deep_clone(seen_tees)),
                metadata: metadata.clone(),
            },
            HydroNode::FilterMap { f, input, metadata } => HydroNode::FilterMap {
                f: f.deep_clone(seen_tees),
                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.deep_clone(seen_tees),
                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.deep_clone(seen_tees),
                acc: acc.deep_clone(seen_tees),
                input: Box::new(input.deep_clone(seen_tees)),
                metadata: metadata.clone(),
            },
            HydroNode::Scan {
                init,
                acc,
                input,
                metadata,
            } => HydroNode::Scan {
                init: init.deep_clone(seen_tees),
                acc: acc.deep_clone(seen_tees),
                input: Box::new(input.deep_clone(seen_tees)),
                metadata: metadata.clone(),
            },
            HydroNode::ScanAsyncBlocking {
                init,
                acc,
                input,
                metadata,
            } => HydroNode::ScanAsyncBlocking {
                init: init.deep_clone(seen_tees),
                acc: acc.deep_clone(seen_tees),
                input: Box::new(input.deep_clone(seen_tees)),
                metadata: metadata.clone(),
            },
            HydroNode::FoldKeyed {
                init,
                acc,
                input,
                metadata,
            } => HydroNode::FoldKeyed {
                init: init.deep_clone(seen_tees),
                acc: acc.deep_clone(seen_tees),
                input: Box::new(input.deep_clone(seen_tees)),
                metadata: metadata.clone(),
            },
            HydroNode::ReduceKeyedWatermark {
                f,
                input,
                watermark,
                metadata,
            } => HydroNode::ReduceKeyedWatermark {
                f: f.deep_clone(seen_tees),
                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.deep_clone(seen_tees),
                input: Box::new(input.deep_clone(seen_tees)),
                metadata: metadata.clone(),
            },
            HydroNode::ReduceKeyed { f, input, metadata } => HydroNode::ReduceKeyed {
                f: f.deep_clone(seen_tees),
                input: Box::new(input.deep_clone(seen_tees)),
                metadata: metadata.clone(),
            },
            HydroNode::Network {
                name,
                networking_info,
                serialize,
                deserialize,
                instantiate_fn,
                input,
                metadata,
            } => HydroNode::Network {
                name: name.clone(),
                networking_info: networking_info.clone(),
                serialize: serialize.clone(),
                deserialize: deserialize.clone(),
                instantiate_fn: instantiate_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(),
            },
            HydroNode::VersionedNetworkFork {
                channel_id,
                channel_name,
                senders,
                metadata,
            } => HydroNode::VersionedNetworkFork {
                channel_id: *channel_id,
                channel_name: channel_name.clone(),
                senders: senders
                    .iter()
                    .map(|(version, sender, serialize)| {
                        (
                            *version,
                            Box::new(sender.deep_clone(seen_tees)),
                            serialize.clone(),
                        )
                    })
                    .collect(),
                metadata: metadata.clone(),
            },
            HydroNode::VersionedNetwork {
                fork,
                version,
                deserialize,
                metadata,
            } => {
                let cloned_fork = if let Some(transformed) = seen_tees.get(&fork.as_ptr()) {
                    SharedNode(transformed.clone())
                } else {
                    let new_rc = Rc::new(RefCell::new(HydroNode::Placeholder));
                    seen_tees.insert(fork.as_ptr(), new_rc.clone());
                    let cloned = fork.0.borrow().deep_clone(seen_tees);
                    *new_rc.borrow_mut() = cloned;
                    SharedNode(new_rc)
                };
                HydroNode::VersionedNetwork {
                    fork: cloned_fork,
                    version: *version,
                    deserialize: deserialize.clone(),
                    metadata: metadata.clone(),
                }
            }
        }
    }

    #[cfg(feature = "build")]
    pub fn emit_core(
        &mut self,
        builders_or_callback: &mut BuildersOrCallback<
            '_,
            impl FnMut(&mut HydroRoot, &mut crate::Counter<StmtId>),
            impl FnMut(&mut HydroNode, &mut crate::Counter<StmtId>),
        >,
        seen_tees: &mut SeenSharedNodes,
        built_tees: &mut HashMap<*const RefCell<HydroNode>, Vec<syn::Ident>>,
        next_stmt_id: &mut crate::Counter<StmtId>,
        fold_hooked_idents: &mut HashSet<String>,
    ) -> 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
                        let _ = next_stmt_id.get_and_increment();
                        match builders_or_callback {
                            BuildersOrCallback::Builders(_) => {}
                            BuildersOrCallback::Callback(_, node_callback) => {
                                node_callback(node, next_stmt_id);
                            }
                        }
                        // input_ident stays on stack as output
                    }

                    HydroNode::UnboundSingleton { .. } => {
                        let inner_ident = ident_stack.pop().unwrap();

                        let stmt_id = next_stmt_id.get_and_increment();
                        let out_ident =
                            syn::Ident::new(&format!("stream_{}", stmt_id), Span::call_site());

                        match builders_or_callback {
                            BuildersOrCallback::Builders(graph_builders) => {
                                if graph_builders.singleton_intermediates() {
                                    graph_builders.add_dfir_at(
                                        &out_location,
                                        parse_quote! {
                                            #out_ident = #inner_ident;
                                        },
                                        None,
                                    );
                                } else {
                                    graph_builders.add_dfir_at(
                                        &out_location,
                                        parse_quote! {
                                            #out_ident = #inner_ident -> persist::<'static>();
                                        },
                                        None,
                                    );
                                }
                            }
                            BuildersOrCallback::Callback(_, node_callback) => {
                                node_callback(node, next_stmt_id);
                            }
                        }

                        ident_stack.push(out_ident);
                    }

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

                        let stmt_id = next_stmt_id.get_and_increment();
                        let out_ident =
                            syn::Ident::new(&format!("stream_{}", stmt_id), Span::call_site());

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

                        ident_stack.push(out_ident);
                    }

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

                        let stmt_id = next_stmt_id.get_and_increment();
                        let observe_ident =
                            syn::Ident::new(&format!("stream_{}", 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);
                            }
                        }

                        ident_stack.push(observe_ident);
                    }

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

                        let stmt_id = next_stmt_id.get_and_increment();
                        let batch_ident =
                            syn::Ident::new(&format!("stream_{}", 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,
                                    fold_hooked_idents,
                                );
                            }
                            BuildersOrCallback::Callback(_, node_callback) => {
                                node_callback(node, next_stmt_id);
                            }
                        }

                        ident_stack.push(batch_ident);
                    }

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

                        let stmt_id = next_stmt_id.get_and_increment();
                        let yield_ident =
                            syn::Ident::new(&format!("stream_{}", 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);
                            }
                        }

                        ident_stack.push(yield_ident);
                    }

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

                        let stmt_id = next_stmt_id.get_and_increment();
                        let begin_ident =
                            syn::Ident::new(&format!("stream_{}", 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);
                            }
                        }

                        ident_stack.push(begin_ident);
                    }

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

                        let stmt_id = next_stmt_id.get_and_increment();
                        let end_ident =
                            syn::Ident::new(&format!("stream_{}", 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);
                            }
                        }

                        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 stmt_id = next_stmt_id.get_and_increment();
                            let source_ident =
                                syn::Ident::new(&format!("stream_{}", stmt_id), Span::call_site());

                            // For tick-located sources, holds the source pipeline (RHS) so that
                            // production codegen can hoist it to the root and window it into the
                            // tick's loop (DFIR sources may not live inside a `loop { ... }`).
                            let mut tick_source_rhs: Option<TokenStream> = None;

                            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_root() {
                                        parse_quote! {
                                            #source_ident = source_iter(#expr);
                                        }
                                    } else {
                                        // Located inside a tick or atomic region (which is fused
                                        // into its tick's loop). DFIR sources may not live inside a
                                        // `loop { ... }`, so hoist the source to the root and window
                                        // it into the loop.
                                        // TODO(shadaj): a more natural semantics would be to re-evaluate the expression on each tick
                                        tick_source_rhs = Some(quote! { source_iter(#expr) });
                                        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) => {
                                    if let Some(source_rhs) = tick_source_rhs {
                                        graph_builders.add_tick_source(
                                            &out_location,
                                            source_rhs,
                                            &source_ident,
                                            true,
                                            Some(&stmt_id.to_string()),
                                        );
                                    } else {
                                        graph_builders.add_dfir_at(
                                            &out_location,
                                            source_stmt,
                                            Some(&stmt_id.to_string()),
                                        );
                                    }
                                }
                                BuildersOrCallback::Callback(_, node_callback) => {
                                    node_callback(node, next_stmt_id);
                                }
                            }

                            ident_stack.push(source_ident);
                        }
                    }

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

                        match builders_or_callback {
                            BuildersOrCallback::Builders(graph_builders) => {
                                if *first_tick_only {
                                    assert!(
                                        !metadata.location_id.is_top_level(),
                                        "first_tick_only SingletonSource must be inside a tick"
                                    );
                                    // Delivered only on the first execution of the tick region.
                                    graph_builders.add_tick_source(
                                        &out_location,
                                        quote! { source_iter([#value]) },
                                        &source_ident,
                                        false,
                                        Some(&stmt_id.to_string()),
                                    );
                                } else if metadata.location_id.is_root() {
                                    if metadata.collection_kind.is_bounded() {
                                        graph_builders.add_dfir_at(
                                            &out_location,
                                            parse_quote! {
                                                #source_ident = source_iter([#value]);
                                            },
                                            Some(&stmt_id.to_string()),
                                        );
                                    } else {
                                        graph_builders.add_dfir_at(
                                            &out_location,
                                            parse_quote! {
                                                #source_ident = source_iter([#value]) -> persist::<'static>();
                                            },
                                            Some(&stmt_id.to_string()),
                                        );
                                    }
                                } else {
                                    // A tick- or atomic-located singleton source (the atomic
                                    // region is fused into its tick's loop) that yields its value
                                    // on every firing. DFIR sources may not live inside a
                                    // `loop { ... }`, so hoist it to the root and window it in.
                                    graph_builders.add_tick_source(
                                        &out_location,
                                        quote! { source_iter([#value]) },
                                        &source_ident,
                                        true,
                                        Some(&stmt_id.to_string()),
                                    );
                                }
                            }
                            BuildersOrCallback::Callback(_, node_callback) => {
                                node_callback(node, next_stmt_id);
                            }
                        }

                        ident_stack.push(source_ident);
                    }

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

                        // consume a stmt id even though we did not emit anything so that we can instrument this
                        let _ = next_stmt_id.get_and_increment();

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

                        ident_stack.push(ident);
                    }

                    HydroNode::Tee { inner, .. } => {
                        // we consume a stmt id regardless of if we emit the tee() operator,
                        // so that during rewrites we touch all recipients of the tee()
                        let stmt_id = next_stmt_id.get_and_increment();

                        let ret_ident = if let Some(built_idents) =
                            built_tees.get(&(std::ptr::from_ref(inner.0.as_ref())))
                        {
                            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_{}", stmt_id), Span::call_site());

                            built_tees.insert(
                                std::ptr::from_ref(inner.0.as_ref()),
                                vec![tee_ident.clone()],
                            );

                            match builders_or_callback {
                                BuildersOrCallback::Builders(graph_builders) => {
                                    // NOTE: With `forward_ref`, the fold codegen may not have
                                    // run yet when we reach this tee, so `fold_hooked_idents`
                                    // might not contain the inner ident. In that case we won't
                                    // propagate the "hooked" status to the tee and the
                                    // downstream singleton batch will use the normal
                                    // `SingletonHook` instead of `PassthroughSingletonHook`.
                                    // This is not a soundness issue: the fallback hook still
                                    // produces correct behavior, just with a redundant decision
                                    // point. TODO(https://github.com/hydro-project/hydro/issues/2856):
                                    // fix ordering so forward_ref folds are always processed
                                    // before their downstream tees.
                                    if fold_hooked_idents.contains(&inner_ident.to_string()) {
                                        fold_hooked_idents.insert(tee_ident.to_string());
                                    }
                                    graph_builders.add_dfir_at(
                                        &out_location,
                                        parse_quote! {
                                            #tee_ident = #inner_ident -> tee();
                                        },
                                        Some(&stmt_id.to_string()),
                                    );
                                }
                                BuildersOrCallback::Callback(_, node_callback) => {
                                    node_callback(node, next_stmt_id);
                                }
                            }

                            tee_ident
                        };

                        ident_stack.push(ret_ident);
                    }

                    HydroNode::Reference { inner, kind, .. } => {
                        // we consume a stmt id regardless of if we emit the operator,
                        // so that during rewrites we touch all recipients
                        let stmt_id = next_stmt_id.get_and_increment();

                        let ret_ident = if let Some(built_idents) =
                            built_tees.get(&(std::ptr::from_ref(inner.0.as_ref())))
                        {
                            built_idents[0].clone()
                        } else {
                            let inner_ident = ident_stack.pop().unwrap();

                            let ref_ident =
                                syn::Ident::new(&format!("stream_{}", stmt_id), Span::call_site());

                            built_tees.insert(
                                std::ptr::from_ref(inner.0.as_ref()),
                                vec![ref_ident.clone()],
                            );

                            match builders_or_callback {
                                BuildersOrCallback::Builders(graph_builders) => {
                                    let op_ident = syn::Ident::new(
                                        match kind {
                                            crate::handoff_ref::HandoffRefKind::Singleton => "singleton",
                                            crate::handoff_ref::HandoffRefKind::Optional => "optional",
                                            crate::handoff_ref::HandoffRefKind::Vec => "handoff",
                                        },
                                        Span::call_site(),
                                    );
                                    graph_builders.add_dfir_at(
                                        &out_location,
                                        parse_quote! {
                                            #ref_ident = #inner_ident -> #op_ident();
                                        },
                                        Some(&stmt_id.to_string()),
                                    );
                                }
                                BuildersOrCallback::Callback(_, node_callback) => {
                                    node_callback(node, next_stmt_id);
                                }
                            }

                            ref_ident
                        };

                        ident_stack.push(ret_ident);
                    }

                    HydroNode::PartitionSide {
                        inner, is_true, metadata: _,
                    } => {
                        let is_true = *is_true; // need to copy early to avoid borrow checking issues with node
                        let ptr = std::ptr::from_ref(inner.0.as_ref());
                        let stmt_id = next_stmt_id.get_and_increment();

                        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 `PartitionShared` node was already processed by transform_bottom_up,
                            // so its ident is on the stack
                            let partition_ident = ident_stack.pop().unwrap();

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

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

                            let stmt_id = next_stmt_id.get_and_increment();
                            match builders_or_callback {
                                BuildersOrCallback::Builders(graph_builders) => {
                                    graph_builders.add_dfir_at(
                                        &out_location,
                                        parse_quote! {
                                            #true_ident = #partition_ident[0];
                                            #false_ident = #partition_ident[1];
                                        },
                                        Some(&stmt_id.to_string()),
                                    );
                                }
                                BuildersOrCallback::Callback(_, node_callback) => {
                                    node_callback(node, next_stmt_id);
                                }
                            }

                            if is_true { true_ident } else { false_ident }
                        };

                        ident_stack.push(ret_ident);
                    }

                    HydroNode::PartitionShared { input, f, metadata } => {
                        // Pop input ident (pushed last by transform_children) before
                        // draining the closure's singleton ref idents below it.
                        let inner_ident = ident_stack.pop().unwrap();
                        let f_tokens = f.emit_tokens(&mut ident_stack);

                        let inner_ident = {
                            maybe_observe_for_mut(
                                f, inner_ident,
                                &input.metadata().location_id,
                                &input.metadata().collection_kind,
                                &metadata.op,
                                builders_or_callback, next_stmt_id,
                            )
                        };

                        let stmt_id = next_stmt_id.get_and_increment();
                        let partition_ident = syn::Ident::new(
                            &format!("stream_{}_partition", stmt_id),
                            Span::call_site(),
                        );

                        let stmt_id = next_stmt_id.get_and_increment();
                        match builders_or_callback {
                            BuildersOrCallback::Builders(graph_builders) => {
                                graph_builders.add_dfir_at(
                                    &out_location,
                                    parse_quote! {
                                        #partition_ident = #inner_ident -> partition(|__item, __num_outputs| if (#f_tokens)(__item) { 0_usize } else { 1_usize });
                                    },
                                    Some(&stmt_id.to_string()),
                                );
                            }
                            BuildersOrCallback::Callback(_, node_callback) => {
                                node_callback(node, next_stmt_id);
                            }
                        }
                        ident_stack.push(partition_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 stmt_id = next_stmt_id.get_and_increment();
                        let chain_ident =
                            syn::Ident::new(&format!("stream_{}", stmt_id), Span::call_site());

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

                        ident_stack.push(chain_ident);
                    }

                    HydroNode::MergeOrdered { first, metadata, .. } => {
                        let second_ident = ident_stack.pop().unwrap();
                        let first_ident = ident_stack.pop().unwrap();

                        let stmt_id = next_stmt_id.get_and_increment();
                        let merge_ident =
                            syn::Ident::new(&format!("stream_{}", stmt_id), Span::call_site());

                        match builders_or_callback {
                            BuildersOrCallback::Builders(graph_builders) => {
                                graph_builders.merge_ordered(
                                    &first.metadata().location_id,
                                    first_ident,
                                    second_ident,
                                    &merge_ident,
                                    &first.metadata().collection_kind,
                                    &metadata.op,
                                    Some(&stmt_id.to_string()),
                                );
                            }
                            BuildersOrCallback::Callback(_, node_callback) => {
                                node_callback(node, next_stmt_id);
                            }
                        }

                        ident_stack.push(merge_ident);
                    }

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

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

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

                        ident_stack.push(chain_ident);
                    }

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

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

                        match builders_or_callback {
                            BuildersOrCallback::Builders(graph_builders) => {
                                if right.metadata().location_id.is_top_level()
                                    && right.metadata().collection_kind.is_bounded()
                                {
                                    let lifetime =
                                        graph_builders.cross_tick_state_lifetime(&out_location);
                                    graph_builders.add_dfir_at(
                                        &out_location,
                                        parse_quote! {
                                            #cross_ident = cross_singleton::<#lifetime>();
                                            #left_ident -> [input]#cross_ident;
                                            #right_ident -> [single]#cross_ident;
                                        },
                                        Some(&stmt_id.to_string()),
                                    );
                                } else {
                                    graph_builders.add_dfir_at(
                                        &out_location,
                                        parse_quote! {
                                            #cross_ident = cross_singleton();
                                            #left_ident -> [input]#cross_ident;
                                            #right_ident -> [single]#cross_ident;
                                        },
                                        Some(&stmt_id.to_string()),
                                    );
                                }
                            }
                            BuildersOrCallback::Callback(_, node_callback) => {
                                node_callback(node, next_stmt_id);
                            }
                        }

                        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_top_level = left.metadata().location_id.is_top_level();
                        let right_top_level = right.metadata().location_id.is_top_level();

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

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

                        match builders_or_callback {
                            BuildersOrCallback::Builders(graph_builders) => {
                                let left_lifetime = if left_top_level {
                                    graph_builders.cross_tick_state_lifetime(&out_location)
                                } else {
                                    graph_builders.tick_state_lifetime(&out_location)
                                };

                                let right_lifetime = if right_top_level {
                                    graph_builders.cross_tick_state_lifetime(&out_location)
                                } else {
                                    graph_builders.tick_state_lifetime(&out_location)
                                };

                                graph_builders.add_dfir_at(
                                    &out_location,
                                    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;
                                        }
                                    },
                                    Some(&stmt_id.to_string()),
                                );
                            }
                            BuildersOrCallback::Callback(_, node_callback) => {
                                node_callback(node, next_stmt_id);
                            }
                        }

                        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_top_level = neg.metadata().location_id.is_top_level();

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

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

                        match builders_or_callback {
                            BuildersOrCallback::Builders(graph_builders) => {
                                let neg_lifetime = if neg_top_level {
                                    graph_builders.cross_tick_state_lifetime(&out_location)
                                } else {
                                    graph_builders.tick_state_lifetime(&out_location)
                                };
                                let pos_lifetime =
                                    graph_builders.tick_state_lifetime(&out_location);

                                graph_builders.add_dfir_at(
                                    &out_location,
                                    parse_quote! {
                                        #stream_ident = #operator::<#pos_lifetime, #neg_lifetime>();
                                        #pos_ident -> [pos]#stream_ident;
                                        #neg_ident -> [neg]#stream_ident;
                                    },
                                    Some(&stmt_id.to_string()),
                                );
                            }
                            BuildersOrCallback::Callback(_, node_callback) => {
                                node_callback(node, next_stmt_id);
                            }
                        }

                        ident_stack.push(stream_ident);
                    }

                    HydroNode::JoinHalf { .. } => {
                        let HydroNode::JoinHalf { right, .. } = node else {
                            unreachable!()
                        };

                        assert!(
                            right.metadata().collection_kind.is_bounded(),
                            "JoinHalf requires the right (build) side to be Bounded, got {:?}",
                            right.metadata().collection_kind
                        );

                        let build_top_level = right.metadata().location_id.is_top_level();

                        let build_ident = ident_stack.pop().unwrap();
                        let probe_ident = ident_stack.pop().unwrap();

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

                        match builders_or_callback {
                            BuildersOrCallback::Builders(graph_builders) => {
                                let build_lifetime = if build_top_level {
                                    graph_builders.cross_tick_state_lifetime(&out_location)
                                } else {
                                    graph_builders.tick_state_lifetime(&out_location)
                                };
                                let probe_lifetime =
                                    graph_builders.tick_state_lifetime(&out_location);

                                graph_builders.add_dfir_at(
                                    &out_location,
                                    parse_quote! {
                                        #stream_ident = join_multiset_half::<#build_lifetime, #probe_lifetime>();
                                        #probe_ident -> [probe]#stream_ident;
                                        #build_ident -> [build]#stream_ident;
                                    },
                                    Some(&stmt_id.to_string()),
                                );
                            }
                            BuildersOrCallback::Callback(_, node_callback) => {
                                node_callback(node, next_stmt_id);
                            }
                        }

                        ident_stack.push(stream_ident);
                    }

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

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

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

                        ident_stack.push(futures_ident);
                    }

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

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

                        match builders_or_callback {
                            BuildersOrCallback::Builders(graph_builders) => {
                                graph_builders.add_dfir_at(
                                    &out_location,
                                    parse_quote! {
                                        #futures_ident = #input_ident -> resolve_futures_blocking();
                                    },
                                    Some(&stmt_id.to_string()),
                                );
                            }
                            BuildersOrCallback::Callback(_, node_callback) => {
                                node_callback(node, next_stmt_id);
                            }
                        }

                        ident_stack.push(futures_ident);
                    }

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

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

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

                        ident_stack.push(futures_ident);
                    }

                    HydroNode::Map {
                        f,
                        input,
                        metadata,
                    } => {
                        // Pop input ident (pushed last by transform_children).
                        let input_ident = ident_stack.pop().unwrap();
                        let f_tokens = f.emit_tokens(&mut ident_stack);

                        let input_ident = maybe_observe_for_mut(
                            f,
                            input_ident,
                            &input.metadata().location_id,
                            &input.metadata().collection_kind,
                            &metadata.op,
                            builders_or_callback,
                            next_stmt_id,
                        );

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

                        match builders_or_callback {
                            BuildersOrCallback::Builders(graph_builders) => {
                                graph_builders.add_dfir_at(
                                    &out_location,
                                    parse_quote! {
                                        #map_ident = #input_ident -> map(#f_tokens);
                                    },
                                    Some(&stmt_id.to_string()),
                                );
                            }
                            BuildersOrCallback::Callback(_, node_callback) => {
                                node_callback(node, next_stmt_id);
                            }
                        }

                        ident_stack.push(map_ident);
                    }

                    HydroNode::FlatMap { f, input, metadata } => {
                        let input_ident = ident_stack.pop().unwrap();
                        let f_tokens = f.emit_tokens(&mut ident_stack);

                        let input_ident = maybe_observe_for_mut(
                            f, input_ident,
                            &input.metadata().location_id,
                            &input.metadata().collection_kind,
                            &metadata.op,
                            builders_or_callback, next_stmt_id,
                        );

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

                        match builders_or_callback {
                            BuildersOrCallback::Builders(graph_builders) => {
                                graph_builders.add_dfir_at(
                                    &out_location,
                                    parse_quote! {
                                        #flat_map_ident = #input_ident -> flat_map(#f_tokens);
                                    },
                                    Some(&stmt_id.to_string()),
                                );
                            }
                            BuildersOrCallback::Callback(_, node_callback) => {
                                node_callback(node, next_stmt_id);
                            }
                        }

                        ident_stack.push(flat_map_ident);
                    }

                    HydroNode::FlatMapStreamBlocking { f, input, metadata } => {
                        let input_ident = ident_stack.pop().unwrap();
                        let f_tokens = f.emit_tokens(&mut ident_stack);

                        let input_ident = maybe_observe_for_mut(
                            f, input_ident,
                            &input.metadata().location_id,
                            &input.metadata().collection_kind,
                            &metadata.op,
                            builders_or_callback, next_stmt_id,
                        );

                        let stmt_id = next_stmt_id.get_and_increment();
                        let flat_map_stream_blocking_ident =
                            syn::Ident::new(&format!("stream_{}", stmt_id), Span::call_site());

                        match builders_or_callback {
                            BuildersOrCallback::Builders(graph_builders) => {
                                graph_builders.add_dfir_at(
                                    &out_location,
                                    parse_quote! {
                                        #flat_map_stream_blocking_ident = #input_ident -> flat_map_stream_blocking(#f_tokens);
                                    },
                                    Some(&stmt_id.to_string()),
                                );
                            }
                            BuildersOrCallback::Callback(_, node_callback) => {
                                node_callback(node, next_stmt_id);
                            }
                        }

                        ident_stack.push(flat_map_stream_blocking_ident);
                    }

                    HydroNode::Filter { f, input, metadata } => {
                        let input_ident = ident_stack.pop().unwrap();
                        let f_tokens = f.emit_tokens(&mut ident_stack);

                        let input_ident = maybe_observe_for_mut(
                            f, input_ident,
                            &input.metadata().location_id,
                            &input.metadata().collection_kind,
                            &metadata.op,
                            builders_or_callback, next_stmt_id,
                        );

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

                        match builders_or_callback {
                            BuildersOrCallback::Builders(graph_builders) => {
                                graph_builders.add_dfir_at(
                                    &out_location,
                                    parse_quote! {
                                        #filter_ident = #input_ident -> filter(#f_tokens);
                                    },
                                    Some(&stmt_id.to_string()),
                                );
                            }
                            BuildersOrCallback::Callback(_, node_callback) => {
                                node_callback(node, next_stmt_id);
                            }
                        }

                        ident_stack.push(filter_ident);
                    }

                    HydroNode::FilterMap { f, input, metadata } => {
                        let input_ident = ident_stack.pop().unwrap();
                        let f_tokens = f.emit_tokens(&mut ident_stack);

                        let input_ident = maybe_observe_for_mut(
                            f, input_ident,
                            &input.metadata().location_id,
                            &input.metadata().collection_kind,
                            &metadata.op,
                            builders_or_callback, next_stmt_id,
                        );

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

                        match builders_or_callback {
                            BuildersOrCallback::Builders(graph_builders) => {
                                graph_builders.add_dfir_at(
                                    &out_location,
                                    parse_quote! {
                                        #filter_map_ident = #input_ident -> filter_map(#f_tokens);
                                    },
                                    Some(&stmt_id.to_string()),
                                );
                            }
                            BuildersOrCallback::Callback(_, node_callback) => {
                                node_callback(node, next_stmt_id);
                            }
                        }

                        ident_stack.push(filter_map_ident);
                    }

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

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

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

                        ident_stack.push(sort_ident);
                    }

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

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

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

                        ident_stack.push(defer_tick_ident);
                    }

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

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

                        match builders_or_callback {
                            BuildersOrCallback::Builders(graph_builders) => {
                                let lifetime = if input.metadata().location_id.is_top_level() {
                                    graph_builders.cross_tick_state_lifetime(&out_location)
                                } else {
                                    graph_builders.tick_state_lifetime(&out_location)
                                };
                                graph_builders.add_dfir_at(
                                    &out_location,
                                    parse_quote! {
                                        #enumerate_ident = #input_ident -> enumerate::<#lifetime>();
                                    },
                                    Some(&stmt_id.to_string()),
                                );
                            }
                            BuildersOrCallback::Callback(_, node_callback) => {
                                node_callback(node, next_stmt_id);
                            }
                        }

                        ident_stack.push(enumerate_ident);
                    }

                    HydroNode::Inspect { f, input, metadata } => {
                        let input_ident = ident_stack.pop().unwrap();
                        let f_tokens = f.emit_tokens(&mut ident_stack);

                        let input_ident = maybe_observe_for_mut(
                            f, input_ident,
                            &input.metadata().location_id,
                            &input.metadata().collection_kind,
                            &metadata.op,
                            builders_or_callback, next_stmt_id,
                        );

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

                        match builders_or_callback {
                            BuildersOrCallback::Builders(graph_builders) => {
                                graph_builders.add_dfir_at(
                                    &out_location,
                                    parse_quote! {
                                        #inspect_ident = #input_ident -> inspect(#f_tokens);
                                    },
                                    Some(&stmt_id.to_string()),
                                );
                            }
                            BuildersOrCallback::Callback(_, node_callback) => {
                                node_callback(node, next_stmt_id);
                            }
                        }

                        ident_stack.push(inspect_ident);
                    }

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

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

                        match builders_or_callback {
                            BuildersOrCallback::Builders(graph_builders) => {
                                let lifetime = if input.metadata().location_id.is_top_level() {
                                    graph_builders.cross_tick_state_lifetime(&out_location)
                                } else {
                                    graph_builders.tick_state_lifetime(&out_location)
                                };

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

                        ident_stack.push(unique_ident);
                    }

                    HydroNode::Fold { .. } | HydroNode::FoldKeyed { .. } | HydroNode::Scan { .. } | HydroNode::ScanAsyncBlocking { .. } => {
                        let operator: syn::Ident = if let HydroNode::Fold { input, .. } = node {
                            if input.metadata().location_id.is_root()
                                && 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 matches!(node, HydroNode::ScanAsyncBlocking { .. }) {
                            parse_quote!(scan_async_blocking)
                        } else if let HydroNode::FoldKeyed { input, .. } = node {
                            if input.metadata().location_id.is_root()
                                && 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, .. }
                        | HydroNode::ScanAsyncBlocking { input, .. }) = node
                        else {
                            unreachable!()
                        };

                        let input_top_level = input.metadata().location_id.is_top_level();

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

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

                        let acc_tokens = acc.emit_tokens(&mut ident_stack);
                        let init_tokens = init.emit_tokens(&mut ident_stack);

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

                        match builders_or_callback {
                            BuildersOrCallback::Builders(graph_builders) => {
                                let lifetime = if input_top_level {
                                    graph_builders.cross_tick_state_lifetime(&out_location)
                                } else {
                                    graph_builders.tick_state_lifetime(&out_location)
                                };

                                if matches!(node, HydroNode::Fold { .. })
                                    && node.metadata().location_id.is_root()
                                    && graph_builders.singleton_intermediates()
                                    && !node.metadata().collection_kind.is_bounded()
                                {
                                    let HydroNode::Fold { input, .. } = &*node else { unreachable!() };
                                    let hooked_input_ident = graph_builders.emit_fold_hook(
                                        &input.metadata().location_id,
                                        &input_ident,
                                        &input.metadata().collection_kind,
                                        &node.metadata().op,
                                    );

                                    let (effective_input, wrapped_acc) = if let Some(ref hooked) = hooked_input_ident {
                                        let acc: syn::Expr = parse_quote!({
                                            let mut __inner = #acc_tokens;
                                            move |__state, __batch: Vec<_>| {
                                                if __batch.is_empty() {
                                                    return None;
                                                }
                                                for __value in __batch {
                                                    __inner(__state, __value);
                                                }
                                                Some(__state.clone())
                                            }
                                        });
                                        (hooked, acc)
                                    } else {
                                        let acc: syn::Expr = parse_quote!({
                                            let mut __inner = #acc_tokens;
                                            move |__state, __value| {
                                                __inner(__state, __value);
                                                Some(__state.clone())
                                            }
                                        });
                                        (&input_ident, acc)
                                    };

                                    graph_builders.add_dfir_at(
                                        &out_location,
                                        parse_quote! {
                                            source_iter([(#init_tokens)()]) -> [0]#fold_ident;
                                            #effective_input -> scan::<#lifetime>(#init_tokens, #wrapped_acc) -> [1]#fold_ident;
                                            #fold_ident = chain();
                                        },
                                        Some(&stmt_id.to_string()),
                                    );

                                    // A *scripted* fold releases exactly one element per
                                    // decision, producing one new version per release, so
                                    // its snapshot takes the ordinary (scriptable) path
                                    // rather than the fuzz-passthrough shortcut.
                                    if hooked_input_ident.is_some()
                                        && node.metadata().op.sim_hook_id.is_none()
                                    {
                                        fold_hooked_idents.insert(fold_ident.to_string());
                                    }
                                } else if matches!(node, HydroNode::FoldKeyed { .. })
                                    && node.metadata().location_id.is_root()
                                    && graph_builders.singleton_intermediates()
                                    && !node.metadata().collection_kind.is_bounded()
                                {
                                    let HydroNode::FoldKeyed { input, .. } = &*node else { unreachable!() };
                                    let hooked_input_ident = graph_builders.emit_fold_hook(
                                        &input.metadata().location_id,
                                        &input_ident,
                                        &input.metadata().collection_kind,
                                        &node.metadata().op,
                                    );

                                    let wrapped_acc: syn::Expr = parse_quote!({
                                        let mut __init = #init_tokens;
                                        let mut __inner = #acc_tokens;
                                        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)))
                                        }
                                    });

                                    if let Some(hooked_input_ident) = hooked_input_ident {
                                        graph_builders.add_dfir_at(
                                            &out_location,
                                            parse_quote! {
                                                #fold_ident = #hooked_input_ident -> flatten() -> scan::<#lifetime>(|| ::std::collections::HashMap::new(), #wrapped_acc);
                                            },
                                            Some(&stmt_id.to_string()),
                                        );

                                        fold_hooked_idents.insert(fold_ident.to_string());
                                    } else {
                                        graph_builders.add_dfir_at(
                                            &out_location,
                                            parse_quote! {
                                                #fold_ident = #input_ident -> scan::<#lifetime>(|| ::std::collections::HashMap::new(), #wrapped_acc);
                                            },
                                            Some(&stmt_id.to_string()),
                                        );
                                    }
                                } else if (matches!(node, HydroNode::Fold { .. })
                                    || matches!(node, HydroNode::FoldKeyed { .. }))
                                    && !node.metadata().location_id.is_top_level()
                                    && graph_builders.singleton_intermediates()
                                {
                                    let input_ref = match &*node {
                                        HydroNode::Fold { input, .. } => input,
                                        HydroNode::FoldKeyed { input, .. } => input,
                                        _ => unreachable!(),
                                    };
                                    let hooked_input_ident = graph_builders.emit_fold_hook(
                                        &input_ref.metadata().location_id,
                                        &input_ident,
                                        &input_ref.metadata().collection_kind,
                                        &node.metadata().op,
                                    );

                                    let actual_input = hooked_input_ident.as_ref().unwrap_or(&input_ident);
                                    graph_builders.add_dfir_at(
                                        &out_location,
                                        parse_quote! {
                                            #fold_ident = #actual_input -> #operator::<#lifetime>(#init_tokens, #acc_tokens);
                                        },
                                        Some(&stmt_id.to_string()),
                                    );
                                } else {
                                    graph_builders.add_dfir_at(
                                        &out_location,
                                        parse_quote! {
                                            #fold_ident = #input_ident -> #operator::<#lifetime>(#init_tokens, #acc_tokens);
                                        },
                                        Some(&stmt_id.to_string()),
                                    );
                                }
                            }
                            BuildersOrCallback::Callback(_, node_callback) => {
                                node_callback(node, next_stmt_id);
                            }
                        }

                        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_root()
                                && 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_root()
                                && 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 input_top_level = input.metadata().location_id.is_top_level();

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

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

                        let f_tokens = f.emit_tokens(&mut ident_stack);

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

                        match builders_or_callback {
                            BuildersOrCallback::Builders(graph_builders) => {
                                let lifetime = if input_top_level {
                                    graph_builders.cross_tick_state_lifetime(&out_location)
                                } else {
                                    graph_builders.tick_state_lifetime(&out_location)
                                };

                                if matches!(node, HydroNode::Reduce { .. })
                                    && node.metadata().location_id.is_root()
                                    && 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_root()
                                    && graph_builders.singleton_intermediates()
                                    && !node.metadata().collection_kind.is_bounded()
                                {
                                    todo!(
                                        "Reduce keyed with optional intermediates is not yet supported in simulator"
                                    );
                                } else {
                                    graph_builders.add_dfir_at(
                                        &out_location,
                                        parse_quote! {
                                            #reduce_ident = #input_ident -> #operator::<#lifetime>(#f_tokens);
                                        },
                                        Some(&stmt_id.to_string()),
                                    );
                                }
                            }
                            BuildersOrCallback::Callback(_, node_callback) => {
                                node_callback(node, next_stmt_id);
                            }
                        }

                        ident_stack.push(reduce_ident);
                    }

                      HydroNode::ReduceKeyedWatermark {
                          f,
                          input,
                          watermark,
                          metadata,
                      } => {

                          // watermark is processed second, so it's on top
                          let watermark_ident = ident_stack.pop().unwrap();
                          let watermark_location = watermark.metadata().location_id.clone();
                          let input_ident = ident_stack.pop().unwrap();
                          let f_tokens = f.emit_tokens(&mut ident_stack);

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

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

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

                          match builders_or_callback {
                              BuildersOrCallback::Builders(graph_builders) => {
                                  // The watermark lives at its own (tick) location; if that is a
                                  // different loop context than where this reduce is emitted
                                  // (`out_location`), un-window it so the edge does not illegally
                                  // exit the loop.
                                  let watermark_ident = graph_builders.unwindow_for_consume(
                                      watermark_ident,
                                      &watermark_location,
                                      &out_location,
                                  );

                                  let lifetime = if input.metadata().location_id.is_top_level() {
                                      graph_builders.cross_tick_state_lifetime(&out_location)
                                  } else {
                                      graph_builders.tick_state_lifetime(&out_location)
                                  };

                                  if metadata.location_id.is_root()
                                      && graph_builders.singleton_intermediates()
                                      && !metadata.collection_kind.is_bounded()
                                  {
                                    todo!(
                                        "Reduce keyed watermarked on a top-level bounded collection is not yet supported"
                                    )
                                } else {
                                    graph_builders.add_dfir_at(
                                        &out_location,
                                        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_tokens;
                                                    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;
                                                                }
                                                            }
                                                            map.retain(|k, _| *k >= watermark);
                                                            *opt_curr_watermark = Some(watermark);
                                                        }
                                                    }
                                                })
                                                -> flat_map(|(map, _curr_watermark)| map);
                                        },
                                        Some(&stmt_id.to_string()),
                                    );
                                }
                            }
                            BuildersOrCallback::Callback(_, node_callback) => {
                                node_callback(node, next_stmt_id);
                            }
                        }

                        ident_stack.push(fold_ident);
                    }

                    HydroNode::Network {
                        networking_info,
                        serialize,
                        deserialize,
                        instantiate_fn,
                        input,
                        ..
                    } => {
                        let input_ident = ident_stack.pop().unwrap();

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

                        // For embedded (external) serialization, this synthesizes only the
                        // member-id tag conversions (if any) and passes the raw payload through.
                        let serialize_pipeline = serialize.pipeline();
                        let deserialize_pipeline = deserialize.pipeline();

                        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(),
                                    serialize.external_element_type(),
                                    stmt_id,
                                    networking_info,
                                );
                            }
                            BuildersOrCallback::Callback(_, node_callback) => {
                                node_callback(node, next_stmt_id);
                            }
                        }

                        ident_stack.push(receiver_stream_ident);
                    }

                    HydroNode::ExternalInput {
                        instantiate_fn,
                        deserialize_fn: deserialize_pipeline,
                        ..
                    } => {
                        let stmt_id = next_stmt_id.get_and_increment();
                        let receiver_stream_ident =
                            syn::Ident::new(&format!("stream_{}", 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(),
                                    stmt_id,
                                );
                            }
                            BuildersOrCallback::Callback(_, node_callback) => {
                                node_callback(node, next_stmt_id);
                            }
                        }

                        ident_stack.push(receiver_stream_ident);
                    }

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

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

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

                        ident_stack.push(counter_ident);
                    }

                    HydroNode::VersionedNetworkFork {
                        channel_id,
                        senders,
                        metadata,
                        ..
                    } => {
                        // sender idents are pushed in order of the 'senders' member.
                        let split_at = ident_stack.len() - senders.len();
                        let sender_idents = ident_stack.split_off(split_at);

                        let stmt_id = next_stmt_id.get_and_increment();

                        // All senders share the channel, so the raw element type (for embedded
                        // serialization) is read from the first sender.
                        let external_element_type =
                            senders.first().and_then(|(_, _, s)| s.external_element_type());

                        match builders_or_callback {
                            BuildersOrCallback::Builders(graph_builders) => {
                                let sender_args: Vec<(LocationId, syn::Ident, Option<DebugExpr>)> =
                                    senders
                                        .iter()
                                        .zip(sender_idents)
                                        .map(|((_version, sender, serialize), ident)| {
                                            (
                                                sender.metadata().location_id.clone(),
                                                ident,
                                                serialize.pipeline(),
                                            )
                                        })
                                        .collect();
                                graph_builders.create_versioned_network_fork(
                                    *channel_id,
                                    &metadata.location_id,
                                    sender_args,
                                    external_element_type,
                                    stmt_id,
                                );
                            }
                            BuildersOrCallback::Callback(_, node_callback) => {
                                node_callback(node, next_stmt_id);
                            }
                        }
                    }

                    HydroNode::VersionedNetwork {
                        fork,
                        deserialize,
                        metadata,
                        ..
                    } => {
                        let stmt_id = next_stmt_id.get_and_increment();
                        let receiver_stream_ident =
                            syn::Ident::new(&format!("stream_{}", stmt_id), Span::call_site());

                        // The wire element type is determined by the channel's *source* kind, which
                        // all senders share; read it from the shared fork's first sender.
                        let (channel_id, source_loc) = {
                            let fork_ref = fork.0.borrow();
                            let HydroNode::VersionedNetworkFork {
                                channel_id,
                                senders,
                                ..
                            } = &*fork_ref
                            else {
                                unreachable!("VersionedNetwork.fork must be a VersionedNetworkFork");
                            };
                            let source_loc = senders
                                .first()
                                .map(|(_v, sender, _s)| sender.metadata().location_id.clone())
                                .expect("a VersionedNetworkFork always has at least one sender");
                            (*channel_id, source_loc)
                        };

                        let deserialize_pipeline = deserialize.pipeline();
                        let external_element_type = deserialize.external_element_type();

                        match builders_or_callback {
                            BuildersOrCallback::Builders(graph_builders) => {
                                graph_builders.create_versioned_network(
                                    channel_id,
                                    &source_loc,
                                    &metadata.location_id,
                                    &receiver_stream_ident,
                                    deserialize_pipeline.as_ref(),
                                    external_element_type,
                                    stmt_id,
                                );
                            }
                            BuildersOrCallback::Callback(_, node_callback) => {
                                node_callback(node, next_stmt_id);
                            }
                        }

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

        let ret = ident_stack
            .pop()
            .expect("ident_stack should have exactly one element after traversal");
        assert!(
            ident_stack.is_empty(),
            "ident_stack should be empty after popping the final ident, but has {} remaining element(s). \
             This indicates a bug in the code gen: some node pushed idents that were never consumed.",
            ident_stack.len()
        );
        ret
    }

    pub fn visit_debug_expr(&mut self, mut transform: impl FnMut(&mut DebugExpr)) {
        match self {
            HydroNode::Placeholder => {
                panic!()
            }
            HydroNode::Cast { .. }
            | HydroNode::ObserveNonDet { .. }
            | HydroNode::UnboundSingleton { .. }
            | HydroNode::AssertIsConsistent { .. } => {}
            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::Reference { .. }
            | HydroNode::YieldConcat { .. }
            | HydroNode::BeginAtomic { .. }
            | HydroNode::EndAtomic { .. }
            | HydroNode::Batch { .. }
            | HydroNode::Chain { .. }
            | HydroNode::MergeOrdered { .. }
            | HydroNode::ChainFirst { .. }
            | HydroNode::CrossProduct { .. }
            | HydroNode::CrossSingleton { .. }
            | HydroNode::ResolveFutures { .. }
            | HydroNode::ResolveFuturesBlocking { .. }
            | HydroNode::ResolveFuturesOrdered { .. }
            | HydroNode::Join { .. }
            | HydroNode::JoinHalf { .. }
            | HydroNode::Difference { .. }
            | HydroNode::AntiJoin { .. }
            | HydroNode::DeferTick { .. }
            | HydroNode::Enumerate { .. }
            | HydroNode::Unique { .. }
            | HydroNode::Sort { .. }
            | HydroNode::PartitionSide { .. }
            | HydroNode::VersionedNetworkFork { .. }
            | HydroNode::VersionedNetwork { .. } => {}
            HydroNode::Map { f, .. }
            | HydroNode::FlatMap { f, .. }
            | HydroNode::FlatMapStreamBlocking { f, .. }
            | HydroNode::Filter { f, .. }
            | HydroNode::FilterMap { f, .. }
            | HydroNode::Inspect { f, .. }
            | HydroNode::PartitionShared { f, .. }
            | HydroNode::Reduce { f, .. }
            | HydroNode::ReduceKeyed { f, .. }
            | HydroNode::ReduceKeyedWatermark { f, .. } => {
                transform(&mut f.expr);
            }
            HydroNode::Fold { init, acc, .. }
            | HydroNode::Scan { init, acc, .. }
            | HydroNode::ScanAsyncBlocking { init, acc, .. }
            | HydroNode::FoldKeyed { init, acc, .. } => {
                transform(&mut init.expr);
                transform(&mut acc.expr);
            }
            HydroNode::Network {
                serialize,
                deserialize,
                ..
            } => {
                if let NetworkSend::Custom {
                    serialize_fn: Some(serialize_fn),
                } = serialize
                {
                    transform(serialize_fn);
                }
                if let NetworkRecv::Custom {
                    deserialize_fn: Some(deserialize_fn),
                } = deserialize
                {
                    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::VersionedNetworkFork { metadata, .. }
            | HydroNode::VersionedNetwork { metadata, .. } => metadata,
            HydroNode::Cast { metadata, .. }
            | HydroNode::ObserveNonDet { metadata, .. }
            | HydroNode::AssertIsConsistent { metadata, .. }
            | HydroNode::UnboundSingleton { metadata, .. }
            | HydroNode::Source { metadata, .. }
            | HydroNode::SingletonSource { metadata, .. }
            | HydroNode::CycleSource { metadata, .. }
            | HydroNode::Tee { metadata, .. }
            | HydroNode::Reference { metadata, .. }
            | HydroNode::PartitionSide { metadata, .. }
            | HydroNode::PartitionShared { metadata, .. }
            | HydroNode::YieldConcat { metadata, .. }
            | HydroNode::BeginAtomic { metadata, .. }
            | HydroNode::EndAtomic { metadata, .. }
            | HydroNode::Batch { metadata, .. }
            | HydroNode::Chain { metadata, .. }
            | HydroNode::MergeOrdered { metadata, .. }
            | HydroNode::ChainFirst { metadata, .. }
            | HydroNode::CrossProduct { metadata, .. }
            | HydroNode::CrossSingleton { metadata, .. }
            | HydroNode::Join { metadata, .. }
            | HydroNode::JoinHalf { metadata, .. }
            | HydroNode::Difference { metadata, .. }
            | HydroNode::AntiJoin { metadata, .. }
            | HydroNode::ResolveFutures { metadata, .. }
            | HydroNode::ResolveFuturesBlocking { metadata, .. }
            | HydroNode::ResolveFuturesOrdered { metadata, .. }
            | HydroNode::Map { metadata, .. }
            | HydroNode::FlatMap { metadata, .. }
            | HydroNode::FlatMapStreamBlocking { metadata, .. }
            | HydroNode::Filter { metadata, .. }
            | HydroNode::FilterMap { metadata, .. }
            | HydroNode::DeferTick { metadata, .. }
            | HydroNode::Enumerate { metadata, .. }
            | HydroNode::Inspect { metadata, .. }
            | HydroNode::Unique { metadata, .. }
            | HydroNode::Sort { metadata, .. }
            | HydroNode::Scan { metadata, .. }
            | HydroNode::ScanAsyncBlocking { metadata, .. }
            | HydroNode::Fold { metadata, .. }
            | HydroNode::FoldKeyed { metadata, .. }
            | HydroNode::Reduce { metadata, .. }
            | HydroNode::ReduceKeyed { metadata, .. }
            | HydroNode::ReduceKeyedWatermark { metadata, .. }
            | HydroNode::ExternalInput { metadata, .. }
            | HydroNode::Network { 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::VersionedNetworkFork { metadata, .. }
            | HydroNode::VersionedNetwork { metadata, .. } => metadata,
            HydroNode::Cast { metadata, .. }
            | HydroNode::ObserveNonDet { metadata, .. }
            | HydroNode::AssertIsConsistent { metadata, .. }
            | HydroNode::UnboundSingleton { metadata, .. }
            | HydroNode::Source { metadata, .. }
            | HydroNode::SingletonSource { metadata, .. }
            | HydroNode::CycleSource { metadata, .. }
            | HydroNode::Tee { metadata, .. }
            | HydroNode::Reference { metadata, .. }
            | HydroNode::PartitionSide { metadata, .. }
            | HydroNode::PartitionShared { metadata, .. }
            | HydroNode::YieldConcat { metadata, .. }
            | HydroNode::BeginAtomic { metadata, .. }
            | HydroNode::EndAtomic { metadata, .. }
            | HydroNode::Batch { metadata, .. }
            | HydroNode::Chain { metadata, .. }
            | HydroNode::MergeOrdered { metadata, .. }
            | HydroNode::ChainFirst { metadata, .. }
            | HydroNode::CrossProduct { metadata, .. }
            | HydroNode::CrossSingleton { metadata, .. }
            | HydroNode::Join { metadata, .. }
            | HydroNode::JoinHalf { metadata, .. }
            | HydroNode::Difference { metadata, .. }
            | HydroNode::AntiJoin { metadata, .. }
            | HydroNode::ResolveFutures { metadata, .. }
            | HydroNode::ResolveFuturesBlocking { metadata, .. }
            | HydroNode::ResolveFuturesOrdered { metadata, .. }
            | HydroNode::Map { metadata, .. }
            | HydroNode::FlatMap { metadata, .. }
            | HydroNode::FlatMapStreamBlocking { metadata, .. }
            | HydroNode::Filter { metadata, .. }
            | HydroNode::FilterMap { metadata, .. }
            | HydroNode::DeferTick { metadata, .. }
            | HydroNode::Enumerate { metadata, .. }
            | HydroNode::Inspect { metadata, .. }
            | HydroNode::Unique { metadata, .. }
            | HydroNode::Sort { metadata, .. }
            | HydroNode::Scan { metadata, .. }
            | HydroNode::ScanAsyncBlocking { metadata, .. }
            | HydroNode::Fold { metadata, .. }
            | HydroNode::FoldKeyed { metadata, .. }
            | HydroNode::Reduce { metadata, .. }
            | HydroNode::ReduceKeyed { metadata, .. }
            | HydroNode::ReduceKeyedWatermark { metadata, .. }
            | HydroNode::ExternalInput { metadata, .. }
            | HydroNode::Network { 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::Reference { .. }
            | HydroNode::PartitionSide { .. }
            | HydroNode::VersionedNetwork { .. } => {
                // Tee/PartitionSide/VersionedNetwork 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, .. }
            | HydroNode::UnboundSingleton { inner, .. }
            | HydroNode::AssertIsConsistent { inner, .. } => {
                vec![inner]
            }
            HydroNode::Chain { first, second, .. }
            | HydroNode::MergeOrdered { first, second, .. }
            | HydroNode::ChainFirst { first, second, .. } => {
                vec![first, second]
            }
            HydroNode::CrossProduct { left, right, .. }
            | HydroNode::CrossSingleton { left, right, .. }
            | HydroNode::Join { left, right, .. }
            | HydroNode::JoinHalf { left, right, .. } => {
                vec![left, right]
            }
            HydroNode::Difference { pos, neg, .. } | HydroNode::AntiJoin { pos, neg, .. } => {
                vec![pos, neg]
            }
            HydroNode::Counter { input, .. }
            | HydroNode::DeferTick { input, .. }
            | HydroNode::Enumerate { input, .. }
            | HydroNode::Filter { input, .. }
            | HydroNode::FilterMap { input, .. }
            | HydroNode::FlatMap { input, .. }
            | HydroNode::FlatMapStreamBlocking { input, .. }
            | HydroNode::Fold { input, .. }
            | HydroNode::FoldKeyed { input, .. }
            | HydroNode::Inspect { input, .. }
            | HydroNode::Map { input, .. }
            | HydroNode::Network { input, .. }
            | HydroNode::PartitionShared { input, .. }
            | HydroNode::Reduce { input, .. }
            | HydroNode::ReduceKeyed { input, .. }
            | HydroNode::ResolveFutures { input, .. }
            | HydroNode::ResolveFuturesBlocking { input, .. }
            | HydroNode::ResolveFuturesOrdered { input, .. }
            | HydroNode::Scan { input, .. }
            | HydroNode::ScanAsyncBlocking { input, .. }
            | HydroNode::Sort { input, .. }
            | HydroNode::Unique { input, .. } => {
                vec![input]
            }
            HydroNode::ReduceKeyedWatermark {
                input, watermark, ..
            } => {
                vec![input, watermark]
            }
            HydroNode::VersionedNetworkFork { senders, .. } => senders
                .iter()
                .map(|(_version, sender, _serialize)| sender.as_ref())
                .collect(),
        }
    }

    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::PartitionSide { inner, .. } => {
                Rc::strong_count(&inner.0) > 1
            }
            // A zero-output reference node is valid in DFIR (it drains itself at
            // end of tick), so it doesn't need to be driven by another consumer.
            HydroNode::Reference { .. } => false,
            _ => false,
        }
    }

    pub fn print_root(&self) -> String {
        match self {
            HydroNode::Placeholder => {
                panic!()
            }
            HydroNode::Cast { .. } => "Cast()".to_owned(),
            HydroNode::UnboundSingleton { .. } => "UnboundSingleton()".to_owned(),
            HydroNode::ObserveNonDet { .. } => "ObserveNonDet()".to_owned(),
            HydroNode::AssertIsConsistent { .. } => "AssertIsConsistent()".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::Reference { inner, kind, .. } => {
                format!("Reference({:?}, {})", kind, inner.0.borrow().print_root())
            }
            HydroNode::PartitionSide { inner, is_true, .. } => {
                format!(
                    "PartitionSide(is_true={}, {})",
                    is_true,
                    inner.0.borrow().print_root(),
                )
            }
            HydroNode::PartitionShared { f, .. } => format!("PartitionShared({:?})", f),
            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::MergeOrdered { first, second, .. } => {
                format!(
                    "MergeOrdered({}, {})",
                    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::JoinHalf { left, right, .. } => {
                format!("JoinHalf({}, {})", 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::ResolveFuturesBlocking { .. } => "ResolveFuturesBlocking()".to_owned(),
            HydroNode::ResolveFuturesOrdered { .. } => "ResolveFuturesOrdered()".to_owned(),
            HydroNode::Map { f, .. } => format!("Map({:?})", f),
            HydroNode::FlatMap { f, .. } => format!("FlatMap({:?})", f),
            HydroNode::FlatMapStreamBlocking { f, .. } => format!("FlatMapStreamBlocking({:?})", 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::ScanAsyncBlocking { init, acc, .. } => {
                format!("ScanAsyncBlocking({:?}, {:?})", 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)
            }
            HydroNode::VersionedNetworkFork {
                channel_name,
                senders,
                ..
            } => {
                let versions: Vec<u32> = senders.iter().map(|(v, _, _)| *v).collect();
                format!(
                    "VersionedNetworkFork({}, senders={:?})",
                    channel_name, versions
                )
            }
            HydroNode::VersionedNetwork { version, .. } => {
                format!("VersionedNetwork(v{})", version)
            }
        }
    }
}

#[cfg(feature = "build")]
#[expect(clippy::too_many_arguments, reason = "networking codegen")]
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,
    external_types: Option<(&syn::Type, &syn::Type)>,
) -> (syn::Expr, syn::Expr, Box<dyn FnOnce()>)
where
    D: Deploy<'a>,
{
    if external_types.is_some() && !D::SUPPORTS_EXTERNAL_SERIALIZATION {
        panic!(
            "`.embedded()` serialization leaves serialization to code outside of Hydro and is \
             only supported by the embedded deployment backend. Use `.bincode()` (or another \
             supported serialization backend) for this deployment target instead."
        );
    }

    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,
                    external_types,
                ),
                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,
                    external_types,
                ),
                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,
                    external_types,
                ),
                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,
                    external_types,
                ),
                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 serde_test;

#[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>(), 280);
    }

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

    #[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());
    }
}