hydro_lang/

handoff_ref.rs

//! Reference handles for capturing singletons, optionals, and streams in `q!()` closures.
//!
//! Each handle type wraps a `&RefCell<HydroNode>` and, when captured inside a `q!()` closure,
//! registers itself with the current capture scope. At codegen time, the IR node is lowered
//! to the corresponding DFIR pseudo-operator (`singleton()`, `optional()`, or `handoff()`),
//! and the reference resolves to the appropriate borrow type.
//!
//! Each handle tracks the [`Location`] **and** the boundedness of the collection it refers to
//! (see [`OperatorContext`]). A handle can only be captured inside closures passed to operators
//! on collections with a matching location and boundedness. This is required for soundness:
//! a bounded collection is only materialized on the first tick, while closures on unbounded
//! collections continue to run on later ticks, where the referenced value no longer exists and
//! accessing it would crash at runtime.

use std::cell::RefCell;
use std::marker::PhantomData;
use std::rc::Rc;

use proc_macro2::Span;
use quote::quote;
use stageleft::runtime_support::{FreeVariableWithContextWithProps, QuoteTokens};

use crate::compile::ir::{AccessCounter, HydroNode, SharedNode};
use crate::live_collections::OperatorContext;
use crate::location::Location;

/// Determines which DFIR pseudo-operator a reference node lowers to.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)]
pub enum HandoffRefKind {
    /// `-> singleton()` — exactly one item, `#var` gives `&T`.
    Singleton,
    /// `-> optional()` — zero or one item, `#var` gives `&Option<T>`.
    Optional,
    /// `-> handoff()` — zero or more items, `#var` gives `&Vec<T>`.
    Vec,
}

// Thread-local storage for handoff references captured during `q!()` expansion.
// Stores the `HydroNode::Reference` and `is_mut: bool` for each reference captured in the current closure.
// The index determines the ident name via `handoff_ref_ident`.
thread_local! {
    static CAPTURED_REFS: RefCell<Option<Vec<(HydroNode, bool)>>> = const { RefCell::new(None) };
}

/// Returns the canonical ident for a captured ref at the given index within a closure.
pub(crate) fn handoff_ref_ident(index: usize) -> syn::Ident {
    syn::Ident::new(
        &format!("__hydro_singleton_ref_{}", index),
        Span::call_site(),
    )
}

/// Activate the reference capture context. Must be called before `q!()` expansion
/// that may capture handoff references. Returns a `ClosureExpr` bundling the expression with any
/// captured references.
pub fn with_ref_capture(
    f: impl FnOnce() -> crate::compile::ir::DebugExpr,
) -> crate::compile::ir::ClosureExpr {
    CAPTURED_REFS.with(|cell| {
        let prev = cell.borrow_mut().replace(Vec::new());
        assert!(
            prev.is_none(),
            "nested handoff reference capture scopes are not supported"
        );
    });
    let expr = (f)();
    let captured_refs = CAPTURED_REFS.with(|cell| cell.borrow_mut().take().unwrap());
    crate::compile::ir::ClosureExpr::new(expr, captured_refs)
}

/// Shared registration logic: wraps the IR node in `HydroNode::Reference` if needed,
/// pushes it to the capture list, and returns the ident to use in the closure body.
fn register_handoff_ref(
    ir_node: &RefCell<HydroNode>,
    is_mut: bool,
    kind: HandoffRefKind,
) -> syn::Ident {
    CAPTURED_REFS.with(|cell| {
        let mut guard = cell.borrow_mut();
        let refs = guard.as_mut().expect(
            "HandoffRef used inside q!() but no reference capture scope is active. \
             This is a bug — reference capture should be set up by the operator that uses q!().",
        );

        let index = refs.len();
        let ident = handoff_ref_ident(index);

        let metadata = ir_node.borrow().metadata().clone();

        // Wrap in HydroNode::Reference for materialization + identity tracking.
        // If already a Reference node, reuse it.
        if !matches!(&*ir_node.borrow(), HydroNode::Reference { .. }) {
            let orig = ir_node.replace(HydroNode::Placeholder);
            *ir_node.borrow_mut() = HydroNode::Reference {
                inner: SharedNode(Rc::new(RefCell::new(orig))),
                kind,
                access_counter: AccessCounter::new(),
                metadata: metadata.clone(),
            };
        }

        let borrow: std::cell::Ref<'_, HydroNode> = ir_node.borrow();
        let HydroNode::Reference {
            inner,
            access_counter,
            ..
        } = &*borrow
        else {
            unreachable!()
        };

        // Compute access group at staging time (code order).
        let group = access_counter.next_group(is_mut);

        refs.push((
            HydroNode::Reference {
                inner: SharedNode(Rc::clone(&inner.0)),
                kind,
                access_counter: group,
                metadata,
            },
            is_mut,
        ));

        ident
    })
}

/// Macro to define a handoff reference struct with all necessary trait impls.
macro_rules! define_handoff_ref {
    (
        $(
            $(#[$meta:meta])*
            $name:ident, $is_mut:expr, $kind:expr, $output:ty
        )+
    ) => {
        $(
            $(#[$meta])*
            pub struct $name<'a, 'slf, T, L, B> {
                pub(crate) ir_node: &'slf RefCell<HydroNode>,
                _phantom: PhantomData<(&'a T, L, B)>,
            }

            impl<'slf, T, L, B> $name<'_, 'slf, T, L, B> {
                /// Creates a new reference handle from an IR node cell.
                pub(crate) fn new(ir_node: &'slf RefCell<HydroNode>) -> Self {
                    Self {
                        ir_node,
                        _phantom: PhantomData,
                    }
                }
            }

            impl<T, L, B> Copy for $name<'_, '_, T, L, B> {}
            impl<T, L, B> Clone for $name<'_, '_, T, L, B> {
                fn clone(&self) -> Self {
                    *self
                }
            }

            impl<'a, 'slf, T: 'a, L, B> FreeVariableWithContextWithProps<OperatorContext<L, B>, ()>
                for $name<'a, 'slf, T, L, B>
            where
                L: Location<'a>,
            {
                type O = $output;

                fn to_tokens(self, _ctx: &OperatorContext<L, B>) -> (QuoteTokens, ()) {
                    let ident = register_handoff_ref(
                        self.ir_node,
                        $is_mut,
                        $kind,
                    );
                    (
                        QuoteTokens {
                            prelude: None,
                            expr: Some(quote!(#ident)),
                        },
                        (),
                    )
                }
            }
        )+
    };
}

#[stageleft::export(
    SingletonRef,
    SingletonMut,
    OptionalRef,
    OptionalMut,
    StreamRef,
    StreamMut
)]
define_handoff_ref!(
    /// A shared reference handle to a singleton, resolves to `&T` at runtime.
    ///
    /// Created via [`Singleton::by_ref()`](crate::live_collections::Singleton::by_ref).
    SingletonRef, false, HandoffRefKind::Singleton, &'a T

    /// A mutable reference handle to a singleton, resolves to `&mut T` at runtime.
    ///
    /// Created via [`Singleton::by_mut()`](crate::live_collections::Singleton::by_mut).
    SingletonMut, true, HandoffRefKind::Singleton, &'a mut T

    /// A shared reference handle to an optional, resolves to `&Option<T>` at runtime.
    ///
    /// Created via [`Optional::by_ref()`](crate::live_collections::Optional::by_ref).
    OptionalRef, false, HandoffRefKind::Optional, &'a Option<T>

    /// A mutable reference handle to an optional, resolves to `&mut Option<T>` at runtime.
    ///
    /// Created via [`Optional::by_mut()`](crate::live_collections::Optional::by_mut).
    OptionalMut, true, HandoffRefKind::Optional, &'a mut Option<T>

    /// A shared reference handle to a stream's handoff buffer, resolves to `&Vec<T>` at runtime.
    ///
    /// Created via [`Stream::by_ref()`](crate::live_collections::Stream::by_ref).
    StreamRef, false, HandoffRefKind::Vec, &'a Vec<T>

    /// A mutable reference handle to a stream's handoff buffer, resolves to `&mut Vec<T>` at runtime.
    ///
    /// Created via [`Stream::by_mut()`](crate::live_collections::Stream::by_mut).
    StreamMut, true, HandoffRefKind::Vec, &'a mut Vec<T>
);

#[cfg(test)]
#[cfg(feature = "build")]
mod tests {
    use stageleft::q;

    use crate::compile::builder::FlowBuilder;
    use crate::location::Location;

    struct P1 {}

    /// Compile-only test: verifies that `by_ref()` + `q!()` produces valid IR.
    #[test]
    fn singleton_by_ref_compiles() {
        let mut flow = FlowBuilder::new();
        let node = flow.process::<P1>();

        let my_count = node
            .source_iter(q!(0..5i32))
            .fold(q!(|| 0i32), q!(|acc: &mut i32, x| *acc += x));
        let count_ref = my_count.by_ref();

        node.source_iter(q!(1..=3i32))
            .map(q!(|x| x + *count_ref))
            .for_each(q!(|_| {}));

        my_count.into_stream().for_each(q!(|_| {}));
        let _built = flow.finalize();
    }

    /// Test with a non-Copy type (Vec) to ensure we're borrowing, not copying.
    #[test]
    fn singleton_by_ref_non_copy() {
        let mut flow = FlowBuilder::new();
        let node = flow.process::<P1>();

        let my_vec = node.source_iter(q!(0..5i32)).fold(
            q!(|| Vec::<i32>::new()),
            q!(|acc: &mut Vec<i32>, x| acc.push(x)),
        );
        let vec_ref = my_vec.by_ref();

        node.source_iter(q!(1..=3i32))
            .map(q!(|x| x + vec_ref.len() as i32))
            .for_each(q!(|_| {}));

        my_vec.into_stream().for_each(q!(|_| {}));
        let _built = flow.finalize();
    }

    /// Compile-only: singleton ref inside filter closure.
    #[test]
    fn singleton_by_ref_filter() {
        let mut flow = FlowBuilder::new();
        let node = flow.process::<P1>();

        let threshold = node
            .source_iter(q!(0..5i32))
            .fold(q!(|| 0i32), q!(|acc: &mut i32, x| *acc += x));
        let threshold_ref = threshold.by_ref();

        node.source_iter(q!(1..=10i32))
            .filter(q!(|x| *x > *threshold_ref))
            .for_each(q!(|_| {}));

        threshold.into_stream().for_each(q!(|_| {}));
        let _built = flow.finalize();
    }

    /// Compile-only: singleton ref inside flat_map closure.
    #[test]
    fn singleton_by_ref_flat_map() {
        let mut flow = FlowBuilder::new();
        let node = flow.process::<P1>();

        let count = node
            .source_iter(q!(0..3i32))
            .fold(q!(|| 0i32), q!(|acc: &mut i32, _| *acc += 1));
        let count_ref = count.by_ref();

        node.source_iter(q!(1..=2i32))
            .flat_map_ordered(q!(|x| (0..*count_ref).map(move |i| x + i)))
            .for_each(q!(|_| {}));

        count.into_stream().for_each(q!(|_| {}));
        let _built = flow.finalize();
    }

    /// Compile-only: singleton ref inside inspect closure.
    #[test]
    fn singleton_by_ref_inspect() {
        let mut flow = FlowBuilder::new();
        let node = flow.process::<P1>();

        let count = node
            .source_iter(q!(0..5i32))
            .fold(q!(|| 0i32), q!(|acc: &mut i32, _| *acc += 1));
        let count_ref = count.by_ref();

        node.source_iter(q!(1..=3i32))
            .inspect(q!(|x| println!("count={}, x={}", *count_ref, x)))
            .for_each(q!(|_| {}));

        count.into_stream().for_each(q!(|_| {}));
        let _built = flow.finalize();
    }

    /// Compile-only: singleton ref inside partition predicate.
    #[test]
    fn singleton_by_ref_partition() {
        let mut flow = FlowBuilder::new();
        let node = flow.process::<P1>();

        let threshold = node
            .source_iter(q!(0..5i32))
            .fold(q!(|| 0i32), q!(|acc: &mut i32, x| *acc += x));
        let threshold_ref = threshold.by_ref();

        let (above, below) = node
            .source_iter(q!(1..=10i32))
            .partition(q!(|x| *x > *threshold_ref));

        above.for_each(q!(|_| {}));
        below.for_each(q!(|_| {}));
        threshold.into_stream().for_each(q!(|_| {}));
        let _built = flow.finalize();
    }

    /// Compile-only: singleton ref inside partition with downstream operators on both branches.
    #[test]
    fn singleton_by_ref_partition_with_downstream_ops() {
        let mut flow = FlowBuilder::new();
        let node = flow.process::<P1>();

        let threshold = node
            .source_iter(q!(0..5i32))
            .fold(q!(|| 0i32), q!(|acc: &mut i32, x| *acc += x));
        let threshold_ref = threshold.by_ref();

        let (above, below) = node
            .source_iter(q!(1..=10i32))
            .partition(q!(|x| *x > *threshold_ref));

        above.map(q!(|x| x * 2)).for_each(q!(|_| {}));
        below.map(q!(|x| x + 100)).for_each(q!(|_| {}));
        threshold.into_stream().for_each(q!(|_| {}));
        let _built = flow.finalize();
    }

    /// Compile-only test: singleton by_mut.
    #[test]
    fn singleton_by_mut_compiles() {
        let mut flow = FlowBuilder::new();
        let node = flow.process::<P1>();

        let my_count = node
            .source_iter(q!(0..5i32))
            .fold(q!(|| 0i32), q!(|acc: &mut i32, x| *acc += x));
        let count_mut = my_count.by_mut();

        node.source_iter(q!(1..=3i32))
            .map(q!(|x| {
                *count_mut += x;
                x
            }))
            .for_each(q!(|_| {}));

        my_count.into_stream().for_each(q!(|_| {}));
        let _built = flow.finalize();
    }

    /// Compile-only test: optional by_ref.
    #[test]
    fn optional_by_ref_compiles() {
        let mut flow = FlowBuilder::new();
        let node = flow.process::<P1>();

        let my_opt = node.source_iter(q!(0..5i32)).reduce(q!(|a, b| *a += b));
        let opt_ref = my_opt.by_ref();

        node.source_iter(q!(1..=3i32))
            .map(q!(|x| x + opt_ref.unwrap_or(0)))
            .for_each(q!(|_| {}));

        my_opt.into_stream().for_each(q!(|_| {}));
        let _built = flow.finalize();
    }

    /// Compile-only test: stream by_ref.
    #[test]
    fn stream_by_ref_compiles() {
        let mut flow = FlowBuilder::new();
        let node = flow.process::<P1>();

        let my_stream = node.source_iter(q!(0..5i32));
        let stream_ref = my_stream.by_ref();

        node.source_iter(q!(1..=3i32))
            .map(q!(|x| x + stream_ref.len() as i32))
            .for_each(q!(|_| {}));

        my_stream.for_each(q!(|_| {}));
        let _built = flow.finalize();
    }

    /// Compile-only test: singleton by_mut in filter (TotalOrder).
    #[test]
    fn singleton_by_mut_filter() {
        let mut flow = FlowBuilder::new();
        let node = flow.process::<P1>();

        let my_count = node
            .source_iter(q!(0..5i32))
            .fold(q!(|| 0i32), q!(|acc: &mut i32, x| *acc += x));
        let count_mut = my_count.by_mut();

        node.source_iter(q!(1..=3i32))
            .filter(q!(|x| {
                *count_mut += *x;
                *count_mut > 0
            }))
            .for_each(q!(|_| {}));

        my_count.into_stream().for_each(q!(|_| {}));
        let _built = flow.finalize();
    }

    /// Compile-only test: singleton by_mut in flat_map_ordered (TotalOrder).
    #[test]
    fn singleton_by_mut_flat_map() {
        let mut flow = FlowBuilder::new();
        let node = flow.process::<P1>();

        let my_count = node
            .source_iter(q!(0..5i32))
            .fold(q!(|| 0i32), q!(|acc: &mut i32, x| *acc += x));
        let count_mut = my_count.by_mut();

        node.source_iter(q!(1..=3i32))
            .flat_map_ordered(q!(|x| {
                *count_mut += x;
                vec![*count_mut]
            }))
            .for_each(q!(|_| {}));

        my_count.into_stream().for_each(q!(|_| {}));
        let _built = flow.finalize();
    }

    /// Compile-only test: singleton by_mut in filter_map (TotalOrder).
    #[test]
    fn singleton_by_mut_filter_map() {
        let mut flow = FlowBuilder::new();
        let node = flow.process::<P1>();

        let my_count = node
            .source_iter(q!(0..5i32))
            .fold(q!(|| 0i32), q!(|acc: &mut i32, x| *acc += x));
        let count_mut = my_count.by_mut();

        node.source_iter(q!(1..=3i32))
            .filter_map(q!(|x| {
                *count_mut += x;
                Some(*count_mut)
            }))
            .for_each(q!(|_| {}));

        my_count.into_stream().for_each(q!(|_| {}));
        let _built = flow.finalize();
    }

    /// Compile-only test: singleton by_mut in inspect (TotalOrder).
    #[test]
    fn singleton_by_mut_inspect() {
        let mut flow = FlowBuilder::new();
        let node = flow.process::<P1>();

        let my_count = node
            .source_iter(q!(0..5i32))
            .fold(q!(|| 0i32), q!(|acc: &mut i32, x| *acc += x));
        let count_mut = my_count.by_mut();

        node.source_iter(q!(1..=3i32))
            .inspect(q!(|x| {
                *count_mut += *x;
            }))
            .for_each(q!(|_| {}));

        my_count.into_stream().for_each(q!(|_| {}));
        let _built = flow.finalize();
    }

    /// Compile-only test: singleton by_ref in for_each.
    #[test]
    fn singleton_by_ref_for_each() {
        let mut flow = FlowBuilder::new();
        let node = flow.process::<P1>();

        let my_count = node
            .source_iter(q!(0..5i32))
            .fold(q!(|| 0i32), q!(|acc: &mut i32, x| *acc += x));
        let count_ref = my_count.by_ref();

        node.source_iter(q!(1..=3i32))
            .for_each(q!(|x| println!("{}", x + *count_ref)));

        my_count.into_stream().for_each(q!(|_| {}));
        let _built = flow.finalize();
    }

    /// Compile-only test: singleton by_mut in for_each.
    #[test]
    fn singleton_by_mut_for_each() {
        let mut flow = FlowBuilder::new();
        let node = flow.process::<P1>();

        let my_count = node
            .source_iter(q!(0..5i32))
            .fold(q!(|| 0i32), q!(|acc: &mut i32, x| *acc += x));
        let count_mut = my_count.by_mut();

        node.source_iter(q!(1..=3i32)).for_each(q!(|x| {
            *count_mut += x;
        }));

        my_count.into_stream().for_each(q!(|_| {}));
        let _built = flow.finalize();
    }

    /// Regression test: a handoff reference whose *only* consumer is a `for_each` closure
    /// must still be materialized during DFIR emission.
    ///
    /// `HydroRoot::ForEach` used to only *look up* captured refs in `built_tees`, assuming
    /// some node-level operator had already emitted them, and panicked with "ForEach
    /// singleton ref not found in built_tees" when the `for_each` closure was the sole
    /// capturer. This test drives the flow through full DFIR emission (which
    /// `flow.finalize()` alone does not) to cover that path.
    #[cfg(feature = "deploy")]
    #[test]
    fn singleton_by_ref_for_each_sole_consumer_emits() {
        use crate::live_collections::sliced::sliced;
        use crate::nondet::nondet;

        let mut flow = FlowBuilder::new();
        let node = flow.process::<P1>();

        let items = node.source_iter(q!(1..=3i32));

        sliced! {
            let items = use::batch(items, nondet!(/** test */));
            let my_count = items
                .location()
                .source_iter(q!(0..5i32))
                .fold(q!(|| 0i32), q!(|acc: &mut i32, x| *acc += x));
            let count_ref = my_count.by_ref();

            // The for_each closure is the ONLY consumer of `my_count` — no other operator
            // emits the reference node before this root is processed.
            items.for_each(q!(|x| println!("{}", x + *count_ref)));
        };

        let _ = flow
            .finalize()
            .with_default_optimize::<crate::deploy::HydroDeploy>()
            .preview_compile();
    }

    /// Regression test: same as [`singleton_by_ref_for_each_sole_consumer_emits`], but for
    /// a mutable reference (`by_mut`) — the common accumulator pattern
    /// `stream.for_each(q!(|x| *acc_mut += x))`.
    #[cfg(feature = "deploy")]
    #[test]
    fn singleton_by_mut_for_each_sole_consumer_emits() {
        use crate::live_collections::sliced::sliced;
        use crate::nondet::nondet;

        let mut flow = FlowBuilder::new();
        let node = flow.process::<P1>();

        let items = node.source_iter(q!(1..=3i32));

        sliced! {
            let items = use::batch(items, nondet!(/** test */));
            let my_count = items
                .location()
                .source_iter(q!(0..5i32))
                .fold(q!(|| 0i32), q!(|acc: &mut i32, x| *acc += x));
            let count_mut = my_count.by_mut();

            // The for_each closure is the ONLY consumer of `my_count`.
            items.for_each(q!(|x| {
                *count_mut += x;
            }));
        };

        let _ = flow
            .finalize()
            .with_default_optimize::<crate::deploy::HydroDeploy>()
            .preview_compile();
    }

    /// Regression test: a singleton ref captured by a partition predicate must be
    /// spliced correctly during DFIR emission. The `PartitionShared` emit previously
    /// drained the ident stack in the wrong order (calling `emit_tokens` before popping
    /// the input ident), so the partition's *input* ident was consumed as if it were the
    /// closure's singleton ref. This test drives the flow through full DFIR emission
    /// (which `flow.finalize()` alone does not) to cover that path.
    #[cfg(feature = "deploy")]
    #[test]
    fn singleton_by_ref_partition_emits() {
        let mut flow = FlowBuilder::new();
        let node = flow.process::<P1>();

        let threshold = node
            .source_iter(q!(0..5i32))
            .fold(q!(|| 0i32), q!(|acc: &mut i32, x| *acc += x));
        let threshold_ref = threshold.by_ref();

        let (above, below) = node
            .source_iter(q!(1..=10i32))
            .partition(q!(|x| *x > *threshold_ref));

        above.for_each(q!(|_| {}));
        below.for_each(q!(|_| {}));
        threshold.into_stream().for_each(q!(|_| {}));

        let _ = flow
            .finalize()
            .with_default_optimize::<crate::deploy::HydroDeploy>()
            .preview_compile();
    }

    /// Compile-only test: singleton by_ref inside scan closures.
    #[test]
    fn singleton_by_ref_scan() {
        let mut flow = FlowBuilder::new();
        let node = flow.process::<P1>();

        let offset = node
            .source_iter(q!(0..5i32))
            .fold(q!(|| 0i32), q!(|acc: &mut i32, x| *acc += x));
        let offset_ref = offset.by_ref();

        node.source_iter(q!(1..=3i32))
            .scan(
                q!(move || *offset_ref),
                q!(move |acc: &mut i32, x| {
                    *acc += x + *offset_ref;
                    Some(*acc)
                }),
            )
            .for_each(q!(|_| {}));

        offset.into_stream().for_each(q!(|_| {}));
        let _built = flow.finalize();
    }

    /// Compile-only test: singleton by_ref inside scan_async_blocking closure.
    #[test]
    fn singleton_by_ref_scan_async_blocking() {
        let mut flow = FlowBuilder::new();
        let node = flow.process::<P1>();

        let offset = node
            .source_iter(q!(0..5i32))
            .fold(q!(|| 0i32), q!(|acc: &mut i32, x| *acc += x));
        let offset_ref = offset.by_ref();

        node.source_iter(q!(1..=3i32))
            .scan_async_blocking(
                q!(|| 0i32),
                q!(move |acc: &mut i32, x| {
                    *acc += x + *offset_ref;
                    let val = *acc;
                    async move { Some(val) }
                }),
            )
            .for_each(q!(|_| {}));

        offset.into_stream().for_each(q!(|_| {}));
        let _built = flow.finalize();
    }

    /// Compile-only test: singleton by_ref inside generator closure.
    #[test]
    fn singleton_by_ref_generator() {
        use crate::live_collections::keyed_stream::Generate;

        let mut flow = FlowBuilder::new();
        let node = flow.process::<P1>();

        let threshold = node
            .source_iter(q!(0..5i32))
            .fold(q!(|| 0i32), q!(|acc: &mut i32, x| *acc += x));
        let threshold_ref = threshold.by_ref();

        node.source_iter(q!(1..=3i32))
            .generator(
                q!(|| 0i32),
                q!(move |acc: &mut i32, x| {
                    *acc += x;
                    if *acc > *threshold_ref {
                        Generate::Return(*acc)
                    } else {
                        Generate::Yield(*acc)
                    }
                }),
            )
            .for_each(q!(|_| {}));

        threshold.into_stream().for_each(q!(|_| {}));
        let _built = flow.finalize();
    }

    /// Compile-only test: singleton by_ref inside keyed scan closure.
    #[test]
    fn singleton_by_ref_keyed_scan() {
        let mut flow = FlowBuilder::new();
        let node = flow.process::<P1>();

        let offset = node
            .source_iter(q!(0..5i32))
            .fold(q!(|| 0i32), q!(|acc: &mut i32, x| *acc += x));
        let offset_ref = offset.by_ref();

        node.source_iter(q!(vec![(0, 1i32), (1, 2i32)]))
            .into_keyed()
            .scan(
                q!(|| 0i32),
                q!(move |acc: &mut i32, x| {
                    *acc += x + *offset_ref;
                    Some(*acc)
                }),
            )
            .entries()
            .assume_ordering::<crate::live_collections::stream::TotalOrder>(
                crate::nondet::nondet!(/** test */),
            )
            .for_each(q!(|_| {}));

        offset.into_stream().for_each(q!(|_| {}));
        let _built = flow.finalize();
    }

    /// Compile-only test: singleton by_ref inside keyed generator closure.
    #[test]
    fn singleton_by_ref_keyed_generator() {
        use crate::live_collections::keyed_stream::Generate;

        let mut flow = FlowBuilder::new();
        let node = flow.process::<P1>();

        let threshold = node
            .source_iter(q!(0..5i32))
            .fold(q!(|| 0i32), q!(|acc: &mut i32, x| *acc += x));
        let threshold_ref = threshold.by_ref();

        node.source_iter(q!(vec![(0, 1i32), (1, 2i32)]))
            .into_keyed()
            .generator(
                q!(|| 0i32),
                q!(move |acc: &mut i32, x| {
                    *acc += x;
                    if *acc > *threshold_ref {
                        Generate::Return(*acc)
                    } else {
                        Generate::Yield(*acc)
                    }
                }),
            )
            .entries()
            .assume_ordering::<crate::live_collections::stream::TotalOrder>(
                crate::nondet::nondet!(/** test */),
            )
            .for_each(q!(|_| {}));

        threshold.into_stream().for_each(q!(|_| {}));
        let _built = flow.finalize();
    }

    /// Regression test for the `Partition` dedup path in `HydroNode::transform_children`.
    ///
    /// A `partition` whose predicate closure captures a `by_ref` singleton produces two output
    /// branches that share the partition's `inner` but each own their own clone of the closure.
    /// When the shared `inner` is rewritten (during `finalize` → `unify_atomic_ticks` →
    /// `transform_bottom_up`), the first branch's closure is transformed and the referenced cell
    /// is emptied to `HydroNode::Placeholder`. The *deduplicated* second branch used to be skipped
    /// entirely, leaving its captured `singleton_ref` dangling at that `Placeholder` cell. This
    /// asserts that no closure `singleton_ref` resolves to a `Placeholder` after finalization.
    #[test]
    fn partition_dedup_does_not_leave_placeholder_refs() {
        use std::cell::RefCell;
        use std::collections::HashSet;

        use crate::compile::ir::{ClosureExpr, HydroNode};

        fn closures(node: &HydroNode) -> Vec<&ClosureExpr> {
            match node {
                HydroNode::Map { f, .. }
                | HydroNode::FlatMap { f, .. }
                | HydroNode::FlatMapStreamBlocking { f, .. }
                | HydroNode::Filter { f, .. }
                | HydroNode::FilterMap { f, .. }
                | HydroNode::Inspect { f, .. }
                | HydroNode::Reduce { f, .. }
                | HydroNode::ReduceKeyed { f, .. }
                | HydroNode::ReduceKeyedWatermark { f, .. }
                | HydroNode::PartitionShared { f, .. } => vec![f],
                HydroNode::Fold { init, acc, .. }
                | HydroNode::FoldKeyed { init, acc, .. }
                | HydroNode::Scan { init, acc, .. }
                | HydroNode::ScanAsyncBlocking { init, acc, .. } => vec![init, acc],
                _ => vec![],
            }
        }

        fn count_placeholder_refs(
            node: &HydroNode,
            visited: &mut HashSet<*const RefCell<HydroNode>>,
            count: &mut usize,
        ) {
            // A captured `singleton_ref` should never resolve to a `Placeholder`.
            for f in closures(node) {
                for (ref_node, _is_mut) in &f.singleton_refs {
                    if let HydroNode::Reference { inner, .. } = ref_node
                        && matches!(&*inner.0.borrow(), HydroNode::Placeholder)
                    {
                        *count += 1;
                    }
                }
            }

            // Recurse the (owned) tree children.
            for child in node.input() {
                count_placeholder_refs(child, visited, count);
            }

            // Recurse the shared inner once, guarding against `Placeholder` / revisits.
            let shared_inner = match node {
                HydroNode::Tee { inner, .. }
                | HydroNode::Reference { inner, .. }
                | HydroNode::PartitionSide { inner, .. } => Some(inner),
                _ => None,
            };
            if let Some(inner) = shared_inner
                && visited.insert(inner.as_ptr())
            {
                let borrowed = inner.0.borrow();
                if !matches!(&*borrowed, HydroNode::Placeholder) {
                    count_placeholder_refs(&borrowed, visited, count);
                }
            }
        }

        let mut flow = FlowBuilder::new();
        let node = flow.process::<P1>();

        let threshold = node
            .source_iter(q!(0..5i32))
            .fold(q!(|| 0i32), q!(|acc: &mut i32, x| *acc += x));
        let threshold_ref = threshold.by_ref();

        let (above, below) = node
            .source_iter(q!(vec![5i32, 8, 10, 11, 15, 3]))
            .partition(q!(|x| *x > *threshold_ref));

        above.for_each(q!(|_| {}));
        below.for_each(q!(|_| {}));
        threshold.into_stream().for_each(q!(|_| {}));

        let built = flow.finalize();

        let mut visited = HashSet::new();
        let mut count = 0usize;
        for root in built.ir() {
            count_placeholder_refs(root.input(), &mut visited, &mut count);
        }

        assert_eq!(
            count, 0,
            "partition's deduplicated branch left {count} closure singleton_ref(s) dangling at a Placeholder cell"
        );
    }
}