Skip to main content

hydro_lang/compile/ir/
mod.rs

1use core::panic;
2use std::cell::{Cell, RefCell};
3use std::collections::HashMap;
4#[cfg(feature = "build")]
5use std::collections::HashSet;
6use std::fmt::{Debug, Display};
7use std::hash::{Hash, Hasher};
8use std::ops::Deref;
9use std::rc::Rc;
10
11#[cfg(feature = "build")]
12use dfir_lang::graph::FlatGraphBuilder;
13#[cfg(feature = "build")]
14use proc_macro2::Span;
15use proc_macro2::TokenStream;
16use quote::ToTokens;
17#[cfg(feature = "build")]
18use quote::quote;
19#[cfg(feature = "build")]
20use slotmap::{SecondaryMap, SparseSecondaryMap};
21#[cfg(feature = "build")]
22use syn::parse_quote;
23
24#[cfg(feature = "build")]
25use crate::compile::builder::ClockId;
26#[cfg(feature = "build")]
27use crate::compile::builder::StmtId;
28use crate::compile::builder::{CycleId, ExternalPortId};
29#[cfg(feature = "build")]
30use crate::compile::deploy_provider::{Deploy, Node, RegisterPort};
31#[cfg(feature = "build")]
32use crate::handoff_ref::handoff_ref_ident;
33use crate::location::dynamic::{ClusterConsistency, LocationId};
34use crate::location::{LocationKey, NetworkHint};
35
36pub mod backtrace;
37use backtrace::Backtrace;
38
39/// A closure expression bundled with any singleton references it captures.
40///
41/// When a `q!()` closure captures a `SingletonRef`, the reference is recorded here
42/// alongside the closure's expression. This allows per-closure tracking of singleton
43/// captures, which is important for nodes with multiple closures (e.g. Fold has `init` and `acc`).
44pub struct ClosureExpr {
45    pub(crate) expr: DebugExpr,
46    /// Each entry is `(HydroNode::Reference, is_mut: bool)`.
47    /// The index in the Vec determines the ident name via [`handoff_ref_ident`].
48    /// The `access_counter` was assigned at staging time in code order.
49    pub(crate) singleton_refs: Vec<(HydroNode, bool)>,
50}
51
52impl Clone for ClosureExpr {
53    fn clone(&self) -> Self {
54        Self {
55            expr: self.expr.clone(),
56            singleton_refs: self
57                .singleton_refs
58                .iter()
59                .map(|(node, is_mut)| {
60                    let HydroNode::Reference {
61                        inner,
62                        kind,
63                        access_counter,
64                        metadata,
65                    } = node
66                    else {
67                        panic!("singleton_refs should only contain HydroNode::Reference");
68                    };
69                    (
70                        HydroNode::Reference {
71                            inner: SharedNode(Rc::clone(&inner.0)),
72                            kind: *kind,
73                            access_counter: access_counter.freeze(),
74                            metadata: metadata.clone(),
75                        },
76                        *is_mut,
77                    )
78                })
79                .collect(),
80        }
81    }
82}
83
84impl Hash for ClosureExpr {
85    fn hash<H: Hasher>(&self, state: &mut H) {
86        self.expr.hash(state);
87        // singleton_refs are structural children (like HydroIrMetadata), not
88        // identity-defining. Two closures with the same expr but different
89        // captured refs are the same closure text — the refs only affect codegen.
90    }
91}
92
93impl serde::Serialize for ClosureExpr {
94    fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
95        use serde::ser::SerializeStruct;
96        let mut s = serializer.serialize_struct("ClosureExpr", 2)?;
97        s.serialize_field("expr", &self.expr)?;
98        s.serialize_field(
99            "singleton_refs",
100            &SerializableSingletonRefs(&self.singleton_refs),
101        )?;
102        s.end()
103    }
104}
105
106struct SerializableSingletonRefs<'a>(&'a [(HydroNode, bool)]);
107
108impl serde::Serialize for SerializableSingletonRefs<'_> {
109    fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
110        use serde::ser::SerializeSeq;
111        let mut seq = serializer.serialize_seq(Some(self.0.len()))?;
112        for (node, is_mut) in self.0.iter() {
113            seq.serialize_element(&(node, is_mut))?;
114        }
115        seq.end()
116    }
117}
118
119impl Debug for ClosureExpr {
120    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
121        Debug::fmt(&self.expr, f)
122    }
123}
124
125impl Display for ClosureExpr {
126    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
127        Display::fmt(&self.expr, f)
128    }
129}
130
131impl From<syn::Expr> for ClosureExpr {
132    fn from(expr: syn::Expr) -> Self {
133        Self {
134            expr: DebugExpr(Box::new(expr)),
135            singleton_refs: Vec::new(),
136        }
137    }
138}
139
140impl From<DebugExpr> for ClosureExpr {
141    fn from(expr: DebugExpr) -> Self {
142        Self {
143            expr,
144            singleton_refs: Vec::new(),
145        }
146    }
147}
148
149impl ClosureExpr {
150    pub fn new(expr: DebugExpr, singleton_refs: Vec<(HydroNode, bool)>) -> Self {
151        Self {
152            expr,
153            singleton_refs,
154        }
155    }
156
157    pub fn has_mut_ref(&self) -> bool {
158        self.singleton_refs.iter().any(|(_, is_mut)| *is_mut)
159    }
160
161    pub fn deep_clone(&self, seen_tees: &mut SeenSharedNodes) -> Self {
162        Self {
163            expr: self.expr.clone(),
164            singleton_refs: self
165                .singleton_refs
166                .iter()
167                .map(|(node, is_mut)| (node.deep_clone(seen_tees), *is_mut))
168                .collect(),
169        }
170    }
171
172    pub fn transform_children(
173        &mut self,
174        transform: &mut impl FnMut(&mut HydroNode, &mut SeenSharedNodes),
175        seen_tees: &mut SeenSharedNodes,
176    ) {
177        for (ref_node, _is_mut) in self.singleton_refs.iter_mut() {
178            transform(ref_node, seen_tees);
179        }
180    }
181
182    /// Pop singleton ref idents from the stack and rewrite the closure's token stream,
183    /// replacing local singleton ref idents with `#{N} dfir_ident` or `#{N} mut dfir_ident` references.
184    #[cfg(feature = "build")]
185    pub fn emit_tokens(&self, ident_stack: &mut Vec<syn::Ident>) -> TokenStream {
186        if self.singleton_refs.is_empty() {
187            self.expr.0.to_token_stream()
188        } else {
189            assert!(
190                ident_stack.len() >= self.singleton_refs.len(),
191                "ident_stack has {} entries but expected at least {} for singleton_refs",
192                ident_stack.len(),
193                self.singleton_refs.len()
194            );
195            let ref_idents = ident_stack.drain(ident_stack.len() - self.singleton_refs.len()..);
196
197            let mut let_bindings = Vec::new();
198            for ((i, (ref_node, is_mut)), ref_ident) in
199                self.singleton_refs.iter().enumerate().zip(ref_idents)
200            {
201                let HydroNode::Reference { access_counter, .. } = ref_node else {
202                    panic!("ClosureExpression expected references to `HydroNode::Reference`");
203                };
204                let group = access_counter.frozen_group();
205                // TODO(mingwei): proper spanning?
206                let local_ident = handoff_ref_ident(i);
207                let hash = proc_macro2::Punct::new('#', proc_macro2::Spacing::Alone);
208                let group_lit = proc_macro2::Literal::u32_unsuffixed(group);
209                let mut_token = is_mut.then(|| quote!(mut));
210                let binding = quote! {
211                    let #local_ident = #hash {#group_lit} #mut_token #ref_ident;
212                };
213                let_bindings.push(binding);
214            }
215
216            let expr = &self.expr.0;
217            quote! {
218                {
219                    #( #let_bindings )*
220                    #expr
221                }
222            }
223        }
224    }
225}
226
227/// Wrapper that displays only the tokens of a parsed expr.
228///
229/// Boxes `syn::Type` which is ~240 bytes.
230#[derive(Clone, Hash)]
231pub struct DebugExpr(pub Box<syn::Expr>);
232
233impl serde::Serialize for DebugExpr {
234    fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
235        serializer.serialize_str(&self.to_string())
236    }
237}
238
239impl From<syn::Expr> for DebugExpr {
240    fn from(expr: syn::Expr) -> Self {
241        Self(Box::new(expr))
242    }
243}
244
245impl Deref for DebugExpr {
246    type Target = syn::Expr;
247
248    fn deref(&self) -> &Self::Target {
249        &self.0
250    }
251}
252
253impl ToTokens for DebugExpr {
254    fn to_tokens(&self, tokens: &mut TokenStream) {
255        self.0.to_tokens(tokens);
256    }
257}
258
259impl Debug for DebugExpr {
260    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
261        write!(f, "{}", self.0.to_token_stream())
262    }
263}
264
265impl Display for DebugExpr {
266    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
267        let original = self.0.as_ref().clone();
268        let simplified = simplify_q_macro(original);
269
270        // For now, just use quote formatting without trying to parse as a statement
271        // This avoids the syn::parse_quote! issues entirely
272        write!(f, "q!({})", quote::quote!(#simplified))
273    }
274}
275
276/// Simplify expanded q! macro calls back to q!(...) syntax for better readability
277fn simplify_q_macro(expr: syn::Expr) -> syn::Expr {
278    if let syn::Expr::Call(ref call) = expr && let syn::Expr::Path(path_expr) = call.func.as_ref()
279        // Look for calls to stageleft::runtime_support::fn*
280        && is_stageleft_runtime_support_call(&path_expr.path)
281        && let syn::Expr::Block(b) = &call.args[0]
282        && b.block.stmts.len() == 3
283        && let Some(syn::Stmt::Expr(e, _)) = b.block.stmts.get(2)
284    // skip the first two, which are imports
285    {
286        let mut e = e.clone();
287        while let syn::Expr::Block(ref mut block) = e
288            && block.block.stmts.len() == 1
289            && let syn::Stmt::Expr(inner_e, _) = block.block.stmts.remove(0)
290        {
291            e = inner_e;
292        }
293
294        e
295    } else {
296        expr
297    }
298}
299
300fn is_stageleft_runtime_support_call(path: &syn::Path) -> bool {
301    // Check if this is a call to stageleft::runtime_support::fn*
302    if let Some(last_segment) = path.segments.last() {
303        let fn_name = last_segment.ident.to_string();
304        path.segments.len() > 2
305            && path.segments[0].ident == "stageleft"
306            && path.segments[1].ident == "runtime_support"
307            && fn_name.contains("_type_hint")
308    } else {
309        false
310    }
311}
312
313/// Debug displays the type's tokens.
314///
315/// Boxes `syn::Type` which is ~320 bytes.
316#[derive(Clone, PartialEq, Eq, Hash)]
317pub struct DebugType(pub Box<syn::Type>);
318
319impl From<syn::Type> for DebugType {
320    fn from(t: syn::Type) -> Self {
321        Self(Box::new(t))
322    }
323}
324
325impl Deref for DebugType {
326    type Target = syn::Type;
327
328    fn deref(&self) -> &Self::Target {
329        &self.0
330    }
331}
332
333impl ToTokens for DebugType {
334    fn to_tokens(&self, tokens: &mut TokenStream) {
335        self.0.to_tokens(tokens);
336    }
337}
338
339impl Debug for DebugType {
340    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
341        write!(f, "{}", self.0.to_token_stream())
342    }
343}
344
345impl serde::Serialize for DebugType {
346    fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
347        serializer.serialize_str(&format!("{}", self.0.to_token_stream()))
348    }
349}
350
351fn serialize_backtrace_as_span<S: serde::Serializer>(
352    backtrace: &Backtrace,
353    serializer: S,
354) -> Result<S::Ok, S::Error> {
355    match backtrace.format_span() {
356        Some(span) => serializer.serialize_some(&span),
357        None => serializer.serialize_none(),
358    }
359}
360
361fn serialize_ident<S: serde::Serializer>(
362    ident: &syn::Ident,
363    serializer: S,
364) -> Result<S::Ok, S::Error> {
365    serializer.serialize_str(&ident.to_string())
366}
367
368pub enum DebugInstantiate {
369    Building,
370    Finalized(Box<DebugInstantiateFinalized>),
371}
372
373impl serde::Serialize for DebugInstantiate {
374    fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
375        match self {
376            DebugInstantiate::Building => {
377                serializer.serialize_unit_variant("DebugInstantiate", 0, "Building")
378            }
379            DebugInstantiate::Finalized(_) => {
380                panic!(
381                    "cannot serialize DebugInstantiate::Finalized: contains non-serializable runtime state (closures)"
382                )
383            }
384        }
385    }
386}
387
388#[cfg_attr(
389    not(feature = "build"),
390    expect(
391        dead_code,
392        reason = "sink, source unused without `feature = \"build\"`."
393    )
394)]
395pub struct DebugInstantiateFinalized {
396    sink: syn::Expr,
397    source: syn::Expr,
398    connect_fn: Option<Box<dyn FnOnce()>>,
399}
400
401impl From<DebugInstantiateFinalized> for DebugInstantiate {
402    fn from(f: DebugInstantiateFinalized) -> Self {
403        Self::Finalized(Box::new(f))
404    }
405}
406
407impl Debug for DebugInstantiate {
408    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
409        write!(f, "<network instantiate>")
410    }
411}
412
413impl Hash for DebugInstantiate {
414    fn hash<H: Hasher>(&self, _state: &mut H) {
415        // Do nothing
416    }
417}
418
419impl Clone for DebugInstantiate {
420    fn clone(&self) -> Self {
421        match self {
422            DebugInstantiate::Building => DebugInstantiate::Building,
423            DebugInstantiate::Finalized(_) => {
424                panic!("DebugInstantiate::Finalized should not be cloned")
425            }
426        }
427    }
428}
429
430/// Tracks the instantiation state of a `ClusterMembers` source.
431///
432/// During `compile_network`, the first `ClusterMembers` node for a given
433/// `(at_location, target_cluster)` pair is promoted to [`Self::Stream`] and
434/// receives the expression returned by `Deploy::cluster_membership_stream`.
435/// All subsequent nodes for the same pair are set to [`Self::Tee`] so that
436/// during code-gen they simply reference the tee output of the first node
437/// instead of creating a redundant `source_stream`.
438#[derive(Debug, Hash, Clone, serde::Serialize)]
439pub enum ClusterMembersState {
440    /// Not yet instantiated.
441    Uninit,
442    /// The primary instance: holds the stream expression and will emit
443    /// `source_stream(expr) -> tee()` during code-gen.
444    Stream(DebugExpr),
445    /// A secondary instance that references the tee output of the primary.
446    /// Stores `(at_location_root, target_cluster_location)` so that `emit_core`
447    /// can derive the deterministic tee ident without extra state.
448    Tee(LocationId, LocationId),
449}
450
451/// A source in a Hydro graph, where data enters the graph.
452#[derive(Debug, Hash, Clone, serde::Serialize)]
453pub enum HydroSource {
454    Stream(DebugExpr),
455    ExternalNetwork(),
456    Iter(DebugExpr),
457    Spin(),
458    ClusterMembers(LocationId, ClusterMembersState),
459    Embedded(#[serde(serialize_with = "serialize_ident")] syn::Ident),
460    EmbeddedSingleton(#[serde(serialize_with = "serialize_ident")] syn::Ident),
461}
462
463#[cfg(feature = "build")]
464/// A trait that abstracts over elements of DFIR code-gen that differ between production deployment
465/// and simulations.
466///
467/// In particular, this lets the simulator fuse together all locations into one DFIR graph, spit
468/// out separate graphs for each tick, and emit hooks for controlling non-deterministic operators.
469pub trait DfirBuilder {
470    /// Whether the representation of singletons should include intermediate states.
471    fn singleton_intermediates(&self) -> bool;
472
473    /// Adds the DFIR statements to the graph for the given location.
474    ///
475    /// The location determines which DFIR graph the statements are placed in (for production,
476    /// the graph of the location's root; for simulation, either the fused async graph or the
477    /// tick's separate graph). In the future (#2902), production codegen will also use the
478    /// location to place tick-located statements inside the tick's `loop { ... }` context.
479    fn add_dfir_at(
480        &mut self,
481        location: &LocationId,
482        dfir: dfir_lang::parse::DfirCode,
483        operator_tag: Option<&str>,
484    );
485
486    /// The DFIR persistence lifetime for operator state scoped to a single tick, for an operator
487    /// at `op_location`.
488    ///
489    /// Returns `'tick`. In the future (#2902), production codegen will emit tick regions as DFIR
490    /// `loop { ... }` blocks, where this must instead be `'none` when `op_location` is a tick.
491    fn tick_state_lifetime(&self, _op_location: &LocationId) -> TokenStream {
492        quote!('tick)
493    }
494
495    /// The DFIR persistence lifetime for operator state that accumulates across ticks, for an
496    /// operator at `op_location`.
497    ///
498    /// Returns `'static`. In the future (#2902), production codegen will emit tick regions as
499    /// DFIR `loop { ... }` blocks, where this must instead be `'loop` when `op_location` is a
500    /// tick.
501    fn cross_tick_state_lifetime(&self, _op_location: &LocationId) -> TokenStream {
502        quote!('static)
503    }
504
505    #[expect(clippy::too_many_arguments, reason = "TODO")]
506    fn batch(
507        &mut self,
508        in_ident: syn::Ident,
509        in_location: &LocationId,
510        in_kind: &CollectionKind,
511        out_ident: &syn::Ident,
512        out_location: &LocationId,
513        op_meta: &HydroIrOpMetadata,
514        fold_hooked_idents: &HashSet<String>,
515    );
516    fn yield_from_tick(
517        &mut self,
518        in_ident: syn::Ident,
519        in_location: &LocationId,
520        in_kind: &CollectionKind,
521        out_ident: &syn::Ident,
522        out_location: &LocationId,
523    );
524
525    fn begin_atomic(
526        &mut self,
527        in_ident: syn::Ident,
528        in_location: &LocationId,
529        in_kind: &CollectionKind,
530        out_ident: &syn::Ident,
531        out_location: &LocationId,
532        op_meta: &HydroIrOpMetadata,
533    );
534    fn end_atomic(
535        &mut self,
536        in_ident: syn::Ident,
537        in_location: &LocationId,
538        in_kind: &CollectionKind,
539        out_ident: &syn::Ident,
540    );
541
542    #[expect(clippy::too_many_arguments, reason = "TODO // internal")]
543    fn observe_nondet(
544        &mut self,
545        trusted: bool,
546        location: &LocationId,
547        in_ident: syn::Ident,
548        in_kind: &CollectionKind,
549        out_ident: &syn::Ident,
550        out_kind: &CollectionKind,
551        op_meta: &HydroIrOpMetadata,
552    );
553
554    #[expect(clippy::too_many_arguments, reason = "TODO")]
555    fn merge_ordered(
556        &mut self,
557        location: &LocationId,
558        first_ident: syn::Ident,
559        second_ident: syn::Ident,
560        out_ident: &syn::Ident,
561        in_kind: &CollectionKind,
562        op_meta: &HydroIrOpMetadata,
563        operator_tag: Option<&str>,
564    );
565
566    #[expect(clippy::too_many_arguments, reason = "TODO")]
567    fn create_network(
568        &mut self,
569        from: &LocationId,
570        to: &LocationId,
571        input_ident: syn::Ident,
572        out_ident: &syn::Ident,
573        serialize: Option<&DebugExpr>,
574        sink: syn::Expr,
575        source: syn::Expr,
576        deserialize: Option<&DebugExpr>,
577        external_element_type: Option<&syn::Type>,
578        tag_id: StmtId,
579        networking_info: &crate::networking::NetworkingInfo,
580    );
581
582    fn create_external_source(
583        &mut self,
584        on: &LocationId,
585        source_expr: syn::Expr,
586        out_ident: &syn::Ident,
587        deserialize: Option<&DebugExpr>,
588        tag_id: StmtId,
589    );
590
591    fn create_external_output(
592        &mut self,
593        on: &LocationId,
594        sink_expr: syn::Expr,
595        input_ident: &syn::Ident,
596        serialize: Option<&DebugExpr>,
597        tag_id: StmtId,
598    );
599
600    /// Optionally emit a fold hook that buffers and permutes inputs before the fold.
601    /// Returns the new input ident to use for the fold if a hook was emitted.
602    fn emit_fold_hook(
603        &mut self,
604        location: &LocationId,
605        in_ident: &syn::Ident,
606        in_kind: &CollectionKind,
607        op_meta: &HydroIrOpMetadata,
608    ) -> Option<syn::Ident>;
609
610    /// Inserts necessary code to validate a manual assertion that at this point the
611    /// input live collection is consistent. In production, this is a no-op, but in simulation
612    /// this will (not yet implemented) inject assertions that validate consistency.
613    fn assert_is_consistent(
614        &mut self,
615        trusted: bool,
616        location: &LocationId,
617        in_ident: syn::Ident,
618        out_ident: &syn::Ident,
619    );
620
621    /// Observes non-determinism introduced by a mut closure operating on a non-strict
622    /// (unordered / at-least-once) input. In production this is identity; in simulation
623    /// it delegates to `observe_nondet` with the strict output kind.
624    fn observe_for_mut(
625        &mut self,
626        location: &LocationId,
627        in_ident: syn::Ident,
628        in_kind: &CollectionKind,
629        out_ident: &syn::Ident,
630        op_meta: &HydroIrOpMetadata,
631    );
632
633    fn create_versioned_network_fork(
634        &mut self,
635        channel_id: u32,
636        dest: &LocationId,
637        senders: Vec<(LocationId, syn::Ident, Option<DebugExpr>)>,
638        external_element_type: Option<&syn::Type>,
639        tag_id: StmtId,
640    );
641
642    #[expect(clippy::too_many_arguments, reason = "networking codegen")]
643    fn create_versioned_network(
644        &mut self,
645        channel_id: u32,
646        source: &LocationId,
647        dest: &LocationId,
648        out_ident: &syn::Ident,
649        deserialize: Option<&DebugExpr>,
650        external_element_type: Option<&syn::Type>,
651        tag_id: StmtId,
652    );
653}
654
655/// The production (deployment) DFIR builder: emits one DFIR graph per root location
656/// (process/cluster).
657///
658/// Tick and atomic locations are collapsed onto their root location's graph. In the future
659/// (#2902), this builder will additionally emit each (unified) tick as a root-level
660/// `loop {{ ... }}` context within its root location's graph.
661#[cfg(feature = "build")]
662#[derive(Default)]
663pub struct ProdDfirBuilder {
664    /// The DFIR graph builder for each root location.
665    pub graphs: SecondaryMap<LocationKey, FlatGraphBuilder>,
666}
667
668#[cfg(feature = "build")]
669impl ProdDfirBuilder {
670    /// Gets the DFIR builder for the given location's root, creating it if necessary.
671    fn graph_mut(&mut self, location: &LocationId) -> &mut FlatGraphBuilder {
672        self.graphs
673            .entry(location.root().key())
674            .expect("location was removed")
675            .or_default()
676    }
677}
678
679#[cfg(feature = "build")]
680impl DfirBuilder for ProdDfirBuilder {
681    fn singleton_intermediates(&self) -> bool {
682        false
683    }
684
685    fn add_dfir_at(
686        &mut self,
687        location: &LocationId,
688        dfir: dfir_lang::parse::DfirCode,
689        operator_tag: Option<&str>,
690    ) {
691        self.graph_mut(location).add_dfir(dfir, None, operator_tag);
692    }
693
694    fn batch(
695        &mut self,
696        in_ident: syn::Ident,
697        in_location: &LocationId,
698        in_kind: &CollectionKind,
699        out_ident: &syn::Ident,
700        _out_location: &LocationId,
701        _op_meta: &HydroIrOpMetadata,
702        _fold_hooked_idents: &HashSet<String>,
703    ) {
704        let builder = self.graph_mut(in_location.root());
705        if in_kind.is_bounded()
706            && matches!(
707                in_kind,
708                CollectionKind::Singleton { .. }
709                    | CollectionKind::Optional { .. }
710                    | CollectionKind::KeyedSingleton { .. }
711            )
712        {
713            assert!(in_location.is_top_level());
714            builder.add_dfir(
715                parse_quote! {
716                    #out_ident = #in_ident -> persist::<'static>();
717                },
718                None,
719                None,
720            );
721        } else {
722            builder.add_dfir(
723                parse_quote! {
724                    #out_ident = #in_ident;
725                },
726                None,
727                None,
728            );
729        }
730    }
731
732    fn yield_from_tick(
733        &mut self,
734        in_ident: syn::Ident,
735        in_location: &LocationId,
736        _in_kind: &CollectionKind,
737        out_ident: &syn::Ident,
738        _out_location: &LocationId,
739    ) {
740        let builder = self.graph_mut(in_location.root());
741        builder.add_dfir(
742            parse_quote! {
743                #out_ident = #in_ident;
744            },
745            None,
746            None,
747        );
748    }
749
750    fn begin_atomic(
751        &mut self,
752        in_ident: syn::Ident,
753        in_location: &LocationId,
754        _in_kind: &CollectionKind,
755        out_ident: &syn::Ident,
756        _out_location: &LocationId,
757        _op_meta: &HydroIrOpMetadata,
758    ) {
759        let builder = self.graph_mut(in_location.root());
760        builder.add_dfir(
761            parse_quote! {
762                #out_ident = #in_ident;
763            },
764            None,
765            None,
766        );
767    }
768
769    fn end_atomic(
770        &mut self,
771        in_ident: syn::Ident,
772        in_location: &LocationId,
773        _in_kind: &CollectionKind,
774        out_ident: &syn::Ident,
775    ) {
776        let builder = self.graph_mut(in_location.root());
777        builder.add_dfir(
778            parse_quote! {
779                #out_ident = #in_ident;
780            },
781            None,
782            None,
783        );
784    }
785
786    fn observe_nondet(
787        &mut self,
788        _trusted: bool,
789        location: &LocationId,
790        in_ident: syn::Ident,
791        _in_kind: &CollectionKind,
792        out_ident: &syn::Ident,
793        _out_kind: &CollectionKind,
794        _op_meta: &HydroIrOpMetadata,
795    ) {
796        let builder = self.graph_mut(location);
797        builder.add_dfir(
798            parse_quote! {
799                #out_ident = #in_ident;
800            },
801            None,
802            None,
803        );
804    }
805
806    fn merge_ordered(
807        &mut self,
808        location: &LocationId,
809        first_ident: syn::Ident,
810        second_ident: syn::Ident,
811        out_ident: &syn::Ident,
812        _in_kind: &CollectionKind,
813        _op_meta: &HydroIrOpMetadata,
814        operator_tag: Option<&str>,
815    ) {
816        let builder = self.graph_mut(location);
817        builder.add_dfir(
818            parse_quote! {
819                #out_ident = union();
820                #first_ident -> [0]#out_ident;
821                #second_ident -> [1]#out_ident;
822            },
823            None,
824            operator_tag,
825        );
826    }
827
828    fn create_network(
829        &mut self,
830        from: &LocationId,
831        to: &LocationId,
832        input_ident: syn::Ident,
833        out_ident: &syn::Ident,
834        serialize: Option<&DebugExpr>,
835        sink: syn::Expr,
836        source: syn::Expr,
837        deserialize: Option<&DebugExpr>,
838        _external_element_type: Option<&syn::Type>,
839        tag_id: StmtId,
840        _networking_info: &crate::networking::NetworkingInfo,
841    ) {
842        let sender_builder = self.graph_mut(from);
843        if let Some(serialize_pipeline) = serialize {
844            sender_builder.add_dfir(
845                parse_quote! {
846                    #input_ident -> map(#serialize_pipeline) -> dest_sink(#sink);
847                },
848                None,
849                // operator tag separates send and receive, which otherwise have the same next_stmt_id
850                Some(&format!("send{}", tag_id)),
851            );
852        } else {
853            sender_builder.add_dfir(
854                parse_quote! {
855                    #input_ident -> dest_sink(#sink);
856                },
857                None,
858                Some(&format!("send{}", tag_id)),
859            );
860        }
861
862        let receiver_builder = self.graph_mut(to);
863        if let Some(deserialize_pipeline) = deserialize {
864            receiver_builder.add_dfir(
865                parse_quote! {
866                    #out_ident = source_stream(#source) -> map(#deserialize_pipeline);
867                },
868                None,
869                Some(&format!("recv{}", tag_id)),
870            );
871        } else {
872            receiver_builder.add_dfir(
873                parse_quote! {
874                    #out_ident = source_stream(#source);
875                },
876                None,
877                Some(&format!("recv{}", tag_id)),
878            );
879        }
880    }
881
882    fn create_external_source(
883        &mut self,
884        on: &LocationId,
885        source_expr: syn::Expr,
886        out_ident: &syn::Ident,
887        deserialize: Option<&DebugExpr>,
888        tag_id: StmtId,
889    ) {
890        let receiver_builder = self.graph_mut(on);
891        if let Some(deserialize_pipeline) = deserialize {
892            receiver_builder.add_dfir(
893                parse_quote! {
894                    #out_ident = source_stream(#source_expr) -> map(#deserialize_pipeline);
895                },
896                None,
897                Some(&format!("recv{}", tag_id)),
898            );
899        } else {
900            receiver_builder.add_dfir(
901                parse_quote! {
902                    #out_ident = source_stream(#source_expr);
903                },
904                None,
905                Some(&format!("recv{}", tag_id)),
906            );
907        }
908    }
909
910    fn create_external_output(
911        &mut self,
912        on: &LocationId,
913        sink_expr: syn::Expr,
914        input_ident: &syn::Ident,
915        serialize: Option<&DebugExpr>,
916        tag_id: StmtId,
917    ) {
918        let sender_builder = self.graph_mut(on);
919        if let Some(serialize_fn) = serialize {
920            sender_builder.add_dfir(
921                parse_quote! {
922                    #input_ident -> map(#serialize_fn) -> dest_sink(#sink_expr);
923                },
924                None,
925                // operator tag separates send and receive, which otherwise have the same next_stmt_id
926                Some(&format!("send{}", tag_id)),
927            );
928        } else {
929            sender_builder.add_dfir(
930                parse_quote! {
931                    #input_ident -> dest_sink(#sink_expr);
932                },
933                None,
934                Some(&format!("send{}", tag_id)),
935            );
936        }
937    }
938
939    fn emit_fold_hook(
940        &mut self,
941        _location: &LocationId,
942        _in_ident: &syn::Ident,
943        _in_kind: &CollectionKind,
944        _op_meta: &HydroIrOpMetadata,
945    ) -> Option<syn::Ident> {
946        None
947    }
948
949    fn assert_is_consistent(
950        &mut self,
951        _trusted: bool,
952        location: &LocationId,
953        in_ident: syn::Ident,
954        out_ident: &syn::Ident,
955    ) {
956        let builder = self.graph_mut(location);
957        builder.add_dfir(
958            parse_quote! {
959                #out_ident = #in_ident;
960            },
961            None,
962            None,
963        );
964    }
965
966    fn observe_for_mut(
967        &mut self,
968        location: &LocationId,
969        in_ident: syn::Ident,
970        _in_kind: &CollectionKind,
971        out_ident: &syn::Ident,
972        _op_meta: &HydroIrOpMetadata,
973    ) {
974        let builder = self.graph_mut(location);
975        builder.add_dfir(
976            parse_quote! {
977                #out_ident = #in_ident;
978            },
979            None,
980            None,
981        );
982    }
983
984    fn create_versioned_network_fork(
985        &mut self,
986        _channel_id: u32,
987        _dest: &LocationId,
988        _senders: Vec<(LocationId, syn::Ident, Option<DebugExpr>)>,
989        _external_element_type: Option<&syn::Type>,
990        _tag_id: StmtId,
991    ) {
992        unreachable!(
993            "HydroNode::VersionedNetworkFork is only produced by the multi-version simulator merge \
994             pass and cannot be emitted by the non-simulation builder"
995        );
996    }
997
998    fn create_versioned_network(
999        &mut self,
1000        _channel_id: u32,
1001        _source: &LocationId,
1002        _dest: &LocationId,
1003        _out_ident: &syn::Ident,
1004        _deserialize: Option<&DebugExpr>,
1005        _external_element_type: Option<&syn::Type>,
1006        _tag_id: StmtId,
1007    ) {
1008        unreachable!(
1009            "HydroNode::VersionedNetwork is only produced by the multi-version simulator merge \
1010             pass and cannot be emitted by the non-simulation builder"
1011        );
1012    }
1013}
1014
1015#[cfg(feature = "build")]
1016pub enum BuildersOrCallback<'a, L, N>
1017where
1018    L: FnMut(&mut HydroRoot, &mut crate::Counter<StmtId>),
1019    N: FnMut(&mut HydroNode, &mut crate::Counter<StmtId>),
1020{
1021    Builders(&'a mut dyn DfirBuilder),
1022    Callback(L, N),
1023}
1024
1025/// An root in a Hydro graph, which is an pipeline that doesn't emit
1026/// any downstream values. Traversals over the dataflow graph and
1027/// generating DFIR IR start from roots.
1028#[derive(Debug, Hash, serde::Serialize)]
1029pub enum HydroRoot {
1030    ForEach {
1031        f: ClosureExpr,
1032        input: Box<HydroNode>,
1033        op_metadata: HydroIrOpMetadata,
1034    },
1035    SendExternal {
1036        to_external_key: LocationKey,
1037        to_port_id: ExternalPortId,
1038        to_many: bool,
1039        unpaired: bool,
1040        serialize_fn: Option<DebugExpr>,
1041        instantiate_fn: DebugInstantiate,
1042        input: Box<HydroNode>,
1043        op_metadata: HydroIrOpMetadata,
1044    },
1045    DestSink {
1046        sink: DebugExpr,
1047        input: Box<HydroNode>,
1048        op_metadata: HydroIrOpMetadata,
1049    },
1050    CycleSink {
1051        cycle_id: CycleId,
1052        input: Box<HydroNode>,
1053        op_metadata: HydroIrOpMetadata,
1054    },
1055    EmbeddedOutput {
1056        #[serde(serialize_with = "serialize_ident")]
1057        ident: syn::Ident,
1058        input: Box<HydroNode>,
1059        op_metadata: HydroIrOpMetadata,
1060    },
1061    Null {
1062        input: Box<HydroNode>,
1063        op_metadata: HydroIrOpMetadata,
1064    },
1065}
1066
1067impl HydroRoot {
1068    #[cfg(feature = "build")]
1069    #[expect(clippy::too_many_arguments, reason = "TODO(internal)")]
1070    pub fn compile_network<'a, D>(
1071        &mut self,
1072        extra_stmts: &mut SparseSecondaryMap<LocationKey, Vec<syn::Stmt>>,
1073        seen_tees: &mut SeenSharedNodes,
1074        seen_cluster_members: &mut HashSet<(LocationId, LocationKey)>,
1075        processes: &SparseSecondaryMap<LocationKey, D::Process>,
1076        clusters: &SparseSecondaryMap<LocationKey, D::Cluster>,
1077        externals: &SparseSecondaryMap<LocationKey, D::External>,
1078        env: &mut D::InstantiateEnv,
1079    ) where
1080        D: Deploy<'a>,
1081    {
1082        let refcell_extra_stmts = RefCell::new(extra_stmts);
1083        let refcell_env = RefCell::new(env);
1084        let refcell_seen_cluster_members = RefCell::new(seen_cluster_members);
1085        self.transform_bottom_up(
1086            &mut |l| {
1087                if let HydroRoot::SendExternal {
1088                    #[cfg(feature = "tokio")]
1089                    input,
1090                    #[cfg(feature = "tokio")]
1091                    to_external_key,
1092                    #[cfg(feature = "tokio")]
1093                    to_port_id,
1094                    #[cfg(feature = "tokio")]
1095                    to_many,
1096                    #[cfg(feature = "tokio")]
1097                    unpaired,
1098                    #[cfg(feature = "tokio")]
1099                    instantiate_fn,
1100                    ..
1101                } = l
1102                {
1103                    #[cfg(feature = "tokio")]
1104                    let ((sink_expr, source_expr), connect_fn) = match instantiate_fn {
1105                        DebugInstantiate::Building => {
1106                            let to_node = externals
1107                                .get(*to_external_key)
1108                                .unwrap_or_else(|| {
1109                                    panic!("A external used in the graph was not instantiated: {}", to_external_key)
1110                                })
1111                                .clone();
1112
1113                            match input.metadata().location_id.root() {
1114                                &LocationId::Process(process_key) => {
1115                                    if *to_many {
1116                                        (
1117                                            (
1118                                                D::e2o_many_sink(format!("{}_{}", *to_external_key, *to_port_id)),
1119                                                parse_quote!(DUMMY),
1120                                            ),
1121                                            Box::new(|| {}) as Box<dyn FnOnce()>,
1122                                        )
1123                                    } else {
1124                                        let from_node = processes
1125                                            .get(process_key)
1126                                            .unwrap_or_else(|| {
1127                                                panic!("A process used in the graph was not instantiated: {}", process_key)
1128                                            })
1129                                            .clone();
1130
1131                                        let sink_port = from_node.next_port();
1132                                        let source_port = to_node.next_port();
1133
1134                                        if *unpaired {
1135                                            use stageleft::quote_type;
1136                                            use tokio_util::codec::LengthDelimitedCodec;
1137
1138                                            to_node.register(*to_port_id, source_port.clone());
1139
1140                                            let _ = D::e2o_source(
1141                                                refcell_extra_stmts.borrow_mut().entry(process_key).expect("location was removed").or_default(),
1142                                                &to_node, &source_port,
1143                                                &from_node, &sink_port,
1144                                                &quote_type::<LengthDelimitedCodec>(),
1145                                                format!("{}_{}", *to_external_key, *to_port_id)
1146                                            );
1147                                        }
1148
1149                                        (
1150                                            (
1151                                                D::o2e_sink(
1152                                                    &from_node,
1153                                                    &sink_port,
1154                                                    &to_node,
1155                                                    &source_port,
1156                                                    format!("{}_{}", *to_external_key, *to_port_id)
1157                                                ),
1158                                                parse_quote!(DUMMY),
1159                                            ),
1160                                            if *unpaired {
1161                                                D::e2o_connect(
1162                                                    &to_node,
1163                                                    &source_port,
1164                                                    &from_node,
1165                                                    &sink_port,
1166                                                    *to_many,
1167                                                    NetworkHint::Auto,
1168                                                )
1169                                            } else {
1170                                                Box::new(|| {}) as Box<dyn FnOnce()>
1171                                            },
1172                                        )
1173                                    }
1174                                }
1175                                LocationId::Cluster(cluster_key) => {
1176                                    let from_node = clusters
1177                                        .get(*cluster_key)
1178                                        .unwrap_or_else(|| {
1179                                            panic!("A cluster used in the graph was not instantiated: {}", cluster_key)
1180                                        })
1181                                        .clone();
1182
1183                                    let sink_port = from_node.next_port();
1184                                    let source_port = to_node.next_port();
1185
1186                                    if *unpaired {
1187                                        to_node.register(*to_port_id, source_port.clone());
1188                                    }
1189
1190                                    (
1191                                        (
1192                                            D::m2e_sink(
1193                                                &from_node,
1194                                                &sink_port,
1195                                                &to_node,
1196                                                &source_port,
1197                                                format!("{}_{}", *to_external_key, *to_port_id)
1198                                            ),
1199                                            parse_quote!(DUMMY),
1200                                        ),
1201                                        Box::new(|| {}) as Box<dyn FnOnce()>,
1202                                    )
1203                                }
1204                                _ => panic!()
1205                            }
1206                        },
1207
1208                        DebugInstantiate::Finalized(_) => panic!("network already finalized"),
1209                    };
1210
1211                    #[cfg(not(feature = "tokio"))]
1212                    {
1213                        panic!("Cannot instantiate external inputs without tokio");
1214                    };
1215
1216                    #[cfg(feature = "tokio")]
1217                    {
1218                        *instantiate_fn = DebugInstantiateFinalized {
1219                            sink: sink_expr,
1220                            source: source_expr,
1221                            connect_fn: Some(connect_fn),
1222                        }
1223                        .into();
1224                    };
1225                } else if let HydroRoot::EmbeddedOutput { ident, input, .. } = l {
1226                    let element_type = match &input.metadata().collection_kind {
1227                        CollectionKind::Stream { element_type, .. } => element_type.0.as_ref().clone(),
1228                        _ => panic!("Embedded output must have Stream collection kind"),
1229                    };
1230                    let location_key = match input.metadata().location_id.root() {
1231                        LocationId::Process(key) | LocationId::Cluster(key) => *key,
1232                        _ => panic!("Embedded output must be on a process or cluster"),
1233                    };
1234                    D::register_embedded_output(
1235                        &mut refcell_env.borrow_mut(),
1236                        location_key,
1237                        ident,
1238                        &element_type,
1239                    );
1240                }
1241            },
1242            &mut |n| {
1243                if let HydroNode::Network {
1244                    name,
1245                    networking_info,
1246                    input,
1247                    instantiate_fn,
1248                    serialize,
1249                    deserialize,
1250                    metadata,
1251                    ..
1252                } = n
1253                {
1254                    let external_types = match (
1255                        serialize.external_element_type(),
1256                        deserialize.external_element_type(),
1257                    ) {
1258                        (Some(input_type), Some(output_type)) => Some((input_type, output_type)),
1259                        _ => None,
1260                    };
1261                    let (sink_expr, source_expr, connect_fn) = match instantiate_fn {
1262                        DebugInstantiate::Building => instantiate_network::<D>(
1263                            &mut refcell_env.borrow_mut(),
1264                            input.metadata().location_id.root(),
1265                            metadata.location_id.root(),
1266                            processes,
1267                            clusters,
1268                            name.as_deref(),
1269                            networking_info,
1270                            external_types,
1271                        ),
1272
1273                        DebugInstantiate::Finalized(_) => panic!("network already finalized"),
1274                    };
1275
1276                    *instantiate_fn = DebugInstantiateFinalized {
1277                        sink: sink_expr,
1278                        source: source_expr,
1279                        connect_fn: Some(connect_fn),
1280                    }
1281                    .into();
1282                } else if let HydroNode::ExternalInput {
1283                    from_external_key,
1284                    from_port_id,
1285                    from_many,
1286                    codec_type,
1287                    port_hint,
1288                    instantiate_fn,
1289                    metadata,
1290                    ..
1291                } = n
1292                {
1293                    let ((sink_expr, source_expr), connect_fn) = match instantiate_fn {
1294                        DebugInstantiate::Building => {
1295                            let from_node = externals
1296                                .get(*from_external_key)
1297                                .unwrap_or_else(|| {
1298                                    panic!(
1299                                        "A external used in the graph was not instantiated: {}",
1300                                        from_external_key,
1301                                    )
1302                                })
1303                                .clone();
1304
1305                            match metadata.location_id.root() {
1306                                &LocationId::Process(process_key) => {
1307                                    let to_node = processes
1308                                        .get(process_key)
1309                                        .unwrap_or_else(|| {
1310                                            panic!("A process used in the graph was not instantiated: {}", process_key)
1311                                        })
1312                                        .clone();
1313
1314                                    let sink_port = from_node.next_port();
1315                                    let source_port = to_node.next_port();
1316
1317                                    from_node.register(*from_port_id, sink_port.clone());
1318
1319                                    (
1320                                        (
1321                                            parse_quote!(DUMMY),
1322                                            if *from_many {
1323                                                D::e2o_many_source(
1324                                                    refcell_extra_stmts.borrow_mut().entry(process_key).expect("location was removed").or_default(),
1325                                                    &to_node, &source_port,
1326                                                    codec_type.0.as_ref(),
1327                                                    format!("{}_{}", *from_external_key, *from_port_id)
1328                                                )
1329                                            } else {
1330                                                D::e2o_source(
1331                                                    refcell_extra_stmts.borrow_mut().entry(process_key).expect("location was removed").or_default(),
1332                                                    &from_node, &sink_port,
1333                                                    &to_node, &source_port,
1334                                                    codec_type.0.as_ref(),
1335                                                    format!("{}_{}", *from_external_key, *from_port_id)
1336                                                )
1337                                            },
1338                                        ),
1339                                        D::e2o_connect(&from_node, &sink_port, &to_node, &source_port, *from_many, *port_hint),
1340                                    )
1341                                }
1342                                LocationId::Cluster(cluster_key) => {
1343                                    let to_node = clusters
1344                                        .get(*cluster_key)
1345                                        .unwrap_or_else(|| {
1346                                            panic!("A cluster used in the graph was not instantiated: {}", cluster_key)
1347                                        })
1348                                        .clone();
1349
1350                                    let sink_port = from_node.next_port();
1351                                    let source_port = to_node.next_port();
1352
1353                                    from_node.register(*from_port_id, sink_port.clone());
1354
1355                                    (
1356                                        (
1357                                            parse_quote!(DUMMY),
1358                                            D::e2m_source(
1359                                                refcell_extra_stmts.borrow_mut().entry(*cluster_key).expect("location was removed").or_default(),
1360                                                &from_node, &sink_port,
1361                                                &to_node, &source_port,
1362                                                codec_type.0.as_ref(),
1363                                                format!("{}_{}", *from_external_key, *from_port_id)
1364                                            ),
1365                                        ),
1366                                        D::e2m_connect(&from_node, &sink_port, &to_node, &source_port, *port_hint),
1367                                    )
1368                                }
1369                                _ => panic!()
1370                            }
1371                        },
1372
1373                        DebugInstantiate::Finalized(_) => panic!("network already finalized"),
1374                    };
1375
1376                    *instantiate_fn = DebugInstantiateFinalized {
1377                        sink: sink_expr,
1378                        source: source_expr,
1379                        connect_fn: Some(connect_fn),
1380                    }
1381                    .into();
1382                } else if let HydroNode::Source { source: HydroSource::Embedded(ident), metadata } = n {
1383                    let element_type = match &metadata.collection_kind {
1384                        CollectionKind::Stream { element_type, .. } => element_type.0.as_ref().clone(),
1385                        _ => panic!("Embedded source must have Stream collection kind"),
1386                    };
1387                    let location_key = match metadata.location_id.root() {
1388                        LocationId::Process(key) | LocationId::Cluster(key) => *key,
1389                        _ => panic!("Embedded source must be on a process or cluster"),
1390                    };
1391                    D::register_embedded_stream_input(
1392                        &mut refcell_env.borrow_mut(),
1393                        location_key,
1394                        ident,
1395                        &element_type,
1396                    );
1397                } else if let HydroNode::Source { source: HydroSource::EmbeddedSingleton(ident), metadata } = n {
1398                    let element_type = match &metadata.collection_kind {
1399                        CollectionKind::Singleton { element_type, .. } => element_type.0.as_ref().clone(),
1400                        _ => panic!("EmbeddedSingleton source must have Singleton collection kind"),
1401                    };
1402                    let location_key = match metadata.location_id.root() {
1403                        LocationId::Process(key) | LocationId::Cluster(key) => *key,
1404                        _ => panic!("EmbeddedSingleton source must be on a process or cluster"),
1405                    };
1406                    D::register_embedded_singleton_input(
1407                        &mut refcell_env.borrow_mut(),
1408                        location_key,
1409                        ident,
1410                        &element_type,
1411                    );
1412                } else if let HydroNode::Source { source: HydroSource::ClusterMembers(location_id, state), metadata } = n {
1413                    match state {
1414                        ClusterMembersState::Uninit => {
1415                            let at_location = metadata.location_id.root().clone();
1416                            let key = (at_location.clone(), location_id.key());
1417                            if refcell_seen_cluster_members.borrow_mut().insert(key) {
1418                                // First occurrence: call cluster_membership_stream and mark as Stream.
1419                                let expr = stageleft::QuotedWithContext::splice_untyped_ctx(
1420                                    D::cluster_membership_stream(&mut refcell_env.borrow_mut(), &at_location, location_id),
1421                                    &(),
1422                                );
1423                                *state = ClusterMembersState::Stream(expr.into());
1424                            } else {
1425                                // Already instantiated for this (at, target) pair: just tee.
1426                                *state = ClusterMembersState::Tee(at_location, location_id.clone());
1427                            }
1428                        }
1429                        ClusterMembersState::Stream(_) | ClusterMembersState::Tee(..) => {
1430                            panic!("cluster members already finalized");
1431                        }
1432                    }
1433                }
1434            },
1435            seen_tees,
1436            false,
1437        );
1438    }
1439
1440    pub fn connect_network(&mut self, seen_tees: &mut SeenSharedNodes) {
1441        self.transform_bottom_up(
1442            &mut |l| {
1443                if let HydroRoot::SendExternal { instantiate_fn, .. } = l {
1444                    match instantiate_fn {
1445                        DebugInstantiate::Building => panic!("network not built"),
1446
1447                        DebugInstantiate::Finalized(finalized) => {
1448                            (finalized.connect_fn.take().unwrap())();
1449                        }
1450                    }
1451                }
1452            },
1453            &mut |n| {
1454                if let HydroNode::Network { instantiate_fn, .. }
1455                | HydroNode::ExternalInput { instantiate_fn, .. } = n
1456                {
1457                    match instantiate_fn {
1458                        DebugInstantiate::Building => panic!("network not built"),
1459
1460                        DebugInstantiate::Finalized(finalized) => {
1461                            (finalized.connect_fn.take().unwrap())();
1462                        }
1463                    }
1464                }
1465            },
1466            seen_tees,
1467            false,
1468        );
1469    }
1470
1471    pub fn transform_bottom_up(
1472        &mut self,
1473        transform_root: &mut impl FnMut(&mut HydroRoot),
1474        transform_node: &mut impl FnMut(&mut HydroNode),
1475        seen_tees: &mut SeenSharedNodes,
1476        check_well_formed: bool,
1477    ) {
1478        self.transform_children(
1479            |n, s| n.transform_bottom_up(transform_node, s, check_well_formed),
1480            seen_tees,
1481        );
1482
1483        transform_root(self);
1484    }
1485
1486    pub fn transform_children(
1487        &mut self,
1488        mut transform: impl FnMut(&mut HydroNode, &mut SeenSharedNodes),
1489        seen_tees: &mut SeenSharedNodes,
1490    ) {
1491        match self {
1492            HydroRoot::ForEach { f, input, .. } => {
1493                f.transform_children(&mut transform, seen_tees);
1494                transform(input, seen_tees);
1495            }
1496            HydroRoot::SendExternal { input, .. }
1497            | HydroRoot::DestSink { input, .. }
1498            | HydroRoot::CycleSink { input, .. }
1499            | HydroRoot::EmbeddedOutput { input, .. }
1500            | HydroRoot::Null { input, .. } => {
1501                transform(input, seen_tees);
1502            }
1503        }
1504    }
1505
1506    pub fn deep_clone(&self, seen_tees: &mut SeenSharedNodes) -> HydroRoot {
1507        match self {
1508            HydroRoot::ForEach {
1509                f,
1510                input,
1511                op_metadata,
1512            } => HydroRoot::ForEach {
1513                f: f.deep_clone(seen_tees),
1514                input: Box::new(input.deep_clone(seen_tees)),
1515                op_metadata: op_metadata.clone(),
1516            },
1517            HydroRoot::SendExternal {
1518                to_external_key,
1519                to_port_id,
1520                to_many,
1521                unpaired,
1522                serialize_fn,
1523                instantiate_fn,
1524                input,
1525                op_metadata,
1526            } => HydroRoot::SendExternal {
1527                to_external_key: *to_external_key,
1528                to_port_id: *to_port_id,
1529                to_many: *to_many,
1530                unpaired: *unpaired,
1531                serialize_fn: serialize_fn.clone(),
1532                instantiate_fn: instantiate_fn.clone(),
1533                input: Box::new(input.deep_clone(seen_tees)),
1534                op_metadata: op_metadata.clone(),
1535            },
1536            HydroRoot::DestSink {
1537                sink,
1538                input,
1539                op_metadata,
1540            } => HydroRoot::DestSink {
1541                sink: sink.clone(),
1542                input: Box::new(input.deep_clone(seen_tees)),
1543                op_metadata: op_metadata.clone(),
1544            },
1545            HydroRoot::CycleSink {
1546                cycle_id,
1547                input,
1548                op_metadata,
1549            } => HydroRoot::CycleSink {
1550                cycle_id: *cycle_id,
1551                input: Box::new(input.deep_clone(seen_tees)),
1552                op_metadata: op_metadata.clone(),
1553            },
1554            HydroRoot::EmbeddedOutput {
1555                ident,
1556                input,
1557                op_metadata,
1558            } => HydroRoot::EmbeddedOutput {
1559                ident: ident.clone(),
1560                input: Box::new(input.deep_clone(seen_tees)),
1561                op_metadata: op_metadata.clone(),
1562            },
1563            HydroRoot::Null { input, op_metadata } => HydroRoot::Null {
1564                input: Box::new(input.deep_clone(seen_tees)),
1565                op_metadata: op_metadata.clone(),
1566            },
1567        }
1568    }
1569
1570    #[cfg(feature = "build")]
1571    pub fn emit(
1572        &mut self,
1573        graph_builders: &mut dyn DfirBuilder,
1574        seen_tees: &mut SeenSharedNodes,
1575        built_tees: &mut HashMap<*const RefCell<HydroNode>, Vec<syn::Ident>>,
1576        next_stmt_id: &mut crate::Counter<StmtId>,
1577        fold_hooked_idents: &mut HashSet<String>,
1578    ) {
1579        self.emit_core(
1580            &mut BuildersOrCallback::<
1581                fn(&mut HydroRoot, &mut crate::Counter<StmtId>),
1582                fn(&mut HydroNode, &mut crate::Counter<StmtId>),
1583            >::Builders(graph_builders),
1584            seen_tees,
1585            built_tees,
1586            next_stmt_id,
1587            fold_hooked_idents,
1588        );
1589    }
1590
1591    #[cfg(feature = "build")]
1592    pub fn emit_core(
1593        &mut self,
1594        builders_or_callback: &mut BuildersOrCallback<
1595            impl FnMut(&mut HydroRoot, &mut crate::Counter<StmtId>),
1596            impl FnMut(&mut HydroNode, &mut crate::Counter<StmtId>),
1597        >,
1598        seen_tees: &mut SeenSharedNodes,
1599        built_tees: &mut HashMap<*const RefCell<HydroNode>, Vec<syn::Ident>>,
1600        next_stmt_id: &mut crate::Counter<StmtId>,
1601        fold_hooked_idents: &mut HashSet<String>,
1602    ) {
1603        match self {
1604            HydroRoot::ForEach { f, input, .. } => {
1605                let input_ident = input.emit_core(
1606                    builders_or_callback,
1607                    seen_tees,
1608                    built_tees,
1609                    next_stmt_id,
1610                    fold_hooked_idents,
1611                );
1612
1613                // for_each is always side-effecting, so we observe non-determinism
1614                // even when the closure does not capture a mut ref (unlike map/filter
1615                // which only observe when they have a mut ref).
1616                let input_ident = if !input.metadata().collection_kind.is_strict() {
1617                    let observe_stmt_id = next_stmt_id.get_and_increment();
1618                    let observe_ident =
1619                        syn::Ident::new(&format!("stream_{}", observe_stmt_id), Span::call_site());
1620                    if let BuildersOrCallback::Builders(graph_builders) = builders_or_callback {
1621                        graph_builders.observe_for_mut(
1622                            &input.metadata().location_id,
1623                            input_ident,
1624                            &input.metadata().collection_kind,
1625                            &observe_ident,
1626                            &input.metadata().op,
1627                        );
1628                    }
1629                    observe_ident
1630                } else {
1631                    input_ident
1632                };
1633
1634                let stmt_id = next_stmt_id.get_and_increment();
1635
1636                match builders_or_callback {
1637                    BuildersOrCallback::Builders(graph_builders) => {
1638                        let mut ident_stack: Vec<syn::Ident> = Vec::new();
1639
1640                        // Look up each captured ref's ident from built_tees
1641                        for (ref_node, _is_mut) in f.singleton_refs.iter() {
1642                            let HydroNode::Reference { inner, .. } = ref_node else {
1643                                panic!("singleton_refs should only contain HydroNode::Reference");
1644                            };
1645                            let ptr = inner.0.as_ref() as *const RefCell<HydroNode>;
1646                            let idents = built_tees.get(&ptr).expect(
1647                                "ForEach singleton ref not found in built_tees — ref node was not emitted",
1648                            );
1649                            ident_stack.push(idents[0].clone());
1650                        }
1651
1652                        let f_tokens = f.emit_tokens(&mut ident_stack);
1653
1654                        graph_builders.add_dfir_at(
1655                            &input.metadata().location_id,
1656                            parse_quote! {
1657                                #input_ident -> for_each(#f_tokens);
1658                            },
1659                            Some(&stmt_id.to_string()),
1660                        );
1661                    }
1662                    BuildersOrCallback::Callback(leaf_callback, _) => {
1663                        leaf_callback(self, next_stmt_id);
1664                    }
1665                }
1666            }
1667
1668            HydroRoot::SendExternal {
1669                serialize_fn,
1670                instantiate_fn,
1671                input,
1672                ..
1673            } => {
1674                let input_ident = input.emit_core(
1675                    builders_or_callback,
1676                    seen_tees,
1677                    built_tees,
1678                    next_stmt_id,
1679                    fold_hooked_idents,
1680                );
1681
1682                let stmt_id = next_stmt_id.get_and_increment();
1683
1684                match builders_or_callback {
1685                    BuildersOrCallback::Builders(graph_builders) => {
1686                        let (sink_expr, _) = match instantiate_fn {
1687                            DebugInstantiate::Building => (
1688                                syn::parse_quote!(DUMMY_SINK),
1689                                syn::parse_quote!(DUMMY_SOURCE),
1690                            ),
1691
1692                            DebugInstantiate::Finalized(finalized) => {
1693                                (finalized.sink.clone(), finalized.source.clone())
1694                            }
1695                        };
1696
1697                        graph_builders.create_external_output(
1698                            &input.metadata().location_id,
1699                            sink_expr,
1700                            &input_ident,
1701                            serialize_fn.as_ref(),
1702                            stmt_id,
1703                        );
1704                    }
1705                    BuildersOrCallback::Callback(leaf_callback, _) => {
1706                        leaf_callback(self, next_stmt_id);
1707                    }
1708                }
1709            }
1710
1711            HydroRoot::DestSink { sink, input, .. } => {
1712                let input_ident = input.emit_core(
1713                    builders_or_callback,
1714                    seen_tees,
1715                    built_tees,
1716                    next_stmt_id,
1717                    fold_hooked_idents,
1718                );
1719
1720                let stmt_id = next_stmt_id.get_and_increment();
1721
1722                match builders_or_callback {
1723                    BuildersOrCallback::Builders(graph_builders) => {
1724                        graph_builders.add_dfir_at(
1725                            &input.metadata().location_id,
1726                            parse_quote! {
1727                                #input_ident -> dest_sink(#sink);
1728                            },
1729                            Some(&stmt_id.to_string()),
1730                        );
1731                    }
1732                    BuildersOrCallback::Callback(leaf_callback, _) => {
1733                        leaf_callback(self, next_stmt_id);
1734                    }
1735                }
1736            }
1737
1738            HydroRoot::CycleSink {
1739                cycle_id, input, ..
1740            } => {
1741                let input_ident = input.emit_core(
1742                    builders_or_callback,
1743                    seen_tees,
1744                    built_tees,
1745                    next_stmt_id,
1746                    fold_hooked_idents,
1747                );
1748
1749                match builders_or_callback {
1750                    BuildersOrCallback::Builders(graph_builders) => {
1751                        let elem_type: syn::Type = match &input.metadata().collection_kind {
1752                            CollectionKind::KeyedSingleton {
1753                                key_type,
1754                                value_type,
1755                                ..
1756                            }
1757                            | CollectionKind::KeyedStream {
1758                                key_type,
1759                                value_type,
1760                                ..
1761                            } => {
1762                                parse_quote!((#key_type, #value_type))
1763                            }
1764                            CollectionKind::Stream { element_type, .. }
1765                            | CollectionKind::Singleton { element_type, .. }
1766                            | CollectionKind::Optional { element_type, .. } => {
1767                                parse_quote!(#element_type)
1768                            }
1769                        };
1770
1771                        let cycle_id_ident = cycle_id.as_ident();
1772                        graph_builders.add_dfir_at(
1773                            &input.metadata().location_id,
1774                            parse_quote! {
1775                                #cycle_id_ident = #input_ident -> identity::<#elem_type>();
1776                            },
1777                            None,
1778                        );
1779                    }
1780                    // No ID, no callback
1781                    BuildersOrCallback::Callback(_, _) => {}
1782                }
1783            }
1784
1785            HydroRoot::EmbeddedOutput { ident, input, .. } => {
1786                let input_ident = input.emit_core(
1787                    builders_or_callback,
1788                    seen_tees,
1789                    built_tees,
1790                    next_stmt_id,
1791                    fold_hooked_idents,
1792                );
1793
1794                let stmt_id = next_stmt_id.get_and_increment();
1795
1796                match builders_or_callback {
1797                    BuildersOrCallback::Builders(graph_builders) => {
1798                        graph_builders.add_dfir_at(
1799                            &input.metadata().location_id,
1800                            parse_quote! {
1801                                #input_ident -> for_each(&mut #ident);
1802                            },
1803                            Some(&stmt_id.to_string()),
1804                        );
1805                    }
1806                    BuildersOrCallback::Callback(leaf_callback, _) => {
1807                        leaf_callback(self, next_stmt_id);
1808                    }
1809                }
1810            }
1811
1812            HydroRoot::Null { input, .. } => {
1813                let input_ident = input.emit_core(
1814                    builders_or_callback,
1815                    seen_tees,
1816                    built_tees,
1817                    next_stmt_id,
1818                    fold_hooked_idents,
1819                );
1820
1821                let stmt_id = next_stmt_id.get_and_increment();
1822
1823                match builders_or_callback {
1824                    BuildersOrCallback::Builders(graph_builders) => {
1825                        graph_builders.add_dfir_at(
1826                            &input.metadata().location_id,
1827                            parse_quote! {
1828                                #input_ident -> for_each(|_| {});
1829                            },
1830                            Some(&stmt_id.to_string()),
1831                        );
1832                    }
1833                    BuildersOrCallback::Callback(leaf_callback, _) => {
1834                        leaf_callback(self, next_stmt_id);
1835                    }
1836                }
1837            }
1838        }
1839    }
1840
1841    pub fn op_metadata(&self) -> &HydroIrOpMetadata {
1842        match self {
1843            HydroRoot::ForEach { op_metadata, .. }
1844            | HydroRoot::SendExternal { op_metadata, .. }
1845            | HydroRoot::DestSink { op_metadata, .. }
1846            | HydroRoot::CycleSink { op_metadata, .. }
1847            | HydroRoot::EmbeddedOutput { op_metadata, .. }
1848            | HydroRoot::Null { op_metadata, .. } => op_metadata,
1849        }
1850    }
1851
1852    pub fn op_metadata_mut(&mut self) -> &mut HydroIrOpMetadata {
1853        match self {
1854            HydroRoot::ForEach { op_metadata, .. }
1855            | HydroRoot::SendExternal { op_metadata, .. }
1856            | HydroRoot::DestSink { op_metadata, .. }
1857            | HydroRoot::CycleSink { op_metadata, .. }
1858            | HydroRoot::EmbeddedOutput { op_metadata, .. }
1859            | HydroRoot::Null { op_metadata, .. } => op_metadata,
1860        }
1861    }
1862
1863    pub fn input(&self) -> &HydroNode {
1864        match self {
1865            HydroRoot::ForEach { input, .. }
1866            | HydroRoot::SendExternal { input, .. }
1867            | HydroRoot::DestSink { input, .. }
1868            | HydroRoot::CycleSink { input, .. }
1869            | HydroRoot::EmbeddedOutput { input, .. }
1870            | HydroRoot::Null { input, .. } => input,
1871        }
1872    }
1873
1874    pub fn input_metadata(&self) -> &HydroIrMetadata {
1875        self.input().metadata()
1876    }
1877
1878    pub fn print_root(&self) -> String {
1879        match self {
1880            HydroRoot::ForEach { f, .. } => format!("ForEach({:?})", f),
1881            HydroRoot::SendExternal { .. } => "SendExternal".to_owned(),
1882            HydroRoot::DestSink { sink, .. } => format!("DestSink({:?})", sink),
1883            HydroRoot::CycleSink { cycle_id, .. } => format!("CycleSink({})", cycle_id),
1884            HydroRoot::EmbeddedOutput { ident, .. } => {
1885                format!("EmbeddedOutput({})", ident)
1886            }
1887            HydroRoot::Null { .. } => "Null".to_owned(),
1888        }
1889    }
1890
1891    pub fn visit_debug_expr(&mut self, mut transform: impl FnMut(&mut DebugExpr)) {
1892        match self {
1893            HydroRoot::ForEach { f, .. } => {
1894                transform(&mut f.expr);
1895            }
1896            HydroRoot::DestSink { sink, .. } => {
1897                transform(sink);
1898            }
1899            HydroRoot::SendExternal { .. }
1900            | HydroRoot::CycleSink { .. }
1901            | HydroRoot::EmbeddedOutput { .. }
1902            | HydroRoot::Null { .. } => {}
1903        }
1904    }
1905}
1906
1907#[cfg(feature = "build")]
1908fn tick_of(loc: &LocationId) -> Option<ClockId> {
1909    match loc {
1910        LocationId::Tick(id, _) => Some(*id),
1911        LocationId::Atomic(inner) => tick_of(inner),
1912        _ => None,
1913    }
1914}
1915
1916#[cfg(feature = "build")]
1917fn remap_location(loc: &mut LocationId, uf: &mut HashMap<ClockId, ClockId>) {
1918    match loc {
1919        LocationId::Tick(id, inner) => {
1920            *id = uf_find(uf, *id);
1921            remap_location(inner, uf);
1922        }
1923        LocationId::Atomic(inner) => {
1924            remap_location(inner, uf);
1925        }
1926        LocationId::Process(_) | LocationId::Cluster(_) => {}
1927    }
1928}
1929
1930#[cfg(feature = "build")]
1931fn uf_find(parent: &mut HashMap<ClockId, ClockId>, x: ClockId) -> ClockId {
1932    let p = *parent.get(&x).unwrap_or(&x);
1933    if p == x {
1934        return x;
1935    }
1936    let root = uf_find(parent, p);
1937    parent.insert(x, root);
1938    root
1939}
1940
1941#[cfg(feature = "build")]
1942fn uf_union(parent: &mut HashMap<ClockId, ClockId>, a: ClockId, b: ClockId) {
1943    let ra = uf_find(parent, a);
1944    let rb = uf_find(parent, b);
1945    if ra != rb {
1946        parent.insert(ra, rb);
1947    }
1948}
1949
1950/// Traverse the IR to build a union-find that unifies tick IDs connected
1951/// through `Batch` and `YieldConcat` nodes at atomic boundaries, then
1952/// rewrite all `LocationId`s to use the representative tick ID.
1953#[cfg(feature = "build")]
1954pub fn unify_atomic_ticks(ir: &mut [HydroRoot]) {
1955    let mut uf: HashMap<ClockId, ClockId> = HashMap::new();
1956
1957    // Pass 1: collect unifications.
1958    transform_bottom_up(
1959        ir,
1960        &mut |_| {},
1961        &mut |node: &mut HydroNode| match node {
1962            HydroNode::Batch { inner, metadata } | HydroNode::YieldConcat { inner, metadata } => {
1963                if let (Some(a), Some(b)) = (
1964                    tick_of(&inner.metadata().location_id),
1965                    tick_of(&metadata.location_id),
1966                ) {
1967                    uf_union(&mut uf, a, b);
1968                }
1969            }
1970            HydroNode::Chain {
1971                first,
1972                second,
1973                metadata,
1974            }
1975            | HydroNode::ChainFirst {
1976                first,
1977                second,
1978                metadata,
1979            }
1980            | HydroNode::MergeOrdered {
1981                first,
1982                second,
1983                metadata,
1984            } => {
1985                if let (Some(a), Some(b)) = (
1986                    tick_of(&first.metadata().location_id),
1987                    tick_of(&metadata.location_id),
1988                ) {
1989                    uf_union(&mut uf, a, b);
1990                }
1991                if let (Some(a), Some(b)) = (
1992                    tick_of(&second.metadata().location_id),
1993                    tick_of(&metadata.location_id),
1994                ) {
1995                    uf_union(&mut uf, a, b);
1996                }
1997            }
1998            _ => {}
1999        },
2000        false,
2001    );
2002
2003    // Pass 2: rewrite all LocationIds.
2004    transform_bottom_up(
2005        ir,
2006        &mut |_| {},
2007        &mut |node: &mut HydroNode| {
2008            remap_location(&mut node.metadata_mut().location_id, &mut uf);
2009        },
2010        false,
2011    );
2012}
2013
2014#[cfg(feature = "build")]
2015pub fn emit(ir: &mut Vec<HydroRoot>) -> SecondaryMap<LocationKey, FlatGraphBuilder> {
2016    let mut builders = ProdDfirBuilder::default();
2017    let mut seen_tees = HashMap::new();
2018    let mut built_tees = HashMap::new();
2019    let mut next_stmt_id = crate::Counter::<StmtId>::default();
2020    let mut fold_hooked_idents = HashSet::new();
2021    for leaf in ir {
2022        leaf.emit(
2023            &mut builders,
2024            &mut seen_tees,
2025            &mut built_tees,
2026            &mut next_stmt_id,
2027            &mut fold_hooked_idents,
2028        );
2029    }
2030    builders.graphs
2031}
2032
2033#[cfg(feature = "build")]
2034pub fn traverse_dfir(
2035    ir: &mut [HydroRoot],
2036    transform_root: impl FnMut(&mut HydroRoot, &mut crate::Counter<StmtId>),
2037    transform_node: impl FnMut(&mut HydroNode, &mut crate::Counter<StmtId>),
2038) {
2039    let mut seen_tees = HashMap::new();
2040    let mut built_tees = HashMap::new();
2041    let mut next_stmt_id = crate::Counter::<StmtId>::default();
2042    let mut fold_hooked_idents = HashSet::new();
2043    let mut callback = BuildersOrCallback::Callback(transform_root, transform_node);
2044    ir.iter_mut().for_each(|leaf| {
2045        leaf.emit_core(
2046            &mut callback,
2047            &mut seen_tees,
2048            &mut built_tees,
2049            &mut next_stmt_id,
2050            &mut fold_hooked_idents,
2051        );
2052    });
2053}
2054
2055pub fn transform_bottom_up(
2056    ir: &mut [HydroRoot],
2057    transform_root: &mut impl FnMut(&mut HydroRoot),
2058    transform_node: &mut impl FnMut(&mut HydroNode),
2059    check_well_formed: bool,
2060) {
2061    let mut seen_tees = HashMap::new();
2062    ir.iter_mut().for_each(|leaf| {
2063        leaf.transform_bottom_up(
2064            transform_root,
2065            transform_node,
2066            &mut seen_tees,
2067            check_well_formed,
2068        );
2069    });
2070}
2071
2072pub fn deep_clone(ir: &[HydroRoot]) -> Vec<HydroRoot> {
2073    let mut seen_tees = HashMap::new();
2074    ir.iter()
2075        .map(|leaf| leaf.deep_clone(&mut seen_tees))
2076        .collect()
2077}
2078
2079type PrintedTees = RefCell<Option<(usize, HashMap<*const RefCell<HydroNode>, usize>)>>;
2080thread_local! {
2081    static PRINTED_TEES: PrintedTees = const { RefCell::new(None) };
2082    /// Tracks shared nodes already serialized so that `SharedNode::serialize`
2083    /// emits the full subtree only once and uses a `"<shared N>"` back-reference
2084    /// on subsequent encounters, preventing infinite loops.
2085    static SERIALIZED_SHARED: PrintedTees
2086        = const { RefCell::new(None) };
2087}
2088
2089pub fn dbg_dedup_tee<T>(f: impl FnOnce() -> T) -> T {
2090    PRINTED_TEES.with(|printed_tees| {
2091        let mut printed_tees_mut = printed_tees.borrow_mut();
2092        *printed_tees_mut = Some((0, HashMap::new()));
2093        drop(printed_tees_mut);
2094
2095        let ret = f();
2096
2097        let mut printed_tees_mut = printed_tees.borrow_mut();
2098        *printed_tees_mut = None;
2099
2100        ret
2101    })
2102}
2103
2104/// Runs `f` with a fresh shared-node deduplication scope for serialization.
2105/// Any `SharedNode` serialized inside `f` will be tracked; the first occurrence
2106/// emits the full subtree while later occurrences emit a `{"$shared_ref": id}`
2107/// back-reference.  The tracking state is restored when `f` returns or panics.
2108pub fn serialize_dedup_shared<T>(f: impl FnOnce() -> T) -> T {
2109    let _guard = SerializedSharedGuard::enter();
2110    f()
2111}
2112
2113/// RAII guard that saves/restores the `SERIALIZED_SHARED` thread-local,
2114/// making `serialize_dedup_shared` re-entrant and panic-safe.
2115struct SerializedSharedGuard {
2116    previous: Option<(usize, HashMap<*const RefCell<HydroNode>, usize>)>,
2117}
2118
2119impl SerializedSharedGuard {
2120    fn enter() -> Self {
2121        let previous = SERIALIZED_SHARED.with(|cell| {
2122            let mut guard = cell.borrow_mut();
2123            guard.replace((0, HashMap::new()))
2124        });
2125        Self { previous }
2126    }
2127}
2128
2129impl Drop for SerializedSharedGuard {
2130    fn drop(&mut self) {
2131        SERIALIZED_SHARED.with(|cell| {
2132            *cell.borrow_mut() = self.previous.take();
2133        });
2134    }
2135}
2136
2137pub struct SharedNode(pub Rc<RefCell<HydroNode>>);
2138
2139impl serde::Serialize for SharedNode {
2140    /// Multiple `SharedNode`s can point to the same underlying `HydroNode` (via
2141    /// `Tee` / `Partition`).  A naïve recursive serialization would revisit the
2142    /// same subtree every time and, if the graph ever contains a cycle, loop
2143    /// forever.
2144    ///
2145    /// We keep a thread-local map (`SERIALIZED_SHARED`) from raw `Rc` pointer →
2146    /// integer id.  The first time we see a pointer we assign it the next id and
2147    /// emit the full subtree as `{"$shared": <id>, "node": …}`.  Every later
2148    /// encounter of the same pointer emits `{"$shared_ref": <id>}`, cutting the
2149    /// recursion.  Requires an active `serialize_dedup_shared` scope.
2150    fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
2151        SERIALIZED_SHARED.with(|cell| {
2152            let mut guard = cell.borrow_mut();
2153            // (next_id, pointer → assigned_id)
2154            let state = guard.as_mut().ok_or_else(|| {
2155                serde::ser::Error::custom(
2156                    "SharedNode serialization requires an active serialize_dedup_shared scope",
2157                )
2158            })?;
2159            let ptr = self.0.as_ptr() as *const RefCell<HydroNode>;
2160
2161            if let Some(&id) = state.1.get(&ptr) {
2162                drop(guard);
2163                use serde::ser::SerializeMap;
2164                let mut map = serializer.serialize_map(Some(1))?;
2165                map.serialize_entry("$shared_ref", &id)?;
2166                map.end()
2167            } else {
2168                let id = state.0;
2169                state.0 += 1;
2170                state.1.insert(ptr, id);
2171                drop(guard);
2172
2173                use serde::ser::SerializeMap;
2174                let mut map = serializer.serialize_map(Some(2))?;
2175                map.serialize_entry("$shared", &id)?;
2176                map.serialize_entry("node", &*self.0.borrow())?;
2177                map.end()
2178            }
2179        })
2180    }
2181}
2182
2183impl SharedNode {
2184    pub fn as_ptr(&self) -> *const RefCell<HydroNode> {
2185        Rc::as_ptr(&self.0)
2186    }
2187}
2188
2189impl Debug for SharedNode {
2190    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
2191        PRINTED_TEES.with(|printed_tees| {
2192            let mut printed_tees_mut_borrow = printed_tees.borrow_mut();
2193            let printed_tees_mut = printed_tees_mut_borrow.as_mut();
2194
2195            if let Some(printed_tees_mut) = printed_tees_mut {
2196                if let Some(existing) = printed_tees_mut
2197                    .1
2198                    .get(&(self.0.as_ref() as *const RefCell<HydroNode>))
2199                {
2200                    write!(f, "<shared {}>", existing)
2201                } else {
2202                    let next_id = printed_tees_mut.0;
2203                    printed_tees_mut.0 += 1;
2204                    printed_tees_mut
2205                        .1
2206                        .insert(self.0.as_ref() as *const RefCell<HydroNode>, next_id);
2207                    drop(printed_tees_mut_borrow);
2208                    write!(f, "<shared {}>: ", next_id)?;
2209                    Debug::fmt(&self.0.borrow(), f)
2210                }
2211            } else {
2212                drop(printed_tees_mut_borrow);
2213                write!(f, "<shared>: ")?;
2214                Debug::fmt(&self.0.borrow(), f)
2215            }
2216        })
2217    }
2218}
2219
2220impl Hash for SharedNode {
2221    fn hash<H: Hasher>(&self, state: &mut H) {
2222        self.0.borrow_mut().hash(state);
2223    }
2224}
2225
2226/// A counter for tracking singleton access groups on a `HydroNode::Reference`.
2227///
2228/// Each mutable access increments the counter (before and after) to isolate itself in its own group;
2229/// immutable accesses share the current group.
2230#[derive(Debug)]
2231pub enum AccessCounter {
2232    Counting(Cell<u32>),
2233    Frozen(u32),
2234}
2235
2236impl AccessCounter {
2237    pub fn new() -> Self {
2238        Self::Counting(Cell::new(0))
2239    }
2240
2241    /// Assign the next access group for this reference.
2242    /// Mutable accesses get an isolated group (counter increments before and after).
2243    /// Immutable accesses share the current group.
2244    pub fn next_group(&self, is_mut: bool) -> Self {
2245        let AccessCounter::Counting(count) = self else {
2246            panic!("Cannot count on `AccessCounter::Frozen`");
2247        };
2248        let c = if is_mut {
2249            let c = count.get() + 1;
2250            count.set(c + 1);
2251            c
2252        } else {
2253            count.get()
2254        };
2255        Self::Frozen(c)
2256    }
2257
2258    /// Creates a frozen counter to prevent further counting.
2259    pub fn freeze(&self) -> Self {
2260        Self::Frozen(match self {
2261            Self::Counting(count) => count.get(),
2262            Self::Frozen(count) => *count,
2263        })
2264    }
2265
2266    pub fn frozen_group(&self) -> u32 {
2267        let Self::Frozen(count) = self else {
2268            panic!("`AccessCounter` not frozen");
2269        };
2270        *count
2271    }
2272}
2273
2274impl Default for AccessCounter {
2275    fn default() -> Self {
2276        Self::new()
2277    }
2278}
2279
2280impl Hash for AccessCounter {
2281    fn hash<H: Hasher>(&self, _state: &mut H) {
2282        // Access counter does not participate in hashing — it is runtime bookkeeping.
2283    }
2284}
2285
2286impl serde::Serialize for AccessCounter {
2287    fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
2288        let count = match self {
2289            AccessCounter::Counting(count) => count.get(),
2290            AccessCounter::Frozen(count) => *count,
2291        };
2292        count.serialize(serializer)
2293    }
2294}
2295
2296#[derive(serde::Serialize, Clone, PartialEq, Eq, Debug)]
2297pub enum BoundKind {
2298    Unbounded,
2299    Bounded,
2300}
2301
2302#[derive(serde::Serialize, Clone, PartialEq, Eq, Debug)]
2303pub enum StreamOrder {
2304    NoOrder,
2305    TotalOrder,
2306}
2307
2308#[derive(serde::Serialize, Clone, PartialEq, Eq, Debug)]
2309pub enum StreamRetry {
2310    AtLeastOnce,
2311    ExactlyOnce,
2312}
2313
2314#[derive(serde::Serialize, Clone, PartialEq, Eq, Debug)]
2315pub enum KeyedSingletonBoundKind {
2316    Unbounded,
2317    MonotonicKeys,
2318    MonotonicValue,
2319    BoundedValue,
2320    Bounded,
2321}
2322
2323#[derive(serde::Serialize, Clone, PartialEq, Eq, Debug)]
2324pub enum SingletonBoundKind {
2325    Unbounded,
2326    Monotonic,
2327    Bounded,
2328}
2329
2330#[derive(Clone, PartialEq, Eq, Debug, serde::Serialize)]
2331pub enum CollectionKind {
2332    Stream {
2333        bound: BoundKind,
2334        order: StreamOrder,
2335        retry: StreamRetry,
2336        element_type: DebugType,
2337    },
2338    Singleton {
2339        bound: SingletonBoundKind,
2340        element_type: DebugType,
2341    },
2342    Optional {
2343        bound: BoundKind,
2344        element_type: DebugType,
2345    },
2346    KeyedStream {
2347        bound: BoundKind,
2348        value_order: StreamOrder,
2349        value_retry: StreamRetry,
2350        key_type: DebugType,
2351        value_type: DebugType,
2352    },
2353    KeyedSingleton {
2354        bound: KeyedSingletonBoundKind,
2355        key_type: DebugType,
2356        value_type: DebugType,
2357    },
2358}
2359
2360impl CollectionKind {
2361    pub fn is_bounded(&self) -> bool {
2362        matches!(
2363            self,
2364            CollectionKind::Stream {
2365                bound: BoundKind::Bounded,
2366                ..
2367            } | CollectionKind::Singleton {
2368                bound: SingletonBoundKind::Bounded,
2369                ..
2370            } | CollectionKind::Optional {
2371                bound: BoundKind::Bounded,
2372                ..
2373            } | CollectionKind::KeyedStream {
2374                bound: BoundKind::Bounded,
2375                ..
2376            } | CollectionKind::KeyedSingleton {
2377                bound: KeyedSingletonBoundKind::Bounded,
2378                ..
2379            }
2380        )
2381    }
2382
2383    /// Returns whether this collection kind is already "strict" (TotalOrder + ExactlyOnce),
2384    /// meaning no non-determinism needs to be observed for mut closures.
2385    pub fn is_strict(&self) -> bool {
2386        match self {
2387            CollectionKind::Stream { order, retry, .. } => {
2388                *order == StreamOrder::TotalOrder && *retry == StreamRetry::ExactlyOnce
2389            }
2390            CollectionKind::KeyedStream {
2391                value_order,
2392                value_retry,
2393                ..
2394            } => {
2395                *value_order == StreamOrder::TotalOrder && *value_retry == StreamRetry::ExactlyOnce
2396            }
2397            // Singletons/Optionals/KeyedSingletons do not have observable
2398            // non-determinism other than snapshots / batching
2399            CollectionKind::Singleton { .. }
2400            | CollectionKind::Optional { .. }
2401            | CollectionKind::KeyedSingleton { .. } => true,
2402        }
2403    }
2404
2405    /// Creates a "strict" version of this kind with TotalOrder and ExactlyOnce.
2406    pub fn strict_kind(&self) -> CollectionKind {
2407        match self {
2408            CollectionKind::Stream {
2409                bound,
2410                element_type,
2411                ..
2412            } => CollectionKind::Stream {
2413                bound: bound.clone(),
2414                order: StreamOrder::TotalOrder,
2415                retry: StreamRetry::ExactlyOnce,
2416                element_type: element_type.clone(),
2417            },
2418            CollectionKind::KeyedStream {
2419                bound,
2420                key_type,
2421                value_type,
2422                ..
2423            } => CollectionKind::KeyedStream {
2424                bound: bound.clone(),
2425                value_order: StreamOrder::TotalOrder,
2426                value_retry: StreamRetry::ExactlyOnce,
2427                key_type: key_type.clone(),
2428                value_type: value_type.clone(),
2429            },
2430            other => other.clone(),
2431        }
2432    }
2433}
2434
2435#[derive(Clone, serde::Serialize)]
2436pub struct HydroIrMetadata {
2437    pub location_id: LocationId,
2438    pub collection_kind: CollectionKind,
2439    pub consistency: Option<ClusterConsistency>,
2440    pub cardinality: Option<usize>,
2441    pub tag: Option<String>,
2442    pub op: HydroIrOpMetadata,
2443}
2444
2445// HydroIrMetadata shouldn't be used to hash or compare
2446impl Hash for HydroIrMetadata {
2447    fn hash<H: Hasher>(&self, _: &mut H) {}
2448}
2449
2450impl PartialEq for HydroIrMetadata {
2451    fn eq(&self, _: &Self) -> bool {
2452        true
2453    }
2454}
2455
2456impl Eq for HydroIrMetadata {}
2457
2458impl Debug for HydroIrMetadata {
2459    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
2460        f.debug_struct("HydroIrMetadata")
2461            .field("location_id", &self.location_id)
2462            .field("collection_kind", &self.collection_kind)
2463            .finish()
2464    }
2465}
2466
2467/// Metadata that is specific to the operator itself, rather than its outputs.
2468/// This is available on _both_ inner nodes and roots.
2469#[derive(Clone, serde::Serialize)]
2470pub struct HydroIrOpMetadata {
2471    #[serde(rename = "span", serialize_with = "serialize_backtrace_as_span")]
2472    pub backtrace: Backtrace,
2473    pub cpu_usage: Option<f64>,
2474    pub network_recv_cpu_usage: Option<f64>,
2475    pub id: Option<usize>,
2476}
2477
2478impl HydroIrOpMetadata {
2479    #[expect(
2480        clippy::new_without_default,
2481        reason = "explicit calls to new ensure correct backtrace bounds"
2482    )]
2483    pub fn new() -> HydroIrOpMetadata {
2484        Self::new_with_skip(1)
2485    }
2486
2487    fn new_with_skip(skip_count: usize) -> HydroIrOpMetadata {
2488        HydroIrOpMetadata {
2489            backtrace: Backtrace::get_backtrace(2 + skip_count),
2490            cpu_usage: None,
2491            network_recv_cpu_usage: None,
2492            id: None,
2493        }
2494    }
2495}
2496
2497impl Debug for HydroIrOpMetadata {
2498    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
2499        f.debug_struct("HydroIrOpMetadata").finish()
2500    }
2501}
2502
2503impl Hash for HydroIrOpMetadata {
2504    fn hash<H: Hasher>(&self, _: &mut H) {}
2505}
2506
2507/// How a network channel's *sender* prepares each message before it is handed to the transport.
2508///
2509/// A channel's serialization is split into a send half ([`NetworkSend`]) and a receive half
2510/// ([`NetworkRecv`]) so that the multi-version simulation merge can reason about each side
2511/// independently (the sender fork and the receiver are separate IR nodes).
2512#[derive(Debug, Clone, Hash, serde::Serialize)]
2513pub enum NetworkSend {
2514    /// Serialization is performed within the Hydro dataflow using the provided serialize
2515    /// expression. This is how channels using [`crate::networking::Bincode`] are lowered.
2516    Custom { serialize_fn: Option<DebugExpr> },
2517    /// Serialization is left to code outside of Hydro (see [`crate::networking::Embedded`]). The
2518    /// raw `element_type` is passed through unserialized; the only transformation is converting a
2519    /// routing [`crate::location::MemberId`] (the destination cluster `tag`, when demuxing) into
2520    /// the raw `TaglessMemberId` used by the transport. Only supported by the embedded backend.
2521    ///
2522    /// Stored as structured info (rather than a pre-baked expression) so that the code can be
2523    /// synthesized in a post-IR codegen pass.
2524    Embedded {
2525        tag: Option<DebugType>,
2526        element_type: DebugType,
2527    },
2528}
2529
2530/// How a network channel's *receiver* recovers each message from the transport. See
2531/// [`NetworkSend`] for the sender half.
2532#[derive(Debug, Clone, Hash, serde::Serialize)]
2533pub enum NetworkRecv {
2534    /// Deserialization is performed within the Hydro dataflow using the provided deserialize
2535    /// expression. This is how channels using [`crate::networking::Bincode`] are lowered.
2536    Custom { deserialize_fn: Option<DebugExpr> },
2537    /// Deserialization is left to code outside of Hydro (see [`crate::networking::Embedded`]). The
2538    /// raw `element_type` is delivered to the receiver directly, with no transport `Result` to
2539    /// unwrap (the external code that produces the stream decides how to handle faults). The only
2540    /// transformation is converting a `TaglessMemberId` back into a typed
2541    /// [`crate::location::MemberId`] (the sender cluster `tag`, when the receiver is keyed by
2542    /// sender). Only supported by the embedded backend.
2543    Embedded {
2544        tag: Option<DebugType>,
2545        element_type: DebugType,
2546    },
2547}
2548
2549impl NetworkSend {
2550    /// The raw payload type flowing across the channel when serialization is left to external code,
2551    /// or [`None`] when the channel serializes internally.
2552    pub(crate) fn external_element_type(&self) -> Option<&syn::Type> {
2553        match self {
2554            NetworkSend::Custom { .. } => None,
2555            NetworkSend::Embedded { element_type, .. } => Some(&element_type.0),
2556        }
2557    }
2558}
2559
2560impl NetworkRecv {
2561    /// See [`NetworkSend::external_element_type`].
2562    pub(crate) fn external_element_type(&self) -> Option<&syn::Type> {
2563        match self {
2564            NetworkRecv::Custom { .. } => None,
2565            NetworkRecv::Embedded { element_type, .. } => Some(&element_type.0),
2566        }
2567    }
2568}
2569
2570#[cfg(feature = "build")]
2571impl NetworkSend {
2572    /// The expression applied on the sender to prepare each message for the transport, if any.
2573    pub(crate) fn pipeline(&self) -> Option<DebugExpr> {
2574        match self {
2575            NetworkSend::Custom { serialize_fn } => serialize_fn.clone(),
2576            NetworkSend::Embedded { tag, element_type } => {
2577                let root = crate::staging_util::get_this_crate();
2578                let element_type = &element_type.0;
2579                let expr: syn::Expr = if let Some(tag) = tag {
2580                    let tag = &tag.0;
2581                    parse_quote! {
2582                        #root::runtime_support::stageleft::runtime_support::fn1_type_hint::<(#root::__staged::location::MemberId<#tag>, #element_type), _>(
2583                            |(id, data)| (id.into_tagless(), data)
2584                        )
2585                    }
2586                } else {
2587                    parse_quote! {
2588                        #root::runtime_support::stageleft::runtime_support::fn1_type_hint::<#element_type, _>(
2589                            |data| data
2590                        )
2591                    }
2592                };
2593                Some(expr.into())
2594            }
2595        }
2596    }
2597}
2598
2599#[cfg(feature = "build")]
2600impl NetworkRecv {
2601    /// The expression applied on the receiver to recover each message from the transport, if any.
2602    pub(crate) fn pipeline(&self) -> Option<DebugExpr> {
2603        match self {
2604            NetworkRecv::Custom { deserialize_fn } => deserialize_fn.clone(),
2605            // Embedded channels hand the raw payload to the receiver directly (no transport
2606            // `Result`), so the developer's external code decides how to handle serialization
2607            // faults. The only transformation is restoring the typed `MemberId` when the receiver
2608            // is keyed by the sender.
2609            NetworkRecv::Embedded { tag, .. } => {
2610                let tag = tag.as_ref()?;
2611                let root = crate::staging_util::get_this_crate();
2612                let tag = &tag.0;
2613                let expr: syn::Expr = parse_quote! {
2614                    |(id, b)| (#root::__staged::location::MemberId::<#tag>::from_tagless(id as #root::__staged::location::TaglessMemberId), b)
2615                };
2616                Some(expr.into())
2617            }
2618        }
2619    }
2620}
2621
2622/// An intermediate node in a Hydro graph, which consumes data
2623/// from upstream nodes and emits data to downstream nodes.
2624#[derive(Debug, Hash, serde::Serialize)]
2625pub enum HydroNode {
2626    Placeholder,
2627
2628    /// Manually "casts" between two different collection kinds.
2629    ///
2630    /// Using this IR node requires special care, since it bypasses many of Hydro's core
2631    /// correctness checks. In particular, the user must ensure that every possible
2632    /// "interpretation" of the input corresponds to a distinct "interpretation" of the output,
2633    /// where an "interpretation" is a possible output of `ObserveNonDet` applied to the
2634    /// collection. This ensures that the simulator does not miss any possible outputs.
2635    Cast {
2636        inner: Box<HydroNode>,
2637        metadata: HydroIrMetadata,
2638    },
2639
2640    /// Strengthens the guarantees of a stream by non-deterministically selecting a possible
2641    /// interpretation of the input stream.
2642    ///
2643    /// In production, this simply passes through the input, but in simulation, this operator
2644    /// explicitly selects a randomized interpretation.
2645    ObserveNonDet {
2646        inner: Box<HydroNode>,
2647        trusted: bool, // if true, we do not need to simulate non-determinism
2648        metadata: HydroIrMetadata,
2649    },
2650
2651    Source {
2652        source: HydroSource,
2653        metadata: HydroIrMetadata,
2654    },
2655
2656    SingletonSource {
2657        value: DebugExpr,
2658        first_tick_only: bool,
2659        metadata: HydroIrMetadata,
2660    },
2661
2662    CycleSource {
2663        cycle_id: CycleId,
2664        metadata: HydroIrMetadata,
2665    },
2666
2667    Tee {
2668        inner: SharedNode,
2669        metadata: HydroIrMetadata,
2670    },
2671
2672    /// A reference materialization point. Wraps a SharedNode so that:
2673    /// - The pipe output delivers data to one consumer
2674    /// - `#var` references can borrow the value from the slot
2675    ///
2676    /// In DFIR codegen, emits `ident = inner_ident -> singleton()` or `-> optional()` or
2677    /// `-> handoff()` depending on `kind`.
2678    ///
2679    /// Uses the same `built_tees` dedup pattern as `Tee`.
2680    Reference {
2681        inner: SharedNode,
2682        kind: crate::handoff_ref::HandoffRefKind,
2683        access_counter: AccessCounter,
2684        metadata: HydroIrMetadata,
2685    },
2686
2687    Partition {
2688        inner: SharedNode,
2689        f: ClosureExpr,
2690        is_true: bool,
2691        metadata: HydroIrMetadata,
2692    },
2693
2694    BeginAtomic {
2695        inner: Box<HydroNode>,
2696        metadata: HydroIrMetadata,
2697    },
2698
2699    EndAtomic {
2700        inner: Box<HydroNode>,
2701        metadata: HydroIrMetadata,
2702    },
2703
2704    Batch {
2705        inner: Box<HydroNode>,
2706        metadata: HydroIrMetadata,
2707    },
2708
2709    YieldConcat {
2710        inner: Box<HydroNode>,
2711        metadata: HydroIrMetadata,
2712    },
2713
2714    Chain {
2715        first: Box<HydroNode>,
2716        second: Box<HydroNode>,
2717        metadata: HydroIrMetadata,
2718    },
2719
2720    MergeOrdered {
2721        first: Box<HydroNode>,
2722        second: Box<HydroNode>,
2723        metadata: HydroIrMetadata,
2724    },
2725
2726    ChainFirst {
2727        first: Box<HydroNode>,
2728        second: Box<HydroNode>,
2729        metadata: HydroIrMetadata,
2730    },
2731
2732    CrossProduct {
2733        left: Box<HydroNode>,
2734        right: Box<HydroNode>,
2735        metadata: HydroIrMetadata,
2736    },
2737
2738    CrossSingleton {
2739        left: Box<HydroNode>,
2740        right: Box<HydroNode>,
2741        metadata: HydroIrMetadata,
2742    },
2743
2744    Join {
2745        left: Box<HydroNode>,
2746        right: Box<HydroNode>,
2747        metadata: HydroIrMetadata,
2748    },
2749
2750    /// Asymmetric join where the right (build) side is bounded.
2751    /// The build side is accumulated (stratum-delayed) into a hash table,
2752    /// then the left (probe) side streams through preserving its ordering.
2753    JoinHalf {
2754        left: Box<HydroNode>,
2755        right: Box<HydroNode>,
2756        metadata: HydroIrMetadata,
2757    },
2758
2759    Difference {
2760        pos: Box<HydroNode>,
2761        neg: Box<HydroNode>,
2762        metadata: HydroIrMetadata,
2763    },
2764
2765    AntiJoin {
2766        pos: Box<HydroNode>,
2767        neg: Box<HydroNode>,
2768        metadata: HydroIrMetadata,
2769    },
2770
2771    ResolveFutures {
2772        input: Box<HydroNode>,
2773        metadata: HydroIrMetadata,
2774    },
2775    ResolveFuturesBlocking {
2776        input: Box<HydroNode>,
2777        metadata: HydroIrMetadata,
2778    },
2779    ResolveFuturesOrdered {
2780        input: Box<HydroNode>,
2781        metadata: HydroIrMetadata,
2782    },
2783
2784    Map {
2785        f: ClosureExpr,
2786        input: Box<HydroNode>,
2787        metadata: HydroIrMetadata,
2788    },
2789    FlatMap {
2790        f: ClosureExpr,
2791        input: Box<HydroNode>,
2792        metadata: HydroIrMetadata,
2793    },
2794    FlatMapStreamBlocking {
2795        f: ClosureExpr,
2796        input: Box<HydroNode>,
2797        metadata: HydroIrMetadata,
2798    },
2799    Filter {
2800        f: ClosureExpr,
2801        input: Box<HydroNode>,
2802        metadata: HydroIrMetadata,
2803    },
2804    FilterMap {
2805        f: ClosureExpr,
2806        input: Box<HydroNode>,
2807        metadata: HydroIrMetadata,
2808    },
2809
2810    DeferTick {
2811        input: Box<HydroNode>,
2812        metadata: HydroIrMetadata,
2813    },
2814    Enumerate {
2815        input: Box<HydroNode>,
2816        metadata: HydroIrMetadata,
2817    },
2818    Inspect {
2819        f: ClosureExpr,
2820        input: Box<HydroNode>,
2821        metadata: HydroIrMetadata,
2822    },
2823
2824    Unique {
2825        input: Box<HydroNode>,
2826        metadata: HydroIrMetadata,
2827    },
2828
2829    Sort {
2830        input: Box<HydroNode>,
2831        metadata: HydroIrMetadata,
2832    },
2833    Fold {
2834        init: ClosureExpr,
2835        acc: ClosureExpr,
2836        input: Box<HydroNode>,
2837        metadata: HydroIrMetadata,
2838    },
2839
2840    Scan {
2841        init: ClosureExpr,
2842        acc: ClosureExpr,
2843        input: Box<HydroNode>,
2844        metadata: HydroIrMetadata,
2845    },
2846    ScanAsyncBlocking {
2847        init: ClosureExpr,
2848        acc: ClosureExpr,
2849        input: Box<HydroNode>,
2850        metadata: HydroIrMetadata,
2851    },
2852    FoldKeyed {
2853        init: ClosureExpr,
2854        acc: ClosureExpr,
2855        input: Box<HydroNode>,
2856        metadata: HydroIrMetadata,
2857    },
2858
2859    Reduce {
2860        f: ClosureExpr,
2861        input: Box<HydroNode>,
2862        metadata: HydroIrMetadata,
2863    },
2864    ReduceKeyed {
2865        f: ClosureExpr,
2866        input: Box<HydroNode>,
2867        metadata: HydroIrMetadata,
2868    },
2869    ReduceKeyedWatermark {
2870        f: ClosureExpr,
2871        input: Box<HydroNode>,
2872        watermark: Box<HydroNode>,
2873        metadata: HydroIrMetadata,
2874    },
2875
2876    Network {
2877        name: Option<String>,
2878        networking_info: crate::networking::NetworkingInfo,
2879        serialize: NetworkSend,
2880        deserialize: NetworkRecv,
2881        instantiate_fn: DebugInstantiate,
2882        input: Box<HydroNode>,
2883        metadata: HydroIrMetadata,
2884    },
2885
2886    VersionedNetworkFork {
2887        channel_id: u32,
2888        channel_name: String,
2889        senders: Vec<(u32, Box<HydroNode>, NetworkSend)>,
2890        metadata: HydroIrMetadata,
2891    },
2892
2893    VersionedNetwork {
2894        fork: SharedNode,
2895        version: u32,
2896        deserialize: NetworkRecv,
2897        metadata: HydroIrMetadata,
2898    },
2899
2900    ExternalInput {
2901        from_external_key: LocationKey,
2902        from_port_id: ExternalPortId,
2903        from_many: bool,
2904        codec_type: DebugType,
2905        #[serde(skip)]
2906        port_hint: NetworkHint,
2907        instantiate_fn: DebugInstantiate,
2908        deserialize_fn: Option<DebugExpr>,
2909        metadata: HydroIrMetadata,
2910    },
2911
2912    Counter {
2913        tag: String,
2914        duration: DebugExpr,
2915        prefix: String,
2916        input: Box<HydroNode>,
2917        metadata: HydroIrMetadata,
2918    },
2919
2920    AssertIsConsistent {
2921        inner: Box<HydroNode>,
2922        trusted: bool,
2923        metadata: HydroIrMetadata,
2924    },
2925
2926    UnboundSingleton {
2927        inner: Box<HydroNode>,
2928        metadata: HydroIrMetadata,
2929    },
2930}
2931
2932pub type SeenSharedNodes = HashMap<*const RefCell<HydroNode>, Rc<RefCell<HydroNode>>>;
2933pub type SeenSharedNodeLocations = HashMap<*const RefCell<HydroNode>, LocationId>;
2934
2935/// If `f` has a mut singleton ref and `in_kind` is non-strict, emits an
2936/// `observe_for_mut` node and returns the new ident. Otherwise returns
2937/// `in_ident` unchanged. Always consumes a stmt_id when applicable.
2938#[cfg(feature = "build")]
2939fn maybe_observe_for_mut(
2940    f: &ClosureExpr,
2941    in_ident: syn::Ident,
2942    in_location: &LocationId,
2943    in_kind: &CollectionKind,
2944    op_meta: &HydroIrOpMetadata,
2945    builders_or_callback: &mut BuildersOrCallback<
2946        impl FnMut(&mut HydroRoot, &mut crate::Counter<StmtId>),
2947        impl FnMut(&mut HydroNode, &mut crate::Counter<StmtId>),
2948    >,
2949    next_stmt_id: &mut crate::Counter<StmtId>,
2950) -> syn::Ident {
2951    if f.has_mut_ref() && !in_kind.is_strict() {
2952        let observe_stmt_id = next_stmt_id.get_and_increment();
2953        let observe_ident =
2954            syn::Ident::new(&format!("stream_{}", observe_stmt_id), Span::call_site());
2955        if let BuildersOrCallback::Builders(graph_builders) = builders_or_callback {
2956            graph_builders.observe_for_mut(in_location, in_ident, in_kind, &observe_ident, op_meta);
2957        }
2958        observe_ident
2959    } else {
2960        in_ident
2961    }
2962}
2963
2964impl HydroNode {
2965    pub fn transform_bottom_up(
2966        &mut self,
2967        transform: &mut impl FnMut(&mut HydroNode),
2968        seen_tees: &mut SeenSharedNodes,
2969        check_well_formed: bool,
2970    ) {
2971        self.transform_children(
2972            |n, s| n.transform_bottom_up(transform, s, check_well_formed),
2973            seen_tees,
2974        );
2975
2976        transform(self);
2977
2978        let self_location = self.metadata().location_id.root();
2979
2980        if check_well_formed {
2981            match &*self {
2982                HydroNode::Network { .. } => {}
2983                _ => {
2984                    self.input_metadata().iter().for_each(|i| {
2985                        if i.location_id.root() != self_location {
2986                            panic!(
2987                                "Mismatching IR locations, child: {:?} ({:?}) of: {:?} ({:?})",
2988                                i,
2989                                i.location_id.root(),
2990                                self,
2991                                self_location
2992                            )
2993                        }
2994                    });
2995                }
2996            }
2997        }
2998    }
2999
3000    #[inline(always)]
3001    pub fn transform_children(
3002        &mut self,
3003        mut transform: impl FnMut(&mut HydroNode, &mut SeenSharedNodes),
3004        seen_tees: &mut SeenSharedNodes,
3005    ) {
3006        match self {
3007            HydroNode::Placeholder => {
3008                panic!();
3009            }
3010
3011            HydroNode::Source { .. }
3012            | HydroNode::SingletonSource { .. }
3013            | HydroNode::CycleSource { .. }
3014            | HydroNode::ExternalInput { .. } => {}
3015
3016            HydroNode::Tee { inner, .. } | HydroNode::Reference { inner, .. } => {
3017                if let Some(transformed) = seen_tees.get(&inner.as_ptr()) {
3018                    *inner = SharedNode(transformed.clone());
3019                } else {
3020                    let transformed_cell = Rc::new(RefCell::new(HydroNode::Placeholder));
3021                    seen_tees.insert(inner.as_ptr(), transformed_cell.clone());
3022                    let mut orig = inner.0.replace(HydroNode::Placeholder);
3023                    transform(&mut orig, seen_tees);
3024                    *transformed_cell.borrow_mut() = orig;
3025                    *inner = SharedNode(transformed_cell);
3026                }
3027            }
3028
3029            HydroNode::Partition { inner, f, .. } => {
3030                if let Some(transformed) = seen_tees.get(&inner.as_ptr()) {
3031                    *inner = SharedNode(transformed.clone());
3032                } else {
3033                    f.transform_children(&mut transform, seen_tees);
3034                    let transformed_cell = Rc::new(RefCell::new(HydroNode::Placeholder));
3035                    seen_tees.insert(inner.as_ptr(), transformed_cell.clone());
3036                    let mut orig = inner.0.replace(HydroNode::Placeholder);
3037                    transform(&mut orig, seen_tees);
3038                    *transformed_cell.borrow_mut() = orig;
3039                    *inner = SharedNode(transformed_cell);
3040                }
3041            }
3042
3043            HydroNode::Cast { inner, .. }
3044            | HydroNode::ObserveNonDet { inner, .. }
3045            | HydroNode::BeginAtomic { inner, .. }
3046            | HydroNode::EndAtomic { inner, .. }
3047            | HydroNode::Batch { inner, .. }
3048            | HydroNode::YieldConcat { inner, .. }
3049            | HydroNode::UnboundSingleton { inner, .. }
3050            | HydroNode::AssertIsConsistent { inner, .. } => {
3051                transform(inner.as_mut(), seen_tees);
3052            }
3053
3054            HydroNode::Chain { first, second, .. } => {
3055                transform(first.as_mut(), seen_tees);
3056                transform(second.as_mut(), seen_tees);
3057            }
3058
3059            HydroNode::MergeOrdered { first, second, .. } => {
3060                transform(first.as_mut(), seen_tees);
3061                transform(second.as_mut(), seen_tees);
3062            }
3063
3064            HydroNode::ChainFirst { first, second, .. } => {
3065                transform(first.as_mut(), seen_tees);
3066                transform(second.as_mut(), seen_tees);
3067            }
3068
3069            HydroNode::CrossSingleton { left, right, .. }
3070            | HydroNode::CrossProduct { left, right, .. }
3071            | HydroNode::Join { left, right, .. }
3072            | HydroNode::JoinHalf { left, right, .. } => {
3073                transform(left.as_mut(), seen_tees);
3074                transform(right.as_mut(), seen_tees);
3075            }
3076
3077            HydroNode::Difference { pos, neg, .. } | HydroNode::AntiJoin { pos, neg, .. } => {
3078                transform(pos.as_mut(), seen_tees);
3079                transform(neg.as_mut(), seen_tees);
3080            }
3081
3082            HydroNode::Map { f, input, .. } => {
3083                f.transform_children(&mut transform, seen_tees);
3084                transform(input.as_mut(), seen_tees);
3085            }
3086            HydroNode::FlatMap { f, input, .. }
3087            | HydroNode::FlatMapStreamBlocking { f, input, .. }
3088            | HydroNode::Filter { f, input, .. }
3089            | HydroNode::FilterMap { f, input, .. }
3090            | HydroNode::Inspect { f, input, .. }
3091            | HydroNode::Reduce { f, input, .. }
3092            | HydroNode::ReduceKeyed { f, input, .. } => {
3093                f.transform_children(&mut transform, seen_tees);
3094                transform(input.as_mut(), seen_tees);
3095            }
3096            HydroNode::ReduceKeyedWatermark {
3097                f,
3098                input,
3099                watermark,
3100                ..
3101            } => {
3102                f.transform_children(&mut transform, seen_tees);
3103                transform(input.as_mut(), seen_tees);
3104                transform(watermark.as_mut(), seen_tees);
3105            }
3106            HydroNode::Fold {
3107                init, acc, input, ..
3108            }
3109            | HydroNode::Scan {
3110                init, acc, input, ..
3111            }
3112            | HydroNode::ScanAsyncBlocking {
3113                init, acc, input, ..
3114            }
3115            | HydroNode::FoldKeyed {
3116                init, acc, input, ..
3117            } => {
3118                init.transform_children(&mut transform, seen_tees);
3119                acc.transform_children(&mut transform, seen_tees);
3120                transform(input.as_mut(), seen_tees);
3121            }
3122            HydroNode::ResolveFutures { input, .. }
3123            | HydroNode::ResolveFuturesBlocking { input, .. }
3124            | HydroNode::ResolveFuturesOrdered { input, .. }
3125            | HydroNode::Sort { input, .. }
3126            | HydroNode::DeferTick { input, .. }
3127            | HydroNode::Enumerate { input, .. }
3128            | HydroNode::Unique { input, .. }
3129            | HydroNode::Network { input, .. }
3130            | HydroNode::Counter { input, .. } => {
3131                transform(input.as_mut(), seen_tees);
3132            }
3133
3134            HydroNode::VersionedNetworkFork { senders, .. } => {
3135                for (_version, sender, _serialize) in senders.iter_mut() {
3136                    transform(sender.as_mut(), seen_tees);
3137                }
3138            }
3139
3140            HydroNode::VersionedNetwork { fork, .. } => {
3141                if let Some(transformed) = seen_tees.get(&fork.as_ptr()) {
3142                    *fork = SharedNode(transformed.clone());
3143                } else {
3144                    let transformed_cell = Rc::new(RefCell::new(HydroNode::Placeholder));
3145                    seen_tees.insert(fork.as_ptr(), transformed_cell.clone());
3146                    let mut orig = fork.0.replace(HydroNode::Placeholder);
3147                    transform(&mut orig, seen_tees);
3148                    *transformed_cell.borrow_mut() = orig;
3149                    *fork = SharedNode(transformed_cell);
3150                }
3151            }
3152        }
3153    }
3154
3155    pub fn deep_clone(&self, seen_tees: &mut SeenSharedNodes) -> HydroNode {
3156        match self {
3157            HydroNode::Placeholder => HydroNode::Placeholder,
3158            HydroNode::Cast { inner, metadata } => HydroNode::Cast {
3159                inner: Box::new(inner.deep_clone(seen_tees)),
3160                metadata: metadata.clone(),
3161            },
3162            HydroNode::UnboundSingleton { inner, metadata } => HydroNode::UnboundSingleton {
3163                inner: Box::new(inner.deep_clone(seen_tees)),
3164                metadata: metadata.clone(),
3165            },
3166            HydroNode::ObserveNonDet {
3167                inner,
3168                trusted,
3169                metadata,
3170            } => HydroNode::ObserveNonDet {
3171                inner: Box::new(inner.deep_clone(seen_tees)),
3172                trusted: *trusted,
3173                metadata: metadata.clone(),
3174            },
3175            HydroNode::AssertIsConsistent {
3176                inner,
3177                trusted,
3178                metadata,
3179            } => HydroNode::AssertIsConsistent {
3180                inner: Box::new(inner.deep_clone(seen_tees)),
3181                trusted: *trusted,
3182                metadata: metadata.clone(),
3183            },
3184            HydroNode::Source { source, metadata } => HydroNode::Source {
3185                source: source.clone(),
3186                metadata: metadata.clone(),
3187            },
3188            HydroNode::SingletonSource {
3189                value,
3190                first_tick_only,
3191                metadata,
3192            } => HydroNode::SingletonSource {
3193                value: value.clone(),
3194                first_tick_only: *first_tick_only,
3195                metadata: metadata.clone(),
3196            },
3197            HydroNode::CycleSource { cycle_id, metadata } => HydroNode::CycleSource {
3198                cycle_id: *cycle_id,
3199                metadata: metadata.clone(),
3200            },
3201            HydroNode::Tee { inner, metadata }
3202            | HydroNode::Reference {
3203                inner, metadata, ..
3204            } => {
3205                let cloned_inner = if let Some(transformed) = seen_tees.get(&inner.as_ptr()) {
3206                    SharedNode(transformed.clone())
3207                } else {
3208                    let new_rc = Rc::new(RefCell::new(HydroNode::Placeholder));
3209                    seen_tees.insert(inner.as_ptr(), new_rc.clone());
3210                    let cloned = inner.0.borrow().deep_clone(seen_tees);
3211                    *new_rc.borrow_mut() = cloned;
3212                    SharedNode(new_rc)
3213                };
3214                if let HydroNode::Reference {
3215                    kind,
3216                    access_counter,
3217                    ..
3218                } = self
3219                {
3220                    HydroNode::Reference {
3221                        inner: cloned_inner,
3222                        kind: *kind,
3223                        access_counter: access_counter.freeze(),
3224                        metadata: metadata.clone(),
3225                    }
3226                } else {
3227                    HydroNode::Tee {
3228                        inner: cloned_inner,
3229                        metadata: metadata.clone(),
3230                    }
3231                }
3232            }
3233            HydroNode::Partition {
3234                inner,
3235                f,
3236                is_true,
3237                metadata,
3238            } => {
3239                if let Some(transformed) = seen_tees.get(&inner.as_ptr()) {
3240                    HydroNode::Partition {
3241                        inner: SharedNode(transformed.clone()),
3242                        f: f.deep_clone(seen_tees),
3243                        is_true: *is_true,
3244                        metadata: metadata.clone(),
3245                    }
3246                } else {
3247                    let new_rc = Rc::new(RefCell::new(HydroNode::Placeholder));
3248                    seen_tees.insert(inner.as_ptr(), new_rc.clone());
3249                    let cloned = inner.0.borrow().deep_clone(seen_tees);
3250                    *new_rc.borrow_mut() = cloned;
3251                    HydroNode::Partition {
3252                        inner: SharedNode(new_rc),
3253                        f: f.deep_clone(seen_tees),
3254                        is_true: *is_true,
3255                        metadata: metadata.clone(),
3256                    }
3257                }
3258            }
3259            HydroNode::YieldConcat { inner, metadata } => HydroNode::YieldConcat {
3260                inner: Box::new(inner.deep_clone(seen_tees)),
3261                metadata: metadata.clone(),
3262            },
3263            HydroNode::BeginAtomic { inner, metadata } => HydroNode::BeginAtomic {
3264                inner: Box::new(inner.deep_clone(seen_tees)),
3265                metadata: metadata.clone(),
3266            },
3267            HydroNode::EndAtomic { inner, metadata } => HydroNode::EndAtomic {
3268                inner: Box::new(inner.deep_clone(seen_tees)),
3269                metadata: metadata.clone(),
3270            },
3271            HydroNode::Batch { inner, metadata } => HydroNode::Batch {
3272                inner: Box::new(inner.deep_clone(seen_tees)),
3273                metadata: metadata.clone(),
3274            },
3275            HydroNode::Chain {
3276                first,
3277                second,
3278                metadata,
3279            } => HydroNode::Chain {
3280                first: Box::new(first.deep_clone(seen_tees)),
3281                second: Box::new(second.deep_clone(seen_tees)),
3282                metadata: metadata.clone(),
3283            },
3284            HydroNode::MergeOrdered {
3285                first,
3286                second,
3287                metadata,
3288            } => HydroNode::MergeOrdered {
3289                first: Box::new(first.deep_clone(seen_tees)),
3290                second: Box::new(second.deep_clone(seen_tees)),
3291                metadata: metadata.clone(),
3292            },
3293            HydroNode::ChainFirst {
3294                first,
3295                second,
3296                metadata,
3297            } => HydroNode::ChainFirst {
3298                first: Box::new(first.deep_clone(seen_tees)),
3299                second: Box::new(second.deep_clone(seen_tees)),
3300                metadata: metadata.clone(),
3301            },
3302            HydroNode::CrossProduct {
3303                left,
3304                right,
3305                metadata,
3306            } => HydroNode::CrossProduct {
3307                left: Box::new(left.deep_clone(seen_tees)),
3308                right: Box::new(right.deep_clone(seen_tees)),
3309                metadata: metadata.clone(),
3310            },
3311            HydroNode::CrossSingleton {
3312                left,
3313                right,
3314                metadata,
3315            } => HydroNode::CrossSingleton {
3316                left: Box::new(left.deep_clone(seen_tees)),
3317                right: Box::new(right.deep_clone(seen_tees)),
3318                metadata: metadata.clone(),
3319            },
3320            HydroNode::Join {
3321                left,
3322                right,
3323                metadata,
3324            } => HydroNode::Join {
3325                left: Box::new(left.deep_clone(seen_tees)),
3326                right: Box::new(right.deep_clone(seen_tees)),
3327                metadata: metadata.clone(),
3328            },
3329            HydroNode::JoinHalf {
3330                left,
3331                right,
3332                metadata,
3333            } => HydroNode::JoinHalf {
3334                left: Box::new(left.deep_clone(seen_tees)),
3335                right: Box::new(right.deep_clone(seen_tees)),
3336                metadata: metadata.clone(),
3337            },
3338            HydroNode::Difference { pos, neg, metadata } => HydroNode::Difference {
3339                pos: Box::new(pos.deep_clone(seen_tees)),
3340                neg: Box::new(neg.deep_clone(seen_tees)),
3341                metadata: metadata.clone(),
3342            },
3343            HydroNode::AntiJoin { pos, neg, metadata } => HydroNode::AntiJoin {
3344                pos: Box::new(pos.deep_clone(seen_tees)),
3345                neg: Box::new(neg.deep_clone(seen_tees)),
3346                metadata: metadata.clone(),
3347            },
3348            HydroNode::ResolveFutures { input, metadata } => HydroNode::ResolveFutures {
3349                input: Box::new(input.deep_clone(seen_tees)),
3350                metadata: metadata.clone(),
3351            },
3352            HydroNode::ResolveFuturesBlocking { input, metadata } => {
3353                HydroNode::ResolveFuturesBlocking {
3354                    input: Box::new(input.deep_clone(seen_tees)),
3355                    metadata: metadata.clone(),
3356                }
3357            }
3358            HydroNode::ResolveFuturesOrdered { input, metadata } => {
3359                HydroNode::ResolveFuturesOrdered {
3360                    input: Box::new(input.deep_clone(seen_tees)),
3361                    metadata: metadata.clone(),
3362                }
3363            }
3364            HydroNode::Map { f, input, metadata } => HydroNode::Map {
3365                f: f.deep_clone(seen_tees),
3366                input: Box::new(input.deep_clone(seen_tees)),
3367                metadata: metadata.clone(),
3368            },
3369            HydroNode::FlatMap { f, input, metadata } => HydroNode::FlatMap {
3370                f: f.deep_clone(seen_tees),
3371                input: Box::new(input.deep_clone(seen_tees)),
3372                metadata: metadata.clone(),
3373            },
3374            HydroNode::FlatMapStreamBlocking { f, input, metadata } => {
3375                HydroNode::FlatMapStreamBlocking {
3376                    f: f.deep_clone(seen_tees),
3377                    input: Box::new(input.deep_clone(seen_tees)),
3378                    metadata: metadata.clone(),
3379                }
3380            }
3381            HydroNode::Filter { f, input, metadata } => HydroNode::Filter {
3382                f: f.deep_clone(seen_tees),
3383                input: Box::new(input.deep_clone(seen_tees)),
3384                metadata: metadata.clone(),
3385            },
3386            HydroNode::FilterMap { f, input, metadata } => HydroNode::FilterMap {
3387                f: f.deep_clone(seen_tees),
3388                input: Box::new(input.deep_clone(seen_tees)),
3389                metadata: metadata.clone(),
3390            },
3391            HydroNode::DeferTick { input, metadata } => HydroNode::DeferTick {
3392                input: Box::new(input.deep_clone(seen_tees)),
3393                metadata: metadata.clone(),
3394            },
3395            HydroNode::Enumerate { input, metadata } => HydroNode::Enumerate {
3396                input: Box::new(input.deep_clone(seen_tees)),
3397                metadata: metadata.clone(),
3398            },
3399            HydroNode::Inspect { f, input, metadata } => HydroNode::Inspect {
3400                f: f.deep_clone(seen_tees),
3401                input: Box::new(input.deep_clone(seen_tees)),
3402                metadata: metadata.clone(),
3403            },
3404            HydroNode::Unique { input, metadata } => HydroNode::Unique {
3405                input: Box::new(input.deep_clone(seen_tees)),
3406                metadata: metadata.clone(),
3407            },
3408            HydroNode::Sort { input, metadata } => HydroNode::Sort {
3409                input: Box::new(input.deep_clone(seen_tees)),
3410                metadata: metadata.clone(),
3411            },
3412            HydroNode::Fold {
3413                init,
3414                acc,
3415                input,
3416                metadata,
3417            } => HydroNode::Fold {
3418                init: init.deep_clone(seen_tees),
3419                acc: acc.deep_clone(seen_tees),
3420                input: Box::new(input.deep_clone(seen_tees)),
3421                metadata: metadata.clone(),
3422            },
3423            HydroNode::Scan {
3424                init,
3425                acc,
3426                input,
3427                metadata,
3428            } => HydroNode::Scan {
3429                init: init.deep_clone(seen_tees),
3430                acc: acc.deep_clone(seen_tees),
3431                input: Box::new(input.deep_clone(seen_tees)),
3432                metadata: metadata.clone(),
3433            },
3434            HydroNode::ScanAsyncBlocking {
3435                init,
3436                acc,
3437                input,
3438                metadata,
3439            } => HydroNode::ScanAsyncBlocking {
3440                init: init.deep_clone(seen_tees),
3441                acc: acc.deep_clone(seen_tees),
3442                input: Box::new(input.deep_clone(seen_tees)),
3443                metadata: metadata.clone(),
3444            },
3445            HydroNode::FoldKeyed {
3446                init,
3447                acc,
3448                input,
3449                metadata,
3450            } => HydroNode::FoldKeyed {
3451                init: init.deep_clone(seen_tees),
3452                acc: acc.deep_clone(seen_tees),
3453                input: Box::new(input.deep_clone(seen_tees)),
3454                metadata: metadata.clone(),
3455            },
3456            HydroNode::ReduceKeyedWatermark {
3457                f,
3458                input,
3459                watermark,
3460                metadata,
3461            } => HydroNode::ReduceKeyedWatermark {
3462                f: f.deep_clone(seen_tees),
3463                input: Box::new(input.deep_clone(seen_tees)),
3464                watermark: Box::new(watermark.deep_clone(seen_tees)),
3465                metadata: metadata.clone(),
3466            },
3467            HydroNode::Reduce { f, input, metadata } => HydroNode::Reduce {
3468                f: f.deep_clone(seen_tees),
3469                input: Box::new(input.deep_clone(seen_tees)),
3470                metadata: metadata.clone(),
3471            },
3472            HydroNode::ReduceKeyed { f, input, metadata } => HydroNode::ReduceKeyed {
3473                f: f.deep_clone(seen_tees),
3474                input: Box::new(input.deep_clone(seen_tees)),
3475                metadata: metadata.clone(),
3476            },
3477            HydroNode::Network {
3478                name,
3479                networking_info,
3480                serialize,
3481                deserialize,
3482                instantiate_fn,
3483                input,
3484                metadata,
3485            } => HydroNode::Network {
3486                name: name.clone(),
3487                networking_info: networking_info.clone(),
3488                serialize: serialize.clone(),
3489                deserialize: deserialize.clone(),
3490                instantiate_fn: instantiate_fn.clone(),
3491                input: Box::new(input.deep_clone(seen_tees)),
3492                metadata: metadata.clone(),
3493            },
3494            HydroNode::ExternalInput {
3495                from_external_key,
3496                from_port_id,
3497                from_many,
3498                codec_type,
3499                port_hint,
3500                instantiate_fn,
3501                deserialize_fn,
3502                metadata,
3503            } => HydroNode::ExternalInput {
3504                from_external_key: *from_external_key,
3505                from_port_id: *from_port_id,
3506                from_many: *from_many,
3507                codec_type: codec_type.clone(),
3508                port_hint: *port_hint,
3509                instantiate_fn: instantiate_fn.clone(),
3510                deserialize_fn: deserialize_fn.clone(),
3511                metadata: metadata.clone(),
3512            },
3513            HydroNode::Counter {
3514                tag,
3515                duration,
3516                prefix,
3517                input,
3518                metadata,
3519            } => HydroNode::Counter {
3520                tag: tag.clone(),
3521                duration: duration.clone(),
3522                prefix: prefix.clone(),
3523                input: Box::new(input.deep_clone(seen_tees)),
3524                metadata: metadata.clone(),
3525            },
3526            HydroNode::VersionedNetworkFork {
3527                channel_id,
3528                channel_name,
3529                senders,
3530                metadata,
3531            } => HydroNode::VersionedNetworkFork {
3532                channel_id: *channel_id,
3533                channel_name: channel_name.clone(),
3534                senders: senders
3535                    .iter()
3536                    .map(|(version, sender, serialize)| {
3537                        (
3538                            *version,
3539                            Box::new(sender.deep_clone(seen_tees)),
3540                            serialize.clone(),
3541                        )
3542                    })
3543                    .collect(),
3544                metadata: metadata.clone(),
3545            },
3546            HydroNode::VersionedNetwork {
3547                fork,
3548                version,
3549                deserialize,
3550                metadata,
3551            } => {
3552                let cloned_fork = if let Some(transformed) = seen_tees.get(&fork.as_ptr()) {
3553                    SharedNode(transformed.clone())
3554                } else {
3555                    let new_rc = Rc::new(RefCell::new(HydroNode::Placeholder));
3556                    seen_tees.insert(fork.as_ptr(), new_rc.clone());
3557                    let cloned = fork.0.borrow().deep_clone(seen_tees);
3558                    *new_rc.borrow_mut() = cloned;
3559                    SharedNode(new_rc)
3560                };
3561                HydroNode::VersionedNetwork {
3562                    fork: cloned_fork,
3563                    version: *version,
3564                    deserialize: deserialize.clone(),
3565                    metadata: metadata.clone(),
3566                }
3567            }
3568        }
3569    }
3570
3571    #[cfg(feature = "build")]
3572    pub fn emit_core(
3573        &mut self,
3574        builders_or_callback: &mut BuildersOrCallback<
3575            impl FnMut(&mut HydroRoot, &mut crate::Counter<StmtId>),
3576            impl FnMut(&mut HydroNode, &mut crate::Counter<StmtId>),
3577        >,
3578        seen_tees: &mut SeenSharedNodes,
3579        built_tees: &mut HashMap<*const RefCell<HydroNode>, Vec<syn::Ident>>,
3580        next_stmt_id: &mut crate::Counter<StmtId>,
3581        fold_hooked_idents: &mut HashSet<String>,
3582    ) -> syn::Ident {
3583        let mut ident_stack: Vec<syn::Ident> = Vec::new();
3584
3585        self.transform_bottom_up(
3586            &mut |node: &mut HydroNode| {
3587                let out_location = node.metadata().location_id.clone();
3588                match node {
3589                    HydroNode::Placeholder => {
3590                        panic!()
3591                    }
3592
3593                    HydroNode::Cast { .. } => {
3594                        // Cast passes through the input ident unchanged
3595                        // The input ident is already on the stack from processing the child
3596                        let _ = next_stmt_id.get_and_increment();
3597                        match builders_or_callback {
3598                            BuildersOrCallback::Builders(_) => {}
3599                            BuildersOrCallback::Callback(_, node_callback) => {
3600                                node_callback(node, next_stmt_id);
3601                            }
3602                        }
3603                        // input_ident stays on stack as output
3604                    }
3605
3606                    HydroNode::UnboundSingleton { .. } => {
3607                        let inner_ident = ident_stack.pop().unwrap();
3608
3609                        let stmt_id = next_stmt_id.get_and_increment();
3610                        let out_ident =
3611                            syn::Ident::new(&format!("stream_{}", stmt_id), Span::call_site());
3612
3613                        match builders_or_callback {
3614                            BuildersOrCallback::Builders(graph_builders) => {
3615                                if graph_builders.singleton_intermediates() {
3616                                    graph_builders.add_dfir_at(
3617                                        &out_location,
3618                                        parse_quote! {
3619                                            #out_ident = #inner_ident;
3620                                        },
3621                                        None,
3622                                    );
3623                                } else {
3624                                    graph_builders.add_dfir_at(
3625                                        &out_location,
3626                                        parse_quote! {
3627                                            #out_ident = #inner_ident -> persist::<'static>();
3628                                        },
3629                                        None,
3630                                    );
3631                                }
3632                            }
3633                            BuildersOrCallback::Callback(_, node_callback) => {
3634                                node_callback(node, next_stmt_id);
3635                            }
3636                        }
3637
3638                        ident_stack.push(out_ident);
3639                    }
3640
3641                    HydroNode::AssertIsConsistent { inner, trusted, .. } => {
3642                        let inner_ident = ident_stack.pop().unwrap();
3643
3644                        let stmt_id = next_stmt_id.get_and_increment();
3645                        let out_ident =
3646                            syn::Ident::new(&format!("stream_{}", stmt_id), Span::call_site());
3647
3648                        match builders_or_callback {
3649                            BuildersOrCallback::Builders(graph_builders) => {
3650                                graph_builders.assert_is_consistent(
3651                                    *trusted,
3652                                    &inner.metadata().location_id,
3653                                    inner_ident,
3654                                    &out_ident,
3655                                );
3656                            }
3657                            BuildersOrCallback::Callback(_, node_callback) => {
3658                                node_callback(node, next_stmt_id);
3659                            }
3660                        }
3661
3662                        ident_stack.push(out_ident);
3663                    }
3664
3665                    HydroNode::ObserveNonDet {
3666                        inner,
3667                        trusted,
3668                        metadata,
3669                        ..
3670                    } => {
3671                        let inner_ident = ident_stack.pop().unwrap();
3672
3673                        let stmt_id = next_stmt_id.get_and_increment();
3674                        let observe_ident =
3675                            syn::Ident::new(&format!("stream_{}", stmt_id), Span::call_site());
3676
3677                        match builders_or_callback {
3678                            BuildersOrCallback::Builders(graph_builders) => {
3679                                graph_builders.observe_nondet(
3680                                    *trusted,
3681                                    &inner.metadata().location_id,
3682                                    inner_ident,
3683                                    &inner.metadata().collection_kind,
3684                                    &observe_ident,
3685                                    &metadata.collection_kind,
3686                                    &metadata.op,
3687                                );
3688                            }
3689                            BuildersOrCallback::Callback(_, node_callback) => {
3690                                node_callback(node, next_stmt_id);
3691                            }
3692                        }
3693
3694                        ident_stack.push(observe_ident);
3695                    }
3696
3697                    HydroNode::Batch {
3698                        inner, metadata, ..
3699                    } => {
3700                        let inner_ident = ident_stack.pop().unwrap();
3701
3702                        let stmt_id = next_stmt_id.get_and_increment();
3703                        let batch_ident =
3704                            syn::Ident::new(&format!("stream_{}", stmt_id), Span::call_site());
3705
3706                        match builders_or_callback {
3707                            BuildersOrCallback::Builders(graph_builders) => {
3708                                graph_builders.batch(
3709                                    inner_ident,
3710                                    &inner.metadata().location_id,
3711                                    &inner.metadata().collection_kind,
3712                                    &batch_ident,
3713                                    &out_location,
3714                                    &metadata.op,
3715                                    fold_hooked_idents,
3716                                );
3717                            }
3718                            BuildersOrCallback::Callback(_, node_callback) => {
3719                                node_callback(node, next_stmt_id);
3720                            }
3721                        }
3722
3723                        ident_stack.push(batch_ident);
3724                    }
3725
3726                    HydroNode::YieldConcat { inner, .. } => {
3727                        let inner_ident = ident_stack.pop().unwrap();
3728
3729                        let stmt_id = next_stmt_id.get_and_increment();
3730                        let yield_ident =
3731                            syn::Ident::new(&format!("stream_{}", stmt_id), Span::call_site());
3732
3733                        match builders_or_callback {
3734                            BuildersOrCallback::Builders(graph_builders) => {
3735                                graph_builders.yield_from_tick(
3736                                    inner_ident,
3737                                    &inner.metadata().location_id,
3738                                    &inner.metadata().collection_kind,
3739                                    &yield_ident,
3740                                    &out_location,
3741                                );
3742                            }
3743                            BuildersOrCallback::Callback(_, node_callback) => {
3744                                node_callback(node, next_stmt_id);
3745                            }
3746                        }
3747
3748                        ident_stack.push(yield_ident);
3749                    }
3750
3751                    HydroNode::BeginAtomic { inner, metadata } => {
3752                        let inner_ident = ident_stack.pop().unwrap();
3753
3754                        let stmt_id = next_stmt_id.get_and_increment();
3755                        let begin_ident =
3756                            syn::Ident::new(&format!("stream_{}", stmt_id), Span::call_site());
3757
3758                        match builders_or_callback {
3759                            BuildersOrCallback::Builders(graph_builders) => {
3760                                graph_builders.begin_atomic(
3761                                    inner_ident,
3762                                    &inner.metadata().location_id,
3763                                    &inner.metadata().collection_kind,
3764                                    &begin_ident,
3765                                    &out_location,
3766                                    &metadata.op,
3767                                );
3768                            }
3769                            BuildersOrCallback::Callback(_, node_callback) => {
3770                                node_callback(node, next_stmt_id);
3771                            }
3772                        }
3773
3774                        ident_stack.push(begin_ident);
3775                    }
3776
3777                    HydroNode::EndAtomic { inner, .. } => {
3778                        let inner_ident = ident_stack.pop().unwrap();
3779
3780                        let stmt_id = next_stmt_id.get_and_increment();
3781                        let end_ident =
3782                            syn::Ident::new(&format!("stream_{}", stmt_id), Span::call_site());
3783
3784                        match builders_or_callback {
3785                            BuildersOrCallback::Builders(graph_builders) => {
3786                                graph_builders.end_atomic(
3787                                    inner_ident,
3788                                    &inner.metadata().location_id,
3789                                    &inner.metadata().collection_kind,
3790                                    &end_ident,
3791                                );
3792                            }
3793                            BuildersOrCallback::Callback(_, node_callback) => {
3794                                node_callback(node, next_stmt_id);
3795                            }
3796                        }
3797
3798                        ident_stack.push(end_ident);
3799                    }
3800
3801                    HydroNode::Source {
3802                        source, metadata, ..
3803                    } => {
3804                        if let HydroSource::ExternalNetwork() = source {
3805                            ident_stack.push(syn::Ident::new("DUMMY", Span::call_site()));
3806                        } else {
3807                            let stmt_id = next_stmt_id.get_and_increment();
3808                            let source_ident =
3809                                syn::Ident::new(&format!("stream_{}", stmt_id), Span::call_site());
3810
3811                            let source_stmt = match source {
3812                                HydroSource::Stream(expr) => {
3813                                    debug_assert!(metadata.location_id.is_top_level());
3814                                    parse_quote! {
3815                                        #source_ident = source_stream(#expr);
3816                                    }
3817                                }
3818
3819                                HydroSource::ExternalNetwork() => {
3820                                    unreachable!()
3821                                }
3822
3823                                HydroSource::Iter(expr) => {
3824                                    if metadata.location_id.is_top_level() {
3825                                        parse_quote! {
3826                                            #source_ident = source_iter(#expr);
3827                                        }
3828                                    } else {
3829                                        // TODO(shadaj): a more natural semantics would be to to re-evaluate the expression on each tick
3830                                        parse_quote! {
3831                                            #source_ident = source_iter(#expr) -> persist::<'static>();
3832                                        }
3833                                    }
3834                                }
3835
3836                                HydroSource::Spin() => {
3837                                    debug_assert!(metadata.location_id.is_top_level());
3838                                    parse_quote! {
3839                                        #source_ident = spin();
3840                                    }
3841                                }
3842
3843                                HydroSource::ClusterMembers(target_loc, state) => {
3844                                    debug_assert!(metadata.location_id.is_top_level());
3845
3846                                    let members_tee_ident = syn::Ident::new(
3847                                        &format!(
3848                                            "__cluster_members_tee_{}_{}",
3849                                            metadata.location_id.root().key(),
3850                                            target_loc.key(),
3851                                        ),
3852                                        Span::call_site(),
3853                                    );
3854
3855                                    match state {
3856                                        ClusterMembersState::Stream(d) => {
3857                                            parse_quote! {
3858                                                #members_tee_ident = source_stream(#d) -> tee();
3859                                                #source_ident = #members_tee_ident;
3860                                            }
3861                                        },
3862                                        ClusterMembersState::Uninit => syn::parse_quote! {
3863                                            #source_ident = source_stream(DUMMY);
3864                                        },
3865                                        ClusterMembersState::Tee(..) => parse_quote! {
3866                                            #source_ident = #members_tee_ident;
3867                                        },
3868                                    }
3869                                }
3870
3871                                HydroSource::Embedded(ident) => {
3872                                    parse_quote! {
3873                                        #source_ident = source_stream(#ident);
3874                                    }
3875                                }
3876
3877                                HydroSource::EmbeddedSingleton(ident) => {
3878                                    parse_quote! {
3879                                        #source_ident = source_iter([#ident]);
3880                                    }
3881                                }
3882                            };
3883
3884                            match builders_or_callback {
3885                                BuildersOrCallback::Builders(graph_builders) => {
3886                                    graph_builders.add_dfir_at(
3887                                        &out_location,
3888                                        source_stmt,
3889                                        Some(&stmt_id.to_string()),
3890                                    );
3891                                }
3892                                BuildersOrCallback::Callback(_, node_callback) => {
3893                                    node_callback(node, next_stmt_id);
3894                                }
3895                            }
3896
3897                            ident_stack.push(source_ident);
3898                        }
3899                    }
3900
3901                    HydroNode::SingletonSource { value, first_tick_only, metadata } => {
3902                        let stmt_id = next_stmt_id.get_and_increment();
3903                        let source_ident =
3904                            syn::Ident::new(&format!("stream_{}", stmt_id), Span::call_site());
3905
3906                        match builders_or_callback {
3907                            BuildersOrCallback::Builders(graph_builders) => {
3908                                if *first_tick_only {
3909                                    assert!(
3910                                        !metadata.location_id.is_top_level(),
3911                                        "first_tick_only SingletonSource must be inside a tick"
3912                                    );
3913                                }
3914
3915                                if *first_tick_only
3916                                    || (metadata.location_id.is_top_level()
3917                                        && metadata.collection_kind.is_bounded())
3918                                {
3919                                    graph_builders.add_dfir_at(
3920                                        &out_location,
3921                                        parse_quote! {
3922                                            #source_ident = source_iter([#value]);
3923                                        },
3924                                        Some(&stmt_id.to_string()),
3925                                    );
3926                                } else {
3927                                    graph_builders.add_dfir_at(
3928                                        &out_location,
3929                                        parse_quote! {
3930                                            #source_ident = source_iter([#value]) -> persist::<'static>();
3931                                        },
3932                                        Some(&stmt_id.to_string()),
3933                                    );
3934                                }
3935                            }
3936                            BuildersOrCallback::Callback(_, node_callback) => {
3937                                node_callback(node, next_stmt_id);
3938                            }
3939                        }
3940
3941                        ident_stack.push(source_ident);
3942                    }
3943
3944                    HydroNode::CycleSource { cycle_id, .. } => {
3945                        let ident = cycle_id.as_ident();
3946
3947                        // consume a stmt id even though we did not emit anything so that we can instrument this
3948                        let _ = next_stmt_id.get_and_increment();
3949
3950                        match builders_or_callback {
3951                            BuildersOrCallback::Builders(_) => {}
3952                            BuildersOrCallback::Callback(_, node_callback) => {
3953                                node_callback(node, next_stmt_id);
3954                            }
3955                        }
3956
3957                        ident_stack.push(ident);
3958                    }
3959
3960                    HydroNode::Tee { inner, .. } => {
3961                        // we consume a stmt id regardless of if we emit the tee() operator,
3962                        // so that during rewrites we touch all recipients of the tee()
3963                        let stmt_id = next_stmt_id.get_and_increment();
3964
3965                        let ret_ident = if let Some(built_idents) =
3966                            built_tees.get(&(inner.0.as_ref() as *const RefCell<HydroNode>))
3967                        {
3968                            match builders_or_callback {
3969                                BuildersOrCallback::Builders(_) => {}
3970                                BuildersOrCallback::Callback(_, node_callback) => {
3971                                    node_callback(node, next_stmt_id);
3972                                }
3973                            }
3974
3975                            built_idents[0].clone()
3976                        } else {
3977                            // The inner node was already processed by transform_bottom_up,
3978                            // so its ident is on the stack
3979                            let inner_ident = ident_stack.pop().unwrap();
3980
3981                            let tee_ident =
3982                                syn::Ident::new(&format!("stream_{}", stmt_id), Span::call_site());
3983
3984                            built_tees.insert(
3985                                inner.0.as_ref() as *const RefCell<HydroNode>,
3986                                vec![tee_ident.clone()],
3987                            );
3988
3989                            match builders_or_callback {
3990                                BuildersOrCallback::Builders(graph_builders) => {
3991                                    // NOTE: With `forward_ref`, the fold codegen may not have
3992                                    // run yet when we reach this tee, so `fold_hooked_idents`
3993                                    // might not contain the inner ident. In that case we won't
3994                                    // propagate the "hooked" status to the tee and the
3995                                    // downstream singleton batch will use the normal
3996                                    // `SingletonHook` instead of `PassthroughSingletonHook`.
3997                                    // This is not a soundness issue: the fallback hook still
3998                                    // produces correct behavior, just with a redundant decision
3999                                    // point. TODO(https://github.com/hydro-project/hydro/issues/2856):
4000                                    // fix ordering so forward_ref folds are always processed
4001                                    // before their downstream tees.
4002                                    if fold_hooked_idents.contains(&inner_ident.to_string()) {
4003                                        fold_hooked_idents.insert(tee_ident.to_string());
4004                                    }
4005                                    graph_builders.add_dfir_at(
4006                                        &out_location,
4007                                        parse_quote! {
4008                                            #tee_ident = #inner_ident -> tee();
4009                                        },
4010                                        Some(&stmt_id.to_string()),
4011                                    );
4012                                }
4013                                BuildersOrCallback::Callback(_, node_callback) => {
4014                                    node_callback(node, next_stmt_id);
4015                                }
4016                            }
4017
4018                            tee_ident
4019                        };
4020
4021                        ident_stack.push(ret_ident);
4022                    }
4023
4024                    HydroNode::Reference { inner, kind, .. } => {
4025                        // we consume a stmt id regardless of if we emit the operator,
4026                        // so that during rewrites we touch all recipients
4027                        let stmt_id = next_stmt_id.get_and_increment();
4028
4029                        let ret_ident = if let Some(built_idents) =
4030                            built_tees.get(&(inner.0.as_ref() as *const RefCell<HydroNode>))
4031                        {
4032                            built_idents[0].clone()
4033                        } else {
4034                            let inner_ident = ident_stack.pop().unwrap();
4035
4036                            let ref_ident =
4037                                syn::Ident::new(&format!("stream_{}", stmt_id), Span::call_site());
4038
4039                            built_tees.insert(
4040                                inner.0.as_ref() as *const RefCell<HydroNode>,
4041                                vec![ref_ident.clone()],
4042                            );
4043
4044                            match builders_or_callback {
4045                                BuildersOrCallback::Builders(graph_builders) => {
4046                                    let op_ident = syn::Ident::new(
4047                                        match kind {
4048                                            crate::handoff_ref::HandoffRefKind::Singleton => "singleton",
4049                                            crate::handoff_ref::HandoffRefKind::Optional => "optional",
4050                                            crate::handoff_ref::HandoffRefKind::Vec => "handoff",
4051                                        },
4052                                        Span::call_site(),
4053                                    );
4054                                    graph_builders.add_dfir_at(
4055                                        &out_location,
4056                                        parse_quote! {
4057                                            #ref_ident = #inner_ident -> #op_ident();
4058                                        },
4059                                        Some(&stmt_id.to_string()),
4060                                    );
4061                                }
4062                                BuildersOrCallback::Callback(_, node_callback) => {
4063                                    node_callback(node, next_stmt_id);
4064                                }
4065                            }
4066
4067                            ref_ident
4068                        };
4069
4070                        ident_stack.push(ret_ident);
4071                    }
4072
4073                    HydroNode::Partition {
4074                        inner, f, is_true, metadata,
4075                    } => {
4076                        let is_true = *is_true; // need to copy early to avoid borrow checking issues with node
4077                        let ptr = inner.0.as_ref() as *const RefCell<HydroNode>;
4078                        let stmt_id = next_stmt_id.get_and_increment();
4079
4080                        let ret_ident = if let Some(built_idents) = built_tees.get(&ptr) {
4081                            match builders_or_callback {
4082                                BuildersOrCallback::Builders(_) => {}
4083                                BuildersOrCallback::Callback(_, node_callback) => {
4084                                    node_callback(node, next_stmt_id);
4085                                }
4086                            }
4087
4088                            let idx = if is_true { 0 } else { 1 };
4089                            built_idents[idx].clone()
4090                        } else {
4091                            // The inner node was already processed by transform_bottom_up,
4092                            // so its ident is on the stack
4093                            let inner_ident = ident_stack.pop().unwrap();
4094                            let f_tokens = f.emit_tokens(&mut ident_stack);
4095
4096                            let inner_ident = {
4097                                let inner_borrow = inner.0.borrow();
4098                                maybe_observe_for_mut(
4099                                    f, inner_ident,
4100                                    &inner_borrow.metadata().location_id,
4101                                    &inner_borrow.metadata().collection_kind,
4102                                    &metadata.op,
4103                                    builders_or_callback, next_stmt_id,
4104                                )
4105                            };
4106
4107                            let partition_ident = syn::Ident::new(
4108                                &format!("stream_{}_partition", stmt_id),
4109                                Span::call_site(),
4110                            );
4111                            let true_ident = syn::Ident::new(
4112                                &format!("stream_{}_true", stmt_id),
4113                                Span::call_site(),
4114                            );
4115                            let false_ident = syn::Ident::new(
4116                                &format!("stream_{}_false", stmt_id),
4117                                Span::call_site(),
4118                            );
4119
4120                            built_tees.insert(
4121                                ptr,
4122                                vec![true_ident.clone(), false_ident.clone()],
4123                            );
4124
4125                            let stmt_id = next_stmt_id.get_and_increment();
4126                            match builders_or_callback {
4127                                BuildersOrCallback::Builders(graph_builders) => {
4128                                    graph_builders.add_dfir_at(
4129                                        &out_location,
4130                                        parse_quote! {
4131                                            #partition_ident = #inner_ident -> partition(|__item, __num_outputs| if (#f_tokens)(__item) { 0_usize } else { 1_usize });
4132                                            #true_ident = #partition_ident[0];
4133                                            #false_ident = #partition_ident[1];
4134                                        },
4135                                        Some(&stmt_id.to_string()),
4136                                    );
4137                                }
4138                                BuildersOrCallback::Callback(_, node_callback) => {
4139                                    node_callback(node, next_stmt_id);
4140                                }
4141                            }
4142
4143                            if is_true { true_ident } else { false_ident }
4144                        };
4145
4146                        ident_stack.push(ret_ident);
4147                    }
4148
4149                    HydroNode::Chain { .. } => {
4150                        // Children are processed left-to-right, so second is on top
4151                        let second_ident = ident_stack.pop().unwrap();
4152                        let first_ident = ident_stack.pop().unwrap();
4153
4154                        let stmt_id = next_stmt_id.get_and_increment();
4155                        let chain_ident =
4156                            syn::Ident::new(&format!("stream_{}", stmt_id), Span::call_site());
4157
4158                        match builders_or_callback {
4159                            BuildersOrCallback::Builders(graph_builders) => {
4160                                graph_builders.add_dfir_at(
4161                                    &out_location,
4162                                    parse_quote! {
4163                                        #chain_ident = chain();
4164                                        #first_ident -> [0]#chain_ident;
4165                                        #second_ident -> [1]#chain_ident;
4166                                    },
4167                                    Some(&stmt_id.to_string()),
4168                                );
4169                            }
4170                            BuildersOrCallback::Callback(_, node_callback) => {
4171                                node_callback(node, next_stmt_id);
4172                            }
4173                        }
4174
4175                        ident_stack.push(chain_ident);
4176                    }
4177
4178                    HydroNode::MergeOrdered { first, metadata, .. } => {
4179                        let second_ident = ident_stack.pop().unwrap();
4180                        let first_ident = ident_stack.pop().unwrap();
4181
4182                        let stmt_id = next_stmt_id.get_and_increment();
4183                        let merge_ident =
4184                            syn::Ident::new(&format!("stream_{}", stmt_id), Span::call_site());
4185
4186                        match builders_or_callback {
4187                            BuildersOrCallback::Builders(graph_builders) => {
4188                                graph_builders.merge_ordered(
4189                                    &first.metadata().location_id,
4190                                    first_ident,
4191                                    second_ident,
4192                                    &merge_ident,
4193                                    &first.metadata().collection_kind,
4194                                    &metadata.op,
4195                                    Some(&stmt_id.to_string()),
4196                                );
4197                            }
4198                            BuildersOrCallback::Callback(_, node_callback) => {
4199                                node_callback(node, next_stmt_id);
4200                            }
4201                        }
4202
4203                        ident_stack.push(merge_ident);
4204                    }
4205
4206                    HydroNode::ChainFirst { .. } => {
4207                        let second_ident = ident_stack.pop().unwrap();
4208                        let first_ident = ident_stack.pop().unwrap();
4209
4210                        let stmt_id = next_stmt_id.get_and_increment();
4211                        let chain_ident =
4212                            syn::Ident::new(&format!("stream_{}", stmt_id), Span::call_site());
4213
4214                        match builders_or_callback {
4215                            BuildersOrCallback::Builders(graph_builders) => {
4216                                graph_builders.add_dfir_at(
4217                                    &out_location,
4218                                    parse_quote! {
4219                                        #chain_ident = chain_first_n(1);
4220                                        #first_ident -> [0]#chain_ident;
4221                                        #second_ident -> [1]#chain_ident;
4222                                    },
4223                                    Some(&stmt_id.to_string()),
4224                                );
4225                            }
4226                            BuildersOrCallback::Callback(_, node_callback) => {
4227                                node_callback(node, next_stmt_id);
4228                            }
4229                        }
4230
4231                        ident_stack.push(chain_ident);
4232                    }
4233
4234                    HydroNode::CrossSingleton { right, .. } => {
4235                        let right_ident = ident_stack.pop().unwrap();
4236                        let left_ident = ident_stack.pop().unwrap();
4237
4238                        let stmt_id = next_stmt_id.get_and_increment();
4239                        let cross_ident =
4240                            syn::Ident::new(&format!("stream_{}", stmt_id), Span::call_site());
4241
4242                        match builders_or_callback {
4243                            BuildersOrCallback::Builders(graph_builders) => {
4244                                if right.metadata().location_id.is_top_level()
4245                                    && right.metadata().collection_kind.is_bounded()
4246                                {
4247                                    let lifetime =
4248                                        graph_builders.cross_tick_state_lifetime(&out_location);
4249                                    graph_builders.add_dfir_at(
4250                                        &out_location,
4251                                        parse_quote! {
4252                                            #cross_ident = cross_singleton::<#lifetime>();
4253                                            #left_ident -> [input]#cross_ident;
4254                                            #right_ident -> [single]#cross_ident;
4255                                        },
4256                                        Some(&stmt_id.to_string()),
4257                                    );
4258                                } else {
4259                                    graph_builders.add_dfir_at(
4260                                        &out_location,
4261                                        parse_quote! {
4262                                            #cross_ident = cross_singleton();
4263                                            #left_ident -> [input]#cross_ident;
4264                                            #right_ident -> [single]#cross_ident;
4265                                        },
4266                                        Some(&stmt_id.to_string()),
4267                                    );
4268                                }
4269                            }
4270                            BuildersOrCallback::Callback(_, node_callback) => {
4271                                node_callback(node, next_stmt_id);
4272                            }
4273                        }
4274
4275                        ident_stack.push(cross_ident);
4276                    }
4277
4278                    HydroNode::CrossProduct { .. } | HydroNode::Join { .. } => {
4279                        let operator: syn::Ident = if matches!(node, HydroNode::CrossProduct { .. }) {
4280                            parse_quote!(cross_join_multiset)
4281                        } else {
4282                            parse_quote!(join_multiset)
4283                        };
4284
4285                        let (HydroNode::CrossProduct { left, right, .. }
4286                        | HydroNode::Join { left, right, .. }) = node
4287                        else {
4288                            unreachable!()
4289                        };
4290
4291                        let is_top_level = left.metadata().location_id.is_top_level()
4292                            && right.metadata().location_id.is_top_level();
4293                        let left_top_level = left.metadata().location_id.is_top_level();
4294                        let right_top_level = right.metadata().location_id.is_top_level();
4295
4296                        let right_ident = ident_stack.pop().unwrap();
4297                        let left_ident = ident_stack.pop().unwrap();
4298
4299                        let stmt_id = next_stmt_id.get_and_increment();
4300                        let stream_ident =
4301                            syn::Ident::new(&format!("stream_{}", stmt_id), Span::call_site());
4302
4303                        match builders_or_callback {
4304                            BuildersOrCallback::Builders(graph_builders) => {
4305                                let left_lifetime = if left_top_level {
4306                                    graph_builders.cross_tick_state_lifetime(&out_location)
4307                                } else {
4308                                    graph_builders.tick_state_lifetime(&out_location)
4309                                };
4310
4311                                let right_lifetime = if right_top_level {
4312                                    graph_builders.cross_tick_state_lifetime(&out_location)
4313                                } else {
4314                                    graph_builders.tick_state_lifetime(&out_location)
4315                                };
4316
4317                                graph_builders.add_dfir_at(
4318                                    &out_location,
4319                                    if is_top_level {
4320                                        // if both inputs are root, the output is expected to have streamy semantics, so we need
4321                                        // a multiset_delta() to negate the replay behavior
4322                                        parse_quote! {
4323                                            #stream_ident = #operator::<#left_lifetime, #right_lifetime>() -> multiset_delta();
4324                                            #left_ident -> [0]#stream_ident;
4325                                            #right_ident -> [1]#stream_ident;
4326                                        }
4327                                    } else {
4328                                        parse_quote! {
4329                                            #stream_ident = #operator::<#left_lifetime, #right_lifetime>();
4330                                            #left_ident -> [0]#stream_ident;
4331                                            #right_ident -> [1]#stream_ident;
4332                                        }
4333                                    },
4334                                    Some(&stmt_id.to_string()),
4335                                );
4336                            }
4337                            BuildersOrCallback::Callback(_, node_callback) => {
4338                                node_callback(node, next_stmt_id);
4339                            }
4340                        }
4341
4342                        ident_stack.push(stream_ident);
4343                    }
4344
4345                    HydroNode::Difference { .. } | HydroNode::AntiJoin { .. } => {
4346                        let operator: syn::Ident = if matches!(node, HydroNode::Difference { .. }) {
4347                            parse_quote!(difference)
4348                        } else {
4349                            parse_quote!(anti_join)
4350                        };
4351
4352                        let (HydroNode::Difference { neg, .. } | HydroNode::AntiJoin { neg, .. }) =
4353                            node
4354                        else {
4355                            unreachable!()
4356                        };
4357
4358                        let neg_top_level = neg.metadata().location_id.is_top_level();
4359
4360                        let neg_ident = ident_stack.pop().unwrap();
4361                        let pos_ident = ident_stack.pop().unwrap();
4362
4363                        let stmt_id = next_stmt_id.get_and_increment();
4364                        let stream_ident =
4365                            syn::Ident::new(&format!("stream_{}", stmt_id), Span::call_site());
4366
4367                        match builders_or_callback {
4368                            BuildersOrCallback::Builders(graph_builders) => {
4369                                let neg_lifetime = if neg_top_level {
4370                                    graph_builders.cross_tick_state_lifetime(&out_location)
4371                                } else {
4372                                    graph_builders.tick_state_lifetime(&out_location)
4373                                };
4374                                let pos_lifetime =
4375                                    graph_builders.tick_state_lifetime(&out_location);
4376
4377                                graph_builders.add_dfir_at(
4378                                    &out_location,
4379                                    parse_quote! {
4380                                        #stream_ident = #operator::<#pos_lifetime, #neg_lifetime>();
4381                                        #pos_ident -> [pos]#stream_ident;
4382                                        #neg_ident -> [neg]#stream_ident;
4383                                    },
4384                                    Some(&stmt_id.to_string()),
4385                                );
4386                            }
4387                            BuildersOrCallback::Callback(_, node_callback) => {
4388                                node_callback(node, next_stmt_id);
4389                            }
4390                        }
4391
4392                        ident_stack.push(stream_ident);
4393                    }
4394
4395                    HydroNode::JoinHalf { .. } => {
4396                        let HydroNode::JoinHalf { right, .. } = node else {
4397                            unreachable!()
4398                        };
4399
4400                        assert!(
4401                            right.metadata().collection_kind.is_bounded(),
4402                            "JoinHalf requires the right (build) side to be Bounded, got {:?}",
4403                            right.metadata().collection_kind
4404                        );
4405
4406                        let build_top_level = right.metadata().location_id.is_top_level();
4407
4408                        let build_ident = ident_stack.pop().unwrap();
4409                        let probe_ident = ident_stack.pop().unwrap();
4410
4411                        let stmt_id = next_stmt_id.get_and_increment();
4412                        let stream_ident =
4413                            syn::Ident::new(&format!("stream_{}", stmt_id), Span::call_site());
4414
4415                        match builders_or_callback {
4416                            BuildersOrCallback::Builders(graph_builders) => {
4417                                let build_lifetime = if build_top_level {
4418                                    graph_builders.cross_tick_state_lifetime(&out_location)
4419                                } else {
4420                                    graph_builders.tick_state_lifetime(&out_location)
4421                                };
4422                                let probe_lifetime =
4423                                    graph_builders.tick_state_lifetime(&out_location);
4424
4425                                graph_builders.add_dfir_at(
4426                                    &out_location,
4427                                    parse_quote! {
4428                                        #stream_ident = join_multiset_half::<#build_lifetime, #probe_lifetime>();
4429                                        #probe_ident -> [probe]#stream_ident;
4430                                        #build_ident -> [build]#stream_ident;
4431                                    },
4432                                    Some(&stmt_id.to_string()),
4433                                );
4434                            }
4435                            BuildersOrCallback::Callback(_, node_callback) => {
4436                                node_callback(node, next_stmt_id);
4437                            }
4438                        }
4439
4440                        ident_stack.push(stream_ident);
4441                    }
4442
4443                    HydroNode::ResolveFutures { .. } => {
4444                        let input_ident = ident_stack.pop().unwrap();
4445
4446                        let stmt_id = next_stmt_id.get_and_increment();
4447                        let futures_ident =
4448                            syn::Ident::new(&format!("stream_{}", stmt_id), Span::call_site());
4449
4450                        match builders_or_callback {
4451                            BuildersOrCallback::Builders(graph_builders) => {
4452                                graph_builders.add_dfir_at(
4453                                    &out_location,
4454                                    parse_quote! {
4455                                        #futures_ident = #input_ident -> resolve_futures();
4456                                    },
4457                                    Some(&stmt_id.to_string()),
4458                                );
4459                            }
4460                            BuildersOrCallback::Callback(_, node_callback) => {
4461                                node_callback(node, next_stmt_id);
4462                            }
4463                        }
4464
4465                        ident_stack.push(futures_ident);
4466                    }
4467
4468                    HydroNode::ResolveFuturesBlocking { .. } => {
4469                        let input_ident = ident_stack.pop().unwrap();
4470
4471                        let stmt_id = next_stmt_id.get_and_increment();
4472                        let futures_ident =
4473                            syn::Ident::new(&format!("stream_{}", stmt_id), Span::call_site());
4474
4475                        match builders_or_callback {
4476                            BuildersOrCallback::Builders(graph_builders) => {
4477                                graph_builders.add_dfir_at(
4478                                    &out_location,
4479                                    parse_quote! {
4480                                        #futures_ident = #input_ident -> resolve_futures_blocking();
4481                                    },
4482                                    Some(&stmt_id.to_string()),
4483                                );
4484                            }
4485                            BuildersOrCallback::Callback(_, node_callback) => {
4486                                node_callback(node, next_stmt_id);
4487                            }
4488                        }
4489
4490                        ident_stack.push(futures_ident);
4491                    }
4492
4493                    HydroNode::ResolveFuturesOrdered { .. } => {
4494                        let input_ident = ident_stack.pop().unwrap();
4495
4496                        let stmt_id = next_stmt_id.get_and_increment();
4497                        let futures_ident =
4498                            syn::Ident::new(&format!("stream_{}", stmt_id), Span::call_site());
4499
4500                        match builders_or_callback {
4501                            BuildersOrCallback::Builders(graph_builders) => {
4502                                graph_builders.add_dfir_at(
4503                                    &out_location,
4504                                    parse_quote! {
4505                                        #futures_ident = #input_ident -> resolve_futures_ordered();
4506                                    },
4507                                    Some(&stmt_id.to_string()),
4508                                );
4509                            }
4510                            BuildersOrCallback::Callback(_, node_callback) => {
4511                                node_callback(node, next_stmt_id);
4512                            }
4513                        }
4514
4515                        ident_stack.push(futures_ident);
4516                    }
4517
4518                    HydroNode::Map {
4519                        f,
4520                        input,
4521                        metadata,
4522                    } => {
4523                        // Pop input ident (pushed last by transform_children).
4524                        let input_ident = ident_stack.pop().unwrap();
4525                        let f_tokens = f.emit_tokens(&mut ident_stack);
4526
4527                        let input_ident = maybe_observe_for_mut(
4528                            f,
4529                            input_ident,
4530                            &input.metadata().location_id,
4531                            &input.metadata().collection_kind,
4532                            &metadata.op,
4533                            builders_or_callback,
4534                            next_stmt_id,
4535                        );
4536
4537                        let stmt_id = next_stmt_id.get_and_increment();
4538                        let map_ident =
4539                            syn::Ident::new(&format!("stream_{}", stmt_id), Span::call_site());
4540
4541                        match builders_or_callback {
4542                            BuildersOrCallback::Builders(graph_builders) => {
4543                                graph_builders.add_dfir_at(
4544                                    &out_location,
4545                                    parse_quote! {
4546                                        #map_ident = #input_ident -> map(#f_tokens);
4547                                    },
4548                                    Some(&stmt_id.to_string()),
4549                                );
4550                            }
4551                            BuildersOrCallback::Callback(_, node_callback) => {
4552                                node_callback(node, next_stmt_id);
4553                            }
4554                        }
4555
4556                        ident_stack.push(map_ident);
4557                    }
4558
4559                    HydroNode::FlatMap { f, input, metadata } => {
4560                        let input_ident = ident_stack.pop().unwrap();
4561                        let f_tokens = f.emit_tokens(&mut ident_stack);
4562
4563                        let input_ident = maybe_observe_for_mut(
4564                            f, input_ident,
4565                            &input.metadata().location_id,
4566                            &input.metadata().collection_kind,
4567                            &metadata.op,
4568                            builders_or_callback, next_stmt_id,
4569                        );
4570
4571                        let stmt_id = next_stmt_id.get_and_increment();
4572                        let flat_map_ident =
4573                            syn::Ident::new(&format!("stream_{}", stmt_id), Span::call_site());
4574
4575                        match builders_or_callback {
4576                            BuildersOrCallback::Builders(graph_builders) => {
4577                                graph_builders.add_dfir_at(
4578                                    &out_location,
4579                                    parse_quote! {
4580                                        #flat_map_ident = #input_ident -> flat_map(#f_tokens);
4581                                    },
4582                                    Some(&stmt_id.to_string()),
4583                                );
4584                            }
4585                            BuildersOrCallback::Callback(_, node_callback) => {
4586                                node_callback(node, next_stmt_id);
4587                            }
4588                        }
4589
4590                        ident_stack.push(flat_map_ident);
4591                    }
4592
4593                    HydroNode::FlatMapStreamBlocking { f, input, metadata } => {
4594                        let input_ident = ident_stack.pop().unwrap();
4595                        let f_tokens = f.emit_tokens(&mut ident_stack);
4596
4597                        let input_ident = maybe_observe_for_mut(
4598                            f, input_ident,
4599                            &input.metadata().location_id,
4600                            &input.metadata().collection_kind,
4601                            &metadata.op,
4602                            builders_or_callback, next_stmt_id,
4603                        );
4604
4605                        let stmt_id = next_stmt_id.get_and_increment();
4606                        let flat_map_stream_blocking_ident =
4607                            syn::Ident::new(&format!("stream_{}", stmt_id), Span::call_site());
4608
4609                        match builders_or_callback {
4610                            BuildersOrCallback::Builders(graph_builders) => {
4611                                graph_builders.add_dfir_at(
4612                                    &out_location,
4613                                    parse_quote! {
4614                                        #flat_map_stream_blocking_ident = #input_ident -> flat_map_stream_blocking(#f_tokens);
4615                                    },
4616                                    Some(&stmt_id.to_string()),
4617                                );
4618                            }
4619                            BuildersOrCallback::Callback(_, node_callback) => {
4620                                node_callback(node, next_stmt_id);
4621                            }
4622                        }
4623
4624                        ident_stack.push(flat_map_stream_blocking_ident);
4625                    }
4626
4627                    HydroNode::Filter { f, input, metadata } => {
4628                        let input_ident = ident_stack.pop().unwrap();
4629                        let f_tokens = f.emit_tokens(&mut ident_stack);
4630
4631                        let input_ident = maybe_observe_for_mut(
4632                            f, input_ident,
4633                            &input.metadata().location_id,
4634                            &input.metadata().collection_kind,
4635                            &metadata.op,
4636                            builders_or_callback, next_stmt_id,
4637                        );
4638
4639                        let stmt_id = next_stmt_id.get_and_increment();
4640                        let filter_ident =
4641                            syn::Ident::new(&format!("stream_{}", stmt_id), Span::call_site());
4642
4643                        match builders_or_callback {
4644                            BuildersOrCallback::Builders(graph_builders) => {
4645                                graph_builders.add_dfir_at(
4646                                    &out_location,
4647                                    parse_quote! {
4648                                        #filter_ident = #input_ident -> filter(#f_tokens);
4649                                    },
4650                                    Some(&stmt_id.to_string()),
4651                                );
4652                            }
4653                            BuildersOrCallback::Callback(_, node_callback) => {
4654                                node_callback(node, next_stmt_id);
4655                            }
4656                        }
4657
4658                        ident_stack.push(filter_ident);
4659                    }
4660
4661                    HydroNode::FilterMap { f, input, metadata } => {
4662                        let input_ident = ident_stack.pop().unwrap();
4663                        let f_tokens = f.emit_tokens(&mut ident_stack);
4664
4665                        let input_ident = maybe_observe_for_mut(
4666                            f, input_ident,
4667                            &input.metadata().location_id,
4668                            &input.metadata().collection_kind,
4669                            &metadata.op,
4670                            builders_or_callback, next_stmt_id,
4671                        );
4672
4673                        let stmt_id = next_stmt_id.get_and_increment();
4674                        let filter_map_ident =
4675                            syn::Ident::new(&format!("stream_{}", stmt_id), Span::call_site());
4676
4677                        match builders_or_callback {
4678                            BuildersOrCallback::Builders(graph_builders) => {
4679                                graph_builders.add_dfir_at(
4680                                    &out_location,
4681                                    parse_quote! {
4682                                        #filter_map_ident = #input_ident -> filter_map(#f_tokens);
4683                                    },
4684                                    Some(&stmt_id.to_string()),
4685                                );
4686                            }
4687                            BuildersOrCallback::Callback(_, node_callback) => {
4688                                node_callback(node, next_stmt_id);
4689                            }
4690                        }
4691
4692                        ident_stack.push(filter_map_ident);
4693                    }
4694
4695                    HydroNode::Sort { .. } => {
4696                        let input_ident = ident_stack.pop().unwrap();
4697
4698                        let stmt_id = next_stmt_id.get_and_increment();
4699                        let sort_ident =
4700                            syn::Ident::new(&format!("stream_{}", stmt_id), Span::call_site());
4701
4702                        match builders_or_callback {
4703                            BuildersOrCallback::Builders(graph_builders) => {
4704                                graph_builders.add_dfir_at(
4705                                    &out_location,
4706                                    parse_quote! {
4707                                        #sort_ident = #input_ident -> sort();
4708                                    },
4709                                    Some(&stmt_id.to_string()),
4710                                );
4711                            }
4712                            BuildersOrCallback::Callback(_, node_callback) => {
4713                                node_callback(node, next_stmt_id);
4714                            }
4715                        }
4716
4717                        ident_stack.push(sort_ident);
4718                    }
4719
4720                    HydroNode::DeferTick { .. } => {
4721                        let input_ident = ident_stack.pop().unwrap();
4722
4723                        let stmt_id = next_stmt_id.get_and_increment();
4724                        let defer_tick_ident =
4725                            syn::Ident::new(&format!("stream_{}", stmt_id), Span::call_site());
4726
4727                        match builders_or_callback {
4728                            BuildersOrCallback::Builders(graph_builders) => {
4729                                graph_builders.add_dfir_at(
4730                                    &out_location,
4731                                    parse_quote! {
4732                                        #defer_tick_ident = #input_ident -> defer_tick_lazy();
4733                                    },
4734                                    Some(&stmt_id.to_string()),
4735                                );
4736                            }
4737                            BuildersOrCallback::Callback(_, node_callback) => {
4738                                node_callback(node, next_stmt_id);
4739                            }
4740                        }
4741
4742                        ident_stack.push(defer_tick_ident);
4743                    }
4744
4745                    HydroNode::Enumerate { input, .. } => {
4746                        let input_ident = ident_stack.pop().unwrap();
4747
4748                        let stmt_id = next_stmt_id.get_and_increment();
4749                        let enumerate_ident =
4750                            syn::Ident::new(&format!("stream_{}", stmt_id), Span::call_site());
4751
4752                        match builders_or_callback {
4753                            BuildersOrCallback::Builders(graph_builders) => {
4754                                let lifetime = if input.metadata().location_id.is_top_level() {
4755                                    graph_builders.cross_tick_state_lifetime(&out_location)
4756                                } else {
4757                                    graph_builders.tick_state_lifetime(&out_location)
4758                                };
4759                                graph_builders.add_dfir_at(
4760                                    &out_location,
4761                                    parse_quote! {
4762                                        #enumerate_ident = #input_ident -> enumerate::<#lifetime>();
4763                                    },
4764                                    Some(&stmt_id.to_string()),
4765                                );
4766                            }
4767                            BuildersOrCallback::Callback(_, node_callback) => {
4768                                node_callback(node, next_stmt_id);
4769                            }
4770                        }
4771
4772                        ident_stack.push(enumerate_ident);
4773                    }
4774
4775                    HydroNode::Inspect { f, input, metadata } => {
4776                        let input_ident = ident_stack.pop().unwrap();
4777                        let f_tokens = f.emit_tokens(&mut ident_stack);
4778
4779                        let input_ident = maybe_observe_for_mut(
4780                            f, input_ident,
4781                            &input.metadata().location_id,
4782                            &input.metadata().collection_kind,
4783                            &metadata.op,
4784                            builders_or_callback, next_stmt_id,
4785                        );
4786
4787                        let stmt_id = next_stmt_id.get_and_increment();
4788                        let inspect_ident =
4789                            syn::Ident::new(&format!("stream_{}", stmt_id), Span::call_site());
4790
4791                        match builders_or_callback {
4792                            BuildersOrCallback::Builders(graph_builders) => {
4793                                graph_builders.add_dfir_at(
4794                                    &out_location,
4795                                    parse_quote! {
4796                                        #inspect_ident = #input_ident -> inspect(#f_tokens);
4797                                    },
4798                                    Some(&stmt_id.to_string()),
4799                                );
4800                            }
4801                            BuildersOrCallback::Callback(_, node_callback) => {
4802                                node_callback(node, next_stmt_id);
4803                            }
4804                        }
4805
4806                        ident_stack.push(inspect_ident);
4807                    }
4808
4809                    HydroNode::Unique { input, .. } => {
4810                        let input_ident = ident_stack.pop().unwrap();
4811
4812                        let stmt_id = next_stmt_id.get_and_increment();
4813                        let unique_ident =
4814                            syn::Ident::new(&format!("stream_{}", stmt_id), Span::call_site());
4815
4816                        match builders_or_callback {
4817                            BuildersOrCallback::Builders(graph_builders) => {
4818                                let lifetime = if input.metadata().location_id.is_top_level() {
4819                                    graph_builders.cross_tick_state_lifetime(&out_location)
4820                                } else {
4821                                    graph_builders.tick_state_lifetime(&out_location)
4822                                };
4823
4824                                graph_builders.add_dfir_at(
4825                                    &out_location,
4826                                    parse_quote! {
4827                                        #unique_ident = #input_ident -> unique::<#lifetime>();
4828                                    },
4829                                    Some(&stmt_id.to_string()),
4830                                );
4831                            }
4832                            BuildersOrCallback::Callback(_, node_callback) => {
4833                                node_callback(node, next_stmt_id);
4834                            }
4835                        }
4836
4837                        ident_stack.push(unique_ident);
4838                    }
4839
4840                    HydroNode::Fold { .. } | HydroNode::FoldKeyed { .. } | HydroNode::Scan { .. } | HydroNode::ScanAsyncBlocking { .. } => {
4841                        let operator: syn::Ident = if let HydroNode::Fold { input, .. } = node {
4842                            if input.metadata().location_id.is_top_level()
4843                                && input.metadata().collection_kind.is_bounded()
4844                            {
4845                                parse_quote!(fold_no_replay)
4846                            } else {
4847                                parse_quote!(fold)
4848                            }
4849                        } else if matches!(node, HydroNode::Scan { .. }) {
4850                            parse_quote!(scan)
4851                        } else if matches!(node, HydroNode::ScanAsyncBlocking { .. }) {
4852                            parse_quote!(scan_async_blocking)
4853                        } else if let HydroNode::FoldKeyed { input, .. } = node {
4854                            if input.metadata().location_id.is_top_level()
4855                                && input.metadata().collection_kind.is_bounded()
4856                            {
4857                                todo!("Fold keyed on a top-level bounded collection is not yet supported")
4858                            } else {
4859                                parse_quote!(fold_keyed)
4860                            }
4861                        } else {
4862                            unreachable!()
4863                        };
4864
4865                        let (HydroNode::Fold { input, .. }
4866                        | HydroNode::FoldKeyed { input, .. }
4867                        | HydroNode::Scan { input, .. }
4868                        | HydroNode::ScanAsyncBlocking { input, .. }) = node
4869                        else {
4870                            unreachable!()
4871                        };
4872
4873                        let input_top_level = input.metadata().location_id.is_top_level();
4874
4875                        let input_ident = ident_stack.pop().unwrap();
4876
4877                        let (HydroNode::Fold { init, acc, .. }
4878                        | HydroNode::FoldKeyed { init, acc, .. }
4879                        | HydroNode::Scan { init, acc, .. }
4880                        | HydroNode::ScanAsyncBlocking { init, acc, .. }) = &*node
4881                        else {
4882                            unreachable!()
4883                        };
4884
4885                        let acc_tokens = acc.emit_tokens(&mut ident_stack);
4886                        let init_tokens = init.emit_tokens(&mut ident_stack);
4887
4888                        let stmt_id = next_stmt_id.get_and_increment();
4889                        let fold_ident =
4890                            syn::Ident::new(&format!("stream_{}", stmt_id), Span::call_site());
4891
4892                        match builders_or_callback {
4893                            BuildersOrCallback::Builders(graph_builders) => {
4894                                let lifetime = if input_top_level {
4895                                    graph_builders.cross_tick_state_lifetime(&out_location)
4896                                } else {
4897                                    graph_builders.tick_state_lifetime(&out_location)
4898                                };
4899
4900                                if matches!(node, HydroNode::Fold { .. })
4901                                    && node.metadata().location_id.is_top_level()
4902                                    && !(matches!(node.metadata().location_id, LocationId::Atomic(_)))
4903                                    && graph_builders.singleton_intermediates()
4904                                    && !node.metadata().collection_kind.is_bounded()
4905                                {
4906                                    let HydroNode::Fold { input, .. } = &*node else { unreachable!() };
4907                                    let hooked_input_ident = graph_builders.emit_fold_hook(
4908                                        &input.metadata().location_id,
4909                                        &input_ident,
4910                                        &input.metadata().collection_kind,
4911                                        &node.metadata().op,
4912                                    );
4913
4914                                    let (effective_input, wrapped_acc) = if let Some(ref hooked) = hooked_input_ident {
4915                                        let acc: syn::Expr = parse_quote!({
4916                                            let mut __inner = #acc_tokens;
4917                                            move |__state, __batch: Vec<_>| {
4918                                                if __batch.is_empty() {
4919                                                    return None;
4920                                                }
4921                                                for __value in __batch {
4922                                                    __inner(__state, __value);
4923                                                }
4924                                                Some(__state.clone())
4925                                            }
4926                                        });
4927                                        (hooked, acc)
4928                                    } else {
4929                                        let acc: syn::Expr = parse_quote!({
4930                                            let mut __inner = #acc_tokens;
4931                                            move |__state, __value| {
4932                                                __inner(__state, __value);
4933                                                Some(__state.clone())
4934                                            }
4935                                        });
4936                                        (&input_ident, acc)
4937                                    };
4938
4939                                    graph_builders.add_dfir_at(
4940                                        &out_location,
4941                                        parse_quote! {
4942                                            source_iter([(#init_tokens)()]) -> [0]#fold_ident;
4943                                            #effective_input -> scan::<#lifetime>(#init_tokens, #wrapped_acc) -> [1]#fold_ident;
4944                                            #fold_ident = chain();
4945                                        },
4946                                        Some(&stmt_id.to_string()),
4947                                    );
4948
4949                                    if hooked_input_ident.is_some() {
4950                                        fold_hooked_idents.insert(fold_ident.to_string());
4951                                    }
4952                                } else if matches!(node, HydroNode::FoldKeyed { .. })
4953                                    && node.metadata().location_id.is_top_level()
4954                                    && !(matches!(node.metadata().location_id, LocationId::Atomic(_)))
4955                                    && graph_builders.singleton_intermediates()
4956                                    && !node.metadata().collection_kind.is_bounded()
4957                                {
4958                                    let HydroNode::FoldKeyed { input, .. } = &*node else { unreachable!() };
4959                                    let hooked_input_ident = graph_builders.emit_fold_hook(
4960                                        &input.metadata().location_id,
4961                                        &input_ident,
4962                                        &input.metadata().collection_kind,
4963                                        &node.metadata().op,
4964                                    );
4965
4966                                    let wrapped_acc: syn::Expr = parse_quote!({
4967                                        let mut __init = #init_tokens;
4968                                        let mut __inner = #acc_tokens;
4969                                        move |__state, __kv: (_, _)| {
4970                                            // TODO(shadaj): we can avoid the clone when the entry exists
4971                                            let __state = __state
4972                                                .entry(::std::clone::Clone::clone(&__kv.0))
4973                                                .or_insert_with(|| (__init)());
4974                                            __inner(__state, __kv.1);
4975                                            Some((__kv.0, ::std::clone::Clone::clone(&*__state)))
4976                                        }
4977                                    });
4978
4979                                    if let Some(hooked_input_ident) = hooked_input_ident {
4980                                        graph_builders.add_dfir_at(
4981                                            &out_location,
4982                                            parse_quote! {
4983                                                #fold_ident = #hooked_input_ident -> flatten() -> scan::<#lifetime>(|| ::std::collections::HashMap::new(), #wrapped_acc);
4984                                            },
4985                                            Some(&stmt_id.to_string()),
4986                                        );
4987
4988                                        fold_hooked_idents.insert(fold_ident.to_string());
4989                                    } else {
4990                                        graph_builders.add_dfir_at(
4991                                            &out_location,
4992                                            parse_quote! {
4993                                                #fold_ident = #input_ident -> scan::<#lifetime>(|| ::std::collections::HashMap::new(), #wrapped_acc);
4994                                            },
4995                                            Some(&stmt_id.to_string()),
4996                                        );
4997                                    }
4998                                } else if (matches!(node, HydroNode::Fold { .. })
4999                                    || matches!(node, HydroNode::FoldKeyed { .. }))
5000                                    && !node.metadata().location_id.is_top_level()
5001                                    && graph_builders.singleton_intermediates()
5002                                {
5003                                    let input_ref = match &*node {
5004                                        HydroNode::Fold { input, .. } => input,
5005                                        HydroNode::FoldKeyed { input, .. } => input,
5006                                        _ => unreachable!(),
5007                                    };
5008                                    let hooked_input_ident = graph_builders.emit_fold_hook(
5009                                        &input_ref.metadata().location_id,
5010                                        &input_ident,
5011                                        &input_ref.metadata().collection_kind,
5012                                        &node.metadata().op,
5013                                    );
5014
5015                                    let actual_input = hooked_input_ident.as_ref().unwrap_or(&input_ident);
5016                                    graph_builders.add_dfir_at(
5017                                        &out_location,
5018                                        parse_quote! {
5019                                            #fold_ident = #actual_input -> #operator::<#lifetime>(#init_tokens, #acc_tokens);
5020                                        },
5021                                        Some(&stmt_id.to_string()),
5022                                    );
5023                                } else {
5024                                    graph_builders.add_dfir_at(
5025                                        &out_location,
5026                                        parse_quote! {
5027                                            #fold_ident = #input_ident -> #operator::<#lifetime>(#init_tokens, #acc_tokens);
5028                                        },
5029                                        Some(&stmt_id.to_string()),
5030                                    );
5031                                }
5032                            }
5033                            BuildersOrCallback::Callback(_, node_callback) => {
5034                                node_callback(node, next_stmt_id);
5035                            }
5036                        }
5037
5038                        ident_stack.push(fold_ident);
5039                    }
5040
5041                    HydroNode::Reduce { .. } | HydroNode::ReduceKeyed { .. } => {
5042                        let operator: syn::Ident = if let HydroNode::Reduce { input, .. } = node {
5043                            if input.metadata().location_id.is_top_level()
5044                                && input.metadata().collection_kind.is_bounded()
5045                            {
5046                                parse_quote!(reduce_no_replay)
5047                            } else {
5048                                parse_quote!(reduce)
5049                            }
5050                        } else if let HydroNode::ReduceKeyed { input, .. } = node {
5051                            if input.metadata().location_id.is_top_level()
5052                                && input.metadata().collection_kind.is_bounded()
5053                            {
5054                                todo!(
5055                                    "Calling keyed reduce on a top-level bounded collection is not supported"
5056                                )
5057                            } else {
5058                                parse_quote!(reduce_keyed)
5059                            }
5060                        } else {
5061                            unreachable!()
5062                        };
5063
5064                        let (HydroNode::Reduce { input, .. } | HydroNode::ReduceKeyed { input, .. }) = node
5065                        else {
5066                            unreachable!()
5067                        };
5068
5069                        let input_top_level = input.metadata().location_id.is_top_level();
5070
5071                        let input_ident = ident_stack.pop().unwrap();
5072
5073                        let (HydroNode::Reduce { f, .. } | HydroNode::ReduceKeyed { f, .. }) = &*node
5074                        else {
5075                            unreachable!()
5076                        };
5077
5078                        let f_tokens = f.emit_tokens(&mut ident_stack);
5079
5080                        let stmt_id = next_stmt_id.get_and_increment();
5081                        let reduce_ident =
5082                            syn::Ident::new(&format!("stream_{}", stmt_id), Span::call_site());
5083
5084                        match builders_or_callback {
5085                            BuildersOrCallback::Builders(graph_builders) => {
5086                                let lifetime = if input_top_level {
5087                                    graph_builders.cross_tick_state_lifetime(&out_location)
5088                                } else {
5089                                    graph_builders.tick_state_lifetime(&out_location)
5090                                };
5091
5092                                if matches!(node, HydroNode::Reduce { .. })
5093                                    && node.metadata().location_id.is_top_level()
5094                                    && !(matches!(node.metadata().location_id, LocationId::Atomic(_)))
5095                                    && graph_builders.singleton_intermediates()
5096                                    && !node.metadata().collection_kind.is_bounded()
5097                                {
5098                                    todo!(
5099                                        "Reduce with optional intermediates is not yet supported in simulator"
5100                                    );
5101                                } else if matches!(node, HydroNode::ReduceKeyed { .. })
5102                                    && node.metadata().location_id.is_top_level()
5103                                    && !(matches!(node.metadata().location_id, LocationId::Atomic(_)))
5104                                    && graph_builders.singleton_intermediates()
5105                                    && !node.metadata().collection_kind.is_bounded()
5106                                {
5107                                    todo!(
5108                                        "Reduce keyed with optional intermediates is not yet supported in simulator"
5109                                    );
5110                                } else {
5111                                    graph_builders.add_dfir_at(
5112                                        &out_location,
5113                                        parse_quote! {
5114                                            #reduce_ident = #input_ident -> #operator::<#lifetime>(#f_tokens);
5115                                        },
5116                                        Some(&stmt_id.to_string()),
5117                                    );
5118                                }
5119                            }
5120                            BuildersOrCallback::Callback(_, node_callback) => {
5121                                node_callback(node, next_stmt_id);
5122                            }
5123                        }
5124
5125                        ident_stack.push(reduce_ident);
5126                    }
5127
5128                    HydroNode::ReduceKeyedWatermark {
5129                        f,
5130                        input,
5131                        metadata,
5132                        ..
5133                    } => {
5134                        let input_top_level = input.metadata().location_id.is_top_level();
5135
5136                        // watermark is processed second, so it's on top
5137                        let watermark_ident = ident_stack.pop().unwrap();
5138                        let input_ident = ident_stack.pop().unwrap();
5139                        let f_tokens = f.emit_tokens(&mut ident_stack);
5140
5141                        let stmt_id = next_stmt_id.get_and_increment();
5142                        let chain_ident = syn::Ident::new(
5143                            &format!("reduce_keyed_watermark_chain_{}", stmt_id),
5144                            Span::call_site(),
5145                        );
5146
5147                        let fold_ident =
5148                            syn::Ident::new(&format!("stream_{}", stmt_id), Span::call_site());
5149
5150                        let agg_operator: syn::Ident = if input.metadata().location_id.is_top_level()
5151                            && input.metadata().collection_kind.is_bounded()
5152                        {
5153                            parse_quote!(fold_no_replay)
5154                        } else {
5155                            parse_quote!(fold)
5156                        };
5157
5158                        match builders_or_callback {
5159                            BuildersOrCallback::Builders(graph_builders) => {
5160                                let lifetime = if input_top_level {
5161                                    graph_builders.cross_tick_state_lifetime(&out_location)
5162                                } else {
5163                                    graph_builders.tick_state_lifetime(&out_location)
5164                                };
5165
5166                                if metadata.location_id.is_top_level()
5167                                    && !(matches!(metadata.location_id, LocationId::Atomic(_)))
5168                                    && graph_builders.singleton_intermediates()
5169                                    && !metadata.collection_kind.is_bounded()
5170                                {
5171                                    todo!(
5172                                        "Reduce keyed watermarked on a top-level bounded collection is not yet supported"
5173                                    )
5174                                } else {
5175                                    graph_builders.add_dfir_at(
5176                                        &out_location,
5177                                        parse_quote! {
5178                                            #chain_ident = chain();
5179                                            #input_ident
5180                                                -> map(|x| (Some(x), None))
5181                                                -> [0]#chain_ident;
5182                                            #watermark_ident
5183                                                -> map(|watermark| (None, Some(watermark)))
5184                                                -> [1]#chain_ident;
5185
5186                                            #fold_ident = #chain_ident
5187                                                -> #agg_operator::<#lifetime>(|| (::std::collections::HashMap::new(), None), {
5188                                                    let __reduce_keyed_fn = #f_tokens;
5189                                                    move |(map, opt_curr_watermark), (opt_payload, opt_watermark)| {
5190                                                        if let Some((k, v)) = opt_payload {
5191                                                            if let Some(curr_watermark) = *opt_curr_watermark {
5192                                                                if k < curr_watermark {
5193                                                                    return;
5194                                                                }
5195                                                            }
5196                                                            match map.entry(k) {
5197                                                                ::std::collections::hash_map::Entry::Vacant(e) => {
5198                                                                    e.insert(v);
5199                                                                }
5200                                                                ::std::collections::hash_map::Entry::Occupied(mut e) => {
5201                                                                    __reduce_keyed_fn(e.get_mut(), v);
5202                                                                }
5203                                                            }
5204                                                        } else {
5205                                                            let watermark = opt_watermark.unwrap();
5206                                                            if let Some(curr_watermark) = *opt_curr_watermark {
5207                                                                if watermark <= curr_watermark {
5208                                                                    return;
5209                                                                }
5210                                                            }
5211                                                            map.retain(|k, _| *k >= watermark);
5212                                                            *opt_curr_watermark = Some(watermark);
5213                                                        }
5214                                                    }
5215                                                })
5216                                                -> flat_map(|(map, _curr_watermark)| map);
5217                                        },
5218                                        Some(&stmt_id.to_string()),
5219                                    );
5220                                }
5221                            }
5222                            BuildersOrCallback::Callback(_, node_callback) => {
5223                                node_callback(node, next_stmt_id);
5224                            }
5225                        }
5226
5227                        ident_stack.push(fold_ident);
5228                    }
5229
5230                    HydroNode::Network {
5231                        networking_info,
5232                        serialize,
5233                        deserialize,
5234                        instantiate_fn,
5235                        input,
5236                        ..
5237                    } => {
5238                        let input_ident = ident_stack.pop().unwrap();
5239
5240                        let stmt_id = next_stmt_id.get_and_increment();
5241                        let receiver_stream_ident =
5242                            syn::Ident::new(&format!("stream_{}", stmt_id), Span::call_site());
5243
5244                        // For embedded (external) serialization, this synthesizes only the
5245                        // member-id tag conversions (if any) and passes the raw payload through.
5246                        let serialize_pipeline = serialize.pipeline();
5247                        let deserialize_pipeline = deserialize.pipeline();
5248
5249                        match builders_or_callback {
5250                            BuildersOrCallback::Builders(graph_builders) => {
5251                                let (sink_expr, source_expr) = match instantiate_fn {
5252                                    DebugInstantiate::Building => (
5253                                        syn::parse_quote!(DUMMY_SINK),
5254                                        syn::parse_quote!(DUMMY_SOURCE),
5255                                    ),
5256
5257                                    DebugInstantiate::Finalized(finalized) => {
5258                                        (finalized.sink.clone(), finalized.source.clone())
5259                                    }
5260                                };
5261
5262                                graph_builders.create_network(
5263                                    &input.metadata().location_id,
5264                                    &out_location,
5265                                    input_ident,
5266                                    &receiver_stream_ident,
5267                                    serialize_pipeline.as_ref(),
5268                                    sink_expr,
5269                                    source_expr,
5270                                    deserialize_pipeline.as_ref(),
5271                                    serialize.external_element_type(),
5272                                    stmt_id,
5273                                    networking_info,
5274                                );
5275                            }
5276                            BuildersOrCallback::Callback(_, node_callback) => {
5277                                node_callback(node, next_stmt_id);
5278                            }
5279                        }
5280
5281                        ident_stack.push(receiver_stream_ident);
5282                    }
5283
5284                    HydroNode::ExternalInput {
5285                        instantiate_fn,
5286                        deserialize_fn: deserialize_pipeline,
5287                        ..
5288                    } => {
5289                        let stmt_id = next_stmt_id.get_and_increment();
5290                        let receiver_stream_ident =
5291                            syn::Ident::new(&format!("stream_{}", stmt_id), Span::call_site());
5292
5293                        match builders_or_callback {
5294                            BuildersOrCallback::Builders(graph_builders) => {
5295                                let (_, source_expr) = match instantiate_fn {
5296                                    DebugInstantiate::Building => (
5297                                        syn::parse_quote!(DUMMY_SINK),
5298                                        syn::parse_quote!(DUMMY_SOURCE),
5299                                    ),
5300
5301                                    DebugInstantiate::Finalized(finalized) => {
5302                                        (finalized.sink.clone(), finalized.source.clone())
5303                                    }
5304                                };
5305
5306                                graph_builders.create_external_source(
5307                                    &out_location,
5308                                    source_expr,
5309                                    &receiver_stream_ident,
5310                                    deserialize_pipeline.as_ref(),
5311                                    stmt_id,
5312                                );
5313                            }
5314                            BuildersOrCallback::Callback(_, node_callback) => {
5315                                node_callback(node, next_stmt_id);
5316                            }
5317                        }
5318
5319                        ident_stack.push(receiver_stream_ident);
5320                    }
5321
5322                    HydroNode::Counter {
5323                        tag,
5324                        duration,
5325                        prefix,
5326                        ..
5327                    } => {
5328                        let input_ident = ident_stack.pop().unwrap();
5329
5330                        let stmt_id = next_stmt_id.get_and_increment();
5331                        let counter_ident =
5332                            syn::Ident::new(&format!("stream_{}", stmt_id), Span::call_site());
5333
5334                        match builders_or_callback {
5335                            BuildersOrCallback::Builders(graph_builders) => {
5336                                let arg = format!("{}({})", prefix, tag);
5337                                graph_builders.add_dfir_at(
5338                                    &out_location,
5339                                    parse_quote! {
5340                                        #counter_ident = #input_ident -> _counter(#arg, #duration);
5341                                    },
5342                                    Some(&stmt_id.to_string()),
5343                                );
5344                            }
5345                            BuildersOrCallback::Callback(_, node_callback) => {
5346                                node_callback(node, next_stmt_id);
5347                            }
5348                        }
5349
5350                        ident_stack.push(counter_ident);
5351                    }
5352
5353                    HydroNode::VersionedNetworkFork {
5354                        channel_id,
5355                        senders,
5356                        metadata,
5357                        ..
5358                    } => {
5359                        // sender idents are pushed in order of the 'senders' member.
5360                        let split_at = ident_stack.len() - senders.len();
5361                        let sender_idents = ident_stack.split_off(split_at);
5362
5363                        let stmt_id = next_stmt_id.get_and_increment();
5364
5365                        // All senders share the channel, so the raw element type (for embedded
5366                        // serialization) is read from the first sender.
5367                        let external_element_type =
5368                            senders.first().and_then(|(_, _, s)| s.external_element_type());
5369
5370                        match builders_or_callback {
5371                            BuildersOrCallback::Builders(graph_builders) => {
5372                                let sender_args: Vec<(LocationId, syn::Ident, Option<DebugExpr>)> =
5373                                    senders
5374                                        .iter()
5375                                        .zip(sender_idents)
5376                                        .map(|((_version, sender, serialize), ident)| {
5377                                            (
5378                                                sender.metadata().location_id.clone(),
5379                                                ident,
5380                                                serialize.pipeline(),
5381                                            )
5382                                        })
5383                                        .collect();
5384                                graph_builders.create_versioned_network_fork(
5385                                    *channel_id,
5386                                    &metadata.location_id,
5387                                    sender_args,
5388                                    external_element_type,
5389                                    stmt_id,
5390                                );
5391                            }
5392                            BuildersOrCallback::Callback(_, node_callback) => {
5393                                node_callback(node, next_stmt_id);
5394                            }
5395                        }
5396                    }
5397
5398                    HydroNode::VersionedNetwork {
5399                        fork,
5400                        deserialize,
5401                        metadata,
5402                        ..
5403                    } => {
5404                        let stmt_id = next_stmt_id.get_and_increment();
5405                        let receiver_stream_ident =
5406                            syn::Ident::new(&format!("stream_{}", stmt_id), Span::call_site());
5407
5408                        // The wire element type is determined by the channel's *source* kind, which
5409                        // all senders share; read it from the shared fork's first sender.
5410                        let (channel_id, source_loc) = {
5411                            let fork_ref = fork.0.borrow();
5412                            let HydroNode::VersionedNetworkFork {
5413                                channel_id,
5414                                senders,
5415                                ..
5416                            } = &*fork_ref
5417                            else {
5418                                unreachable!("VersionedNetwork.fork must be a VersionedNetworkFork");
5419                            };
5420                            let source_loc = senders
5421                                .first()
5422                                .map(|(_v, sender, _s)| sender.metadata().location_id.clone())
5423                                .expect("a VersionedNetworkFork always has at least one sender");
5424                            (*channel_id, source_loc)
5425                        };
5426
5427                        let deserialize_pipeline = deserialize.pipeline();
5428                        let external_element_type = deserialize.external_element_type();
5429
5430                        match builders_or_callback {
5431                            BuildersOrCallback::Builders(graph_builders) => {
5432                                graph_builders.create_versioned_network(
5433                                    channel_id,
5434                                    &source_loc,
5435                                    &metadata.location_id,
5436                                    &receiver_stream_ident,
5437                                    deserialize_pipeline.as_ref(),
5438                                    external_element_type,
5439                                    stmt_id,
5440                                );
5441                            }
5442                            BuildersOrCallback::Callback(_, node_callback) => {
5443                                node_callback(node, next_stmt_id);
5444                            }
5445                        }
5446
5447                        ident_stack.push(receiver_stream_ident);
5448                    }
5449                }
5450            },
5451            seen_tees,
5452            false,
5453        );
5454
5455        let ret = ident_stack
5456            .pop()
5457            .expect("ident_stack should have exactly one element after traversal");
5458        assert!(
5459            ident_stack.is_empty(),
5460            "ident_stack should be empty after popping the final ident, but has {} remaining element(s). \
5461             This indicates a bug in the code gen: some node pushed idents that were never consumed.",
5462            ident_stack.len()
5463        );
5464        ret
5465    }
5466
5467    pub fn visit_debug_expr(&mut self, mut transform: impl FnMut(&mut DebugExpr)) {
5468        match self {
5469            HydroNode::Placeholder => {
5470                panic!()
5471            }
5472            HydroNode::Cast { .. }
5473            | HydroNode::ObserveNonDet { .. }
5474            | HydroNode::UnboundSingleton { .. }
5475            | HydroNode::AssertIsConsistent { .. } => {}
5476            HydroNode::Source { source, .. } => match source {
5477                HydroSource::Stream(expr) | HydroSource::Iter(expr) => transform(expr),
5478                HydroSource::ExternalNetwork()
5479                | HydroSource::Spin()
5480                | HydroSource::ClusterMembers(_, _)
5481                | HydroSource::Embedded(_)
5482                | HydroSource::EmbeddedSingleton(_) => {} // TODO: what goes here?
5483            },
5484            HydroNode::SingletonSource { value, .. } => {
5485                transform(value);
5486            }
5487            HydroNode::CycleSource { .. }
5488            | HydroNode::Tee { .. }
5489            | HydroNode::Reference { .. }
5490            | HydroNode::YieldConcat { .. }
5491            | HydroNode::BeginAtomic { .. }
5492            | HydroNode::EndAtomic { .. }
5493            | HydroNode::Batch { .. }
5494            | HydroNode::Chain { .. }
5495            | HydroNode::MergeOrdered { .. }
5496            | HydroNode::ChainFirst { .. }
5497            | HydroNode::CrossProduct { .. }
5498            | HydroNode::CrossSingleton { .. }
5499            | HydroNode::ResolveFutures { .. }
5500            | HydroNode::ResolveFuturesBlocking { .. }
5501            | HydroNode::ResolveFuturesOrdered { .. }
5502            | HydroNode::Join { .. }
5503            | HydroNode::JoinHalf { .. }
5504            | HydroNode::Difference { .. }
5505            | HydroNode::AntiJoin { .. }
5506            | HydroNode::DeferTick { .. }
5507            | HydroNode::Enumerate { .. }
5508            | HydroNode::Unique { .. }
5509            | HydroNode::Sort { .. }
5510            | HydroNode::VersionedNetworkFork { .. }
5511            | HydroNode::VersionedNetwork { .. } => {}
5512            HydroNode::Map { f, .. }
5513            | HydroNode::FlatMap { f, .. }
5514            | HydroNode::FlatMapStreamBlocking { f, .. }
5515            | HydroNode::Filter { f, .. }
5516            | HydroNode::FilterMap { f, .. }
5517            | HydroNode::Inspect { f, .. }
5518            | HydroNode::Partition { f, .. }
5519            | HydroNode::Reduce { f, .. }
5520            | HydroNode::ReduceKeyed { f, .. }
5521            | HydroNode::ReduceKeyedWatermark { f, .. } => {
5522                transform(&mut f.expr);
5523            }
5524            HydroNode::Fold { init, acc, .. }
5525            | HydroNode::Scan { init, acc, .. }
5526            | HydroNode::ScanAsyncBlocking { init, acc, .. }
5527            | HydroNode::FoldKeyed { init, acc, .. } => {
5528                transform(&mut init.expr);
5529                transform(&mut acc.expr);
5530            }
5531            HydroNode::Network {
5532                serialize,
5533                deserialize,
5534                ..
5535            } => {
5536                if let NetworkSend::Custom {
5537                    serialize_fn: Some(serialize_fn),
5538                } = serialize
5539                {
5540                    transform(serialize_fn);
5541                }
5542                if let NetworkRecv::Custom {
5543                    deserialize_fn: Some(deserialize_fn),
5544                } = deserialize
5545                {
5546                    transform(deserialize_fn);
5547                }
5548            }
5549            HydroNode::ExternalInput { deserialize_fn, .. } => {
5550                if let Some(deserialize_fn) = deserialize_fn {
5551                    transform(deserialize_fn);
5552                }
5553            }
5554            HydroNode::Counter { duration, .. } => {
5555                transform(duration);
5556            }
5557        }
5558    }
5559
5560    pub fn op_metadata(&self) -> &HydroIrOpMetadata {
5561        &self.metadata().op
5562    }
5563
5564    pub fn metadata(&self) -> &HydroIrMetadata {
5565        match self {
5566            HydroNode::Placeholder => {
5567                panic!()
5568            }
5569            HydroNode::VersionedNetworkFork { metadata, .. }
5570            | HydroNode::VersionedNetwork { metadata, .. } => metadata,
5571            HydroNode::Cast { metadata, .. }
5572            | HydroNode::ObserveNonDet { metadata, .. }
5573            | HydroNode::AssertIsConsistent { metadata, .. }
5574            | HydroNode::UnboundSingleton { metadata, .. }
5575            | HydroNode::Source { metadata, .. }
5576            | HydroNode::SingletonSource { metadata, .. }
5577            | HydroNode::CycleSource { metadata, .. }
5578            | HydroNode::Tee { metadata, .. }
5579            | HydroNode::Reference { metadata, .. }
5580            | HydroNode::Partition { metadata, .. }
5581            | HydroNode::YieldConcat { metadata, .. }
5582            | HydroNode::BeginAtomic { metadata, .. }
5583            | HydroNode::EndAtomic { metadata, .. }
5584            | HydroNode::Batch { metadata, .. }
5585            | HydroNode::Chain { metadata, .. }
5586            | HydroNode::MergeOrdered { metadata, .. }
5587            | HydroNode::ChainFirst { metadata, .. }
5588            | HydroNode::CrossProduct { metadata, .. }
5589            | HydroNode::CrossSingleton { metadata, .. }
5590            | HydroNode::Join { metadata, .. }
5591            | HydroNode::JoinHalf { metadata, .. }
5592            | HydroNode::Difference { metadata, .. }
5593            | HydroNode::AntiJoin { metadata, .. }
5594            | HydroNode::ResolveFutures { metadata, .. }
5595            | HydroNode::ResolveFuturesBlocking { metadata, .. }
5596            | HydroNode::ResolveFuturesOrdered { metadata, .. }
5597            | HydroNode::Map { metadata, .. }
5598            | HydroNode::FlatMap { metadata, .. }
5599            | HydroNode::FlatMapStreamBlocking { metadata, .. }
5600            | HydroNode::Filter { metadata, .. }
5601            | HydroNode::FilterMap { metadata, .. }
5602            | HydroNode::DeferTick { metadata, .. }
5603            | HydroNode::Enumerate { metadata, .. }
5604            | HydroNode::Inspect { metadata, .. }
5605            | HydroNode::Unique { metadata, .. }
5606            | HydroNode::Sort { metadata, .. }
5607            | HydroNode::Scan { metadata, .. }
5608            | HydroNode::ScanAsyncBlocking { metadata, .. }
5609            | HydroNode::Fold { metadata, .. }
5610            | HydroNode::FoldKeyed { metadata, .. }
5611            | HydroNode::Reduce { metadata, .. }
5612            | HydroNode::ReduceKeyed { metadata, .. }
5613            | HydroNode::ReduceKeyedWatermark { metadata, .. }
5614            | HydroNode::ExternalInput { metadata, .. }
5615            | HydroNode::Network { metadata, .. }
5616            | HydroNode::Counter { metadata, .. } => metadata,
5617        }
5618    }
5619
5620    pub fn op_metadata_mut(&mut self) -> &mut HydroIrOpMetadata {
5621        &mut self.metadata_mut().op
5622    }
5623
5624    pub fn metadata_mut(&mut self) -> &mut HydroIrMetadata {
5625        match self {
5626            HydroNode::Placeholder => {
5627                panic!()
5628            }
5629            HydroNode::VersionedNetworkFork { metadata, .. }
5630            | HydroNode::VersionedNetwork { metadata, .. } => metadata,
5631            HydroNode::Cast { metadata, .. }
5632            | HydroNode::ObserveNonDet { metadata, .. }
5633            | HydroNode::AssertIsConsistent { metadata, .. }
5634            | HydroNode::UnboundSingleton { metadata, .. }
5635            | HydroNode::Source { metadata, .. }
5636            | HydroNode::SingletonSource { metadata, .. }
5637            | HydroNode::CycleSource { metadata, .. }
5638            | HydroNode::Tee { metadata, .. }
5639            | HydroNode::Reference { metadata, .. }
5640            | HydroNode::Partition { metadata, .. }
5641            | HydroNode::YieldConcat { metadata, .. }
5642            | HydroNode::BeginAtomic { metadata, .. }
5643            | HydroNode::EndAtomic { metadata, .. }
5644            | HydroNode::Batch { metadata, .. }
5645            | HydroNode::Chain { metadata, .. }
5646            | HydroNode::MergeOrdered { metadata, .. }
5647            | HydroNode::ChainFirst { metadata, .. }
5648            | HydroNode::CrossProduct { metadata, .. }
5649            | HydroNode::CrossSingleton { metadata, .. }
5650            | HydroNode::Join { metadata, .. }
5651            | HydroNode::JoinHalf { metadata, .. }
5652            | HydroNode::Difference { metadata, .. }
5653            | HydroNode::AntiJoin { metadata, .. }
5654            | HydroNode::ResolveFutures { metadata, .. }
5655            | HydroNode::ResolveFuturesBlocking { metadata, .. }
5656            | HydroNode::ResolveFuturesOrdered { metadata, .. }
5657            | HydroNode::Map { metadata, .. }
5658            | HydroNode::FlatMap { metadata, .. }
5659            | HydroNode::FlatMapStreamBlocking { metadata, .. }
5660            | HydroNode::Filter { metadata, .. }
5661            | HydroNode::FilterMap { metadata, .. }
5662            | HydroNode::DeferTick { metadata, .. }
5663            | HydroNode::Enumerate { metadata, .. }
5664            | HydroNode::Inspect { metadata, .. }
5665            | HydroNode::Unique { metadata, .. }
5666            | HydroNode::Sort { metadata, .. }
5667            | HydroNode::Scan { metadata, .. }
5668            | HydroNode::ScanAsyncBlocking { metadata, .. }
5669            | HydroNode::Fold { metadata, .. }
5670            | HydroNode::FoldKeyed { metadata, .. }
5671            | HydroNode::Reduce { metadata, .. }
5672            | HydroNode::ReduceKeyed { metadata, .. }
5673            | HydroNode::ReduceKeyedWatermark { metadata, .. }
5674            | HydroNode::ExternalInput { metadata, .. }
5675            | HydroNode::Network { metadata, .. }
5676            | HydroNode::Counter { metadata, .. } => metadata,
5677        }
5678    }
5679
5680    pub fn input(&self) -> Vec<&HydroNode> {
5681        match self {
5682            HydroNode::Placeholder => {
5683                panic!()
5684            }
5685            HydroNode::Source { .. }
5686            | HydroNode::SingletonSource { .. }
5687            | HydroNode::ExternalInput { .. }
5688            | HydroNode::CycleSource { .. }
5689            | HydroNode::Tee { .. }
5690            | HydroNode::Reference { .. }
5691            | HydroNode::Partition { .. }
5692            | HydroNode::VersionedNetwork { .. } => {
5693                // Tee/Partition/VersionedNetwork find their input in separate special ways
5694                vec![]
5695            }
5696            HydroNode::Cast { inner, .. }
5697            | HydroNode::ObserveNonDet { inner, .. }
5698            | HydroNode::YieldConcat { inner, .. }
5699            | HydroNode::BeginAtomic { inner, .. }
5700            | HydroNode::EndAtomic { inner, .. }
5701            | HydroNode::Batch { inner, .. }
5702            | HydroNode::UnboundSingleton { inner, .. }
5703            | HydroNode::AssertIsConsistent { inner, .. } => {
5704                vec![inner]
5705            }
5706            HydroNode::Chain { first, second, .. } => {
5707                vec![first, second]
5708            }
5709            HydroNode::MergeOrdered { first, second, .. } => {
5710                vec![first, second]
5711            }
5712            HydroNode::ChainFirst { first, second, .. } => {
5713                vec![first, second]
5714            }
5715            HydroNode::CrossProduct { left, right, .. }
5716            | HydroNode::CrossSingleton { left, right, .. }
5717            | HydroNode::Join { left, right, .. }
5718            | HydroNode::JoinHalf { left, right, .. } => {
5719                vec![left, right]
5720            }
5721            HydroNode::Difference { pos, neg, .. } | HydroNode::AntiJoin { pos, neg, .. } => {
5722                vec![pos, neg]
5723            }
5724            HydroNode::Map { input, .. }
5725            | HydroNode::FlatMap { input, .. }
5726            | HydroNode::FlatMapStreamBlocking { input, .. }
5727            | HydroNode::Filter { input, .. }
5728            | HydroNode::FilterMap { input, .. }
5729            | HydroNode::Sort { input, .. }
5730            | HydroNode::DeferTick { input, .. }
5731            | HydroNode::Enumerate { input, .. }
5732            | HydroNode::Inspect { input, .. }
5733            | HydroNode::Unique { input, .. }
5734            | HydroNode::Network { input, .. }
5735            | HydroNode::Counter { input, .. }
5736            | HydroNode::ResolveFutures { input, .. }
5737            | HydroNode::ResolveFuturesBlocking { input, .. }
5738            | HydroNode::ResolveFuturesOrdered { input, .. }
5739            | HydroNode::Fold { input, .. }
5740            | HydroNode::FoldKeyed { input, .. }
5741            | HydroNode::Reduce { input, .. }
5742            | HydroNode::ReduceKeyed { input, .. }
5743            | HydroNode::Scan { input, .. }
5744            | HydroNode::ScanAsyncBlocking { input, .. } => {
5745                vec![input]
5746            }
5747            HydroNode::ReduceKeyedWatermark {
5748                input, watermark, ..
5749            } => {
5750                vec![input, watermark]
5751            }
5752            HydroNode::VersionedNetworkFork { senders, .. } => senders
5753                .iter()
5754                .map(|(_version, sender, _serialize)| sender.as_ref())
5755                .collect(),
5756        }
5757    }
5758
5759    pub fn input_metadata(&self) -> Vec<&HydroIrMetadata> {
5760        self.input()
5761            .iter()
5762            .map(|input_node| input_node.metadata())
5763            .collect()
5764    }
5765
5766    /// Returns `true` if this node is a Tee or Partition whose inner Rc
5767    /// has other live references, meaning the upstream is already driven
5768    /// by another consumer and does not need a Null sink.
5769    pub fn is_shared_with_others(&self) -> bool {
5770        match self {
5771            HydroNode::Tee { inner, .. } | HydroNode::Partition { inner, .. } => {
5772                Rc::strong_count(&inner.0) > 1
5773            }
5774            // A zero-output reference node is valid in DFIR (it drains itself at
5775            // end of tick), so it doesn't need to be driven by another consumer.
5776            HydroNode::Reference { .. } => false,
5777            _ => false,
5778        }
5779    }
5780
5781    pub fn print_root(&self) -> String {
5782        match self {
5783            HydroNode::Placeholder => {
5784                panic!()
5785            }
5786            HydroNode::Cast { .. } => "Cast()".to_owned(),
5787            HydroNode::UnboundSingleton { .. } => "UnboundSingleton()".to_owned(),
5788            HydroNode::ObserveNonDet { .. } => "ObserveNonDet()".to_owned(),
5789            HydroNode::AssertIsConsistent { .. } => "AssertIsConsistent()".to_owned(),
5790            HydroNode::Source { source, .. } => format!("Source({:?})", source),
5791            HydroNode::SingletonSource {
5792                value,
5793                first_tick_only,
5794                ..
5795            } => format!(
5796                "SingletonSource({:?}, first_tick_only={})",
5797                value, first_tick_only
5798            ),
5799            HydroNode::CycleSource { cycle_id, .. } => format!("CycleSource({})", cycle_id),
5800            HydroNode::Tee { inner, .. } => {
5801                format!("Tee({})", inner.0.borrow().print_root())
5802            }
5803            HydroNode::Reference { inner, kind, .. } => {
5804                format!("Reference({:?}, {})", kind, inner.0.borrow().print_root())
5805            }
5806            HydroNode::Partition { f, is_true, .. } => {
5807                format!("Partition({:?}, is_true={})", f, is_true)
5808            }
5809            HydroNode::YieldConcat { .. } => "YieldConcat()".to_owned(),
5810            HydroNode::BeginAtomic { .. } => "BeginAtomic()".to_owned(),
5811            HydroNode::EndAtomic { .. } => "EndAtomic()".to_owned(),
5812            HydroNode::Batch { .. } => "Batch()".to_owned(),
5813            HydroNode::Chain { first, second, .. } => {
5814                format!("Chain({}, {})", first.print_root(), second.print_root())
5815            }
5816            HydroNode::MergeOrdered { first, second, .. } => {
5817                format!(
5818                    "MergeOrdered({}, {})",
5819                    first.print_root(),
5820                    second.print_root()
5821                )
5822            }
5823            HydroNode::ChainFirst { first, second, .. } => {
5824                format!(
5825                    "ChainFirst({}, {})",
5826                    first.print_root(),
5827                    second.print_root()
5828                )
5829            }
5830            HydroNode::CrossProduct { left, right, .. } => {
5831                format!(
5832                    "CrossProduct({}, {})",
5833                    left.print_root(),
5834                    right.print_root()
5835                )
5836            }
5837            HydroNode::CrossSingleton { left, right, .. } => {
5838                format!(
5839                    "CrossSingleton({}, {})",
5840                    left.print_root(),
5841                    right.print_root()
5842                )
5843            }
5844            HydroNode::Join { left, right, .. } => {
5845                format!("Join({}, {})", left.print_root(), right.print_root())
5846            }
5847            HydroNode::JoinHalf { left, right, .. } => {
5848                format!("JoinHalf({}, {})", left.print_root(), right.print_root())
5849            }
5850            HydroNode::Difference { pos, neg, .. } => {
5851                format!("Difference({}, {})", pos.print_root(), neg.print_root())
5852            }
5853            HydroNode::AntiJoin { pos, neg, .. } => {
5854                format!("AntiJoin({}, {})", pos.print_root(), neg.print_root())
5855            }
5856            HydroNode::ResolveFutures { .. } => "ResolveFutures()".to_owned(),
5857            HydroNode::ResolveFuturesBlocking { .. } => "ResolveFuturesBlocking()".to_owned(),
5858            HydroNode::ResolveFuturesOrdered { .. } => "ResolveFuturesOrdered()".to_owned(),
5859            HydroNode::Map { f, .. } => format!("Map({:?})", f),
5860            HydroNode::FlatMap { f, .. } => format!("FlatMap({:?})", f),
5861            HydroNode::FlatMapStreamBlocking { f, .. } => format!("FlatMapStreamBlocking({:?})", f),
5862            HydroNode::Filter { f, .. } => format!("Filter({:?})", f),
5863            HydroNode::FilterMap { f, .. } => format!("FilterMap({:?})", f),
5864            HydroNode::DeferTick { .. } => "DeferTick()".to_owned(),
5865            HydroNode::Enumerate { .. } => "Enumerate()".to_owned(),
5866            HydroNode::Inspect { f, .. } => format!("Inspect({:?})", f),
5867            HydroNode::Unique { .. } => "Unique()".to_owned(),
5868            HydroNode::Sort { .. } => "Sort()".to_owned(),
5869            HydroNode::Fold { init, acc, .. } => format!("Fold({:?}, {:?})", init, acc),
5870            HydroNode::Scan { init, acc, .. } => format!("Scan({:?}, {:?})", init, acc),
5871            HydroNode::ScanAsyncBlocking { init, acc, .. } => {
5872                format!("ScanAsyncBlocking({:?}, {:?})", init, acc)
5873            }
5874            HydroNode::FoldKeyed { init, acc, .. } => format!("FoldKeyed({:?}, {:?})", init, acc),
5875            HydroNode::Reduce { f, .. } => format!("Reduce({:?})", f),
5876            HydroNode::ReduceKeyed { f, .. } => format!("ReduceKeyed({:?})", f),
5877            HydroNode::ReduceKeyedWatermark { f, .. } => format!("ReduceKeyedWatermark({:?})", f),
5878            HydroNode::Network { .. } => "Network()".to_owned(),
5879            HydroNode::ExternalInput { .. } => "ExternalInput()".to_owned(),
5880            HydroNode::Counter { tag, duration, .. } => {
5881                format!("Counter({:?}, {:?})", tag, duration)
5882            }
5883            HydroNode::VersionedNetworkFork {
5884                channel_name,
5885                senders,
5886                ..
5887            } => {
5888                let versions: Vec<u32> = senders.iter().map(|(v, _, _)| *v).collect();
5889                format!(
5890                    "VersionedNetworkFork({}, senders={:?})",
5891                    channel_name, versions
5892                )
5893            }
5894            HydroNode::VersionedNetwork { version, .. } => {
5895                format!("VersionedNetwork(v{})", version)
5896            }
5897        }
5898    }
5899}
5900
5901#[cfg(feature = "build")]
5902#[expect(clippy::too_many_arguments, reason = "networking codegen")]
5903fn instantiate_network<'a, D>(
5904    env: &mut D::InstantiateEnv,
5905    from_location: &LocationId,
5906    to_location: &LocationId,
5907    processes: &SparseSecondaryMap<LocationKey, D::Process>,
5908    clusters: &SparseSecondaryMap<LocationKey, D::Cluster>,
5909    name: Option<&str>,
5910    networking_info: &crate::networking::NetworkingInfo,
5911    external_types: Option<(&syn::Type, &syn::Type)>,
5912) -> (syn::Expr, syn::Expr, Box<dyn FnOnce()>)
5913where
5914    D: Deploy<'a>,
5915{
5916    if external_types.is_some() && !D::SUPPORTS_EXTERNAL_SERIALIZATION {
5917        panic!(
5918            "`.embedded()` serialization leaves serialization to code outside of Hydro and is \
5919             only supported by the embedded deployment backend. Use `.bincode()` (or another \
5920             supported serialization backend) for this deployment target instead."
5921        );
5922    }
5923
5924    let ((sink, source), connect_fn) = match (from_location, to_location) {
5925        (&LocationId::Process(from), &LocationId::Process(to)) => {
5926            let from_node = processes
5927                .get(from)
5928                .unwrap_or_else(|| {
5929                    panic!("A process used in the graph was not instantiated: {}", from)
5930                })
5931                .clone();
5932            let to_node = processes
5933                .get(to)
5934                .unwrap_or_else(|| {
5935                    panic!("A process used in the graph was not instantiated: {}", to)
5936                })
5937                .clone();
5938
5939            let sink_port = from_node.next_port();
5940            let source_port = to_node.next_port();
5941
5942            (
5943                D::o2o_sink_source(
5944                    env,
5945                    &from_node,
5946                    &sink_port,
5947                    &to_node,
5948                    &source_port,
5949                    name,
5950                    networking_info,
5951                    external_types,
5952                ),
5953                D::o2o_connect(&from_node, &sink_port, &to_node, &source_port),
5954            )
5955        }
5956        (&LocationId::Process(from), &LocationId::Cluster(to)) => {
5957            let from_node = processes
5958                .get(from)
5959                .unwrap_or_else(|| {
5960                    panic!("A process used in the graph was not instantiated: {}", from)
5961                })
5962                .clone();
5963            let to_node = clusters
5964                .get(to)
5965                .unwrap_or_else(|| {
5966                    panic!("A cluster used in the graph was not instantiated: {}", to)
5967                })
5968                .clone();
5969
5970            let sink_port = from_node.next_port();
5971            let source_port = to_node.next_port();
5972
5973            (
5974                D::o2m_sink_source(
5975                    env,
5976                    &from_node,
5977                    &sink_port,
5978                    &to_node,
5979                    &source_port,
5980                    name,
5981                    networking_info,
5982                    external_types,
5983                ),
5984                D::o2m_connect(&from_node, &sink_port, &to_node, &source_port),
5985            )
5986        }
5987        (&LocationId::Cluster(from), &LocationId::Process(to)) => {
5988            let from_node = clusters
5989                .get(from)
5990                .unwrap_or_else(|| {
5991                    panic!("A cluster used in the graph was not instantiated: {}", from)
5992                })
5993                .clone();
5994            let to_node = processes
5995                .get(to)
5996                .unwrap_or_else(|| {
5997                    panic!("A process used in the graph was not instantiated: {}", to)
5998                })
5999                .clone();
6000
6001            let sink_port = from_node.next_port();
6002            let source_port = to_node.next_port();
6003
6004            (
6005                D::m2o_sink_source(
6006                    env,
6007                    &from_node,
6008                    &sink_port,
6009                    &to_node,
6010                    &source_port,
6011                    name,
6012                    networking_info,
6013                    external_types,
6014                ),
6015                D::m2o_connect(&from_node, &sink_port, &to_node, &source_port),
6016            )
6017        }
6018        (&LocationId::Cluster(from), &LocationId::Cluster(to)) => {
6019            let from_node = clusters
6020                .get(from)
6021                .unwrap_or_else(|| {
6022                    panic!("A cluster used in the graph was not instantiated: {}", from)
6023                })
6024                .clone();
6025            let to_node = clusters
6026                .get(to)
6027                .unwrap_or_else(|| {
6028                    panic!("A cluster used in the graph was not instantiated: {}", to)
6029                })
6030                .clone();
6031
6032            let sink_port = from_node.next_port();
6033            let source_port = to_node.next_port();
6034
6035            (
6036                D::m2m_sink_source(
6037                    env,
6038                    &from_node,
6039                    &sink_port,
6040                    &to_node,
6041                    &source_port,
6042                    name,
6043                    networking_info,
6044                    external_types,
6045                ),
6046                D::m2m_connect(&from_node, &sink_port, &to_node, &source_port),
6047            )
6048        }
6049        (LocationId::Tick(_, _), _) => panic!(),
6050        (_, LocationId::Tick(_, _)) => panic!(),
6051        (LocationId::Atomic(_), _) => panic!(),
6052        (_, LocationId::Atomic(_)) => panic!(),
6053    };
6054    (sink, source, connect_fn)
6055}
6056
6057#[cfg(test)]
6058mod serde_test;
6059
6060#[cfg(test)]
6061mod test {
6062    use std::mem::size_of;
6063
6064    use stageleft::{QuotedWithContext, q};
6065
6066    use super::*;
6067
6068    #[test]
6069    #[cfg_attr(
6070        not(feature = "build"),
6071        ignore = "expects inclusion of feature-gated fields"
6072    )]
6073    fn hydro_node_size() {
6074        assert_eq!(size_of::<HydroNode>(), 264);
6075    }
6076
6077    #[test]
6078    #[cfg_attr(
6079        not(feature = "build"),
6080        ignore = "expects inclusion of feature-gated fields"
6081    )]
6082    fn hydro_root_size() {
6083        assert_eq!(size_of::<HydroRoot>(), 136);
6084    }
6085
6086    #[test]
6087    fn test_simplify_q_macro_basic() {
6088        // Test basic non-q! expression
6089        let simple_expr: syn::Expr = syn::parse_str("x + y").unwrap();
6090        let result = simplify_q_macro(simple_expr.clone());
6091        assert_eq!(result, simple_expr);
6092    }
6093
6094    #[test]
6095    fn test_simplify_q_macro_actual_stageleft_call() {
6096        // Test a simplified version of what a real stageleft call might look like
6097        let stageleft_call = q!(|x: usize| x + 1).splice_fn1_ctx(&());
6098        let result = simplify_q_macro(stageleft_call);
6099        // This should be processed by our visitor and simplified to q!(...)
6100        // since we detect the stageleft::runtime_support::fn_* pattern
6101        hydro_build_utils::assert_snapshot!(result.to_token_stream().to_string());
6102    }
6103
6104    #[test]
6105    fn test_closure_no_pipe_at_start() {
6106        // Test a closure that does not start with a pipe
6107        let stageleft_call = q!({
6108            let foo = 123;
6109            move |b: usize| b + foo
6110        })
6111        .splice_fn1_ctx(&());
6112        let result = simplify_q_macro(stageleft_call);
6113        hydro_build_utils::assert_snapshot!(result.to_token_stream().to_string());
6114    }
6115}