hydro_lang/live_collections/keyed_stream/

mod.rs

//! Definitions for the [`KeyedStream`] live collection.

use std::cell::RefCell;
use std::collections::HashMap;
use std::hash::Hash;
use std::marker::PhantomData;
use std::ops::Deref;
use std::rc::Rc;

use stageleft::{IntoQuotedMut, QuotedWithContext, QuotedWithContextWithProps, q};

use super::boundedness::{Bounded, Boundedness, IsBounded, Unbounded};
use super::keyed_singleton::KeyedSingleton;
use super::optional::Optional;
use super::stream::{
    ExactlyOnce, IsExactlyOnce, IsOrdered, MinOrder, MinRetries, NoOrder, Stream, TotalOrder,
};
use crate::compile::builder::{CycleId, FlowState};
use crate::compile::ir::{
    CollectionKind, HydroIrOpMetadata, HydroNode, HydroRoot, SharedNode, StreamOrder, StreamRetry,
};
#[cfg(stageleft_runtime)]
use crate::forward_handle::{CycleCollection, ReceiverComplete};
use crate::forward_handle::{ForwardRef, TickCycle};
use crate::live_collections::batch_atomic::BatchAtomic;
use crate::live_collections::stream::{
    AtLeastOnce, Ordering, Retries, WeakerOrderingThan, WeakerRetryThan,
};
#[cfg(stageleft_runtime)]
use crate::location::dynamic::{DynLocation, LocationId};
use crate::location::tick::DeferTick;
use crate::location::{Atomic, Location, NoTick, Tick, check_matching_location};
use crate::manual_expr::ManualExpr;
use crate::nondet::{NonDet, nondet};
use crate::properties::{AggFuncAlgebra, ValidCommutativityFor, ValidIdempotenceFor};

pub mod networking;

/// Streaming elements of type `V` grouped by a key of type `K`.
///
/// Keyed Streams capture streaming elements of type `V` grouped by a key of type `K`, where the
/// order of keys is non-deterministic but the order *within* each group may be deterministic.
///
/// Although keyed streams are conceptually grouped by keys, values are not immediately grouped
/// into buckets when constructing a keyed stream. Instead, keyed streams defer grouping until an
/// operator such as [`KeyedStream::fold`] is called, which requires `K: Hash + Eq`.
///
/// Type Parameters:
/// - `K`: the type of the key for each group
/// - `V`: the type of the elements inside each group
/// - `Loc`: the [`Location`] where the keyed stream is materialized
/// - `Bound`: tracks whether the entries are [`Bounded`] (local and finite) or [`Unbounded`] (asynchronous and possibly infinite)
/// - `Order`: tracks whether the elements within each group have deterministic order
///   ([`TotalOrder`]) or not ([`NoOrder`])
/// - `Retries`: tracks whether the elements within each group have deterministic cardinality
///   ([`ExactlyOnce`]) or may have non-deterministic retries ([`crate::live_collections::stream::AtLeastOnce`])
pub struct KeyedStream<
    K,
    V,
    Loc,
    Bound: Boundedness = Unbounded,
    Order: Ordering = TotalOrder,
    Retry: Retries = ExactlyOnce,
> {
    pub(crate) location: Loc,
    pub(crate) ir_node: RefCell<HydroNode>,
    pub(crate) flow_state: FlowState,

    _phantom: PhantomData<(K, V, Loc, Bound, Order, Retry)>,
}

impl<K, V, L, B: Boundedness, O: Ordering, R: Retries> Drop for KeyedStream<K, V, L, B, O, R> {
    fn drop(&mut self) {
        let ir_node = self.ir_node.replace(HydroNode::Placeholder);
        if !matches!(ir_node, HydroNode::Placeholder) && !ir_node.is_shared_with_others() {
            self.flow_state.borrow_mut().try_push_root(HydroRoot::Null {
                input: Box::new(ir_node),
                op_metadata: HydroIrOpMetadata::new(),
            });
        }
    }
}

impl<'a, K, V, L, O: Ordering, R: Retries> From<KeyedStream<K, V, L, Bounded, O, R>>
    for KeyedStream<K, V, L, Unbounded, O, R>
where
    L: Location<'a>,
{
    fn from(stream: KeyedStream<K, V, L, Bounded, O, R>) -> KeyedStream<K, V, L, Unbounded, O, R> {
        let new_meta = stream
            .location
            .new_node_metadata(KeyedStream::<K, V, L, Unbounded, O, R>::collection_kind());

        KeyedStream {
            location: stream.location.clone(),
            flow_state: stream.flow_state.clone(),
            ir_node: RefCell::new(HydroNode::Cast {
                inner: Box::new(stream.ir_node.replace(HydroNode::Placeholder)),
                metadata: new_meta,
            }),
            _phantom: PhantomData,
        }
    }
}

impl<'a, K, V, L, B: Boundedness, R: Retries> From<KeyedStream<K, V, L, B, TotalOrder, R>>
    for KeyedStream<K, V, L, B, NoOrder, R>
where
    L: Location<'a>,
{
    fn from(stream: KeyedStream<K, V, L, B, TotalOrder, R>) -> KeyedStream<K, V, L, B, NoOrder, R> {
        stream.weaken_ordering()
    }
}

impl<'a, K, V, L, O: Ordering, R: Retries> DeferTick for KeyedStream<K, V, Tick<L>, Bounded, O, R>
where
    L: Location<'a>,
{
    fn defer_tick(self) -> Self {
        KeyedStream::defer_tick(self)
    }
}

impl<'a, K, V, L, O: Ordering, R: Retries> CycleCollection<'a, TickCycle>
    for KeyedStream<K, V, Tick<L>, Bounded, O, R>
where
    L: Location<'a>,
{
    type Location = Tick<L>;

    fn create_source(cycle_id: CycleId, location: Tick<L>) -> Self {
        KeyedStream {
            flow_state: location.flow_state().clone(),
            location: location.clone(),
            ir_node: RefCell::new(HydroNode::CycleSource {
                cycle_id,
                metadata: location.new_node_metadata(
                    KeyedStream::<K, V, Tick<L>, Bounded, O, R>::collection_kind(),
                ),
            }),
            _phantom: PhantomData,
        }
    }
}

impl<'a, K, V, L, O: Ordering, R: Retries> ReceiverComplete<'a, TickCycle>
    for KeyedStream<K, V, Tick<L>, Bounded, O, R>
where
    L: Location<'a>,
{
    fn complete(self, cycle_id: CycleId, expected_location: LocationId) {
        assert_eq!(
            Location::id(&self.location),
            expected_location,
            "locations do not match"
        );

        self.location
            .flow_state()
            .borrow_mut()
            .push_root(HydroRoot::CycleSink {
                cycle_id,
                input: Box::new(self.ir_node.replace(HydroNode::Placeholder)),
                op_metadata: HydroIrOpMetadata::new(),
            });
    }
}

impl<'a, K, V, L, B: Boundedness, O: Ordering, R: Retries> CycleCollection<'a, ForwardRef>
    for KeyedStream<K, V, L, B, O, R>
where
    L: Location<'a> + NoTick,
{
    type Location = L;

    fn create_source(cycle_id: CycleId, location: L) -> Self {
        KeyedStream {
            flow_state: location.flow_state().clone(),
            location: location.clone(),
            ir_node: RefCell::new(HydroNode::CycleSource {
                cycle_id,
                metadata: location
                    .new_node_metadata(KeyedStream::<K, V, L, B, O, R>::collection_kind()),
            }),
            _phantom: PhantomData,
        }
    }
}

impl<'a, K, V, L, B: Boundedness, O: Ordering, R: Retries> ReceiverComplete<'a, ForwardRef>
    for KeyedStream<K, V, L, B, O, R>
where
    L: Location<'a> + NoTick,
{
    fn complete(self, cycle_id: CycleId, expected_location: LocationId) {
        assert_eq!(
            Location::id(&self.location),
            expected_location,
            "locations do not match"
        );
        self.location
            .flow_state()
            .borrow_mut()
            .push_root(HydroRoot::CycleSink {
                cycle_id,
                input: Box::new(self.ir_node.replace(HydroNode::Placeholder)),
                op_metadata: HydroIrOpMetadata::new(),
            });
    }
}

impl<'a, K: Clone, V: Clone, Loc: Location<'a>, Bound: Boundedness, Order: Ordering, R: Retries>
    Clone for KeyedStream<K, V, Loc, Bound, Order, R>
{
    fn clone(&self) -> Self {
        if !matches!(self.ir_node.borrow().deref(), HydroNode::Tee { .. }) {
            let orig_ir_node = self.ir_node.replace(HydroNode::Placeholder);
            *self.ir_node.borrow_mut() = HydroNode::Tee {
                inner: SharedNode(Rc::new(RefCell::new(orig_ir_node))),
                metadata: self.location.new_node_metadata(Self::collection_kind()),
            };
        }

        if let HydroNode::Tee { inner, metadata } = self.ir_node.borrow().deref() {
            KeyedStream {
                location: self.location.clone(),
                flow_state: self.flow_state.clone(),
                ir_node: HydroNode::Tee {
                    inner: SharedNode(inner.0.clone()),
                    metadata: metadata.clone(),
                }
                .into(),
                _phantom: PhantomData,
            }
        } else {
            unreachable!()
        }
    }
}

/// The output of a Hydro generator created with [`KeyedStream::generator`], which can yield elements and
/// control the processing of future elements.
pub enum Generate<T> {
    /// Emit the provided element, and keep processing future inputs.
    Yield(T),
    /// Emit the provided element as the _final_ element, do not process future inputs.
    Return(T),
    /// Do not emit anything, but continue processing future inputs.
    Continue,
    /// Do not emit anything, and do not process further inputs.
    Break,
}

impl<'a, K, V, L: Location<'a>, B: Boundedness, O: Ordering, R: Retries>
    KeyedStream<K, V, L, B, O, R>
{
    pub(crate) fn new(location: L, ir_node: HydroNode) -> Self {
        debug_assert_eq!(ir_node.metadata().location_id, Location::id(&location));
        debug_assert_eq!(ir_node.metadata().collection_kind, Self::collection_kind());

        let flow_state = location.flow_state().clone();
        KeyedStream {
            location,
            flow_state,
            ir_node: RefCell::new(ir_node),
            _phantom: PhantomData,
        }
    }

    /// Returns the [`CollectionKind`] corresponding to this type.
    pub fn collection_kind() -> CollectionKind {
        CollectionKind::KeyedStream {
            bound: B::BOUND_KIND,
            value_order: O::ORDERING_KIND,
            value_retry: R::RETRIES_KIND,
            key_type: stageleft::quote_type::<K>().into(),
            value_type: stageleft::quote_type::<V>().into(),
        }
    }

    /// Returns the [`Location`] where this keyed stream is being materialized.
    pub fn location(&self) -> &L {
        &self.location
    }

    /// Explicitly "casts" the keyed stream to a type with a different ordering
    /// guarantee for each group. Useful in unsafe code where the ordering cannot be proven
    /// by the type-system.
    ///
    /// # Non-Determinism
    /// This function is used as an escape hatch, and any mistakes in the
    /// provided ordering guarantee will propagate into the guarantees
    /// for the rest of the program.
    pub fn assume_ordering<O2: Ordering>(self, _nondet: NonDet) -> KeyedStream<K, V, L, B, O2, R> {
        if O::ORDERING_KIND == O2::ORDERING_KIND {
            KeyedStream::new(
                self.location.clone(),
                self.ir_node.replace(HydroNode::Placeholder),
            )
        } else if O2::ORDERING_KIND == StreamOrder::NoOrder {
            // We can always weaken the ordering guarantee
            KeyedStream::new(
                self.location.clone(),
                HydroNode::Cast {
                    inner: Box::new(self.ir_node.replace(HydroNode::Placeholder)),
                    metadata: self
                        .location
                        .new_node_metadata(KeyedStream::<K, V, L, B, O2, R>::collection_kind()),
                },
            )
        } else {
            KeyedStream::new(
                self.location.clone(),
                HydroNode::ObserveNonDet {
                    inner: Box::new(self.ir_node.replace(HydroNode::Placeholder)),
                    trusted: false,
                    metadata: self
                        .location
                        .new_node_metadata(KeyedStream::<K, V, L, B, O2, R>::collection_kind()),
                },
            )
        }
    }

    fn assume_ordering_trusted<O2: Ordering>(
        self,
        _nondet: NonDet,
    ) -> KeyedStream<K, V, L, B, O2, R> {
        if O::ORDERING_KIND == O2::ORDERING_KIND {
            KeyedStream::new(
                self.location.clone(),
                self.ir_node.replace(HydroNode::Placeholder),
            )
        } else if O2::ORDERING_KIND == StreamOrder::NoOrder {
            // We can always weaken the ordering guarantee
            KeyedStream::new(
                self.location.clone(),
                HydroNode::Cast {
                    inner: Box::new(self.ir_node.replace(HydroNode::Placeholder)),
                    metadata: self
                        .location
                        .new_node_metadata(KeyedStream::<K, V, L, B, O2, R>::collection_kind()),
                },
            )
        } else {
            KeyedStream::new(
                self.location.clone(),
                HydroNode::ObserveNonDet {
                    inner: Box::new(self.ir_node.replace(HydroNode::Placeholder)),
                    trusted: true,
                    metadata: self
                        .location
                        .new_node_metadata(KeyedStream::<K, V, L, B, O2, R>::collection_kind()),
                },
            )
        }
    }

    #[deprecated = "use `weaken_ordering::<NoOrder>()` instead"]
    /// Weakens the ordering guarantee provided by the stream to [`NoOrder`],
    /// which is always safe because that is the weakest possible guarantee.
    pub fn weakest_ordering(self) -> KeyedStream<K, V, L, B, NoOrder, R> {
        self.weaken_ordering::<NoOrder>()
    }

    /// Weakens the ordering guarantee provided by the stream to `O2`, with the type-system
    /// enforcing that `O2` is weaker than the input ordering guarantee.
    pub fn weaken_ordering<O2: WeakerOrderingThan<O>>(self) -> KeyedStream<K, V, L, B, O2, R> {
        let nondet = nondet!(/** this is a weaker ordering guarantee, so it is safe to assume */);
        self.assume_ordering::<O2>(nondet)
    }

    /// Explicitly "casts" the keyed stream to a type with a different retries
    /// guarantee for each group. Useful in unsafe code where the lack of retries cannot
    /// be proven by the type-system.
    ///
    /// # Non-Determinism
    /// This function is used as an escape hatch, and any mistakes in the
    /// provided retries guarantee will propagate into the guarantees
    /// for the rest of the program.
    pub fn assume_retries<R2: Retries>(self, _nondet: NonDet) -> KeyedStream<K, V, L, B, O, R2> {
        if R::RETRIES_KIND == R2::RETRIES_KIND {
            KeyedStream::new(
                self.location.clone(),
                self.ir_node.replace(HydroNode::Placeholder),
            )
        } else if R2::RETRIES_KIND == StreamRetry::AtLeastOnce {
            // We can always weaken the retries guarantee
            KeyedStream::new(
                self.location.clone(),
                HydroNode::Cast {
                    inner: Box::new(self.ir_node.replace(HydroNode::Placeholder)),
                    metadata: self
                        .location
                        .new_node_metadata(KeyedStream::<K, V, L, B, O, R2>::collection_kind()),
                },
            )
        } else {
            KeyedStream::new(
                self.location.clone(),
                HydroNode::ObserveNonDet {
                    inner: Box::new(self.ir_node.replace(HydroNode::Placeholder)),
                    trusted: false,
                    metadata: self
                        .location
                        .new_node_metadata(KeyedStream::<K, V, L, B, O, R2>::collection_kind()),
                },
            )
        }
    }

    #[deprecated = "use `weaken_retries::<AtLeastOnce>()` instead"]
    /// Weakens the retries guarantee provided by the stream to [`AtLeastOnce`],
    /// which is always safe because that is the weakest possible guarantee.
    pub fn weakest_retries(self) -> KeyedStream<K, V, L, B, O, AtLeastOnce> {
        self.weaken_retries::<AtLeastOnce>()
    }

    /// Weakens the retries guarantee provided by the stream to `R2`, with the type-system
    /// enforcing that `R2` is weaker than the input retries guarantee.
    pub fn weaken_retries<R2: WeakerRetryThan<R>>(self) -> KeyedStream<K, V, L, B, O, R2> {
        let nondet = nondet!(/** this is a weaker retries guarantee, so it is safe to assume */);
        self.assume_retries::<R2>(nondet)
    }

    /// Strengthens the ordering guarantee to `TotalOrder`, given that `O: IsOrdered`, which
    /// implies that `O == TotalOrder`.
    pub fn make_totally_ordered(self) -> KeyedStream<K, V, L, B, TotalOrder, R>
    where
        O: IsOrdered,
    {
        self.assume_ordering(nondet!(/** no-op */))
    }

    /// Strengthens the retry guarantee to `ExactlyOnce`, given that `R: IsExactlyOnce`, which
    /// implies that `R == ExactlyOnce`.
    pub fn make_exactly_once(self) -> KeyedStream<K, V, L, B, O, ExactlyOnce>
    where
        R: IsExactlyOnce,
    {
        self.assume_retries(nondet!(/** no-op */))
    }

    /// Strengthens the boundedness guarantee to `Bounded`, given that `B: IsBounded`, which
    /// implies that `B == Bounded`.
    pub fn make_bounded(self) -> KeyedStream<K, V, L, Bounded, O, R>
    where
        B: IsBounded,
    {
        KeyedStream::new(
            self.location.clone(),
            self.ir_node.replace(HydroNode::Placeholder),
        )
    }

    /// Flattens the keyed stream into an unordered stream of key-value pairs.
    ///
    /// # Example
    /// ```rust
    /// # #[cfg(feature = "deploy")] {
    /// # use hydro_lang::prelude::*;
    /// # use futures::StreamExt;
    /// # tokio_test::block_on(hydro_lang::test_util::stream_transform_test(|process| {
    /// process
    ///     .source_iter(q!(vec![(1, 2), (1, 3), (2, 4)]))
    ///     .into_keyed()
    ///     .entries()
    /// # }, |mut stream| async move {
    /// // (1, 2), (1, 3), (2, 4) in any order
    /// # let mut results = Vec::new();
    /// # for _ in 0..3 {
    /// #     results.push(stream.next().await.unwrap());
    /// # }
    /// # results.sort();
    /// # assert_eq!(results, vec![(1, 2), (1, 3), (2, 4)]);
    /// # }));
    /// # }
    /// ```
    pub fn entries(self) -> Stream<(K, V), L, B, NoOrder, R> {
        Stream::new(
            self.location.clone(),
            HydroNode::Cast {
                inner: Box::new(self.ir_node.replace(HydroNode::Placeholder)),
                metadata: self
                    .location
                    .new_node_metadata(Stream::<(K, V), L, B, NoOrder, R>::collection_kind()),
            },
        )
    }

    /// Flattens the keyed stream into an unordered stream of only the values.
    ///
    /// # Example
    /// ```rust
    /// # #[cfg(feature = "deploy")] {
    /// # use hydro_lang::prelude::*;
    /// # use futures::StreamExt;
    /// # tokio_test::block_on(hydro_lang::test_util::stream_transform_test(|process| {
    /// process
    ///     .source_iter(q!(vec![(1, 2), (1, 3), (2, 4)]))
    ///     .into_keyed()
    ///     .values()
    /// # }, |mut stream| async move {
    /// // 2, 3, 4 in any order
    /// # let mut results = Vec::new();
    /// # for _ in 0..3 {
    /// #     results.push(stream.next().await.unwrap());
    /// # }
    /// # results.sort();
    /// # assert_eq!(results, vec![2, 3, 4]);
    /// # }));
    /// # }
    /// ```
    pub fn values(self) -> Stream<V, L, B, NoOrder, R> {
        self.entries().map(q!(|(_, v)| v))
    }

    /// Flattens the keyed stream into an unordered stream of just the keys.
    ///
    /// # Example
    /// ```rust
    /// # #[cfg(feature = "deploy")] {
    /// # use hydro_lang::prelude::*;
    /// # use futures::StreamExt;
    /// # tokio_test::block_on(hydro_lang::test_util::stream_transform_test(|process| {
    /// # process
    /// #     .source_iter(q!(vec![(1, 2), (2, 4), (1, 5)]))
    /// #     .into_keyed()
    /// #     .keys()
    /// # }, |mut stream| async move {
    /// // 1, 2 in any order
    /// # let mut results = Vec::new();
    /// # for _ in 0..2 {
    /// #     results.push(stream.next().await.unwrap());
    /// # }
    /// # results.sort();
    /// # assert_eq!(results, vec![1, 2]);
    /// # }));
    /// # }
    /// ```
    pub fn keys(self) -> Stream<K, L, B, NoOrder, ExactlyOnce>
    where
        K: Eq + Hash,
    {
        self.entries().map(q!(|(k, _)| k)).unique()
    }

    /// Transforms each value by invoking `f` on each element, with keys staying the same
    /// after transformation. If you need access to the key, see [`KeyedStream::map_with_key`].
    ///
    /// If you do not want to modify the stream and instead only want to view
    /// each item use [`KeyedStream::inspect`] instead.
    ///
    /// # Example
    /// ```rust
    /// # #[cfg(feature = "deploy")] {
    /// # use hydro_lang::prelude::*;
    /// # use futures::StreamExt;
    /// # tokio_test::block_on(hydro_lang::test_util::stream_transform_test(|process| {
    /// process
    ///     .source_iter(q!(vec![(1, 2), (1, 3), (2, 4)]))
    ///     .into_keyed()
    ///     .map(q!(|v| v + 1))
    /// #   .entries()
    /// # }, |mut stream| async move {
    /// // { 1: [3, 4], 2: [5] }
    /// # let mut results = Vec::new();
    /// # for _ in 0..3 {
    /// #     results.push(stream.next().await.unwrap());
    /// # }
    /// # results.sort();
    /// # assert_eq!(results, vec![(1, 3), (1, 4), (2, 5)]);
    /// # }));
    /// # }
    /// ```
    pub fn map<U, F>(self, f: impl IntoQuotedMut<'a, F, L> + Copy) -> KeyedStream<K, U, L, B, O, R>
    where
        F: Fn(V) -> U + 'a,
    {
        let f: ManualExpr<F, _> = ManualExpr::new(move |ctx: &L| f.splice_fn1_ctx(ctx));
        let map_f = q!({
            let orig = f;
            move |(k, v)| (k, orig(v))
        })
        .splice_fn1_ctx::<(K, V), (K, U)>(&self.location)
        .into();

        KeyedStream::new(
            self.location.clone(),
            HydroNode::Map {
                f: map_f,
                input: Box::new(self.ir_node.replace(HydroNode::Placeholder)),
                metadata: self
                    .location
                    .new_node_metadata(KeyedStream::<K, U, L, B, O, R>::collection_kind()),
            },
        )
    }

    /// Transforms each value by invoking `f` on each key-value pair. The resulting values are **not**
    /// re-grouped even they are tuples; instead they will be grouped under the original key.
    ///
    /// If you do not want to modify the stream and instead only want to view
    /// each item use [`KeyedStream::inspect_with_key`] instead.
    ///
    /// # Example
    /// ```rust
    /// # #[cfg(feature = "deploy")] {
    /// # use hydro_lang::prelude::*;
    /// # use futures::StreamExt;
    /// # tokio_test::block_on(hydro_lang::test_util::stream_transform_test(|process| {
    /// process
    ///     .source_iter(q!(vec![(1, 2), (1, 3), (2, 4)]))
    ///     .into_keyed()
    ///     .map_with_key(q!(|(k, v)| k + v))
    /// #   .entries()
    /// # }, |mut stream| async move {
    /// // { 1: [3, 4], 2: [6] }
    /// # let mut results = Vec::new();
    /// # for _ in 0..3 {
    /// #     results.push(stream.next().await.unwrap());
    /// # }
    /// # results.sort();
    /// # assert_eq!(results, vec![(1, 3), (1, 4), (2, 6)]);
    /// # }));
    /// # }
    /// ```
    pub fn map_with_key<U, F>(
        self,
        f: impl IntoQuotedMut<'a, F, L> + Copy,
    ) -> KeyedStream<K, U, L, B, O, R>
    where
        F: Fn((K, V)) -> U + 'a,
        K: Clone,
    {
        let f: ManualExpr<F, _> = ManualExpr::new(move |ctx: &L| f.splice_fn1_ctx(ctx));
        let map_f = q!({
            let orig = f;
            move |(k, v)| {
                let out = orig((Clone::clone(&k), v));
                (k, out)
            }
        })
        .splice_fn1_ctx::<(K, V), (K, U)>(&self.location)
        .into();

        KeyedStream::new(
            self.location.clone(),
            HydroNode::Map {
                f: map_f,
                input: Box::new(self.ir_node.replace(HydroNode::Placeholder)),
                metadata: self
                    .location
                    .new_node_metadata(KeyedStream::<K, U, L, B, O, R>::collection_kind()),
            },
        )
    }

    /// Prepends a new value to the key of each element in the stream, producing a new
    /// keyed stream with compound keys. Because the original key is preserved, no re-grouping
    /// occurs and the elements in each group preserve their original order.
    ///
    /// # Example
    /// ```rust
    /// # #[cfg(feature = "deploy")] {
    /// # use hydro_lang::prelude::*;
    /// # use futures::StreamExt;
    /// # tokio_test::block_on(hydro_lang::test_util::stream_transform_test(|process| {
    /// process
    ///     .source_iter(q!(vec![(1, 2), (1, 3), (2, 4)]))
    ///     .into_keyed()
    ///     .prefix_key(q!(|&(k, _)| k % 2))
    /// #   .entries()
    /// # }, |mut stream| async move {
    /// // { (1, 1): [2, 3], (0, 2): [4] }
    /// # let mut results = Vec::new();
    /// # for _ in 0..3 {
    /// #     results.push(stream.next().await.unwrap());
    /// # }
    /// # results.sort();
    /// # assert_eq!(results, vec![((0, 2), 4), ((1, 1), 2), ((1, 1), 3)]);
    /// # }));
    /// # }
    /// ```
    pub fn prefix_key<K2, F>(
        self,
        f: impl IntoQuotedMut<'a, F, L> + Copy,
    ) -> KeyedStream<(K2, K), V, L, B, O, R>
    where
        F: Fn(&(K, V)) -> K2 + 'a,
    {
        let f: ManualExpr<F, _> = ManualExpr::new(move |ctx: &L| f.splice_fn1_borrow_ctx(ctx));
        let map_f = q!({
            let orig = f;
            move |kv| {
                let out = orig(&kv);
                ((out, kv.0), kv.1)
            }
        })
        .splice_fn1_ctx::<(K, V), ((K2, K), V)>(&self.location)
        .into();

        KeyedStream::new(
            self.location.clone(),
            HydroNode::Map {
                f: map_f,
                input: Box::new(self.ir_node.replace(HydroNode::Placeholder)),
                metadata: self
                    .location
                    .new_node_metadata(KeyedStream::<(K2, K), V, L, B, O, R>::collection_kind()),
            },
        )
    }

    /// Creates a stream containing only the elements of each group stream that satisfy a predicate
    /// `f`, preserving the order of the elements within the group.
    ///
    /// The closure `f` receives a reference `&V` rather than an owned value `v` because filtering does
    /// not modify or take ownership of the values. If you need to modify the values while filtering
    /// use [`KeyedStream::filter_map`] instead.
    ///
    /// # Example
    /// ```rust
    /// # #[cfg(feature = "deploy")] {
    /// # use hydro_lang::prelude::*;
    /// # use futures::StreamExt;
    /// # tokio_test::block_on(hydro_lang::test_util::stream_transform_test(|process| {
    /// process
    ///     .source_iter(q!(vec![(1, 2), (1, 3), (2, 4)]))
    ///     .into_keyed()
    ///     .filter(q!(|&x| x > 2))
    /// #   .entries()
    /// # }, |mut stream| async move {
    /// // { 1: [3], 2: [4] }
    /// # let mut results = Vec::new();
    /// # for _ in 0..2 {
    /// #     results.push(stream.next().await.unwrap());
    /// # }
    /// # results.sort();
    /// # assert_eq!(results, vec![(1, 3), (2, 4)]);
    /// # }));
    /// # }
    /// ```
    pub fn filter<F>(self, f: impl IntoQuotedMut<'a, F, L> + Copy) -> KeyedStream<K, V, L, B, O, R>
    where
        F: Fn(&V) -> bool + 'a,
    {
        let f: ManualExpr<F, _> = ManualExpr::new(move |ctx: &L| f.splice_fn1_borrow_ctx(ctx));
        let filter_f = q!({
            let orig = f;
            move |t: &(_, _)| orig(&t.1)
        })
        .splice_fn1_borrow_ctx::<(K, V), bool>(&self.location)
        .into();

        KeyedStream::new(
            self.location.clone(),
            HydroNode::Filter {
                f: filter_f,
                input: Box::new(self.ir_node.replace(HydroNode::Placeholder)),
                metadata: self.location.new_node_metadata(Self::collection_kind()),
            },
        )
    }

    /// Creates a stream containing only the elements of each group stream that satisfy a predicate
    /// `f` (which receives the key-value tuple), preserving the order of the elements within the group.
    ///
    /// The closure `f` receives a reference `&(K, V)` rather than an owned value `(K, V)` because filtering does
    /// not modify or take ownership of the values. If you need to modify the values while filtering
    /// use [`KeyedStream::filter_map_with_key`] instead.
    ///
    /// # Example
    /// ```rust
    /// # #[cfg(feature = "deploy")] {
    /// # use hydro_lang::prelude::*;
    /// # use futures::StreamExt;
    /// # tokio_test::block_on(hydro_lang::test_util::stream_transform_test(|process| {
    /// process
    ///     .source_iter(q!(vec![(1, 2), (1, 3), (2, 4)]))
    ///     .into_keyed()
    ///     .filter_with_key(q!(|&(k, v)| v - k == 2))
    /// #   .entries()
    /// # }, |mut stream| async move {
    /// // { 1: [3], 2: [4] }
    /// # let mut results = Vec::new();
    /// # for _ in 0..2 {
    /// #     results.push(stream.next().await.unwrap());
    /// # }
    /// # results.sort();
    /// # assert_eq!(results, vec![(1, 3), (2, 4)]);
    /// # }));
    /// # }
    /// ```
    pub fn filter_with_key<F>(
        self,
        f: impl IntoQuotedMut<'a, F, L> + Copy,
    ) -> KeyedStream<K, V, L, B, O, R>
    where
        F: Fn(&(K, V)) -> bool + 'a,
    {
        let filter_f = f
            .splice_fn1_borrow_ctx::<(K, V), bool>(&self.location)
            .into();

        KeyedStream::new(
            self.location.clone(),
            HydroNode::Filter {
                f: filter_f,
                input: Box::new(self.ir_node.replace(HydroNode::Placeholder)),
                metadata: self.location.new_node_metadata(Self::collection_kind()),
            },
        )
    }

    /// An operator that both filters and maps each value, with keys staying the same.
    /// It yields only the items for which the supplied closure `f` returns `Some(value)`.
    /// If you need access to the key, see [`KeyedStream::filter_map_with_key`].
    ///
    /// # Example
    /// ```rust
    /// # #[cfg(feature = "deploy")] {
    /// # use hydro_lang::prelude::*;
    /// # use futures::StreamExt;
    /// # tokio_test::block_on(hydro_lang::test_util::stream_transform_test(|process| {
    /// process
    ///     .source_iter(q!(vec![(1, "2"), (1, "hello"), (2, "4")]))
    ///     .into_keyed()
    ///     .filter_map(q!(|s| s.parse::<usize>().ok()))
    /// #   .entries()
    /// # }, |mut stream| async move {
    /// // { 1: [2], 2: [4] }
    /// # let mut results = Vec::new();
    /// # for _ in 0..2 {
    /// #     results.push(stream.next().await.unwrap());
    /// # }
    /// # results.sort();
    /// # assert_eq!(results, vec![(1, 2), (2, 4)]);
    /// # }));
    /// # }
    /// ```
    pub fn filter_map<U, F>(
        self,
        f: impl IntoQuotedMut<'a, F, L> + Copy,
    ) -> KeyedStream<K, U, L, B, O, R>
    where
        F: Fn(V) -> Option<U> + 'a,
    {
        let f: ManualExpr<F, _> = ManualExpr::new(move |ctx: &L| f.splice_fn1_ctx(ctx));
        let filter_map_f = q!({
            let orig = f;
            move |(k, v)| orig(v).map(|o| (k, o))
        })
        .splice_fn1_ctx::<(K, V), Option<(K, U)>>(&self.location)
        .into();

        KeyedStream::new(
            self.location.clone(),
            HydroNode::FilterMap {
                f: filter_map_f,
                input: Box::new(self.ir_node.replace(HydroNode::Placeholder)),
                metadata: self
                    .location
                    .new_node_metadata(KeyedStream::<K, U, L, B, O, R>::collection_kind()),
            },
        )
    }

    /// An operator that both filters and maps each key-value pair. The resulting values are **not**
    /// re-grouped even they are tuples; instead they will be grouped under the original key.
    /// It yields only the items for which the supplied closure `f` returns `Some(value)`.
    ///
    /// # Example
    /// ```rust
    /// # #[cfg(feature = "deploy")] {
    /// # use hydro_lang::prelude::*;
    /// # use futures::StreamExt;
    /// # tokio_test::block_on(hydro_lang::test_util::stream_transform_test(|process| {
    /// process
    ///     .source_iter(q!(vec![(1, "2"), (1, "hello"), (2, "2")]))
    ///     .into_keyed()
    ///     .filter_map_with_key(q!(|(k, s)| s.parse::<usize>().ok().filter(|v| v == &k)))
    /// #   .entries()
    /// # }, |mut stream| async move {
    /// // { 2: [2] }
    /// # let mut results = Vec::new();
    /// # for _ in 0..1 {
    /// #     results.push(stream.next().await.unwrap());
    /// # }
    /// # results.sort();
    /// # assert_eq!(results, vec![(2, 2)]);
    /// # }));
    /// # }
    /// ```
    pub fn filter_map_with_key<U, F>(
        self,
        f: impl IntoQuotedMut<'a, F, L> + Copy,
    ) -> KeyedStream<K, U, L, B, O, R>
    where
        F: Fn((K, V)) -> Option<U> + 'a,
        K: Clone,
    {
        let f: ManualExpr<F, _> = ManualExpr::new(move |ctx: &L| f.splice_fn1_ctx(ctx));
        let filter_map_f = q!({
            let orig = f;
            move |(k, v)| {
                let out = orig((Clone::clone(&k), v));
                out.map(|o| (k, o))
            }
        })
        .splice_fn1_ctx::<(K, V), Option<(K, U)>>(&self.location)
        .into();

        KeyedStream::new(
            self.location.clone(),
            HydroNode::FilterMap {
                f: filter_map_f,
                input: Box::new(self.ir_node.replace(HydroNode::Placeholder)),
                metadata: self
                    .location
                    .new_node_metadata(KeyedStream::<K, U, L, B, O, R>::collection_kind()),
            },
        )
    }

    /// Generates a keyed stream that maps each value `v` to a tuple `(v, x)`,
    /// where `v` is the value of `other`, a bounded [`super::singleton::Singleton`] or
    /// [`Optional`]. If `other` is an empty [`Optional`], no values will be produced.
    ///
    /// # Example
    /// ```rust
    /// # #[cfg(feature = "deploy")] {
    /// # use hydro_lang::prelude::*;
    /// # use futures::StreamExt;
    /// # tokio_test::block_on(hydro_lang::test_util::stream_transform_test(|process| {
    /// let tick = process.tick();
    /// let batch = process
    ///   .source_iter(q!(vec![(1, 123), (1, 456), (2, 123)]))
    ///   .into_keyed()
    ///   .batch(&tick, nondet!(/** test */));
    /// let count = batch.clone().entries().count(); // `count()` returns a singleton
    /// batch.cross_singleton(count).all_ticks().entries()
    /// # }, |mut stream| async move {
    /// // { 1: [(123, 3), (456, 3)], 2: [(123, 3)] }
    /// # let mut results = Vec::new();
    /// # for _ in 0..3 {
    /// #     results.push(stream.next().await.unwrap());
    /// # }
    /// # results.sort();
    /// # assert_eq!(results, vec![(1, (123, 3)), (1, (456, 3)), (2, (123, 3))]);
    /// # }));
    /// # }
    /// ```
    pub fn cross_singleton<O2>(
        self,
        other: impl Into<Optional<O2, L, Bounded>>,
    ) -> KeyedStream<K, (V, O2), L, B, O, R>
    where
        O2: Clone,
    {
        let other: Optional<O2, L, Bounded> = other.into();
        check_matching_location(&self.location, &other.location);

        Stream::new(
            self.location.clone(),
            HydroNode::CrossSingleton {
                left: Box::new(self.ir_node.replace(HydroNode::Placeholder)),
                right: Box::new(other.ir_node.replace(HydroNode::Placeholder)),
                metadata: self
                    .location
                    .new_node_metadata(Stream::<((K, V), O2), L, B, O, R>::collection_kind()),
            },
        )
        .map(q!(|((k, v), o2)| (k, (v, o2))))
        .into_keyed()
    }

    /// For each value `v` in each group, transform `v` using `f` and then treat the
    /// result as an [`Iterator`] to produce values one by one within the same group.
    /// The implementation for [`Iterator`] for the output type `I` must produce items
    /// in a **deterministic** order.
    ///
    /// For example, `I` could be a `Vec`, but not a `HashSet`. If the order of the items in `I` is
    /// not deterministic, use [`KeyedStream::flat_map_unordered`] instead.
    ///
    /// # Example
    /// ```rust
    /// # #[cfg(feature = "deploy")] {
    /// # use hydro_lang::prelude::*;
    /// # use futures::StreamExt;
    /// # tokio_test::block_on(hydro_lang::test_util::stream_transform_test(|process| {
    /// process
    ///     .source_iter(q!(vec![(1, vec![2, 3]), (1, vec![4]), (2, vec![5, 6])]))
    ///     .into_keyed()
    ///     .flat_map_ordered(q!(|x| x))
    /// #   .entries()
    /// # }, |mut stream| async move {
    /// // { 1: [2, 3, 4], 2: [5, 6] }
    /// # let mut results = Vec::new();
    /// # for _ in 0..5 {
    /// #     results.push(stream.next().await.unwrap());
    /// # }
    /// # results.sort();
    /// # assert_eq!(results, vec![(1, 2), (1, 3), (1, 4), (2, 5), (2, 6)]);
    /// # }));
    /// # }
    /// ```
    pub fn flat_map_ordered<U, I, F>(
        self,
        f: impl IntoQuotedMut<'a, F, L> + Copy,
    ) -> KeyedStream<K, U, L, B, O, R>
    where
        I: IntoIterator<Item = U>,
        F: Fn(V) -> I + 'a,
        K: Clone,
    {
        let f: ManualExpr<F, _> = ManualExpr::new(move |ctx: &L| f.splice_fn1_ctx(ctx));
        let flat_map_f = q!({
            let orig = f;
            move |(k, v)| orig(v).into_iter().map(move |u| (Clone::clone(&k), u))
        })
        .splice_fn1_ctx::<(K, V), _>(&self.location)
        .into();

        KeyedStream::new(
            self.location.clone(),
            HydroNode::FlatMap {
                f: flat_map_f,
                input: Box::new(self.ir_node.replace(HydroNode::Placeholder)),
                metadata: self
                    .location
                    .new_node_metadata(KeyedStream::<K, U, L, B, O, R>::collection_kind()),
            },
        )
    }

    /// Like [`KeyedStream::flat_map_ordered`], but allows the implementation of [`Iterator`]
    /// for the output type `I` to produce items in any order.
    ///
    /// # Example
    /// ```rust
    /// # #[cfg(feature = "deploy")] {
    /// # use hydro_lang::{prelude::*, live_collections::stream::{NoOrder, ExactlyOnce}};
    /// # use futures::StreamExt;
    /// # tokio_test::block_on(hydro_lang::test_util::stream_transform_test::<_, _, _, NoOrder, ExactlyOnce>(|process| {
    /// process
    ///     .source_iter(q!(vec![
    ///         (1, std::collections::HashSet::<i32>::from_iter(vec![2, 3])),
    ///         (2, std::collections::HashSet::from_iter(vec![4, 5]))
    ///     ]))
    ///     .into_keyed()
    ///     .flat_map_unordered(q!(|x| x))
    /// #   .entries()
    /// # }, |mut stream| async move {
    /// // { 1: [2, 3], 2: [4, 5] } with values in each group in unknown order
    /// # let mut results = Vec::new();
    /// # for _ in 0..4 {
    /// #     results.push(stream.next().await.unwrap());
    /// # }
    /// # results.sort();
    /// # assert_eq!(results, vec![(1, 2), (1, 3), (2, 4), (2, 5)]);
    /// # }));
    /// # }
    /// ```
    pub fn flat_map_unordered<U, I, F>(
        self,
        f: impl IntoQuotedMut<'a, F, L> + Copy,
    ) -> KeyedStream<K, U, L, B, NoOrder, R>
    where
        I: IntoIterator<Item = U>,
        F: Fn(V) -> I + 'a,
        K: Clone,
    {
        let f: ManualExpr<F, _> = ManualExpr::new(move |ctx: &L| f.splice_fn1_ctx(ctx));
        let flat_map_f = q!({
            let orig = f;
            move |(k, v)| orig(v).into_iter().map(move |u| (Clone::clone(&k), u))
        })
        .splice_fn1_ctx::<(K, V), _>(&self.location)
        .into();

        KeyedStream::new(
            self.location.clone(),
            HydroNode::FlatMap {
                f: flat_map_f,
                input: Box::new(self.ir_node.replace(HydroNode::Placeholder)),
                metadata: self
                    .location
                    .new_node_metadata(KeyedStream::<K, U, L, B, NoOrder, R>::collection_kind()),
            },
        )
    }

    /// For each value `v` in each group, treat `v` as an [`Iterator`] and produce its items one by one
    /// within the same group. The implementation for [`Iterator`] for the value type `V` must produce
    /// items in a **deterministic** order.
    ///
    /// For example, `V` could be a `Vec`, but not a `HashSet`. If the order of the items in `V` is
    /// not deterministic, use [`KeyedStream::flatten_unordered`] instead.
    ///
    /// # Example
    /// ```rust
    /// # #[cfg(feature = "deploy")] {
    /// # use hydro_lang::prelude::*;
    /// # use futures::StreamExt;
    /// # tokio_test::block_on(hydro_lang::test_util::stream_transform_test(|process| {
    /// process
    ///     .source_iter(q!(vec![(1, vec![2, 3]), (1, vec![4]), (2, vec![5, 6])]))
    ///     .into_keyed()
    ///     .flatten_ordered()
    /// #   .entries()
    /// # }, |mut stream| async move {
    /// // { 1: [2, 3, 4], 2: [5, 6] }
    /// # let mut results = Vec::new();
    /// # for _ in 0..5 {
    /// #     results.push(stream.next().await.unwrap());
    /// # }
    /// # results.sort();
    /// # assert_eq!(results, vec![(1, 2), (1, 3), (1, 4), (2, 5), (2, 6)]);
    /// # }));
    /// # }
    /// ```
    pub fn flatten_ordered<U>(self) -> KeyedStream<K, U, L, B, O, R>
    where
        V: IntoIterator<Item = U>,
        K: Clone,
    {
        self.flat_map_ordered(q!(|d| d))
    }

    /// Like [`KeyedStream::flatten_ordered`], but allows the implementation of [`Iterator`]
    /// for the value type `V` to produce items in any order.
    ///
    /// # Example
    /// ```rust
    /// # #[cfg(feature = "deploy")] {
    /// # use hydro_lang::{prelude::*, live_collections::stream::{NoOrder, ExactlyOnce}};
    /// # use futures::StreamExt;
    /// # tokio_test::block_on(hydro_lang::test_util::stream_transform_test::<_, _, _, NoOrder, ExactlyOnce>(|process| {
    /// process
    ///     .source_iter(q!(vec![
    ///         (1, std::collections::HashSet::<i32>::from_iter(vec![2, 3])),
    ///         (2, std::collections::HashSet::from_iter(vec![4, 5]))
    ///     ]))
    ///     .into_keyed()
    ///     .flatten_unordered()
    /// #   .entries()
    /// # }, |mut stream| async move {
    /// // { 1: [2, 3], 2: [4, 5] } with values in each group in unknown order
    /// # let mut results = Vec::new();
    /// # for _ in 0..4 {
    /// #     results.push(stream.next().await.unwrap());
    /// # }
    /// # results.sort();
    /// # assert_eq!(results, vec![(1, 2), (1, 3), (2, 4), (2, 5)]);
    /// # }));
    /// # }
    /// ```
    pub fn flatten_unordered<U>(self) -> KeyedStream<K, U, L, B, NoOrder, R>
    where
        V: IntoIterator<Item = U>,
        K: Clone,
    {
        self.flat_map_unordered(q!(|d| d))
    }

    /// An operator which allows you to "inspect" each element of a stream without
    /// modifying it. The closure `f` is called on a reference to each value. This is
    /// mainly useful for debugging, and should not be used to generate side-effects.
    ///
    /// # Example
    /// ```rust
    /// # #[cfg(feature = "deploy")] {
    /// # use hydro_lang::prelude::*;
    /// # use futures::StreamExt;
    /// # tokio_test::block_on(hydro_lang::test_util::stream_transform_test(|process| {
    /// process
    ///     .source_iter(q!(vec![(1, 2), (1, 3), (2, 4)]))
    ///     .into_keyed()
    ///     .inspect(q!(|v| println!("{}", v)))
    /// #   .entries()
    /// # }, |mut stream| async move {
    /// # let mut results = Vec::new();
    /// # for _ in 0..3 {
    /// #     results.push(stream.next().await.unwrap());
    /// # }
    /// # results.sort();
    /// # assert_eq!(results, vec![(1, 2), (1, 3), (2, 4)]);
    /// # }));
    /// # }
    /// ```
    pub fn inspect<F>(self, f: impl IntoQuotedMut<'a, F, L> + Copy) -> Self
    where
        F: Fn(&V) + 'a,
    {
        let f: ManualExpr<F, _> = ManualExpr::new(move |ctx: &L| f.splice_fn1_borrow_ctx(ctx));
        let inspect_f = q!({
            let orig = f;
            move |t: &(_, _)| orig(&t.1)
        })
        .splice_fn1_borrow_ctx::<(K, V), ()>(&self.location)
        .into();

        KeyedStream::new(
            self.location.clone(),
            HydroNode::Inspect {
                f: inspect_f,
                input: Box::new(self.ir_node.replace(HydroNode::Placeholder)),
                metadata: self.location.new_node_metadata(Self::collection_kind()),
            },
        )
    }

    /// An operator which allows you to "inspect" each element of a stream without
    /// modifying it. The closure `f` is called on a reference to each key-value pair. This is
    /// mainly useful for debugging, and should not be used to generate side-effects.
    ///
    /// # Example
    /// ```rust
    /// # #[cfg(feature = "deploy")] {
    /// # use hydro_lang::prelude::*;
    /// # use futures::StreamExt;
    /// # tokio_test::block_on(hydro_lang::test_util::stream_transform_test(|process| {
    /// process
    ///     .source_iter(q!(vec![(1, 2), (1, 3), (2, 4)]))
    ///     .into_keyed()
    ///     .inspect_with_key(q!(|(k, v)| println!("{}: {}", k, v)))
    /// #   .entries()
    /// # }, |mut stream| async move {
    /// # let mut results = Vec::new();
    /// # for _ in 0..3 {
    /// #     results.push(stream.next().await.unwrap());
    /// # }
    /// # results.sort();
    /// # assert_eq!(results, vec![(1, 2), (1, 3), (2, 4)]);
    /// # }));
    /// # }
    /// ```
    pub fn inspect_with_key<F>(self, f: impl IntoQuotedMut<'a, F, L>) -> Self
    where
        F: Fn(&(K, V)) + 'a,
    {
        let inspect_f = f.splice_fn1_borrow_ctx::<(K, V), ()>(&self.location).into();

        KeyedStream::new(
            self.location.clone(),
            HydroNode::Inspect {
                f: inspect_f,
                input: Box::new(self.ir_node.replace(HydroNode::Placeholder)),
                metadata: self.location.new_node_metadata(Self::collection_kind()),
            },
        )
    }

    /// An operator which allows you to "name" a `HydroNode`.
    /// This is only used for testing, to correlate certain `HydroNode`s with IDs.
    pub fn ir_node_named(self, name: &str) -> KeyedStream<K, V, L, B, O, R> {
        {
            let mut node = self.ir_node.borrow_mut();
            let metadata = node.metadata_mut();
            metadata.tag = Some(name.to_owned());
        }
        self
    }

    /// A special case of [`Stream::scan`] for keyed streams. For each key group the values are transformed via the `f` combinator.
    ///
    /// Unlike [`KeyedStream::fold`] which only returns the final accumulated value, `scan` produces a new stream
    /// containing all intermediate accumulated values paired with the key. The scan operation can also terminate
    /// early by returning `None`.
    ///
    /// The function takes a mutable reference to the accumulator and the current element, and returns
    /// an `Option<U>`. If the function returns `Some(value)`, `value` is emitted to the output stream.
    /// If the function returns `None`, the stream is terminated and no more elements are processed.
    ///
    /// # Example
    /// ```rust
    /// # #[cfg(feature = "deploy")] {
    /// # use hydro_lang::prelude::*;
    /// # use futures::StreamExt;
    /// # tokio_test::block_on(hydro_lang::test_util::stream_transform_test(|process| {
    /// process
    ///     .source_iter(q!(vec![(0, 1), (0, 3), (1, 3), (1, 4)]))
    ///     .into_keyed()
    ///     .scan(
    ///         q!(|| 0),
    ///         q!(|acc, x| {
    ///             *acc += x;
    ///             if *acc % 2 == 0 { None } else { Some(*acc) }
    ///         }),
    ///     )
    /// #   .entries()
    /// # }, |mut stream| async move {
    /// // Output: { 0: [1], 1: [3, 7] }
    /// # let mut results = Vec::new();
    /// # for _ in 0..3 {
    /// #     results.push(stream.next().await.unwrap());
    /// # }
    /// # results.sort();
    /// # assert_eq!(results, vec![(0, 1), (1, 3), (1, 7)]);
    /// # }));
    /// # }
    /// ```
    pub fn scan<A, U, I, F>(
        self,
        init: impl IntoQuotedMut<'a, I, L> + Copy,
        f: impl IntoQuotedMut<'a, F, L> + Copy,
    ) -> KeyedStream<K, U, L, B, TotalOrder, ExactlyOnce>
    where
        O: IsOrdered,
        R: IsExactlyOnce,
        K: Clone + Eq + Hash,
        I: Fn() -> A + 'a,
        F: Fn(&mut A, V) -> Option<U> + 'a,
    {
        let f: ManualExpr<F, _> = ManualExpr::new(move |ctx: &L| f.splice_fn2_borrow_mut_ctx(ctx));
        self.make_totally_ordered().make_exactly_once().generator(
            init,
            q!({
                let orig = f;
                move |state, v| {
                    if let Some(out) = orig(state, v) {
                        Generate::Yield(out)
                    } else {
                        Generate::Break
                    }
                }
            }),
        )
    }

    /// Iteratively processes the elements in each group using a state machine that can yield
    /// elements as it processes its inputs. This is designed to mirror the unstable generator
    /// syntax in Rust, without requiring special syntax.
    ///
    /// Like [`KeyedStream::scan`], this function takes in an initializer that emits the initial
    /// state for each group. The second argument defines the processing logic, taking in a
    /// mutable reference to the group's state and the value to be processed. It emits a
    /// [`Generate`] value, whose variants define what is emitted and whether further inputs
    /// should be processed.
    ///
    /// # Example
    /// ```rust
    /// # #[cfg(feature = "deploy")] {
    /// # use hydro_lang::prelude::*;
    /// # use futures::StreamExt;
    /// # tokio_test::block_on(hydro_lang::test_util::stream_transform_test(|process| {
    /// process
    ///     .source_iter(q!(vec![(0, 1), (0, 3), (0, 100), (0, 10), (1, 3), (1, 4), (1, 3)]))
    ///     .into_keyed()
    ///     .generator(
    ///         q!(|| 0),
    ///         q!(|acc, x| {
    ///             *acc += x;
    ///             if *acc > 100 {
    ///                 hydro_lang::live_collections::keyed_stream::Generate::Return(
    ///                     "done!".to_owned()
    ///                 )
    ///             } else if *acc % 2 == 0 {
    ///                 hydro_lang::live_collections::keyed_stream::Generate::Yield(
    ///                     "even".to_owned()
    ///                 )
    ///             } else {
    ///                 hydro_lang::live_collections::keyed_stream::Generate::Continue
    ///             }
    ///         }),
    ///     )
    /// #   .entries()
    /// # }, |mut stream| async move {
    /// // Output: { 0: ["even", "done!"], 1: ["even"] }
    /// # let mut results = Vec::new();
    /// # for _ in 0..3 {
    /// #     results.push(stream.next().await.unwrap());
    /// # }
    /// # results.sort();
    /// # assert_eq!(results, vec![(0, "done!".to_owned()), (0, "even".to_owned()), (1, "even".to_owned())]);
    /// # }));
    /// # }
    /// ```
    pub fn generator<A, U, I, F>(
        self,
        init: impl IntoQuotedMut<'a, I, L> + Copy,
        f: impl IntoQuotedMut<'a, F, L> + Copy,
    ) -> KeyedStream<K, U, L, B, TotalOrder, ExactlyOnce>
    where
        O: IsOrdered,
        R: IsExactlyOnce,
        K: Clone + Eq + Hash,
        I: Fn() -> A + 'a,
        F: Fn(&mut A, V) -> Generate<U> + 'a,
    {
        let init: ManualExpr<I, _> = ManualExpr::new(move |ctx: &L| init.splice_fn0_ctx(ctx));
        let f: ManualExpr<F, _> = ManualExpr::new(move |ctx: &L| f.splice_fn2_borrow_mut_ctx(ctx));

        let this = self.make_totally_ordered().make_exactly_once();

        let scan_init = q!(|| HashMap::new())
            .splice_fn0_ctx::<HashMap<K, Option<A>>>(&this.location)
            .into();
        let scan_f = q!(move |acc: &mut HashMap<_, _>, (k, v)| {
            let existing_state = acc.entry(Clone::clone(&k)).or_insert_with(|| Some(init()));
            if let Some(existing_state_value) = existing_state {
                match f(existing_state_value, v) {
                    Generate::Yield(out) => Some(Some((k, out))),
                    Generate::Return(out) => {
                        let _ = existing_state.take(); // TODO(shadaj): garbage collect with termination markers
                        Some(Some((k, out)))
                    }
                    Generate::Break => {
                        let _ = existing_state.take(); // TODO(shadaj): garbage collect with termination markers
                        Some(None)
                    }
                    Generate::Continue => Some(None),
                }
            } else {
                Some(None)
            }
        })
        .splice_fn2_borrow_mut_ctx::<HashMap<K, Option<A>>, (K, V), _>(&this.location)
        .into();

        let scan_node = HydroNode::Scan {
            init: scan_init,
            acc: scan_f,
            input: Box::new(this.ir_node.replace(HydroNode::Placeholder)),
            metadata: this.location.new_node_metadata(Stream::<
                Option<(K, U)>,
                L,
                B,
                TotalOrder,
                ExactlyOnce,
            >::collection_kind()),
        };

        let flatten_f = q!(|d| d)
            .splice_fn1_ctx::<Option<(K, U)>, _>(&this.location)
            .into();
        let flatten_node = HydroNode::FlatMap {
            f: flatten_f,
            input: Box::new(scan_node),
            metadata: this.location.new_node_metadata(KeyedStream::<
                K,
                U,
                L,
                B,
                TotalOrder,
                ExactlyOnce,
            >::collection_kind()),
        };

        KeyedStream::new(this.location.clone(), flatten_node)
    }

    /// A variant of [`Stream::fold`], intended for keyed streams. The aggregation is executed
    /// in-order across the values in each group. But the aggregation function returns a boolean,
    /// which when true indicates that the aggregated result is complete and can be released to
    /// downstream computation. Unlike [`KeyedStream::fold`], this means that even if the input
    /// stream is [`super::boundedness::Unbounded`], the outputs of the fold can be processed like
    /// normal stream elements.
    ///
    /// # Example
    /// ```rust
    /// # #[cfg(feature = "deploy")] {
    /// # use hydro_lang::prelude::*;
    /// # use futures::StreamExt;
    /// # tokio_test::block_on(hydro_lang::test_util::stream_transform_test(|process| {
    /// process
    ///     .source_iter(q!(vec![(0, 2), (0, 3), (1, 3), (1, 6)]))
    ///     .into_keyed()
    ///     .fold_early_stop(
    ///         q!(|| 0),
    ///         q!(|acc, x| {
    ///             *acc += x;
    ///             x % 2 == 0
    ///         }),
    ///     )
    /// #   .entries()
    /// # }, |mut stream| async move {
    /// // Output: { 0: 2, 1: 9 }
    /// # let mut results = Vec::new();
    /// # for _ in 0..2 {
    /// #     results.push(stream.next().await.unwrap());
    /// # }
    /// # results.sort();
    /// # assert_eq!(results, vec![(0, 2), (1, 9)]);
    /// # }));
    /// # }
    /// ```
    pub fn fold_early_stop<A, I, F>(
        self,
        init: impl IntoQuotedMut<'a, I, L> + Copy,
        f: impl IntoQuotedMut<'a, F, L> + Copy,
    ) -> KeyedSingleton<K, A, L, B::WhenValueBounded>
    where
        O: IsOrdered,
        R: IsExactlyOnce,
        K: Clone + Eq + Hash,
        I: Fn() -> A + 'a,
        F: Fn(&mut A, V) -> bool + 'a,
    {
        let init: ManualExpr<I, _> = ManualExpr::new(move |ctx: &L| init.splice_fn0_ctx(ctx));
        let f: ManualExpr<F, _> = ManualExpr::new(move |ctx: &L| f.splice_fn2_borrow_mut_ctx(ctx));
        let out_without_bound_cast = self.generator(
            q!(move || Some(init())),
            q!(move |key_state, v| {
                if let Some(key_state_value) = key_state.as_mut() {
                    if f(key_state_value, v) {
                        Generate::Return(key_state.take().unwrap())
                    } else {
                        Generate::Continue
                    }
                } else {
                    unreachable!()
                }
            }),
        );

        KeyedSingleton::new(
            out_without_bound_cast.location.clone(),
            HydroNode::Cast {
                inner: Box::new(
                    out_without_bound_cast
                        .ir_node
                        .replace(HydroNode::Placeholder),
                ),
                metadata: out_without_bound_cast
                    .location
                    .new_node_metadata(
                        KeyedSingleton::<K, A, L, B::WhenValueBounded>::collection_kind(),
                    ),
            },
        )
    }

    /// Gets the first element inside each group of values as a [`KeyedSingleton`] that preserves
    /// the original group keys. Requires the input stream to have [`TotalOrder`] guarantees,
    /// otherwise the first element would be non-deterministic.
    ///
    /// # Example
    /// ```rust
    /// # #[cfg(feature = "deploy")] {
    /// # use hydro_lang::prelude::*;
    /// # use futures::StreamExt;
    /// # tokio_test::block_on(hydro_lang::test_util::stream_transform_test(|process| {
    /// process
    ///     .source_iter(q!(vec![(0, 2), (0, 3), (1, 3), (1, 6)]))
    ///     .into_keyed()
    ///     .first()
    /// #   .entries()
    /// # }, |mut stream| async move {
    /// // Output: { 0: 2, 1: 3 }
    /// # let mut results = Vec::new();
    /// # for _ in 0..2 {
    /// #     results.push(stream.next().await.unwrap());
    /// # }
    /// # results.sort();
    /// # assert_eq!(results, vec![(0, 2), (1, 3)]);
    /// # }));
    /// # }
    /// ```
    pub fn first(self) -> KeyedSingleton<K, V, L, B::WhenValueBounded>
    where
        O: IsOrdered,
        R: IsExactlyOnce,
        K: Clone + Eq + Hash,
    {
        self.fold_early_stop(
            q!(|| None),
            q!(|acc, v| {
                *acc = Some(v);
                true
            }),
        )
        .map(q!(|v| v.unwrap()))
    }

    /// Assigns a zero-based index to each value within each key group, emitting
    /// `(K, (index, V))` tuples with per-key sequential indices.
    ///
    /// The output keyed stream has [`TotalOrder`] and [`ExactlyOnce`] guarantees.
    /// This is a streaming operator that processes elements as they arrive.
    ///
    /// # Example
    /// ```rust
    /// # #[cfg(feature = "deploy")] {
    /// # use hydro_lang::prelude::*;
    /// # use futures::StreamExt;
    /// # tokio_test::block_on(hydro_lang::test_util::stream_transform_test(|process| {
    /// process
    ///     .source_iter(q!(vec![(1, 10), (2, 20), (1, 30)]))
    ///     .into_keyed()
    ///     .enumerate()
    /// # .entries()
    /// # }, |mut stream| async move {
    /// // per-key indices: { 1: [(0, 10), (1, 30)], 2: [(0, 20)] }
    /// # let mut results = Vec::new();
    /// # for _ in 0..3 {
    /// #     results.push(stream.next().await.unwrap());
    /// # }
    /// # let key1: Vec<_> = results.iter().filter(|(k, _)| *k == 1).map(|(_, v)| *v).collect();
    /// # let key2: Vec<_> = results.iter().filter(|(k, _)| *k == 2).map(|(_, v)| *v).collect();
    /// # assert_eq!(key1, vec![(0, 10), (1, 30)]);
    /// # assert_eq!(key2, vec![(0, 20)]);
    /// # }));
    /// # }
    /// ```
    pub fn enumerate(self) -> KeyedStream<K, (usize, V), L, B, TotalOrder, ExactlyOnce>
    where
        O: IsOrdered,
        R: IsExactlyOnce,
        K: Eq + Hash + Clone,
    {
        self.scan(
            q!(|| 0),
            q!(|acc, next| {
                let curr = *acc;
                *acc += 1;
                Some((curr, next))
            }),
        )
    }

    /// Counts the number of elements in each group, producing a [`KeyedSingleton`] with the counts.
    ///
    /// # Example
    /// ```rust
    /// # #[cfg(feature = "deploy")] {
    /// # use hydro_lang::prelude::*;
    /// # use futures::StreamExt;
    /// # tokio_test::block_on(hydro_lang::test_util::stream_transform_test(|process| {
    /// let tick = process.tick();
    /// let numbers = process
    ///     .source_iter(q!(vec![(1, 2), (2, 3), (1, 3), (2, 4), (1, 5)]))
    ///     .into_keyed();
    /// let batch = numbers.batch(&tick, nondet!(/** test */));
    /// batch
    ///     .value_counts()
    ///     .entries()
    ///     .all_ticks()
    /// # }, |mut stream| async move {
    /// // (1, 3), (2, 2)
    /// # let mut results = Vec::new();
    /// # for _ in 0..2 {
    /// #     results.push(stream.next().await.unwrap());
    /// # }
    /// # results.sort();
    /// # assert_eq!(results, vec![(1, 3), (2, 2)]);
    /// # }));
    /// # }
    /// ```
    pub fn value_counts(self) -> KeyedSingleton<K, usize, L, B::WhenValueUnbounded>
    where
        R: IsExactlyOnce,
        K: Eq + Hash,
    {
        self.make_exactly_once()
            .assume_ordering_trusted(
                nondet!(/** ordering within each group affects neither result nor intermediates */),
            )
            .fold(q!(|| 0), q!(|acc, _| *acc += 1))
    }

    /// Like [`Stream::fold`] but in the spirit of SQL `GROUP BY`, aggregates the values in each
    /// group via the `comb` closure.
    ///
    /// Depending on the input stream guarantees, the closure may need to be commutative
    /// (for unordered streams) or idempotent (for streams with non-deterministic duplicates).
    ///
    /// If the input and output value types are the same and do not require initialization then use
    /// [`KeyedStream::reduce`].
    ///
    /// # Example
    /// ```rust
    /// # #[cfg(feature = "deploy")] {
    /// # use hydro_lang::prelude::*;
    /// # use futures::StreamExt;
    /// # tokio_test::block_on(hydro_lang::test_util::stream_transform_test(|process| {
    /// let tick = process.tick();
    /// let numbers = process
    ///     .source_iter(q!(vec![(1, false), (2, true), (1, false), (2, false)]))
    ///     .into_keyed();
    /// let batch = numbers.batch(&tick, nondet!(/** test */));
    /// batch
    ///     .fold(q!(|| false), q!(|acc, x| *acc |= x))
    ///     .entries()
    ///     .all_ticks()
    /// # }, |mut stream| async move {
    /// // (1, false), (2, true)
    /// # let mut results = Vec::new();
    /// # for _ in 0..2 {
    /// #     results.push(stream.next().await.unwrap());
    /// # }
    /// # results.sort();
    /// # assert_eq!(results, vec![(1, false), (2, true)]);
    /// # }));
    /// # }
    /// ```
    pub fn fold<A, I: Fn() -> A + 'a, F: Fn(&mut A, V), C, Idemp>(
        self,
        init: impl IntoQuotedMut<'a, I, L>,
        comb: impl IntoQuotedMut<'a, F, L, AggFuncAlgebra<C, Idemp>>,
    ) -> KeyedSingleton<K, A, L, B::WhenValueUnbounded>
    where
        K: Eq + Hash,
        C: ValidCommutativityFor<O>,
        Idemp: ValidIdempotenceFor<R>,
    {
        let init = init.splice_fn0_ctx(&self.location).into();
        let (comb, proof) = comb.splice_fn2_borrow_mut_ctx_props(&self.location);
        proof.register_proof(&comb);

        let ordered = self
            .assume_retries::<ExactlyOnce>(nondet!(/** the combinator function is idempotent */))
            .assume_ordering::<TotalOrder>(nondet!(/** the combinator function is commutative */));

        KeyedSingleton::new(
            ordered.location.clone(),
            HydroNode::FoldKeyed {
                init,
                acc: comb.into(),
                input: Box::new(ordered.ir_node.replace(HydroNode::Placeholder)),
                metadata: ordered.location.new_node_metadata(KeyedSingleton::<
                    K,
                    A,
                    L,
                    B::WhenValueUnbounded,
                >::collection_kind()),
            },
        )
    }

    /// Like [`Stream::reduce`] but in the spirit of SQL `GROUP BY`, aggregates the values in each
    /// group via the `comb` closure.
    ///
    /// Depending on the input stream guarantees, the closure may need to be commutative
    /// (for unordered streams) or idempotent (for streams with non-deterministic duplicates).
    ///
    /// If you need the accumulated value to have a different type than the input, use [`KeyedStream::fold`].
    ///
    /// # Example
    /// ```rust
    /// # #[cfg(feature = "deploy")] {
    /// # use hydro_lang::prelude::*;
    /// # use futures::StreamExt;
    /// # tokio_test::block_on(hydro_lang::test_util::stream_transform_test(|process| {
    /// let tick = process.tick();
    /// let numbers = process
    ///     .source_iter(q!(vec![(1, false), (2, true), (1, false), (2, false)]))
    ///     .into_keyed();
    /// let batch = numbers.batch(&tick, nondet!(/** test */));
    /// batch
    ///     .reduce(q!(|acc, x| *acc |= x))
    ///     .entries()
    ///     .all_ticks()
    /// # }, |mut stream| async move {
    /// // (1, false), (2, true)
    /// # let mut results = Vec::new();
    /// # for _ in 0..2 {
    /// #     results.push(stream.next().await.unwrap());
    /// # }
    /// # results.sort();
    /// # assert_eq!(results, vec![(1, false), (2, true)]);
    /// # }));
    /// # }
    /// ```
    pub fn reduce<F: Fn(&mut V, V) + 'a, C, Idemp>(
        self,
        comb: impl IntoQuotedMut<'a, F, L, AggFuncAlgebra<C, Idemp>>,
    ) -> KeyedSingleton<K, V, L, B::WhenValueUnbounded>
    where
        K: Eq + Hash,
        C: ValidCommutativityFor<O>,
        Idemp: ValidIdempotenceFor<R>,
    {
        let (f, proof) = comb.splice_fn2_borrow_mut_ctx_props(&self.location);
        proof.register_proof(&f);

        let ordered = self
            .assume_retries::<ExactlyOnce>(nondet!(/** the combinator function is idempotent */))
            .assume_ordering::<TotalOrder>(nondet!(/** the combinator function is commutative */));

        KeyedSingleton::new(
            ordered.location.clone(),
            HydroNode::ReduceKeyed {
                f: f.into(),
                input: Box::new(ordered.ir_node.replace(HydroNode::Placeholder)),
                metadata: ordered.location.new_node_metadata(KeyedSingleton::<
                    K,
                    V,
                    L,
                    B::WhenValueUnbounded,
                >::collection_kind()),
            },
        )
    }

    /// A special case of [`KeyedStream::reduce`] where tuples with keys less than the watermark
    /// are automatically deleted.
    ///
    /// Depending on the input stream guarantees, the closure may need to be commutative
    /// (for unordered streams) or idempotent (for streams with non-deterministic duplicates).
    ///
    /// # Example
    /// ```rust
    /// # #[cfg(feature = "deploy")] {
    /// # use hydro_lang::prelude::*;
    /// # use futures::StreamExt;
    /// # tokio_test::block_on(hydro_lang::test_util::stream_transform_test(|process| {
    /// let tick = process.tick();
    /// let watermark = tick.singleton(q!(2));
    /// let numbers = process
    ///     .source_iter(q!([(0, false), (1, false), (2, false), (2, true)]))
    ///     .into_keyed();
    /// let batch = numbers.batch(&tick, nondet!(/** test */));
    /// batch
    ///     .reduce_watermark(watermark, q!(|acc, x| *acc |= x))
    ///     .entries()
    ///     .all_ticks()
    /// # }, |mut stream| async move {
    /// // (2, true)
    /// # assert_eq!(stream.next().await.unwrap(), (2, true));
    /// # }));
    /// # }
    /// ```
    pub fn reduce_watermark<O2, F, C, Idemp>(
        self,
        other: impl Into<Optional<O2, Tick<L::Root>, Bounded>>,
        comb: impl IntoQuotedMut<'a, F, L, AggFuncAlgebra<C, Idemp>>,
    ) -> KeyedSingleton<K, V, L, B::WhenValueUnbounded>
    where
        K: Eq + Hash,
        O2: Clone,
        F: Fn(&mut V, V) + 'a,
        C: ValidCommutativityFor<O>,
        Idemp: ValidIdempotenceFor<R>,
    {
        let other: Optional<O2, Tick<L::Root>, Bounded> = other.into();
        check_matching_location(&self.location.root(), other.location.outer());
        let (f, proof) = comb.splice_fn2_borrow_mut_ctx_props(&self.location);
        proof.register_proof(&f);

        let ordered = self
            .assume_retries::<ExactlyOnce>(nondet!(/** the combinator function is idempotent */))
            .assume_ordering::<TotalOrder>(nondet!(/** the combinator function is commutative */));

        KeyedSingleton::new(
            ordered.location.clone(),
            HydroNode::ReduceKeyedWatermark {
                f: f.into(),
                input: Box::new(ordered.ir_node.replace(HydroNode::Placeholder)),
                watermark: Box::new(other.ir_node.replace(HydroNode::Placeholder)),
                metadata: ordered.location.new_node_metadata(KeyedSingleton::<
                    K,
                    V,
                    L,
                    B::WhenValueUnbounded,
                >::collection_kind()),
            },
        )
    }

    /// Given a bounded stream of keys `K`, returns a new keyed stream containing only the groups
    /// whose keys are not in the bounded stream.
    ///
    /// # Example
    /// ```rust
    /// # #[cfg(feature = "deploy")] {
    /// # use hydro_lang::prelude::*;
    /// # use futures::StreamExt;
    /// # tokio_test::block_on(hydro_lang::test_util::stream_transform_test(|process| {
    /// let tick = process.tick();
    /// let keyed_stream = process
    ///     .source_iter(q!(vec![ (1, 'a'), (2, 'b'), (3, 'c'), (4, 'd') ]))
    ///     .batch(&tick, nondet!(/** test */))
    ///     .into_keyed();
    /// let keys_to_remove = process
    ///     .source_iter(q!(vec![1, 2]))
    ///     .batch(&tick, nondet!(/** test */));
    /// keyed_stream.filter_key_not_in(keys_to_remove).all_ticks()
    /// #   .entries()
    /// # }, |mut stream| async move {
    /// // { 3: ['c'], 4: ['d'] }
    /// # let mut results = Vec::new();
    /// # for _ in 0..2 {
    /// #     results.push(stream.next().await.unwrap());
    /// # }
    /// # results.sort();
    /// # assert_eq!(results, vec![(3, 'c'), (4, 'd')]);
    /// # }));
    /// # }
    /// ```
    pub fn filter_key_not_in<O2: Ordering, R2: Retries>(
        self,
        other: Stream<K, L, Bounded, O2, R2>,
    ) -> Self
    where
        K: Eq + Hash,
    {
        check_matching_location(&self.location, &other.location);

        KeyedStream::new(
            self.location.clone(),
            HydroNode::AntiJoin {
                pos: Box::new(self.ir_node.replace(HydroNode::Placeholder)),
                neg: Box::new(other.ir_node.replace(HydroNode::Placeholder)),
                metadata: self.location.new_node_metadata(Self::collection_kind()),
            },
        )
    }

    /// Emit a keyed stream containing keys shared between two keyed streams,
    /// where each value in the output keyed stream is a tuple of
    /// (self's value, other's value).
    /// If there are multiple values for the same key, this performs a cross product
    /// for each matching key.
    ///
    /// # Example
    /// ```rust
    /// # #[cfg(feature = "deploy")] {
    /// # use hydro_lang::prelude::*;
    /// # use futures::StreamExt;
    /// # tokio_test::block_on(hydro_lang::test_util::stream_transform_test(|process| {
    /// let tick = process.tick();
    /// let keyed_data = process
    ///     .source_iter(q!(vec![(1, 10), (1, 11), (2, 20)]))
    ///     .into_keyed()
    ///     .batch(&tick, nondet!(/** test */));
    /// let other_data = process
    ///     .source_iter(q!(vec![(1, 100), (2, 200), (2, 201)]))
    ///     .into_keyed()
    ///     .batch(&tick, nondet!(/** test */));
    /// keyed_data.join_keyed_stream(other_data).entries().all_ticks()
    /// # }, |mut stream| async move {
    /// // { 1: [(10, 100), (11, 100)], 2: [(20, 200), (20, 201)] } in any order
    /// # let mut results = vec![];
    /// # for _ in 0..4 {
    /// #     results.push(stream.next().await.unwrap());
    /// # }
    /// # results.sort();
    /// # assert_eq!(results, vec![(1, (10, 100)), (1, (11, 100)), (2, (20, 200)), (2, (20, 201))]);
    /// # }));
    /// # }
    /// ```
    pub fn join_keyed_stream<V2, O2: Ordering, R2: Retries>(
        self,
        other: KeyedStream<K, V2, L, B, O2, R2>,
    ) -> KeyedStream<K, (V, V2), L, B, NoOrder, <R as MinRetries<R2>>::Min>
    where
        K: Eq + Hash,
        R: MinRetries<R2>,
    {
        self.entries().join(other.entries()).into_keyed()
    }

    /// Deduplicates values within each key group, emitting each unique value per key
    /// exactly once.
    ///
    /// # Example
    /// ```rust
    /// # #[cfg(feature = "deploy")] {
    /// # use hydro_lang::prelude::*;
    /// # use futures::StreamExt;
    /// # tokio_test::block_on(hydro_lang::test_util::stream_transform_test(|process| {
    /// process
    ///     .source_iter(q!(vec![(1, 10), (2, 20), (1, 10), (2, 30), (1, 20)]))
    ///     .into_keyed()
    ///     .unique()
    /// # .entries()
    /// # }, |mut stream| async move {
    /// // unique values per key: { 1: [10, 20], 2: [20, 30] }
    /// # let mut results = Vec::new();
    /// # for _ in 0..4 {
    /// #     results.push(stream.next().await.unwrap());
    /// # }
    /// # let mut key1: Vec<_> = results.iter().filter(|(k, _)| *k == 1).map(|(_, v)| *v).collect();
    /// # let mut key2: Vec<_> = results.iter().filter(|(k, _)| *k == 2).map(|(_, v)| *v).collect();
    /// # key1.sort();
    /// # key2.sort();
    /// # assert_eq!(key1, vec![10, 20]);
    /// # assert_eq!(key2, vec![20, 30]);
    /// # }));
    /// # }
    /// ```
    pub fn unique(self) -> KeyedStream<K, V, L, B, NoOrder, ExactlyOnce>
    where
        K: Eq + Hash + Clone,
        V: Eq + Hash + Clone,
    {
        self.entries().unique().into_keyed()
    }

    /// Sorts the values within each key group in ascending order.
    ///
    /// The output keyed stream has a [`TotalOrder`] guarantee on the values within
    /// each group. This operator will block until all elements in the input stream
    /// are available, so it requires the input stream to be [`Bounded`].
    ///
    /// # Example
    /// ```rust
    /// # #[cfg(feature = "deploy")] {
    /// # use hydro_lang::prelude::*;
    /// # use futures::StreamExt;
    /// # tokio_test::block_on(hydro_lang::test_util::stream_transform_test(|process| {
    /// let tick = process.tick();
    /// let numbers = process
    ///     .source_iter(q!(vec![(1, 3), (2, 1), (1, 1), (2, 2)]))
    ///     .into_keyed();
    /// let batch = numbers.batch(&tick, nondet!(/** test */));
    /// batch.sort().all_ticks()
    /// # .entries()
    /// # }, |mut stream| async move {
    /// // values sorted within each key: { 1: [1, 3], 2: [1, 2] }
    /// # let mut results = Vec::new();
    /// # for _ in 0..4 {
    /// #     results.push(stream.next().await.unwrap());
    /// # }
    /// # let key1_vals: Vec<_> = results.iter().filter(|(k, _)| *k == 1).map(|(_, v)| *v).collect();
    /// # let key2_vals: Vec<_> = results.iter().filter(|(k, _)| *k == 2).map(|(_, v)| *v).collect();
    /// # assert_eq!(key1_vals, vec![1, 3]);
    /// # assert_eq!(key2_vals, vec![1, 2]);
    /// # }));
    /// # }
    /// ```
    pub fn sort(self) -> KeyedStream<K, V, L, Bounded, TotalOrder, R>
    where
        B: IsBounded,
        K: Ord,
        V: Ord,
    {
        self.entries().sort().into_keyed()
    }

    /// Produces a new keyed stream that combines the groups of the inputs by first emitting the
    /// elements of the `self` stream, and then emits the elements of the `other` stream (if a key
    /// is only present in one of the inputs, its values are passed through as-is). The output has
    /// a [`TotalOrder`] guarantee if and only if both inputs have a [`TotalOrder`] guarantee.
    ///
    /// Currently, both input streams must be [`Bounded`]. This operator will block
    /// on the first stream until all its elements are available. In a future version,
    /// we will relax the requirement on the `other` stream.
    ///
    /// # Example
    /// ```rust
    /// # #[cfg(feature = "deploy")] {
    /// # use hydro_lang::prelude::*;
    /// # use futures::StreamExt;
    /// # tokio_test::block_on(hydro_lang::test_util::stream_transform_test(|process| {
    /// let tick = process.tick();
    /// let numbers = process.source_iter(q!(vec![(0, 1), (1, 3)])).into_keyed();
    /// let batch = numbers.batch(&tick, nondet!(/** test */));
    /// batch.clone().map(q!(|x| x + 1)).chain(batch).all_ticks()
    /// # .entries()
    /// # }, |mut stream| async move {
    /// // { 0: [2, 1], 1: [4, 3] }
    /// # let mut results = Vec::new();
    /// # for _ in 0..4 {
    /// #     results.push(stream.next().await.unwrap());
    /// # }
    /// # results.sort();
    /// # assert_eq!(results, vec![(0, 1), (0, 2), (1, 3), (1, 4)]);
    /// # }));
    /// # }
    /// ```
    pub fn chain<O2: Ordering, R2: Retries>(
        self,
        other: KeyedStream<K, V, L, Bounded, O2, R2>,
    ) -> KeyedStream<K, V, L, Bounded, <O as MinOrder<O2>>::Min, <R as MinRetries<R2>>::Min>
    where
        B: IsBounded,
        O: MinOrder<O2>,
        R: MinRetries<R2>,
    {
        let this = self.make_bounded();
        check_matching_location(&this.location, &other.location);

        KeyedStream::new(
            this.location.clone(),
            HydroNode::Chain {
                first: Box::new(this.ir_node.replace(HydroNode::Placeholder)),
                second: Box::new(other.ir_node.replace(HydroNode::Placeholder)),
                metadata: this.location.new_node_metadata(KeyedStream::<
                    K,
                    V,
                    L,
                    Bounded,
                    <O as MinOrder<O2>>::Min,
                    <R as MinRetries<R2>>::Min,
                >::collection_kind()),
            },
        )
    }

    /// Emit a keyed stream containing keys shared between the keyed stream and the
    /// keyed singleton, where each value in the output keyed stream is a tuple of
    /// (the keyed stream's value, the keyed singleton's value).
    ///
    /// # Example
    /// ```rust
    /// # #[cfg(feature = "deploy")] {
    /// # use hydro_lang::prelude::*;
    /// # use futures::StreamExt;
    /// # tokio_test::block_on(hydro_lang::test_util::stream_transform_test(|process| {
    /// let tick = process.tick();
    /// let keyed_data = process
    ///     .source_iter(q!(vec![(1, 10), (1, 11), (2, 20)]))
    ///     .into_keyed()
    ///     .batch(&tick, nondet!(/** test */));
    /// let singleton_data = process
    ///     .source_iter(q!(vec![(1, 100), (2, 200)]))
    ///     .into_keyed()
    ///     .batch(&tick, nondet!(/** test */))
    ///     .first();
    /// keyed_data.join_keyed_singleton(singleton_data).entries().all_ticks()
    /// # }, |mut stream| async move {
    /// // { 1: [(10, 100), (11, 100)], 2: [(20, 200)] } in any order
    /// # let mut results = vec![];
    /// # for _ in 0..3 {
    /// #     results.push(stream.next().await.unwrap());
    /// # }
    /// # results.sort();
    /// # assert_eq!(results, vec![(1, (10, 100)), (1, (11, 100)), (2, (20, 200))]);
    /// # }));
    /// # }
    /// ```
    pub fn join_keyed_singleton<V2: Clone>(
        self,
        keyed_singleton: KeyedSingleton<K, V2, L, Bounded>,
    ) -> KeyedStream<K, (V, V2), L, Bounded, NoOrder, R>
    where
        B: IsBounded,
        K: Eq + Hash,
    {
        keyed_singleton
            .join_keyed_stream(self.make_bounded())
            .map(q!(|(v2, v)| (v, v2)))
    }

    /// Gets the values associated with a specific key from the keyed stream.
    /// Returns an empty stream if the key is `None` or there are no associated values.
    ///
    /// # Example
    /// ```rust
    /// # #[cfg(feature = "deploy")] {
    /// # use hydro_lang::prelude::*;
    /// # use futures::StreamExt;
    /// # tokio_test::block_on(hydro_lang::test_util::stream_transform_test(|process| {
    /// let tick = process.tick();
    /// let keyed_data = process
    ///     .source_iter(q!(vec![(1, 10), (1, 11), (2, 20)]))
    ///     .into_keyed()
    ///     .batch(&tick, nondet!(/** test */));
    /// let key = tick.singleton(q!(1));
    /// keyed_data.get(key).all_ticks()
    /// # }, |mut stream| async move {
    /// // 10, 11 in any order
    /// # let mut results = vec![];
    /// # for _ in 0..2 {
    /// #     results.push(stream.next().await.unwrap());
    /// # }
    /// # results.sort();
    /// # assert_eq!(results, vec![10, 11]);
    /// # }));
    /// # }
    /// ```
    pub fn get(self, key: impl Into<Optional<K, L, Bounded>>) -> Stream<V, L, Bounded, NoOrder, R>
    where
        B: IsBounded,
        K: Eq + Hash,
    {
        self.make_bounded()
            .entries()
            .join(key.into().into_stream().map(q!(|k| (k, ()))))
            .map(q!(|(_, (v, _))| v))
    }

    /// For each value in `self`, find the matching key in `lookup`.
    /// The output is a keyed stream with the key from `self`, and a value
    /// that is a tuple of (`self`'s value, Option<`lookup`'s value>).
    /// If the key is not present in `lookup`, the option will be [`None`].
    ///
    /// # Example
    /// ```rust
    /// # #[cfg(feature = "deploy")] {
    /// # use hydro_lang::prelude::*;
    /// # use futures::StreamExt;
    /// # tokio_test::block_on(hydro_lang::test_util::stream_transform_test(|process| {
    /// # let tick = process.tick();
    /// let requests = // { 1: [10, 11], 2: 20 }
    /// # process
    /// #     .source_iter(q!(vec![(1, 10), (1, 11), (2, 20)]))
    /// #     .into_keyed()
    /// #     .batch(&tick, nondet!(/** test */));
    /// let other_data = // { 10: 100, 11: 110 }
    /// # process
    /// #     .source_iter(q!(vec![(10, 100), (11, 110)]))
    /// #     .into_keyed()
    /// #     .batch(&tick, nondet!(/** test */))
    /// #     .first();
    /// requests.lookup_keyed_singleton(other_data)
    /// # .entries().all_ticks()
    /// # }, |mut stream| async move {
    /// // { 1: [(10, Some(100)), (11, Some(110))], 2: (20, None) }
    /// # let mut results = vec![];
    /// # for _ in 0..3 {
    /// #     results.push(stream.next().await.unwrap());
    /// # }
    /// # results.sort();
    /// # assert_eq!(results, vec![(1, (10, Some(100))), (1, (11, Some(110))), (2, (20, None))]);
    /// # }));
    /// # }
    /// ```
    pub fn lookup_keyed_singleton<V2>(
        self,
        lookup: KeyedSingleton<V, V2, L, Bounded>,
    ) -> KeyedStream<K, (V, Option<V2>), L, Bounded, NoOrder, R>
    where
        B: IsBounded,
        K: Eq + Hash + Clone,
        V: Eq + Hash + Clone,
        V2: Clone,
    {
        self.lookup_keyed_stream(
            lookup
                .into_keyed_stream()
                .assume_retries::<R>(nondet!(/** Retries are irrelevant for keyed singletons */)),
        )
    }

    /// For each value in `self`, find the matching key in `lookup`.
    /// The output is a keyed stream with the key from `self`, and a value
    /// that is a tuple of (`self`'s value, Option<`lookup`'s value>).
    /// If the key is not present in `lookup`, the option will be [`None`].
    ///
    /// # Example
    /// ```rust
    /// # #[cfg(feature = "deploy")] {
    /// # use hydro_lang::prelude::*;
    /// # use futures::StreamExt;
    /// # tokio_test::block_on(hydro_lang::test_util::stream_transform_test(|process| {
    /// # let tick = process.tick();
    /// let requests = // { 1: [10, 11], 2: 20 }
    /// # process
    /// #     .source_iter(q!(vec![(1, 10), (1, 11), (2, 20)]))
    /// #     .into_keyed()
    /// #     .batch(&tick, nondet!(/** test */));
    /// let other_data = // { 10: [100, 101], 11: 110 }
    /// # process
    /// #     .source_iter(q!(vec![(10, 100), (10, 101), (11, 110)]))
    /// #     .into_keyed()
    /// #     .batch(&tick, nondet!(/** test */));
    /// requests.lookup_keyed_stream(other_data)
    /// # .entries().all_ticks()
    /// # }, |mut stream| async move {
    /// // { 1: [(10, Some(100)), (10, Some(101)), (11, Some(110))], 2: (20, None) }
    /// # let mut results = vec![];
    /// # for _ in 0..4 {
    /// #     results.push(stream.next().await.unwrap());
    /// # }
    /// # results.sort();
    /// # assert_eq!(results, vec![(1, (10, Some(100))), (1, (10, Some(101))), (1, (11, Some(110))), (2, (20, None))]);
    /// # }));
    /// # }
    /// ```
    #[expect(clippy::type_complexity, reason = "retries propagation")]
    pub fn lookup_keyed_stream<V2, O2: Ordering, R2: Retries>(
        self,
        lookup: KeyedStream<V, V2, L, Bounded, O2, R2>,
    ) -> KeyedStream<K, (V, Option<V2>), L, Bounded, NoOrder, <R as MinRetries<R2>>::Min>
    where
        B: IsBounded,
        K: Eq + Hash + Clone,
        V: Eq + Hash + Clone,
        V2: Clone,
        R: MinRetries<R2>,
    {
        let inverted = self
            .make_bounded()
            .entries()
            .map(q!(|(key, lookup_value)| (lookup_value, key)))
            .into_keyed();
        let found = inverted
            .clone()
            .join_keyed_stream(lookup.clone())
            .entries()
            .map(q!(|(lookup_value, (key, value))| (
                key,
                (lookup_value, Some(value))
            )))
            .into_keyed();
        let not_found = inverted
            .filter_key_not_in(lookup.keys())
            .entries()
            .map(q!(|(lookup_value, key)| (key, (lookup_value, None))))
            .into_keyed();

        found.chain(not_found.weaken_retries::<<R as MinRetries<R2>>::Min>())
    }

    /// Shifts this keyed stream into an atomic context, which guarantees that any downstream logic
    /// will all be executed synchronously before any outputs are yielded (in [`KeyedStream::end_atomic`]).
    ///
    /// This is useful to enforce local consistency constraints, such as ensuring that a write is
    /// processed before an acknowledgement is emitted. Entering an atomic section requires a [`Tick`]
    /// argument that declares where the stream will be atomically processed. Batching a stream into
    /// the _same_ [`Tick`] will preserve the synchronous execution, while batching into a different
    /// [`Tick`] will introduce asynchrony.
    pub fn atomic(self, tick: &Tick<L>) -> KeyedStream<K, V, Atomic<L>, B, O, R> {
        let out_location = Atomic { tick: tick.clone() };
        KeyedStream::new(
            out_location.clone(),
            HydroNode::BeginAtomic {
                inner: Box::new(self.ir_node.replace(HydroNode::Placeholder)),
                metadata: out_location
                    .new_node_metadata(KeyedStream::<K, V, Atomic<L>, B, O, R>::collection_kind()),
            },
        )
    }

    /// Given a tick, returns a keyed stream corresponding to a batch of elements segmented by
    /// that tick. These batches are guaranteed to be contiguous across ticks and preserve
    /// the order of the input.
    ///
    /// # Non-Determinism
    /// The batch boundaries are non-deterministic and may change across executions.
    pub fn batch(
        self,
        tick: &Tick<L>,
        nondet: NonDet,
    ) -> KeyedStream<K, V, Tick<L>, Bounded, O, R> {
        let _ = nondet;
        assert_eq!(Location::id(tick.outer()), Location::id(&self.location));
        KeyedStream::new(
            tick.clone(),
            HydroNode::Batch {
                inner: Box::new(self.ir_node.replace(HydroNode::Placeholder)),
                metadata: tick.new_node_metadata(
                    KeyedStream::<K, V, Tick<L>, Bounded, O, R>::collection_kind(),
                ),
            },
        )
    }
}

impl<'a, K1, K2, V, L: Location<'a>, B: Boundedness, O: Ordering, R: Retries>
    KeyedStream<(K1, K2), V, L, B, O, R>
{
    /// Produces a new keyed stream by dropping the first element of the compound key.
    ///
    /// Because multiple keys may share the same suffix, this operation results in re-grouping
    /// of the values under the new keys. The values across groups with the same new key
    /// will be interleaved, so the resulting stream has [`NoOrder`] within each group.
    ///
    /// # Example
    /// ```rust
    /// # #[cfg(feature = "deploy")] {
    /// # use hydro_lang::prelude::*;
    /// # use futures::StreamExt;
    /// # tokio_test::block_on(hydro_lang::test_util::stream_transform_test(|process| {
    /// process
    ///     .source_iter(q!(vec![((1, 10), 2), ((1, 10), 3), ((2, 20), 4)]))
    ///     .into_keyed()
    ///     .drop_key_prefix()
    /// #   .entries()
    /// # }, |mut stream| async move {
    /// // { 10: [2, 3], 20: [4] }
    /// # let mut results = Vec::new();
    /// # for _ in 0..3 {
    /// #     results.push(stream.next().await.unwrap());
    /// # }
    /// # results.sort();
    /// # assert_eq!(results, vec![(10, 2), (10, 3), (20, 4)]);
    /// # }));
    /// # }
    /// ```
    pub fn drop_key_prefix(self) -> KeyedStream<K2, V, L, B, NoOrder, R> {
        self.entries()
            .map(q!(|((_k1, k2), v)| (k2, v)))
            .into_keyed()
    }
}

impl<'a, K, V, L: Location<'a> + NoTick, O: Ordering, R: Retries>
    KeyedStream<K, V, L, Unbounded, O, R>
{
    /// Produces a new keyed stream that "merges" the inputs by interleaving the elements
    /// of any overlapping groups. The result has [`NoOrder`] on each group because the
    /// order of interleaving is not guaranteed. If the keys across both inputs do not overlap,
    /// the ordering will be deterministic and you can safely use [`Self::assume_ordering`].
    ///
    /// Currently, both input streams must be [`Unbounded`].
    ///
    /// # Example
    /// ```rust
    /// # #[cfg(feature = "deploy")] {
    /// # use hydro_lang::prelude::*;
    /// # use futures::StreamExt;
    /// # tokio_test::block_on(hydro_lang::test_util::stream_transform_test(|process| {
    /// let numbers1: KeyedStream<i32, i32, _> = // { 1: [2], 3: [4] }
    /// # process.source_iter(q!(vec![(1, 2), (3, 4)])).into_keyed().into();
    /// let numbers2: KeyedStream<i32, i32, _> = // { 1: [3], 3: [5] }
    /// # process.source_iter(q!(vec![(1, 3), (3, 5)])).into_keyed().into();
    /// numbers1.merge_unordered(numbers2)
    /// #   .entries()
    /// # }, |mut stream| async move {
    /// // { 1: [2, 3], 3: [4, 5] } with each group in unknown order
    /// # let mut results = Vec::new();
    /// # for _ in 0..4 {
    /// #     results.push(stream.next().await.unwrap());
    /// # }
    /// # results.sort();
    /// # assert_eq!(results, vec![(1, 2), (1, 3), (3, 4), (3, 5)]);
    /// # }));
    /// # }
    /// ```
    pub fn merge_unordered<O2: Ordering, R2: Retries>(
        self,
        other: KeyedStream<K, V, L, Unbounded, O2, R2>,
    ) -> KeyedStream<K, V, L, Unbounded, NoOrder, <R as MinRetries<R2>>::Min>
    where
        R: MinRetries<R2>,
    {
        KeyedStream::new(
            self.location.clone(),
            HydroNode::Chain {
                first: Box::new(self.ir_node.replace(HydroNode::Placeholder)),
                second: Box::new(other.ir_node.replace(HydroNode::Placeholder)),
                metadata: self.location.new_node_metadata(KeyedStream::<
                    K,
                    V,
                    L,
                    Unbounded,
                    NoOrder,
                    <R as MinRetries<R2>>::Min,
                >::collection_kind()),
            },
        )
    }

    /// Deprecated: use [`KeyedStream::merge_unordered`] instead.
    #[deprecated(note = "use `merge_unordered` instead")]
    pub fn interleave<O2: Ordering, R2: Retries>(
        self,
        other: KeyedStream<K, V, L, Unbounded, O2, R2>,
    ) -> KeyedStream<K, V, L, Unbounded, NoOrder, <R as MinRetries<R2>>::Min>
    where
        R: MinRetries<R2>,
    {
        self.merge_unordered(other)
    }
}

impl<'a, K, V, L, B: Boundedness, O: Ordering, R: Retries> KeyedStream<K, V, Atomic<L>, B, O, R>
where
    L: Location<'a> + NoTick,
{
    /// Returns a keyed stream corresponding to the latest batch of elements being atomically
    /// processed. These batches are guaranteed to be contiguous across ticks and preserve
    /// the order of the input. The output keyed stream will execute in the [`Tick`] that was
    /// used to create the atomic section.
    ///
    /// # Non-Determinism
    /// The batch boundaries are non-deterministic and may change across executions.
    pub fn batch_atomic(self, nondet: NonDet) -> KeyedStream<K, V, Tick<L>, Bounded, O, R> {
        let _ = nondet;
        KeyedStream::new(
            self.location.clone().tick,
            HydroNode::Batch {
                inner: Box::new(self.ir_node.replace(HydroNode::Placeholder)),
                metadata: self.location.tick.new_node_metadata(KeyedStream::<
                    K,
                    V,
                    Tick<L>,
                    Bounded,
                    O,
                    R,
                >::collection_kind(
                )),
            },
        )
    }

    /// Yields the elements of this keyed stream back into a top-level, asynchronous execution context.
    /// See [`KeyedStream::atomic`] for more details.
    pub fn end_atomic(self) -> KeyedStream<K, V, L, B, O, R> {
        KeyedStream::new(
            self.location.tick.l.clone(),
            HydroNode::EndAtomic {
                inner: Box::new(self.ir_node.replace(HydroNode::Placeholder)),
                metadata: self
                    .location
                    .tick
                    .l
                    .new_node_metadata(KeyedStream::<K, V, L, B, O, R>::collection_kind()),
            },
        )
    }
}

impl<'a, K, V, L, O: Ordering, R: Retries> KeyedStream<K, V, Tick<L>, Bounded, O, R>
where
    L: Location<'a>,
{
    /// Asynchronously yields this batch of keyed elements outside the tick as an unbounded keyed stream,
    /// which will stream all the elements across _all_ tick iterations by concatenating the batches for
    /// each key.
    pub fn all_ticks(self) -> KeyedStream<K, V, L, Unbounded, O, R> {
        KeyedStream::new(
            self.location.outer().clone(),
            HydroNode::YieldConcat {
                inner: Box::new(self.ir_node.replace(HydroNode::Placeholder)),
                metadata: self.location.outer().new_node_metadata(KeyedStream::<
                    K,
                    V,
                    L,
                    Unbounded,
                    O,
                    R,
                >::collection_kind(
                )),
            },
        )
    }

    /// Synchronously yields this batch of keyed elements outside the tick as an unbounded keyed stream,
    /// which will stream all the elements across _all_ tick iterations by concatenating the batches for
    /// each key.
    ///
    /// Unlike [`KeyedStream::all_ticks`], this preserves synchronous execution, as the output stream
    /// is emitted in an [`Atomic`] context that will process elements synchronously with the input
    /// stream's [`Tick`] context.
    pub fn all_ticks_atomic(self) -> KeyedStream<K, V, Atomic<L>, Unbounded, O, R> {
        let out_location = Atomic {
            tick: self.location.clone(),
        };

        KeyedStream::new(
            out_location.clone(),
            HydroNode::YieldConcat {
                inner: Box::new(self.ir_node.replace(HydroNode::Placeholder)),
                metadata: out_location.new_node_metadata(KeyedStream::<
                    K,
                    V,
                    Atomic<L>,
                    Unbounded,
                    O,
                    R,
                >::collection_kind()),
            },
        )
    }

    /// Transforms the keyed stream using the given closure in "stateful" mode, where stateful operators
    /// such as `fold` retrain their memory for each key across ticks rather than resetting across batches of each key.
    ///
    /// This API is particularly useful for stateful computation on batches of data, such as
    /// maintaining an accumulated state that is up to date with the current batch.
    ///
    /// # Example
    /// ```rust
    /// # #[cfg(feature = "deploy")] {
    /// # use hydro_lang::prelude::*;
    /// # use futures::StreamExt;
    /// # tokio_test::block_on(hydro_lang::test_util::stream_transform_test(|process| {
    /// let tick = process.tick();
    /// # // ticks are lazy by default, forces the second tick to run
    /// # tick.spin_batch(q!(1)).all_ticks().for_each(q!(|_| {}));
    /// # let batch_first_tick = process
    /// #   .source_iter(q!(vec![(0, 1), (1, 2), (2, 3), (3, 4)]))
    /// #   .into_keyed()
    /// #   .batch(&tick, nondet!(/** test */));
    /// # let batch_second_tick = process
    /// #   .source_iter(q!(vec![(0, 5), (1, 6), (2, 7)]))
    /// #   .into_keyed()
    /// #   .batch(&tick, nondet!(/** test */))
    /// #   .defer_tick(); // appears on the second tick
    /// let input = batch_first_tick.chain(batch_second_tick).all_ticks();
    ///
    /// input.batch(&tick, nondet!(/** test */))
    ///     .across_ticks(|s| s.reduce(q!(|sum, new| {
    ///         *sum += new;
    ///     }))).entries().all_ticks()
    /// # }, |mut stream| async move {
    /// // First tick: [(0, 1), (1, 2), (2, 3), (3, 4)]
    /// # let mut results = Vec::new();
    /// # for _ in 0..4 {
    /// #     results.push(stream.next().await.unwrap());
    /// # }
    /// # results.sort();
    /// # assert_eq!(results, vec![(0, 1), (1, 2), (2, 3), (3, 4)]);
    /// // Second tick: [(0, 6), (1, 8), (2, 10), (3, 4)]
    /// # results.clear();
    /// # for _ in 0..4 {
    /// #     results.push(stream.next().await.unwrap());
    /// # }
    /// # results.sort();
    /// # assert_eq!(results, vec![(0, 6), (1, 8), (2, 10), (3, 4)]);
    /// # }));
    /// # }
    /// ```
    pub fn across_ticks<Out: BatchAtomic>(
        self,
        thunk: impl FnOnce(KeyedStream<K, V, Atomic<L>, Unbounded, O, R>) -> Out,
    ) -> Out::Batched {
        thunk(self.all_ticks_atomic()).batched_atomic()
    }

    /// Shifts the entries in `self` to the **next tick**, so that the returned keyed stream at
    /// tick `T` always has the entries of `self` at tick `T - 1`.
    ///
    /// At tick `0`, the output keyed stream is empty, since there is no previous tick.
    ///
    /// This operator enables stateful iterative processing with ticks, by sending data from one
    /// tick to the next. For example, you can use it to combine inputs across consecutive batches.
    ///
    /// # Example
    /// ```rust
    /// # #[cfg(feature = "deploy")] {
    /// # use hydro_lang::prelude::*;
    /// # use futures::StreamExt;
    /// # tokio_test::block_on(hydro_lang::test_util::stream_transform_test(|process| {
    /// let tick = process.tick();
    /// # // ticks are lazy by default, forces the second tick to run
    /// # tick.spin_batch(q!(1)).all_ticks().for_each(q!(|_| {}));
    /// # let batch_first_tick = process
    /// #   .source_iter(q!(vec![(1, 2), (1, 3)]))
    /// #   .batch(&tick, nondet!(/** test */))
    /// #   .into_keyed();
    /// # let batch_second_tick = process
    /// #   .source_iter(q!(vec![(1, 4), (2, 5)]))
    /// #   .batch(&tick, nondet!(/** test */))
    /// #   .defer_tick()
    /// #   .into_keyed(); // appears on the second tick
    /// let changes_across_ticks = // { 1: [2, 3] } (first tick), { 1: [4], 2: [5] } (second tick)
    /// # batch_first_tick.chain(batch_second_tick);
    /// changes_across_ticks.clone().defer_tick().chain( // from the previous tick
    ///     changes_across_ticks // from the current tick
    /// )
    /// # .entries().all_ticks()
    /// # }, |mut stream| async move {
    /// // First tick: { 1: [2, 3] }
    /// # let mut results = Vec::new();
    /// # for _ in 0..2 {
    /// #     results.push(stream.next().await.unwrap());
    /// # }
    /// # results.sort();
    /// # assert_eq!(results, vec![(1, 2), (1, 3)]);
    /// // Second tick: { 1: [2, 3, 4], 2: [5] }
    /// # results.clear();
    /// # for _ in 0..4 {
    /// #     results.push(stream.next().await.unwrap());
    /// # }
    /// # results.sort();
    /// # assert_eq!(results, vec![(1, 2), (1, 3), (1, 4), (2, 5)]);
    /// // Third tick: { 1: [4], 2: [5] }
    /// # results.clear();
    /// # for _ in 0..2 {
    /// #     results.push(stream.next().await.unwrap());
    /// # }
    /// # results.sort();
    /// # assert_eq!(results, vec![(1, 4), (2, 5)]);
    /// # }));
    /// # }
    /// ```
    pub fn defer_tick(self) -> KeyedStream<K, V, Tick<L>, Bounded, O, R> {
        KeyedStream::new(
            self.location.clone(),
            HydroNode::DeferTick {
                input: Box::new(self.ir_node.replace(HydroNode::Placeholder)),
                metadata: self.location.new_node_metadata(KeyedStream::<
                    K,
                    V,
                    Tick<L>,
                    Bounded,
                    O,
                    R,
                >::collection_kind()),
            },
        )
    }
}

#[cfg(test)]
mod tests {
    #[cfg(feature = "deploy")]
    use futures::{SinkExt, StreamExt};
    #[cfg(feature = "deploy")]
    use hydro_deploy::Deployment;
    #[cfg(any(feature = "deploy", feature = "sim"))]
    use stageleft::q;

    #[cfg(any(feature = "deploy", feature = "sim"))]
    use crate::compile::builder::FlowBuilder;
    #[cfg(feature = "deploy")]
    use crate::live_collections::stream::ExactlyOnce;
    #[cfg(feature = "sim")]
    use crate::live_collections::stream::{NoOrder, TotalOrder};
    #[cfg(any(feature = "deploy", feature = "sim"))]
    use crate::location::Location;
    #[cfg(any(feature = "deploy", feature = "sim"))]
    use crate::nondet::nondet;
    #[cfg(feature = "deploy")]
    use crate::properties::manual_proof;

    #[cfg(feature = "deploy")]
    #[tokio::test]
    async fn reduce_watermark_filter() {
        let mut deployment = Deployment::new();

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

        let node_tick = node.tick();
        let watermark = node_tick.singleton(q!(2));

        let sum = node
            .source_stream(q!(tokio_stream::iter([
                (0, 100),
                (1, 101),
                (2, 102),
                (2, 102)
            ])))
            .into_keyed()
            .reduce_watermark(
                watermark,
                q!(|acc, v| {
                    *acc += v;
                }),
            )
            .snapshot(&node_tick, nondet!(/** test */))
            .entries()
            .all_ticks()
            .send_bincode_external(&external);

        let nodes = flow
            .with_process(&node, deployment.Localhost())
            .with_external(&external, deployment.Localhost())
            .deploy(&mut deployment);

        deployment.deploy().await.unwrap();

        let mut out = nodes.connect(sum).await;

        deployment.start().await.unwrap();

        assert_eq!(out.next().await.unwrap(), (2, 204));
    }

    #[cfg(feature = "deploy")]
    #[tokio::test]
    async fn reduce_watermark_bounded() {
        let mut deployment = Deployment::new();

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

        let node_tick = node.tick();
        let watermark = node_tick.singleton(q!(2));

        let sum = node
            .source_iter(q!([(0, 100), (1, 101), (2, 102), (2, 102)]))
            .into_keyed()
            .reduce_watermark(
                watermark,
                q!(|acc, v| {
                    *acc += v;
                }),
            )
            .entries()
            .send_bincode_external(&external);

        let nodes = flow
            .with_process(&node, deployment.Localhost())
            .with_external(&external, deployment.Localhost())
            .deploy(&mut deployment);

        deployment.deploy().await.unwrap();

        let mut out = nodes.connect(sum).await;

        deployment.start().await.unwrap();

        assert_eq!(out.next().await.unwrap(), (2, 204));
    }

    #[cfg(feature = "deploy")]
    #[tokio::test]
    async fn reduce_watermark_garbage_collect() {
        let mut deployment = Deployment::new();

        let mut flow = FlowBuilder::new();
        let node = flow.process::<()>();
        let external = flow.external::<()>();
        let (tick_send, tick_trigger) =
            node.source_external_bincode::<_, _, _, ExactlyOnce>(&external);

        let node_tick = node.tick();
        let (watermark_complete_cycle, watermark) =
            node_tick.cycle_with_initial(node_tick.singleton(q!(2)));
        let next_watermark = watermark.clone().map(q!(|v| v + 1));
        watermark_complete_cycle.complete_next_tick(next_watermark);

        let tick_triggered_input = node_tick
            .singleton(q!((3, 103)))
            .into_stream()
            .filter_if(
                tick_trigger
                    .clone()
                    .batch(&node_tick, nondet!(/** test */))
                    .first()
                    .is_some(),
            )
            .all_ticks();

        let sum = node
            .source_stream(q!(tokio_stream::iter([
                (0, 100),
                (1, 101),
                (2, 102),
                (2, 102)
            ])))
            .merge_unordered(tick_triggered_input)
            .into_keyed()
            .reduce_watermark(
                watermark,
                q!(
                    |acc, v| {
                        *acc += v;
                    },
                    commutative = manual_proof!(/** integer addition is commutative */)
                ),
            )
            .snapshot(&node_tick, nondet!(/** test */))
            .entries()
            .all_ticks()
            .send_bincode_external(&external);

        let nodes = flow
            .with_default_optimize()
            .with_process(&node, deployment.Localhost())
            .with_external(&external, deployment.Localhost())
            .deploy(&mut deployment);

        deployment.deploy().await.unwrap();

        let mut tick_send = nodes.connect(tick_send).await;
        let mut out_recv = nodes.connect(sum).await;

        deployment.start().await.unwrap();

        assert_eq!(out_recv.next().await.unwrap(), (2, 204));

        tick_send.send(()).await.unwrap();

        assert_eq!(out_recv.next().await.unwrap(), (3, 103));
    }

    #[cfg(feature = "sim")]
    #[test]
    #[should_panic]
    fn sim_batch_nondet_size() {
        let mut flow = FlowBuilder::new();
        let node = flow.process::<()>();

        let input = node.source_iter(q!([(1, 1), (1, 2), (2, 3)])).into_keyed();

        let tick = node.tick();
        let out_recv = input
            .batch(&tick, nondet!(/** test */))
            .fold(q!(|| vec![]), q!(|acc, v| acc.push(v)))
            .entries()
            .all_ticks()
            .sim_output();

        flow.sim().exhaustive(async || {
            out_recv
                .assert_yields_only_unordered([(1, vec![1, 2])])
                .await;
        });
    }

    #[cfg(feature = "sim")]
    #[test]
    fn sim_batch_preserves_group_order() {
        let mut flow = FlowBuilder::new();
        let node = flow.process::<()>();

        let input = node.source_iter(q!([(1, 1), (1, 2), (2, 3)])).into_keyed();

        let tick = node.tick();
        let out_recv = input
            .batch(&tick, nondet!(/** test */))
            .all_ticks()
            .fold_early_stop(
                q!(|| 0),
                q!(|acc, v| {
                    *acc = std::cmp::max(v, *acc);
                    *acc >= 2
                }),
            )
            .entries()
            .sim_output();

        let instances = flow.sim().exhaustive(async || {
            out_recv
                .assert_yields_only_unordered([(1, 2), (2, 3)])
                .await;
        });

        assert_eq!(instances, 8);
        // - three cases: all three in a separate tick (pick where (2, 3) is)
        // - two cases: (1, 1) and (1, 2) together, (2, 3) before or after
        // - two cases: (1, 1) and (1, 2) separate, (2, 3) grouped with one of them
        // - one case: all three together
    }

    #[cfg(feature = "sim")]
    #[test]
    fn sim_batch_unordered_shuffles() {
        let mut flow = FlowBuilder::new();
        let node = flow.process::<()>();

        let input = node
            .source_iter(q!([(1, 1), (1, 2), (2, 3)]))
            .into_keyed()
            .weaken_ordering::<NoOrder>();

        let tick = node.tick();
        let out_recv = input
            .batch(&tick, nondet!(/** test */))
            .all_ticks()
            .entries()
            .sim_output();

        let instances = flow.sim().exhaustive(async || {
            out_recv
                .assert_yields_only_unordered([(1, 1), (1, 2), (2, 3)])
                .await;
        });

        assert_eq!(instances, 13);
        // - 6 (3 * 2) cases: all three in a separate tick (pick where (2, 3) is), and order of (1, 1), (1, 2)
        // - two cases: (1, 1) and (1, 2) together, (2, 3) before or after (order of (1, 1), (1, 2) doesn't matter because batched is still unordered)
        // - 4 (2 * 2) cases: (1, 1) and (1, 2) separate, (2, 3) grouped with one of them, and order of (1, 1), (1, 2)
        // - one case: all three together (order of (1, 1), (1, 2) doesn't matter because batched is still unordered)
    }

    #[cfg(feature = "sim")]
    #[test]
    #[should_panic]
    fn sim_observe_order_batched() {
        let mut flow = FlowBuilder::new();
        let node = flow.process::<()>();

        let (port, input) = node.sim_input::<_, NoOrder, _>();

        let tick = node.tick();
        let batch = input.into_keyed().batch(&tick, nondet!(/** test */));
        let out_recv = batch
            .assume_ordering::<TotalOrder>(nondet!(/** test */))
            .all_ticks()
            .first()
            .entries()
            .sim_output();

        flow.sim().exhaustive(async || {
            port.send_many_unordered([(1, 1), (1, 2), (2, 1), (2, 2)]);
            out_recv
                .assert_yields_only_unordered([(1, 1), (2, 1)])
                .await; // fails with assume_ordering
        });
    }

    #[cfg(feature = "sim")]
    #[test]
    fn sim_observe_order_batched_count() {
        let mut flow = FlowBuilder::new();
        let node = flow.process::<()>();

        let (port, input) = node.sim_input::<_, NoOrder, _>();

        let tick = node.tick();
        let batch = input.into_keyed().batch(&tick, nondet!(/** test */));
        let out_recv = batch
            .assume_ordering::<TotalOrder>(nondet!(/** test */))
            .all_ticks()
            .entries()
            .sim_output();

        let instance_count = flow.sim().exhaustive(async || {
            port.send_many_unordered([(1, 1), (1, 2), (2, 1), (2, 2)]);
            let _ = out_recv.collect_sorted::<Vec<_>>().await;
        });

        assert_eq!(instance_count, 104); // too complicated to enumerate here, but less than stream equivalent
    }

    #[cfg(feature = "sim")]
    #[test]
    fn sim_top_level_assume_ordering() {
        use std::collections::HashMap;

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

        let (in_send, input) = node.sim_input::<_, NoOrder, _>();

        let out_recv = input
            .into_keyed()
            .assume_ordering::<TotalOrder>(nondet!(/** test */))
            .fold_early_stop(
                q!(|| Vec::new()),
                q!(|acc, v| {
                    acc.push(v);
                    acc.len() >= 2
                }),
            )
            .entries()
            .sim_output();

        let instance_count = flow.sim().exhaustive(async || {
            in_send.send_many_unordered([(1, 'a'), (1, 'b'), (2, 'c'), (2, 'd')]);
            let out: HashMap<_, _> = out_recv
                .collect_sorted::<Vec<_>>()
                .await
                .into_iter()
                .collect();
            // Each key accumulates its values; we get one entry per key
            assert_eq!(out.len(), 2);
        });

        assert_eq!(instance_count, 24)
    }

    #[cfg(feature = "sim")]
    #[test]
    fn sim_top_level_assume_ordering_cycle_back() {
        use std::collections::HashMap;

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

        let (in_send, input) = node.sim_input::<_, NoOrder, _>();

        let (complete_cycle_back, cycle_back) =
            node.forward_ref::<super::KeyedStream<_, _, _, _, NoOrder>>();
        let ordered = input
            .into_keyed()
            .merge_unordered(cycle_back)
            .assume_ordering::<TotalOrder>(nondet!(/** test */));
        complete_cycle_back.complete(
            ordered
                .clone()
                .map(q!(|v| v + 1))
                .filter(q!(|v| v % 2 == 1)),
        );

        let out_recv = ordered
            .fold_early_stop(
                q!(|| Vec::new()),
                q!(|acc, v| {
                    acc.push(v);
                    acc.len() >= 2
                }),
            )
            .entries()
            .sim_output();

        let mut saw = false;
        let instance_count = flow.sim().exhaustive(async || {
            // Send (1, 0) and (1, 2). 0+1=1 is odd so cycles back.
            // We want to see [0, 1] - the cycled back value interleaved
            in_send.send_many_unordered([(1, 0), (1, 2)]);
            let out: HashMap<_, _> = out_recv
                .collect_sorted::<Vec<_>>()
                .await
                .into_iter()
                .collect();

            // We want to see an instance where key 1 gets: 0, then 1 (cycled back from 0+1)
            if let Some(values) = out.get(&1)
                && *values == vec![0, 1]
            {
                saw = true;
            }
        });

        assert!(
            saw,
            "did not see an instance with key 1 having [0, 1] in order"
        );
        assert_eq!(instance_count, 6);
    }

    #[cfg(feature = "sim")]
    #[test]
    fn sim_top_level_assume_ordering_cross_key_cycle() {
        use std::collections::HashMap;

        // This test demonstrates why releasing one entry at a time is important:
        // When one key's observed order cycles back into a different key, we need
        // to be able to interleave the cycled-back entry with pending items for
        // that other key.
        let mut flow = FlowBuilder::new();
        let node = flow.process::<()>();

        let (in_send, input) = node.sim_input::<_, NoOrder, _>();

        let (complete_cycle_back, cycle_back) =
            node.forward_ref::<super::KeyedStream<_, _, _, _, NoOrder>>();
        let ordered = input
            .into_keyed()
            .merge_unordered(cycle_back)
            .assume_ordering::<TotalOrder>(nondet!(/** test */));

        // Cycle back: when we see (1, 10), emit (2, 100) to key 2
        complete_cycle_back.complete(
            ordered
                .clone()
                .filter(q!(|v| *v == 10))
                .map(q!(|_| 100))
                .entries()
                .map(q!(|(_, v)| (2, v))) // Change key from 1 to 2
                .into_keyed(),
        );

        let out_recv = ordered
            .fold_early_stop(
                q!(|| Vec::new()),
                q!(|acc, v| {
                    acc.push(v);
                    acc.len() >= 2
                }),
            )
            .entries()
            .sim_output();

        // We want to see an instance where:
        // - (1, 10) is released first
        // - This causes (2, 100) to be cycled back
        // - (2, 100) is released BEFORE (2, 20) which was already pending
        let mut saw_cross_key_interleave = false;
        let instance_count = flow.sim().exhaustive(async || {
            // Send (1, 10), (1, 11) for key 1, and (2, 20), (2, 21) for key 2
            in_send.send_many_unordered([(1, 10), (1, 11), (2, 20), (2, 21)]);
            let out: HashMap<_, _> = out_recv
                .collect_sorted::<Vec<_>>()
                .await
                .into_iter()
                .collect();

            // Check if we see the cross-key interleaving:
            // key 2 should have [100, 20] or [100, 21] - cycled back 100 before a pending item
            if let Some(values) = out.get(&2)
                && values.len() >= 2
                && values[0] == 100
            {
                saw_cross_key_interleave = true;
            }
        });

        assert!(
            saw_cross_key_interleave,
            "did not see an instance where cycled-back 100 was released before pending items for key 2"
        );
        assert_eq!(instance_count, 60);
    }

    #[cfg(feature = "sim")]
    #[test]
    fn sim_top_level_assume_ordering_cycle_back_tick() {
        use std::collections::HashMap;

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

        let (in_send, input) = node.sim_input::<_, NoOrder, _>();

        let (complete_cycle_back, cycle_back) =
            node.forward_ref::<super::KeyedStream<_, _, _, _, NoOrder>>();
        let ordered = input
            .into_keyed()
            .merge_unordered(cycle_back)
            .assume_ordering::<TotalOrder>(nondet!(/** test */));
        complete_cycle_back.complete(
            ordered
                .clone()
                .batch(&node.tick(), nondet!(/** test */))
                .all_ticks()
                .map(q!(|v| v + 1))
                .filter(q!(|v| v % 2 == 1)),
        );

        let out_recv = ordered
            .fold_early_stop(
                q!(|| Vec::new()),
                q!(|acc, v| {
                    acc.push(v);
                    acc.len() >= 2
                }),
            )
            .entries()
            .sim_output();

        let mut saw = false;
        let instance_count = flow.sim().exhaustive(async || {
            in_send.send_many_unordered([(1, 0), (1, 2)]);
            let out: HashMap<_, _> = out_recv
                .collect_sorted::<Vec<_>>()
                .await
                .into_iter()
                .collect();

            if let Some(values) = out.get(&1)
                && *values == vec![0, 1]
            {
                saw = true;
            }
        });

        assert!(
            saw,
            "did not see an instance with key 1 having [0, 1] in order"
        );
        assert_eq!(instance_count, 58);
    }
}