hydro_lang/live_collections/

keyed_singleton.rs

//! Definitions for the [`KeyedSingleton`] 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 sealed::sealed;
use stageleft::{IntoQuotedMut, QuotedWithContext, q};

use super::OperatorContext;
use super::boundedness::{Bounded, Boundedness, IsBounded, Unbounded};
use super::keyed_stream::KeyedStream;
use super::optional::Optional;
use super::singleton::Singleton;
use super::sliced::sliced;
use super::stream::{ExactlyOnce, NoOrder, Stream, TotalOrder};
use crate::compile::builder::{CycleId, FlowState};
use crate::compile::ir::{
    CollectionKind, HydroIrOpMetadata, HydroNode, HydroRoot, KeyedSingletonBoundKind, SharedNode,
};
#[cfg(stageleft_runtime)]
use crate::forward_handle::{CycleCollection, ReceiverComplete};
use crate::forward_handle::{ForwardRef, TickCycle};
use crate::live_collections::stream::{Ordering, Retries};
#[cfg(stageleft_runtime)]
use crate::location::dynamic::{DynLocation, LocationId};
use crate::location::tick::DeferTick;
use crate::location::{Atomic, Location, Tick, TopLevel, check_matching_location};
use crate::manual_expr::ManualExpr;
use crate::nondet::{NonDet, nondet};
use crate::properties::manual_proof;

/// A marker trait indicating which components of a [`KeyedSingleton`] may change.
///
/// In addition to [`Bounded`] (all entries are fixed) and [`Unbounded`] (entries may be added /
/// changed, but not removed), this also includes an additional variant [`BoundedValue`], which
/// indicates that entries may be added over time, but once an entry is added it will never be
/// removed and its value will never change.
pub trait KeyedSingletonBound {
    /// The [`Boundedness`] of the [`Stream`] underlying the keyed singleton.
    type UnderlyingBound: Boundedness;
    /// The [`Boundedness`] of each entry's value; [`Bounded`] means it is immutable.
    type ValueBound: Boundedness;

    /// The type of the keyed singleton if the value for each key is immutable.
    type WithBoundedValue: KeyedSingletonBound<
            UnderlyingBound = Self::UnderlyingBound,
            ValueBound = Bounded,
            EraseMonotonic = Self::WithBoundedValue,
        >;

    /// The [`Boundedness`] of this [`Singleton`] if it is produced from a [`KeyedStream`] with [`Self`] boundedness.
    type KeyedStreamToMonotone: KeyedSingletonBound<UnderlyingBound = Self::UnderlyingBound, ValueBound = Self::ValueBound>;

    /// The [`Boundedness`] of the keyed singleton produced by folding a [`KeyedStream`] with
    /// [`Self`] boundedness when the aggregation does *not* have a monotonicity proof.
    ///
    /// Without a monotonicity proof, the per-key values may change arbitrarily, so an unbounded
    /// input collapses to [`MonotonicKeys`] (keys are still only added, never removed).
    type KeyedStreamToNonMonotone: KeyedSingletonBound<UnderlyingBound = Self::UnderlyingBound, ValueBound = Self::ValueBound>;

    /// The type of the keyed singleton if the value for each key is no longer monotonic.
    type EraseMonotonic: KeyedSingletonBound<UnderlyingBound = Self::UnderlyingBound, ValueBound = Self::ValueBound>;

    /// Returns the [`KeyedSingletonBoundKind`] corresponding to this type.
    fn bound_kind() -> KeyedSingletonBoundKind;
}

impl KeyedSingletonBound for Unbounded {
    type UnderlyingBound = Unbounded;
    type ValueBound = Unbounded;
    type WithBoundedValue = BoundedValue;
    type KeyedStreamToMonotone = MonotonicValue;
    type KeyedStreamToNonMonotone = MonotonicKeys;
    type EraseMonotonic = Unbounded;

    fn bound_kind() -> KeyedSingletonBoundKind {
        KeyedSingletonBoundKind::Unbounded
    }
}

impl KeyedSingletonBound for Bounded {
    type UnderlyingBound = Bounded;
    type ValueBound = Bounded;
    type WithBoundedValue = Bounded;
    type KeyedStreamToMonotone = Bounded;
    type KeyedStreamToNonMonotone = Bounded;
    type EraseMonotonic = Bounded;

    fn bound_kind() -> KeyedSingletonBoundKind {
        KeyedSingletonBoundKind::Bounded
    }
}

/// A variation of boundedness specific to [`KeyedSingleton`], which indicates that once a key appears,
/// its value is bounded and will never change, but new entries may appear asynchronously
pub enum BoundedValue {}

impl KeyedSingletonBound for BoundedValue {
    type UnderlyingBound = Unbounded;
    type ValueBound = Bounded;
    type WithBoundedValue = BoundedValue;
    type KeyedStreamToMonotone = BoundedValue;
    type KeyedStreamToNonMonotone = BoundedValue;
    type EraseMonotonic = BoundedValue;

    fn bound_kind() -> KeyedSingletonBoundKind {
        KeyedSingletonBoundKind::BoundedValue
    }
}

/// A variation of boundedness specific to [`KeyedSingleton`], which indicates that once a key appears,
/// it will never be removed, and the corresponding value will only increase monotonically.
pub enum MonotonicValue {}

impl KeyedSingletonBound for MonotonicValue {
    type UnderlyingBound = Unbounded;
    type ValueBound = Unbounded;
    type WithBoundedValue = BoundedValue;
    type KeyedStreamToMonotone = MonotonicValue;
    type KeyedStreamToNonMonotone = MonotonicKeys;
    type EraseMonotonic = MonotonicKeys;

    fn bound_kind() -> KeyedSingletonBoundKind {
        KeyedSingletonBoundKind::MonotonicValue
    }
}

/// A variation of boundedness specific to [`KeyedSingleton`], which indicates that once a key
/// appears, it will never be removed, but the corresponding value may change arbitrarily.
pub enum MonotonicKeys {}

impl KeyedSingletonBound for MonotonicKeys {
    type UnderlyingBound = Unbounded;
    type ValueBound = Unbounded;
    type WithBoundedValue = BoundedValue;
    type KeyedStreamToMonotone = MonotonicKeys;
    type KeyedStreamToNonMonotone = MonotonicKeys;
    type EraseMonotonic = MonotonicKeys;

    fn bound_kind() -> KeyedSingletonBoundKind {
        KeyedSingletonBoundKind::MonotonicKeys
    }
}

#[sealed]
#[diagnostic::on_unimplemented(
    message = "The keyed singleton must have monotonic values (`MonotonicValue`) or be bounded (`Bounded`), but has bound `{Self}`. Strengthen the monotonicity upstream or consider a different API.",
    label = "required here",
    note = "To intentionally process a non-deterministic snapshot or batch, you may want to use a `sliced!` region. This introduces non-determinism so avoid unless necessary."
)]
/// Marker trait that is implemented for [`KeyedSingletonBound`] types whose per-key values
/// are monotonically non-decreasing (or bounded).
pub trait IsKeyedMonotonic: KeyedSingletonBound {}

#[sealed]
#[diagnostic::do_not_recommend]
impl IsKeyedMonotonic for MonotonicValue {}

#[sealed]
#[diagnostic::do_not_recommend]
impl IsKeyedMonotonic for BoundedValue {}

#[sealed]
#[diagnostic::do_not_recommend]
impl<B: IsBounded + KeyedSingletonBound> IsKeyedMonotonic for B {}

/// Mapping from keys of type `K` to values of type `V`.
///
/// Keyed Singletons capture an asynchronously updated mapping from keys of the `K` to values of
/// type `V`, where the order of keys is non-deterministic. In addition to the standard boundedness
/// variants ([`Bounded`] for finite and immutable, [`Unbounded`] for asynchronously changing),
/// keyed singletons can use [`BoundedValue`] to declare that new keys may be added over time, but
/// keys cannot be removed and the value for each key is immutable.
///
/// Type Parameters:
/// - `K`: the type of the key for each entry
/// - `V`: the type of the value for each entry
/// - `Loc`: the [`Location`] where the keyed singleton is materialized
/// - `Bound`: tracks whether the entries are:
///     - [`Bounded`] (local and finite)
///     - [`Unbounded`] (asynchronous with entries added / removed / changed over time)
///     - [`BoundedValue`] (asynchronous with immutable values for each key and no removals)
pub struct KeyedSingleton<K, V, Loc, Bound: KeyedSingletonBound> {
    pub(crate) location: Loc,
    pub(crate) ir_node: Rc<RefCell<HydroNode>>,
    pub(crate) flow_state: FlowState,

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

impl<K, V, L, B: KeyedSingletonBound> Drop for KeyedSingleton<K, V, L, B> {
    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: Clone, V: Clone, Loc: Location<'a>, Bound: KeyedSingletonBound> Clone
    for KeyedSingleton<K, V, Loc, Bound>
{
    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() {
            KeyedSingleton {
                location: self.location.clone(),
                flow_state: self.flow_state.clone(),
                ir_node: super::tracked_ir_node(
                    &self.flow_state,
                    HydroNode::Tee {
                        inner: SharedNode(inner.0.clone()),
                        metadata: metadata.clone(),
                    },
                ),
                _phantom: PhantomData,
            }
        } else {
            unreachable!()
        }
    }
}

impl<'a, K, V, L, B: KeyedSingletonBound> CycleCollection<'a, ForwardRef>
    for KeyedSingleton<K, V, L, B>
where
    L: Location<'a>,
{
    type Location = L;

    fn create_source(cycle_id: CycleId, location: L) -> Self {
        let flow_state = location.flow_state().clone();
        KeyedSingleton {
            ir_node: super::tracked_ir_node(
                &flow_state,
                HydroNode::CycleSource {
                    cycle_id,
                    metadata: location.new_node_metadata(Self::collection_kind()),
                },
            ),
            flow_state,
            location,
            _phantom: PhantomData,
        }
    }
}

impl<'a, K, V, L> CycleCollection<'a, TickCycle> for KeyedSingleton<K, V, Tick<L>, Bounded>
where
    L: Location<'a>,
{
    type Location = Tick<L>;

    fn create_source(cycle_id: CycleId, location: Tick<L>) -> Self {
        KeyedSingleton::new(
            location.clone(),
            HydroNode::CycleSource {
                cycle_id,
                metadata: location.new_node_metadata(Self::collection_kind()),
            },
        )
    }
}

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

impl<'a, K, V, L, B: KeyedSingletonBound> ReceiverComplete<'a, ForwardRef>
    for KeyedSingleton<K, V, L, B>
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> ReceiverComplete<'a, TickCycle> for KeyedSingleton<K, V, Tick<L>, Bounded>
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: Location<'a>, B: KeyedSingletonBound> KeyedSingleton<K, V, L, B> {
    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();
        let ir_node = super::tracked_ir_node(&flow_state, ir_node);
        KeyedSingleton {
            location,
            flow_state,
            ir_node,
            _phantom: PhantomData,
        }
    }

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

    /// Weakens the consistency of this live collection to not guarantee any consistency across
    /// cluster members (if this collection is on a cluster).
    pub fn weaken_consistency(self) -> KeyedSingleton<K, V, L::DropConsistency, B>
    where
        L: Location<'a>,
    {
        if L::consistency()
            .is_none_or(|c| c == crate::location::dynamic::ClusterConsistency::NoConsistency)
        {
            // already no consistency
            KeyedSingleton::new(
                self.location.drop_consistency(),
                self.ir_node.replace(HydroNode::Placeholder),
            )
        } else {
            KeyedSingleton::new(
                self.location.drop_consistency(),
                HydroNode::Cast {
                    inner: Box::new(self.ir_node.replace(HydroNode::Placeholder)),
                    metadata: self
                        .location
                        .drop_consistency()
                        .new_node_metadata(
                            KeyedSingleton::<K, V, L::DropConsistency, B>::collection_kind(),
                        ),
                },
            )
        }
    }

    /// Casts this live collection to have the consistency guarantees specified in the given
    /// location type parameter. The developer must ensure that the strengthened consistency
    /// is actually guaranteed, via the proof field (see [`crate::prelude::manual_proof`]).
    pub fn assert_has_consistency_of<L2: Location<'a, DropConsistency = L::DropConsistency>>(
        self,
        _proof: impl crate::properties::ConsistencyProof,
    ) -> KeyedSingleton<K, V, L2, B>
    where
        L: Location<'a>,
    {
        if L::consistency() == L2::consistency() {
            // already consistent
            KeyedSingleton::new(
                self.location.with_consistency_of(),
                self.ir_node.replace(HydroNode::Placeholder),
            )
        } else {
            KeyedSingleton::new(
                self.location.with_consistency_of(),
                HydroNode::AssertIsConsistent {
                    inner: Box::new(self.ir_node.replace(HydroNode::Placeholder)),
                    trusted: false,
                    metadata: self
                        .location
                        .clone()
                        .with_consistency_of::<L2>()
                        .new_node_metadata(KeyedSingleton::<K, V, L2, B>::collection_kind()),
                },
            )
        }
    }
}

#[cfg(stageleft_runtime)]
fn key_count_inside_tick<'a, K, V, L: Location<'a>>(
    me: KeyedSingleton<K, V, L, Bounded>,
) -> Singleton<usize, L, Bounded> {
    me.entries().count()
}

#[cfg(stageleft_runtime)]
fn into_singleton_inside_tick<'a, K, V, L: Location<'a>>(
    me: KeyedSingleton<K, V, L, Bounded>,
) -> Singleton<HashMap<K, V>, L, Bounded>
where
    K: Eq + Hash,
{
    me.entries()
        .assume_ordering_trusted(nondet!(
            /// There is only one element associated with each key. The closure technically
            /// isn't commutative in the case where both passed entries have the same key
            /// but different values.
            ///
            /// In the future, we may want to have an `assume!(...)` statement in the UDF that
            /// the key is never already present in the map.
        ))
        .fold(
            q!(|| HashMap::new()),
            q!(|map, (k, v)| {
                map.insert(k, v);
            }),
        )
}

impl<'a, K, V, L: Location<'a>, B: KeyedSingletonBound> KeyedSingleton<K, V, L, B> {
    pub(crate) fn collection_kind() -> CollectionKind {
        CollectionKind::KeyedSingleton {
            bound: B::bound_kind(),
            key_type: stageleft::quote_type::<K>().into(),
            value_type: stageleft::quote_type::<V>().into(),
        }
    }

    /// Transforms each value by invoking `f` on each element, with keys staying the same
    /// after transformation. If you need access to the key, see [`KeyedSingleton::map_with_key`].
    ///
    /// If you do not want to modify the stream and instead only want to view
    /// each item use [`KeyedSingleton::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| {
    /// let keyed_singleton = // { 1: 2, 2: 4 }
    /// # process
    /// #     .source_iter(q!(vec![(1, 2), (2, 4)]))
    /// #     .into_keyed()
    /// #     .first();
    /// keyed_singleton.map(q!(|v| v + 1))
    /// #   .entries()
    /// # }, |mut stream| async move {
    /// // { 1: 3, 2: 5 }
    /// # let mut results = Vec::new();
    /// # for _ in 0..2 {
    /// #     results.push(stream.next().await.unwrap());
    /// # }
    /// # results.sort();
    /// # assert_eq!(results, vec![(1, 3), (2, 5)]);
    /// # }));
    /// # }
    /// ```
    pub fn map<U, F>(
        self,
        f: impl IntoQuotedMut<'a, F, OperatorContext<L, B::UnderlyingBound>> + Copy,
    ) -> KeyedSingleton<K, U, L, B::EraseMonotonic>
    where
        F: Fn(V) -> U + 'a,
    {
        let f: ManualExpr<F, _> =
            ManualExpr::new(move |ctx: &OperatorContext<L, B::UnderlyingBound>| {
                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)>(&OperatorContext::<L, B::UnderlyingBound>::new(
            &self.location,
        ))
        .into();

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

    /// Transforms each value by invoking `f` on each key-value pair, with keys staying the same
    /// after transformation. Unlike [`KeyedSingleton::map`], this gives access to both the key and value.
    ///
    /// The closure `f` receives a tuple `(K, V)` containing both the key and value, and returns
    /// the new value `U`. The key remains unchanged in the output.
    ///
    /// # 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 keyed_singleton = // { 1: 2, 2: 4 }
    /// # process
    /// #     .source_iter(q!(vec![(1, 2), (2, 4)]))
    /// #     .into_keyed()
    /// #     .first();
    /// keyed_singleton.map_with_key(q!(|(k, v)| k + v))
    /// #   .entries()
    /// # }, |mut stream| async move {
    /// // { 1: 3, 2: 6 }
    /// # let mut results = Vec::new();
    /// # for _ in 0..2 {
    /// #     results.push(stream.next().await.unwrap());
    /// # }
    /// # results.sort();
    /// # assert_eq!(results, vec![(1, 3), (2, 6)]);
    /// # }));
    /// # }
    /// ```
    pub fn map_with_key<U, F>(
        self,
        f: impl IntoQuotedMut<'a, F, OperatorContext<L, B::UnderlyingBound>> + Copy,
    ) -> KeyedSingleton<K, U, L, B::EraseMonotonic>
    where
        F: Fn((K, V)) -> U + 'a,
        K: Clone,
    {
        let f: ManualExpr<F, _> =
            ManualExpr::new(move |ctx: &OperatorContext<L, B::UnderlyingBound>| {
                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)>(&OperatorContext::<L, B::UnderlyingBound>::new(
            &self.location,
        ))
        .into();

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

    /// Gets the number of keys in the keyed singleton.
    ///
    /// The output singleton will be unbounded if the input is [`Unbounded`] or [`BoundedValue`],
    /// since keys may be added / removed over time. When the set of keys changes, the count will
    /// be asynchronously updated.
    ///
    /// # 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_singleton = // { 1: "a", 2: "b", 3: "c" }
    /// # process
    /// #     .source_iter(q!(vec![(1, "a"), (2, "b"), (3, "c")]))
    /// #     .into_keyed()
    /// #     .batch(&tick, nondet!(/** test */))
    /// #     .first();
    /// keyed_singleton.key_count()
    /// # .all_ticks()
    /// # }, |mut stream| async move {
    /// // 3
    /// # assert_eq!(stream.next().await.unwrap(), 3);
    /// # }));
    /// # }
    /// ```
    pub fn key_count(self) -> Singleton<usize, L, B::UnderlyingBound> {
        if B::ValueBound::BOUNDED {
            let me: KeyedSingleton<K, V, L, B::WithBoundedValue> = KeyedSingleton {
                location: self.location.clone(),
                flow_state: self.flow_state.clone(),
                ir_node: super::tracked_ir_node(
                    &self.flow_state,
                    self.ir_node.replace(HydroNode::Placeholder),
                ),
                _phantom: PhantomData,
            };

            me.entries().count().ignore_monotonic()
        } else if L::is_top_level()
            && let Some(tick) = self.location.try_tick()
            && (B::bound_kind() == KeyedSingletonBoundKind::Unbounded
                || B::bound_kind() == KeyedSingletonBoundKind::MonotonicKeys
                || B::bound_kind() == KeyedSingletonBoundKind::MonotonicValue)
        {
            let location = self.location.clone();
            let ir_node = self.ir_node.replace(HydroNode::Placeholder);
            let me: KeyedSingleton<K, V, L, MonotonicKeys> =
                KeyedSingleton::new(location.clone(), ir_node);

            let out = key_count_inside_tick(me.snapshot(&tick, nondet!(/** eventually stabilizes */)))
                    .latest()
                    // The key count is folded with an initial value, so it is always present
                    // (0 when there are no keys). `latest()` is null until the producing tick
                    // first runs; fill that prefix with 0 to recover an always-present count.
                    .unwrap_or(location.singleton(q!(0usize)).into());
            // Re-tag the node from the concrete `Unbounded` singleton to the `B::UnderlyingBound`
            // that this method returns (equal at runtime for this branch).
            Singleton::new(location, out.ir_node.replace(HydroNode::Placeholder))
        } else {
            panic!("BoundedValue or Unbounded KeyedSingleton inside a tick, not supported");
        }
    }

    /// Converts this keyed singleton into a [`Singleton`] containing a `HashMap` from keys to values.
    ///
    /// As the values for each key are updated asynchronously, the `HashMap` will be updated
    /// asynchronously as well.
    ///
    /// # 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 keyed_singleton = // { 1: "a", 2: "b", 3: "c" }
    /// # process
    /// #     .source_iter(q!(vec![(1, "a".to_owned()), (2, "b".to_owned()), (3, "c".to_owned())]))
    /// #     .into_keyed()
    /// #     .batch(&process.tick(), nondet!(/** test */))
    /// #     .first();
    /// keyed_singleton.into_singleton()
    /// # .all_ticks()
    /// # }, |mut stream| async move {
    /// // { 1: "a", 2: "b", 3: "c" }
    /// # assert_eq!(stream.next().await.unwrap(), vec![(1, "a".to_owned()), (2, "b".to_owned()), (3, "c".to_owned())].into_iter().collect());
    /// # }));
    /// # }
    /// ```
    pub fn into_singleton(self) -> Singleton<HashMap<K, V>, L, B::UnderlyingBound>
    where
        K: Eq + Hash + Clone + 'a,
        V: Clone + 'a,
    {
        if B::ValueBound::BOUNDED {
            let me: KeyedSingleton<K, V, L, B::WithBoundedValue> = KeyedSingleton {
                location: self.location.clone(),
                flow_state: self.flow_state.clone(),
                ir_node: super::tracked_ir_node(
                    &self.flow_state,
                    self.ir_node.replace(HydroNode::Placeholder),
                ),
                _phantom: PhantomData,
            };

            me.entries()
                .assume_ordering_trusted(nondet!(
                    /// There is only one element associated with each key. The closure technically
                    /// isn't commutative in the case where both passed entries have the same key
                    /// but different values.
                    ///
                    /// In the future, we may want to have an `assume!(...)` statement in the UDF that
                    /// the key is never already present in the map.
                ))
                .fold(
                    q!(|| HashMap::new()),
                    q!(|map, (k, v)| {
                        // TODO(shadaj): make this commutative but really-debug-assert that there is no key overlap
                        map.insert(k, v);
                    }),
                )
        } else if L::is_top_level()
            && let Some(tick) = self.location.try_tick()
            && (B::bound_kind() == KeyedSingletonBoundKind::Unbounded
                || B::bound_kind() == KeyedSingletonBoundKind::MonotonicKeys
                || B::bound_kind() == KeyedSingletonBoundKind::MonotonicValue)
        {
            let location = self.location.clone();
            let ir_node = self.ir_node.replace(HydroNode::Placeholder);
            let me: KeyedSingleton<K, V, L, MonotonicKeys> =
                KeyedSingleton::new(location.clone(), ir_node);

            let out = into_singleton_inside_tick(
                me.snapshot(&tick, nondet!(/** eventually stabilizes */)),
            )
            .latest()
            // The map is folded with an initial value, so it is always present (empty when
            // there are no keys). `latest()` is null until the producing tick first runs; fill
            // that prefix with an empty map to recover an always-present map.
            .unwrap_or(location.singleton(q!(HashMap::new())).into());
            // Re-tag the node from the concrete `Unbounded` singleton to the `B::UnderlyingBound`
            // that this method returns (equal at runtime for this branch).
            Singleton::new(location, out.ir_node.replace(HydroNode::Placeholder))
        } else {
            panic!("BoundedValue or Unbounded KeyedSingleton inside a tick, not supported");
        }
    }

    /// 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) -> KeyedSingleton<K, V, L, B> {
        {
            let mut node = self.ir_node.borrow_mut();
            let metadata = node.metadata_mut();
            metadata.tag = Some(name.to_owned());
        }
        self
    }

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

    /// Gets the value associated with a specific key from the keyed singleton.
    /// Returns `None` if the key is `None` or there is no associated 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, 2), (2, 3)]))
    ///     .into_keyed()
    ///     .batch(&tick, nondet!(/** test */))
    ///     .first();
    /// let key = tick.singleton(q!(1));
    /// keyed_data.get(key).all_ticks()
    /// # }, |mut stream| async move {
    /// // 2
    /// # assert_eq!(stream.next().await.unwrap(), 2);
    /// # }));
    /// # }
    /// ```
    pub fn get(self, key: impl Into<Optional<K, L, Bounded>>) -> Optional<V, L, Bounded>
    where
        B: IsBounded,
        K: Hash + Eq + Clone,
        V: Clone,
    {
        self.make_bounded()
            .into_keyed_stream()
            .get(key)
            .cast_at_most_one_element()
    }

    /// Emit a keyed stream containing keys shared between the keyed singleton and the
    /// keyed stream, where each value in the output keyed stream is a tuple of
    /// (the keyed singleton's value, the keyed stream'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), (2, 20)]))
    ///     .into_keyed()
    ///     .batch(&tick, nondet!(/** test */))
    ///     .first();
    /// let other_data = process
    ///     .source_iter(q!(vec![(1, 100), (2, 200), (1, 101)]))
    ///     .into_keyed()
    ///     .batch(&tick, nondet!(/** test */));
    /// keyed_data.join_keyed_stream(other_data).entries().all_ticks()
    /// # }, |mut stream| async move {
    /// // { 1: [(10, 100), (10, 101)], 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, (10, 101)), (2, (20, 200))]);
    /// # }));
    /// # }
    /// ```
    pub fn join_keyed_stream<O2: Ordering, R2: Retries, V2, B2: Boundedness>(
        self,
        other: KeyedStream<K, V2, L, B2, O2, R2>,
    ) -> KeyedStream<K, (V, V2), L, B2, O2, R2>
    where
        B: IsBounded,
        K: Eq + Hash + Clone,
        V: Clone,
        V2: Clone,
    {
        // TODO(shadaj): if DFIR guarantees that joining unbounded keyed stream x bounded keyed stream
        // always produces deterministic order per key (nested loop join), this could just use
        // `join_keyed_stream` without constructing IRs manually
        KeyedStream::new(
            self.location.clone(),
            HydroNode::Join {
                left: Box::new(self.ir_node.replace(HydroNode::Placeholder)),
                right: Box::new(other.ir_node.replace(HydroNode::Placeholder)),
                metadata: self
                    .location
                    .new_node_metadata(KeyedStream::<K, (V, V2), L, B2, O2, R2>::collection_kind()),
            },
        )
    }

    /// Emit a keyed singleton containing all keys shared between two keyed singletons,
    /// where each value in the output keyed singleton is a tuple of
    /// (self.value, other.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 requests = // { 1: 10, 2: 20, 3: 30 }
    /// # process
    /// #     .source_iter(q!(vec![(1, 10), (2, 20), (3, 30)]))
    /// #     .into_keyed()
    /// #     .batch(&tick, nondet!(/** test */))
    /// #     .first();
    /// let other = // { 1: 100, 2: 200, 4: 400 }
    /// # process
    /// #     .source_iter(q!(vec![(1, 100), (2, 200), (4, 400)]))
    /// #     .into_keyed()
    /// #     .batch(&tick, nondet!(/** test */))
    /// #     .first();
    /// requests.join_keyed_singleton(other)
    /// # .entries().all_ticks()
    /// # }, |mut stream| async move {
    /// // { 1: (10, 100), 2: (20, 200) }
    /// # let mut results = vec![];
    /// # for _ in 0..2 {
    /// #     results.push(stream.next().await.unwrap());
    /// # }
    /// # results.sort();
    /// # assert_eq!(results, vec![(1, (10, 100)), (2, (20, 200))]);
    /// # }));
    /// # }
    /// ```
    pub fn join_keyed_singleton<V2: Clone>(
        self,
        other: KeyedSingleton<K, V2, L, Bounded>,
    ) -> KeyedSingleton<K, (V, V2), L, Bounded>
    where
        B: IsBounded,
        K: Eq + Hash + Clone,
        V: Clone,
    {
        let result_stream = self
            .make_bounded()
            .entries()
            .join(other.entries())
            .into_keyed();

        // The cast is guaranteed to succeed, since each key (in both `self` and `other`) has at most one value.
        result_stream.cast_at_most_one_entry_per_key()
    }

    /// For each value in `self`, find the matching key in `lookup`.
    /// The output is a keyed singleton 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, 2: 20 }
    /// # process
    /// #     .source_iter(q!(vec![(1, 10), (2, 20)]))
    /// #     .into_keyed()
    /// #     .batch(&tick, nondet!(/** test */))
    /// #     .first();
    /// 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)), 2: (20, None) }
    /// # let mut results = vec![];
    /// # for _ in 0..2 {
    /// #     results.push(stream.next().await.unwrap());
    /// # }
    /// # results.sort();
    /// # assert_eq!(results, vec![(1, (10, Some(100))), (2, (20, None))]);
    /// # }));
    /// # }
    /// ```
    pub fn lookup_keyed_singleton<V2>(
        self,
        lookup: KeyedSingleton<V, V2, L, Bounded>,
    ) -> KeyedSingleton<K, (V, Option<V2>), L, Bounded>
    where
        B: IsBounded,
        K: Eq + Hash + Clone,
        V: Eq + Hash + Clone,
        V2: Clone,
    {
        let result_stream = self
            .make_bounded()
            .into_keyed_stream()
            .lookup_keyed_stream(lookup.into_keyed_stream());

        // The cast is guaranteed to succeed since both lookup and self contain at most 1 value per key
        result_stream.cast_at_most_one_entry_per_key()
    }

    /// 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, 2: 20 }
    /// # process
    /// #     .source_iter(q!(vec![(1, 10), (2, 20)]))
    /// #     .into_keyed()
    /// #     .batch(&tick, nondet!(/** test */))
    /// #     .first();
    /// let other_data = // { 10: 100, 10: 110 }
    /// # process
    /// #     .source_iter(q!(vec![(10, 100), (10, 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(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, (10, Some(110))), (2, (20, None))]);
    /// # }));
    /// # }
    /// ```
    pub fn lookup_keyed_stream<V2, O: Ordering, R: Retries>(
        self,
        lookup: KeyedStream<V, V2, L, Bounded, O, R>,
    ) -> KeyedStream<K, (V, Option<V2>), L, Bounded, NoOrder, R>
    where
        B: IsBounded,
        K: Eq + Hash + Clone,
        V: Eq + Hash + Clone,
        V2: Clone,
    {
        self.make_bounded()
            .entries()
            .weaken_retries::<R>() // TODO: Once weaken_retries() is implemented for KeyedSingleton, remove entries() and into_keyed()
            .into_keyed()
            .lookup_keyed_stream(lookup)
    }

    /// For each key present in both `self` and `thresholds`, emits a [`KeyedStream`] event the first
    /// time that key's value becomes greater than or equal to the corresponding threshold value.
    /// The emitted value for each key is the threshold value itself.
    ///
    /// This requires the keyed singleton to have monotonic values ([`MonotonicValue`] or [`Bounded`]),
    /// because otherwise the threshold detection would be non-deterministic.
    ///
    /// The `thresholds` parameter is a [`BoundedValue`] keyed singleton mapping each key to its
    /// threshold. Thresholds may arrive asynchronously (new keys appear over time), but once set
    /// for a key, the threshold value is fixed. Late-arriving thresholds are checked against the
    /// current snapshot value immediately.
    ///
    /// # Example
    /// ```rust,ignore
    /// use hydro_lang::prelude::*;
    ///
    /// // Given a monotonically increasing keyed singleton (e.g. from fold with monotone proof)
    /// let counts: KeyedSingleton<u32, usize, _, MonotonicValue> = events.into_keyed()
    ///     .fold(q!(|| 0), q!(|acc, _| *acc += 1, monotone = manual_proof!(/** +1 is monotone */)));
    ///
    /// // BoundedValue keyed singleton of thresholds (from .first())
    /// let thresholds = threshold_source.into_keyed().first();
    ///
    /// // Emits (key, threshold_value) the first time each key's value >= threshold
    /// let crossed = counts.threshold_greater_or_equal(thresholds);
    /// ```
    pub fn threshold_greater_or_equal(
        self,
        thresholds: KeyedSingleton<K, V, L, BoundedValue>,
    ) -> KeyedStream<K, V, L, B::UnderlyingBound, NoOrder, ExactlyOnce>
    where
        K: Clone + Eq + Hash,
        V: Clone + PartialOrd,
        B: IsKeyedMonotonic,
    {
        let self_location = self.location.clone();
        match B::bound_kind() {
            KeyedSingletonBoundKind::Bounded => {
                // Bounded case: self is already fixed, just join and filter
                let me: KeyedSingleton<K, V, L, Bounded> = KeyedSingleton::new(
                    self.location.clone(),
                    self.ir_node.replace(HydroNode::Placeholder),
                );
                let result = me
                    .entries()
                    .join(thresholds.entries())
                    .filter_map(q!(|(k, (val, thresh))| {
                        if val >= thresh {
                            Some((k, thresh))
                        } else {
                            None
                        }
                    }))
                    .into_keyed();
                KeyedStream::new(
                    result.location.clone(),
                    result.ir_node.replace(HydroNode::Placeholder),
                )
            }
            KeyedSingletonBoundKind::MonotonicValue => {
                let me: KeyedSingleton<K, V, L, MonotonicValue> = KeyedSingleton::new(
                    self.location.clone(),
                    self.ir_node.replace(HydroNode::Placeholder),
                );

                let result = sliced! {
                    let snapshot = use::snapshot(me, nondet!(/** thresholds are deterministic */));
                    let thresh_snapshot =
                        use::batch(thresholds, nondet!(/** thresholds are deterministic */));
                    let mut already_crossed =
                        use::state_null::<Stream<K, Tick<_>, Bounded, NoOrder>>();

                    let joined = thresh_snapshot.entries().join(snapshot.entries());
                    let passed = joined
                        .filter(q!(|(_, (thresh, val))| *val >= *thresh))
                        .map(q!(|(k, (thresh, _))| (k, thresh)));

                    let newly_crossed = passed.anti_join(already_crossed.clone());
                    already_crossed =
                        already_crossed.chain(newly_crossed.clone().map(q!(|(k, _)| k)));

                    newly_crossed.into_keyed()
                };

                KeyedStream::new(
                    self_location,
                    result.ir_node.replace(HydroNode::Placeholder),
                )
            }
            KeyedSingletonBoundKind::BoundedValue => {
                let me: KeyedSingleton<K, V, L, BoundedValue> = KeyedSingleton::new(
                    self.location.clone(),
                    self.ir_node.replace(HydroNode::Placeholder),
                );

                let result = sliced! {
                    let snapshot = use::batch(me, nondet!(/** thresholds are deterministic */));
                    let thresh_snapshot =
                        use::batch(thresholds, nondet!(/** thresholds are deterministic */));
                    let mut already_crossed =
                        use::state_null::<Stream<K, Tick<_>, Bounded, NoOrder>>();

                    let joined = thresh_snapshot.entries().join(snapshot.entries());
                    let passed = joined
                        .filter(q!(|(_, (thresh, val))| *val >= *thresh))
                        .map(q!(|(k, (thresh, _))| (k, thresh)));

                    let newly_crossed = passed.anti_join(already_crossed.clone());
                    already_crossed =
                        already_crossed.chain(newly_crossed.clone().map(q!(|(k, _)| k)));

                    newly_crossed.into_keyed()
                };

                KeyedStream::new(
                    self_location,
                    result.ir_node.replace(HydroNode::Placeholder),
                )
            }
            _ => {
                unreachable!(
                    "IsKeyedMonotonic is only implemented for Bounded, BoundedValue, and MonotonicValue"
                )
            }
        }
    }

    /// Like [`Self::threshold_greater_or_equal`], but uses a single [`Singleton`] threshold
    /// shared across all keys. Emits a `(K, V)` event for each key the first time that key's
    /// value becomes >= the threshold. The emitted value is the threshold itself.
    ///
    /// Because the threshold is a [`Bounded`] singleton, it is a compile-time constant and
    /// does not carry ongoing memory cost.
    ///
    /// # Example
    /// ```rust
    /// # #[cfg(feature = "deploy")] {
    /// # use hydro_lang::prelude::*;
    /// # use futures::StreamExt;
    /// # tokio_test::block_on(hydro_lang::test_util::stream_transform_test(|process| {
    /// // A keyed singleton of per-key values (in practice often a monotone counter): { 1: 6, 2: 4 }
    /// let counts = process
    ///     .source_iter(q!(vec![(1, 6), (2, 4)]))
    ///     .into_keyed()
    ///     .first();
    ///
    /// // A single threshold value shared across all keys
    /// let threshold = process.singleton(q!(5));
    ///
    /// // Emits (key, threshold) the first time each key's value >= threshold
    /// counts.threshold_greater_or_equal_uniform(threshold)
    /// #   .entries()
    /// # }, |mut stream| async move {
    /// // { 1: 5 } -- key 1's value 6 >= 5, but key 2's value 4 < 5
    /// # assert_eq!(stream.next().await.unwrap(), (1, 5));
    /// # }));
    /// # }
    /// ```
    pub fn threshold_greater_or_equal_uniform(
        self,
        threshold: Singleton<V, L, Bounded>,
    ) -> KeyedStream<K, V, L, B::UnderlyingBound, NoOrder, ExactlyOnce>
    where
        K: Clone + Eq + Hash,
        V: Clone + PartialOrd,
        B: IsKeyedMonotonic,
    {
        let self_location = self.location.clone();
        match B::bound_kind() {
            KeyedSingletonBoundKind::Bounded => {
                let me: KeyedSingleton<K, V, L, Bounded> = KeyedSingleton::new(
                    self.location.clone(),
                    self.ir_node.replace(HydroNode::Placeholder),
                );
                let result = me
                    .entries()
                    .cross_singleton(threshold)
                    .filter_map(q!(|((k, val), thresh)| {
                        if val >= thresh {
                            Some((k, thresh))
                        } else {
                            None
                        }
                    }))
                    .into_keyed();
                KeyedStream::new(
                    result.location.clone(),
                    result.ir_node.replace(HydroNode::Placeholder),
                )
            }
            KeyedSingletonBoundKind::MonotonicValue => {
                let me: KeyedSingleton<K, V, L, MonotonicValue> = KeyedSingleton::new(
                    self.location.clone(),
                    self.ir_node.replace(HydroNode::Placeholder),
                );

                let result = sliced! {
                    let snapshot = use::snapshot(me, nondet!(/** thresholds are deterministic */));
                    let mut already_crossed =
                        use::state_null::<Stream<K, Tick<_>, Bounded, NoOrder>>();

                    let tick = snapshot.location().clone();
                    let thresh_in_tick = threshold.clone_into_tick(&tick);

                    let crossing = snapshot
                        .entries()
                        .cross_singleton(thresh_in_tick)
                        .filter_map(q!(|((k, val), thresh)| {
                            if val >= thresh {
                                Some((k, thresh))
                            } else {
                                None
                            }
                        }));

                    let newly_crossed = crossing.anti_join(already_crossed.clone());
                    already_crossed =
                        already_crossed.chain(newly_crossed.clone().map(q!(|(k, _)| k)));

                    newly_crossed.into_keyed()
                };

                KeyedStream::new(
                    self_location,
                    result.ir_node.replace(HydroNode::Placeholder),
                )
            }
            KeyedSingletonBoundKind::BoundedValue => {
                let me: KeyedSingleton<K, V, L, BoundedValue> = KeyedSingleton::new(
                    self.location.clone(),
                    self.ir_node.replace(HydroNode::Placeholder),
                );

                let result = sliced! {
                    let snapshot = use::batch(me, nondet!(/** thresholds are deterministic */));
                    let mut already_crossed =
                        use::state_null::<Stream<K, Tick<_>, Bounded, NoOrder>>();

                    let tick = snapshot.location().clone();
                    let thresh_in_tick = threshold.clone_into_tick(&tick);

                    let crossing = snapshot
                        .entries()
                        .cross_singleton(thresh_in_tick)
                        .filter_map(q!(|((k, val), thresh)| {
                            if val >= thresh {
                                Some((k, thresh))
                            } else {
                                None
                            }
                        }));

                    let newly_crossed = crossing.anti_join(already_crossed.clone());
                    already_crossed =
                        already_crossed.chain(newly_crossed.clone().map(q!(|(k, _)| k)));

                    newly_crossed.into_keyed()
                };

                KeyedStream::new(
                    self_location,
                    result.ir_node.replace(HydroNode::Placeholder),
                )
            }
            _ => {
                unreachable!(
                    "IsKeyedMonotonic is only implemented for Bounded, BoundedValue, and MonotonicValue"
                )
            }
        }
    }
}

impl<'a, K, V, L: Location<'a>, B: KeyedSingletonBound<ValueBound = Bounded>>
    KeyedSingleton<K, V, L, B>
{
    /// Flattens the keyed singleton into an unordered stream of key-value pairs.
    ///
    /// The value for each key must be bounded, otherwise the resulting stream elements would be
    /// non-deterministic. As new entries are added to the keyed singleton, they will be streamed
    /// into the output.
    ///
    /// # 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 keyed_singleton = // { 1: 2, 2: 4 }
    /// # process
    /// #     .source_iter(q!(vec![(1, 2), (2, 4)]))
    /// #     .into_keyed()
    /// #     .first();
    /// keyed_singleton.entries()
    /// # }, |mut stream| async move {
    /// // (1, 2), (2, 4) 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), (2, 4)]);
    /// # }));
    /// # }
    /// ```
    pub fn entries(self) -> Stream<(K, V), L, B::UnderlyingBound, NoOrder, ExactlyOnce> {
        self.into_keyed_stream().entries()
    }

    /// Flattens the keyed singleton into an unordered stream of just the values.
    ///
    /// The value for each key must be bounded, otherwise the resulting stream elements would be
    /// non-deterministic. As new entries are added to the keyed singleton, they will be streamed
    /// into the output.
    ///
    /// # 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 keyed_singleton = // { 1: 2, 2: 4 }
    /// # process
    /// #     .source_iter(q!(vec![(1, 2), (2, 4)]))
    /// #     .into_keyed()
    /// #     .first();
    /// keyed_singleton.values()
    /// # }, |mut stream| async move {
    /// // 2, 4 in any order
    /// # let mut results = Vec::new();
    /// # for _ in 0..2 {
    /// #     results.push(stream.next().await.unwrap());
    /// # }
    /// # results.sort();
    /// # assert_eq!(results, vec![2, 4]);
    /// # }));
    /// # }
    /// ```
    pub fn values(self) -> Stream<V, L, B::UnderlyingBound, NoOrder, ExactlyOnce> {
        let map_f = q!(|(_, v)| v)
            .splice_fn1_ctx::<(K, V), V>(&OperatorContext::<L, B::UnderlyingBound>::new(
                &self.location,
            ))
            .into();

        Stream::new(
            self.location.clone(),
            HydroNode::Map {
                f: map_f,
                input: Box::new(self.ir_node.replace(HydroNode::Placeholder)),
                metadata: self.location.new_node_metadata(Stream::<
                    V,
                    L,
                    B::UnderlyingBound,
                    NoOrder,
                    ExactlyOnce,
                >::collection_kind()),
            },
        )
    }

    /// Flattens the keyed singleton into an unordered stream of just the keys.
    ///
    /// The value for each key must be bounded, otherwise the removal of keys would result in
    /// non-determinism. As new entries are added to the keyed singleton, they will be streamed
    /// into the output.
    ///
    /// # 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 keyed_singleton = // { 1: 2, 2: 4 }
    /// # process
    /// #     .source_iter(q!(vec![(1, 2), (2, 4)]))
    /// #     .into_keyed()
    /// #     .first();
    /// keyed_singleton.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::UnderlyingBound, NoOrder, ExactlyOnce> {
        self.entries().map(q!(|(k, _)| k))
    }

    /// Given a bounded stream of keys `K`, returns a new keyed singleton containing only the
    /// entries whose keys are not in the provided 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_singleton = // { 1: 2, 2: 4 }
    /// # process
    /// #     .source_iter(q!(vec![(1, 2), (2, 4)]))
    /// #     .into_keyed()
    /// #     .first()
    /// #     .batch(&tick, nondet!(/** test */));
    /// let keys_to_remove = process
    ///     .source_iter(q!(vec![1]))
    ///     .batch(&tick, nondet!(/** test */));
    /// keyed_singleton.filter_key_not_in(keys_to_remove)
    /// #   .entries().all_ticks()
    /// # }, |mut stream| async move {
    /// // { 2: 4 }
    /// # for w in vec![(2, 4)] {
    /// #     assert_eq!(stream.next().await.unwrap(), w);
    /// # }
    /// # }));
    /// # }
    /// ```
    pub fn filter_key_not_in<O2: Ordering, R2: Retries>(
        self,
        other: Stream<K, L, Bounded, O2, R2>,
    ) -> Self
    where
        K: Hash + Eq,
    {
        check_matching_location(&self.location, &other.location);

        KeyedSingleton::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()),
            },
        )
    }

    /// An operator which allows you to "inspect" each value of a keyed singleton 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| {
    /// let keyed_singleton = // { 1: 2, 2: 4 }
    /// # process
    /// #     .source_iter(q!(vec![(1, 2), (2, 4)]))
    /// #     .into_keyed()
    /// #     .first();
    /// keyed_singleton
    ///     .inspect(q!(|v| println!("{}", v)))
    /// #   .entries()
    /// # }, |mut stream| async move {
    /// // { 1: 2, 2: 4 }
    /// # for w in vec![(1, 2), (2, 4)] {
    /// #     assert_eq!(stream.next().await.unwrap(), w);
    /// # }
    /// # }));
    /// # }
    /// ```
    pub fn inspect<F>(
        self,
        f: impl IntoQuotedMut<'a, F, OperatorContext<L, B::UnderlyingBound>> + Copy,
    ) -> Self
    where
        F: Fn(&V) + 'a,
    {
        let f: ManualExpr<F, _> =
            ManualExpr::new(move |ctx: &OperatorContext<L, B::UnderlyingBound>| {
                f.splice_fn1_borrow_ctx(ctx)
            });
        let inspect_f = q!({
            let orig = f;
            move |t: &(_, _)| orig(&t.1)
        })
        .splice_fn1_borrow_ctx::<(K, V), ()>(&OperatorContext::<L, B::UnderlyingBound>::new(
            &self.location,
        ))
        .into();

        KeyedSingleton::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 entry of a keyed singleton 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| {
    /// let keyed_singleton = // { 1: 2, 2: 4 }
    /// # process
    /// #     .source_iter(q!(vec![(1, 2), (2, 4)]))
    /// #     .into_keyed()
    /// #     .first();
    /// keyed_singleton
    ///     .inspect_with_key(q!(|(k, v)| println!("{}: {}", k, v)))
    /// #   .entries()
    /// # }, |mut stream| async move {
    /// // { 1: 2, 2: 4 }
    /// # for w in vec![(1, 2), (2, 4)] {
    /// #     assert_eq!(stream.next().await.unwrap(), w);
    /// # }
    /// # }));
    /// # }
    /// ```
    pub fn inspect_with_key<F>(
        self,
        f: impl IntoQuotedMut<'a, F, OperatorContext<L, B::UnderlyingBound>>,
    ) -> Self
    where
        F: Fn(&(K, V)) + 'a,
    {
        let inspect_f = f
            .splice_fn1_borrow_ctx::<(K, V), ()>(&OperatorContext::<L, B::UnderlyingBound>::new(
                &self.location,
            ))
            .into();

        KeyedSingleton::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()),
            },
        )
    }

    /// Gets the key-value tuple with the largest key among all entries in this [`KeyedSingleton`].
    ///
    /// Because this method requires values to be bounded, the output [`Optional`] will only be
    /// asynchronously updated if a new key is added that is higher than the previous max 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_singleton = // { 1: 123, 2: 456, 0: 789 }
    /// # Stream::<_, _>::from(process.source_iter(q!(vec![(1, 123), (2, 456), (0, 789)])))
    /// #     .into_keyed()
    /// #     .first();
    /// keyed_singleton.get_max_key()
    /// # .sample_eager(nondet!(/** test */))
    /// # }, |mut stream| async move {
    /// // (2, 456)
    /// # assert_eq!(stream.next().await.unwrap(), (2, 456));
    /// # }));
    /// # }
    /// ```
    pub fn get_max_key(
        self,
    ) -> Optional<(K, V), L, <B::UnderlyingBound as Boundedness>::AggregatedOptional>
    where
        K: Ord,
    {
        self.entries()
            .assume_ordering_trusted(nondet!(
                /// There is only one element associated with each key, and the keys are totallly
                /// ordered so we will produce a deterministic value. The closure technically
                /// isn't commutative in the case where both passed entries have the same key
                /// but different values.
                ///
                /// In the future, we may want to have an `assume!(...)` statement in the UDF that
                /// the two inputs do not have the same key.
            ))
            .reduce(q!(
                move |curr, new| {
                    if new.0 > curr.0 {
                        *curr = new;
                    }
                },
                idempotent = manual_proof!(/** repeated elements are ignored */)
            ))
    }

    /// Converts this keyed singleton into a [`KeyedStream`] with each group having a single
    /// element, the value.
    ///
    /// This is the equivalent of [`Singleton::into_stream`] but keyed.
    ///
    /// # 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 keyed_singleton = // { 1: 2, 2: 4 }
    /// # Stream::<_, _>::from(process.source_iter(q!(vec![(1, 2), (2, 4)])))
    /// #     .into_keyed()
    /// #     .first();
    /// keyed_singleton
    ///     .clone()
    ///     .into_keyed_stream()
    ///     .merge_unordered(
    ///         keyed_singleton.into_keyed_stream()
    ///     )
    /// #   .entries()
    /// # }, |mut stream| async move {
    /// /// // { 1: [2, 2], 2: [4, 4] }
    /// # for w in vec![(1, 2), (2, 4), (1, 2), (2, 4)] {
    /// #     assert_eq!(stream.next().await.unwrap(), w);
    /// # }
    /// # }));
    /// # }
    /// ```
    pub fn into_keyed_stream(
        self,
    ) -> KeyedStream<K, V, L, B::UnderlyingBound, TotalOrder, ExactlyOnce> {
        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::UnderlyingBound,
                    TotalOrder,
                    ExactlyOnce,
                >::collection_kind()),
            },
        )
    }
}

impl<'a, K, V, L, B: KeyedSingletonBound> KeyedSingleton<K, V, L, B>
where
    L: Location<'a>,
    B: KeyedSingletonBound<ValueBound = Bounded>,
{
    /// Shifts this bounded-value keyed singleton into an atomic context, which guarantees that any downstream logic
    /// will all be executed synchronously before any outputs are yielded (in [`KeyedSingleton::end_atomic`]).
    ///
    /// This is useful to enforce local consistency constraints, such as ensuring that a write is
    /// processed before an acknowledgement is emitted.
    pub fn atomic(self) -> KeyedSingleton<K, V, Atomic<L>, B>
    where
        L: TopLevel<'a>,
    {
        let out_location = Atomic {
            tick: self.location.tick(),
        };
        KeyedSingleton::new(
            out_location.clone(),
            HydroNode::BeginAtomic {
                inner: Box::new(self.ir_node.replace(HydroNode::Placeholder)),
                metadata: out_location
                    .new_node_metadata(KeyedSingleton::<K, V, Atomic<L>, B>::collection_kind()),
            },
        )
    }
}

impl<'a, K, V, L, B: KeyedSingletonBound> KeyedSingleton<K, V, Atomic<L>, B>
where
    L: Location<'a>,
{
    /// Yields the elements of this keyed singleton back into a top-level, asynchronous execution context.
    /// See [`KeyedSingleton::atomic`] for more details.
    pub fn end_atomic(self) -> KeyedSingleton<K, V, L, B> {
        KeyedSingleton::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(KeyedSingleton::<K, V, L, B>::collection_kind()),
            },
        )
    }
}

impl<'a, K, V, L: Location<'a>> KeyedSingleton<K, V, Tick<L>, Bounded> {
    /// Shifts the state in `self` to the **next tick**, so that the returned keyed singleton at
    /// tick `T` always has the entries of `self` at tick `T - 1`.
    ///
    /// At tick `0`, the output has no entries, 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 compare state 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), (2, 3)]))
    /// #   .batch(&tick, nondet!(/** test */))
    /// #   .into_keyed();
    /// # let batch_second_tick = process
    /// #   .source_iter(q!(vec![(2, 4), (3, 5)]))
    /// #   .batch(&tick, nondet!(/** test */))
    /// #   .into_keyed()
    /// #   .defer_tick(); // appears on the second tick
    /// let input_batch = // first tick: { 1: 2, 2: 3 }, second tick: { 2: 4, 3: 5 }
    /// # batch_first_tick.chain(batch_second_tick).first();
    /// input_batch.clone().filter_key_not_in(
    ///     input_batch.defer_tick().keys() // keys present in the previous tick
    /// )
    /// # .entries().all_ticks()
    /// # }, |mut stream| async move {
    /// // { 1: 2, 2: 3 } (first tick), { 3: 5 } (second tick)
    /// # for w in vec![(1, 2), (2, 3), (3, 5)] {
    /// #     assert_eq!(stream.next().await.unwrap(), w);
    /// # }
    /// # }));
    /// # }
    /// ```
    pub fn defer_tick(self) -> KeyedSingleton<K, V, Tick<L>, Bounded> {
        KeyedSingleton::new(
            self.location.clone(),
            HydroNode::DeferTick {
                input: Box::new(self.ir_node.replace(HydroNode::Placeholder)),
                metadata: self
                    .location
                    .new_node_metadata(KeyedSingleton::<K, V, Tick<L>, Bounded>::collection_kind()),
            },
        )
    }
}

impl<'a, K, V, L, B: KeyedSingletonBound<ValueBound = Unbounded>> KeyedSingleton<K, V, L, B>
where
    L: Location<'a>,
{
    /// Returns a keyed singleton with a snapshot of each key-value entry at a non-deterministic
    /// point in time.
    ///
    /// # Non-Determinism
    /// Because this picks a snapshot of each entry, which is continuously changing, each output has a
    /// non-deterministic set of entries since each snapshot can be at an arbitrary point in time.
    pub fn snapshot<L2: Location<'a, DropConsistency = L::DropConsistency>>(
        self,
        tick: &Tick<L2>,
        _nondet: NonDet,
    ) -> KeyedSingleton<K, V, Tick<L::DropConsistency>, Bounded> {
        assert_eq!(
            Location::id(tick.parent_location()),
            Location::id(&self.location)
        );
        KeyedSingleton::new(
            tick.drop_consistency(),
            HydroNode::Batch {
                inner: Box::new(self.ir_node.replace(HydroNode::Placeholder)),
                metadata: tick
                    .new_node_metadata(KeyedSingleton::<K, V, Tick<L>, Bounded>::collection_kind()),
            },
        )
    }
}

impl<'a, K, V, L, B: KeyedSingletonBound<ValueBound = Unbounded>> KeyedSingleton<K, V, Atomic<L>, B>
where
    L: Location<'a>,
{
    /// Returns a keyed singleton with a snapshot of each key-value entry, consistent with the
    /// state of the keyed singleton being atomically processed.
    ///
    /// # Non-Determinism
    /// Because this picks a snapshot of each entry, which is continuously changing, each output has a
    /// non-deterministic set of entries since each snapshot can be at an arbitrary point in time.
    pub fn snapshot_atomic<L2: Location<'a, DropConsistency = L::DropConsistency>>(
        self,
        tick: &Tick<L2>,
        _nondet: NonDet,
    ) -> KeyedSingleton<K, V, Tick<L::DropConsistency>, Bounded> {
        assert_eq!(
            Location::id(tick.parent_location()),
            Location::id(self.location.tick.parent_location())
        );
        KeyedSingleton::new(
            tick.drop_consistency(),
            HydroNode::Batch {
                inner: Box::new(self.ir_node.replace(HydroNode::Placeholder)),
                metadata: tick
                    .new_node_metadata(KeyedSingleton::<K, V, Tick<L>, Bounded>::collection_kind()),
            },
        )
    }
}

impl<'a, K, V, L, B: KeyedSingletonBound<ValueBound = Bounded>> KeyedSingleton<K, V, L, B>
where
    L: Location<'a>,
{
    /// Creates a keyed singleton containing only the key-value pairs where the value satisfies a predicate `f`.
    ///
    /// The closure `f` receives a reference `&V` to each value and returns a boolean. If the predicate
    /// returns `true`, the key-value pair is included in the output. If it returns `false`, the pair
    /// is filtered out.
    ///
    /// 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 [`KeyedSingleton::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| {
    /// let keyed_singleton = // { 1: 2, 2: 4, 3: 1 }
    /// # process
    /// #     .source_iter(q!(vec![(1, 2), (2, 4), (3, 1)]))
    /// #     .into_keyed()
    /// #     .first();
    /// keyed_singleton.filter(q!(|&v| v > 1))
    /// #   .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<F>(
        self,
        f: impl IntoQuotedMut<'a, F, OperatorContext<L, B::UnderlyingBound>> + Copy,
    ) -> KeyedSingleton<K, V, L, B>
    where
        F: Fn(&V) -> bool + 'a,
    {
        let f: ManualExpr<F, _> =
            ManualExpr::new(move |ctx: &OperatorContext<L, B::UnderlyingBound>| {
                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>(&OperatorContext::<L, B::UnderlyingBound>::new(
            &self.location,
        ))
        .into();

        KeyedSingleton::new(
            self.location.clone(),
            HydroNode::Filter {
                f: filter_f,
                input: Box::new(self.ir_node.replace(HydroNode::Placeholder)),
                metadata: self
                    .location
                    .new_node_metadata(KeyedSingleton::<K, V, L, B>::collection_kind()),
            },
        )
    }

    /// An operator that both filters and maps values. It yields only the key-value pairs where
    /// the supplied closure `f` returns `Some(value)`.
    ///
    /// The closure `f` receives each value `V` and returns `Option<U>`. If the closure returns
    /// `Some(new_value)`, the key-value pair `(key, new_value)` is included in the output.
    /// If it returns `None`, the key-value pair is filtered out.
    ///
    /// # 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 keyed_singleton = // { 1: "42", 2: "hello", 3: "100" }
    /// # process
    /// #     .source_iter(q!(vec![(1, "42"), (2, "hello"), (3, "100")]))
    /// #     .into_keyed()
    /// #     .first();
    /// keyed_singleton.filter_map(q!(|s| s.parse::<i32>().ok()))
    /// #   .entries()
    /// # }, |mut stream| async move {
    /// // { 1: 42, 3: 100 }
    /// # let mut results = Vec::new();
    /// # for _ in 0..2 {
    /// #     results.push(stream.next().await.unwrap());
    /// # }
    /// # results.sort();
    /// # assert_eq!(results, vec![(1, 42), (3, 100)]);
    /// # }));
    /// # }
    /// ```
    pub fn filter_map<F, U>(
        self,
        f: impl IntoQuotedMut<'a, F, OperatorContext<L, B::UnderlyingBound>> + Copy,
    ) -> KeyedSingleton<K, U, L, B::EraseMonotonic>
    where
        F: Fn(V) -> Option<U> + 'a,
    {
        let f: ManualExpr<F, _> =
            ManualExpr::new(move |ctx: &OperatorContext<L, B::UnderlyingBound>| {
                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)>>(&OperatorContext::<L, B::UnderlyingBound>::new(
            &self.location,
        ))
        .into();

        KeyedSingleton::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(KeyedSingleton::<
                    K,
                    U,
                    L,
                    B::EraseMonotonic,
                >::collection_kind()),
            },
        )
    }

    /// Returns a keyed singleton with entries consisting of _new_ key-value pairs that have
    /// arrived since the previous batch was released.
    ///
    /// Currently, there is no `all_ticks` dual on [`KeyedSingleton`], instead you may want to use
    /// [`KeyedSingleton::into_keyed_stream`] then yield with [`KeyedStream::all_ticks`].
    ///
    /// # Non-Determinism
    /// Because this picks a batch of asynchronously added entries, each output keyed singleton
    /// has a non-deterministic set of key-value pairs.
    pub fn batch<L2: Location<'a, DropConsistency = L::DropConsistency>>(
        self,
        tick: &Tick<L2>,
        _nondet: NonDet,
    ) -> KeyedSingleton<K, V, Tick<L::DropConsistency>, Bounded> {
        assert_eq!(
            Location::id(tick.parent_location()),
            Location::id(&self.location)
        );
        KeyedSingleton::new(
            tick.drop_consistency(),
            HydroNode::Batch {
                inner: Box::new(self.ir_node.replace(HydroNode::Placeholder)),
                metadata: tick
                    .new_node_metadata(KeyedSingleton::<K, V, Tick<L>, Bounded>::collection_kind()),
            },
        )
    }
}

impl<'a, K, V, L, B: KeyedSingletonBound<ValueBound = Bounded>> KeyedSingleton<K, V, Atomic<L>, B>
where
    L: Location<'a>,
{
    /// Returns a keyed singleton with entries consisting of _new_ key-value pairs that are being
    /// atomically processed.
    ///
    /// Currently, there is no dual to asynchronously yield back outside the tick, instead you
    /// should use [`KeyedSingleton::into_keyed_stream`] and yield a [`KeyedStream`].
    ///
    /// # Non-Determinism
    /// Because this picks a batch of asynchronously added entries, each output keyed singleton
    /// has a non-deterministic set of key-value pairs.
    pub fn batch_atomic<L2: Location<'a, DropConsistency = L::DropConsistency>>(
        self,
        tick: &Tick<L2>,
        nondet: NonDet,
    ) -> KeyedSingleton<K, V, Tick<L::DropConsistency>, Bounded> {
        let _ = nondet;
        assert_eq!(
            Location::id(tick.parent_location()),
            Location::id(self.location.tick.parent_location())
        );
        KeyedSingleton::new(
            tick.drop_consistency(),
            HydroNode::Batch {
                inner: Box::new(self.ir_node.replace(HydroNode::Placeholder)),
                metadata: tick
                    .new_node_metadata(KeyedSingleton::<K, V, Tick<L>, Bounded>::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(any(feature = "deploy", feature = "sim"))]
    use crate::location::Location;
    #[cfg(any(feature = "deploy", feature = "sim"))]
    use crate::nondet::nondet;

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

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

        let (input_port, input) = node.source_external_bincode(&external);
        let out = input
            .into_keyed()
            .first()
            .key_count()
            .sample_eager(nondet!(/** test */))
            .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 external_in = nodes.connect(input_port).await;
        let mut external_out = nodes.connect(out).await;

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

        assert_eq!(external_out.next().await.unwrap(), 0);

        external_in.send((1, 1)).await.unwrap();
        assert_eq!(external_out.next().await.unwrap(), 1);

        external_in.send((2, 2)).await.unwrap();
        assert_eq!(external_out.next().await.unwrap(), 2);
    }

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

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

        let (input_port, input) = node.source_external_bincode(&external);
        let out = input
            .into_keyed()
            .fold(q!(|| 0), q!(|acc, _| *acc += 1))
            .key_count()
            .sample_eager(nondet!(/** test */))
            .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 external_in = nodes.connect(input_port).await;
        let mut external_out = nodes.connect(out).await;

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

        assert_eq!(external_out.next().await.unwrap(), 0);

        external_in.send((1, 1)).await.unwrap();
        assert_eq!(external_out.next().await.unwrap(), 1);

        external_in.send((1, 2)).await.unwrap();
        assert_eq!(external_out.next().await.unwrap(), 1);

        external_in.send((2, 2)).await.unwrap();
        assert_eq!(external_out.next().await.unwrap(), 2);

        external_in.send((1, 1)).await.unwrap();
        assert_eq!(external_out.next().await.unwrap(), 2);

        external_in.send((3, 1)).await.unwrap();
        assert_eq!(external_out.next().await.unwrap(), 3);
    }

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

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

        let (input_port, input) = node.source_external_bincode(&external);
        let out = input
            .into_keyed()
            .first()
            .into_singleton()
            .sample_eager(nondet!(/** test */))
            .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 external_in = nodes.connect(input_port).await;
        let mut external_out = nodes.connect(out).await;

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

        assert_eq!(
            external_out.next().await.unwrap(),
            std::collections::HashMap::new()
        );

        external_in.send((1, 1)).await.unwrap();
        assert_eq!(
            external_out.next().await.unwrap(),
            vec![(1, 1)].into_iter().collect()
        );

        external_in.send((2, 2)).await.unwrap();
        assert_eq!(
            external_out.next().await.unwrap(),
            vec![(1, 1), (2, 2)].into_iter().collect()
        );
    }

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

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

        let (input_port, input) = node.source_external_bincode(&external);
        let out = input
            .into_keyed()
            .fold(q!(|| 0), q!(|acc, _| *acc += 1))
            .into_singleton()
            .sample_eager(nondet!(/** test */))
            .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 external_in = nodes.connect(input_port).await;
        let mut external_out = nodes.connect(out).await;

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

        assert_eq!(
            external_out.next().await.unwrap(),
            std::collections::HashMap::new()
        );

        external_in.send((1, 1)).await.unwrap();
        assert_eq!(
            external_out.next().await.unwrap(),
            vec![(1, 1)].into_iter().collect()
        );

        external_in.send((1, 2)).await.unwrap();
        assert_eq!(
            external_out.next().await.unwrap(),
            vec![(1, 2)].into_iter().collect()
        );

        external_in.send((2, 2)).await.unwrap();
        assert_eq!(
            external_out.next().await.unwrap(),
            vec![(1, 2), (2, 1)].into_iter().collect()
        );

        external_in.send((1, 1)).await.unwrap();
        assert_eq!(
            external_out.next().await.unwrap(),
            vec![(1, 3), (2, 1)].into_iter().collect()
        );

        external_in.send((3, 1)).await.unwrap();
        assert_eq!(
            external_out.next().await.unwrap(),
            vec![(1, 3), (2, 1), (3, 1)].into_iter().collect()
        );
    }

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

        let (input_port, input) = node.sim_input();
        let output = input
            .into_keyed()
            .fold(q!(|| 0), q!(|acc, _| *acc += 1))
            .snapshot(&node.tick(), nondet!(/** test */))
            .entries()
            .all_ticks()
            .sim_output();

        let count = flow.sim().exhaustive(async || {
            input_port.send((1, 123));
            input_port.send((1, 456));
            input_port.send((2, 123));

            let all = output.collect_sorted::<Vec<_>>().await;
            assert_eq!(all.last().unwrap(), &(2, 1));
        });

        assert_eq!(count, 8);
    }

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

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

        let tick = node.tick();
        let keyed_data = node
            .source_iter(q!(vec![(1, 10), (2, 20)]))
            .into_keyed()
            .batch(&tick, nondet!(/** test */))
            .first();
        let requests = node
            .source_iter(q!(vec![(1, 100), (2, 200), (3, 300)]))
            .into_keyed()
            .batch(&tick, nondet!(/** test */));

        let out = keyed_data
            .join_keyed_stream(requests)
            .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 external_out = nodes.connect(out).await;

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

        let mut results = vec![];
        for _ in 0..2 {
            results.push(external_out.next().await.unwrap());
        }
        results.sort();

        assert_eq!(results, vec![(1, (10, 100)), (2, (20, 200))]);
    }

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

        let (input_port, input) = node.sim_input::<(u32, usize), _, _>();
        let (thresh_port, thresh_input) = node.sim_input::<(u32, usize), _, _>();

        // Create a monotonically increasing keyed singleton via fold with monotone proof
        let counts: super::KeyedSingleton<u32, usize, _, super::MonotonicValue> =
            input.into_keyed().fold(
                q!(|| 0usize),
                q!(
                    |acc, v| *acc += v,
                    monotone = crate::properties::manual_proof!(/** += is monotonic */)
                ),
            );

        // BoundedValue keyed singleton of thresholds (from .first() on unbounded stream)
        let thresholds = thresh_input.into_keyed().first();

        let output = counts
            .threshold_greater_or_equal(thresholds)
            .entries()
            .sim_output();

        let count = flow.sim().exhaustive(async || {
            // Set thresholds: key 1 needs value >= 5, key 2 needs value >= 10
            thresh_port.send((1, 5));
            thresh_port.send((2, 10));

            // key 1 gets increments: 3 + 3 = 6, which is >= 5 ✓
            input_port.send((1, 3));
            input_port.send((1, 3));
            // key 2 gets increments: 3 + 3 = 6, which is < 10 ✗
            input_port.send((2, 3));
            input_port.send((2, 3));

            let results = output.collect_sorted::<Vec<_>>().await;
            assert_eq!(results, vec![(1, 5)]);
        });

        assert!(count > 0);
    }

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

        let (input_port, input) = node.sim_input::<(u32, usize), _, _>();

        let counts: super::KeyedSingleton<u32, usize, _, super::MonotonicValue> =
            input.into_keyed().fold(
                q!(|| 0usize),
                q!(
                    |acc, v| *acc += v,
                    monotone = crate::properties::manual_proof!(/** += is monotonic */)
                ),
            );

        // Uniform threshold: all keys need value >= 5
        let threshold = node.singleton(q!(5usize));

        let output = counts
            .threshold_greater_or_equal_uniform(threshold)
            .entries()
            .sim_output();

        let count = flow.sim().exhaustive(async || {
            // key 1: 3 + 3 = 6 >= 5 ✓
            input_port.send((1, 3));
            input_port.send((1, 3));
            // key 2: 2 + 2 = 4 < 5 ✗
            input_port.send((2, 2));
            input_port.send((2, 2));

            let results = output.collect_sorted::<Vec<_>>().await;
            assert_eq!(results, vec![(1, 5)]);
        });

        assert!(count > 0);
    }

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

        let (input_port, input) = node.sim_input::<(u32, usize), _, _>();
        let (thresh_port, thresh_input) = node.sim_input::<(u32, usize), _, _>();

        // BoundedValue keyed singleton (values fixed once per key via .first())
        let values = input.into_keyed().first();

        // BoundedValue keyed singleton of thresholds
        let thresholds = thresh_input.into_keyed().first();

        let output = values
            .threshold_greater_or_equal(thresholds)
            .entries()
            .sim_output();

        let count = flow.sim().exhaustive(async || {
            // Set thresholds: key 1 needs >= 3, key 2 needs >= 10
            thresh_port.send((1, 3));
            thresh_port.send((2, 10));

            // key 1 gets value 5 >= 3 ✓, key 2 gets value 4 < 10 ✗
            input_port.send((1, 5));
            input_port.send((2, 4));

            let results = output.collect_sorted::<Vec<_>>().await;
            assert_eq!(results, vec![(1, 3)]);
        });

        assert!(count > 0);
    }

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

        let (input_port, input) = node.sim_input::<(u32, usize), _, _>();

        // BoundedValue keyed singleton (values fixed once per key via .first())
        let values = input.into_keyed().first();

        // Uniform threshold: all keys need value >= 5
        let threshold = node.singleton(q!(5usize));

        let output = values
            .threshold_greater_or_equal_uniform(threshold)
            .entries()
            .sim_output();

        let count = flow.sim().exhaustive(async || {
            // key 1 gets value 7 >= 5 ✓, key 2 gets value 3 < 5 ✗
            input_port.send((1, 7));
            input_port.send((2, 3));

            let results = output.collect_sorted::<Vec<_>>().await;
            assert_eq!(results, vec![(1, 5)]);
        });

        assert!(count > 0);
    }

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

        // Bounded keyed singleton (fully known upfront)
        let values = node
            .source_iter(q!(vec![(1, 6usize), (2, 4usize)]))
            .into_keyed()
            .first();

        // BoundedValue thresholds (from async source)
        let (thresh_port, thresh_input) = node.sim_input::<(u32, usize), _, _>();
        let thresholds = thresh_input.into_keyed().first();

        let output = values
            .threshold_greater_or_equal(thresholds)
            .entries()
            .sim_output();

        let count = flow.sim().exhaustive(async || {
            thresh_port.send((1, 5));
            thresh_port.send((2, 10));

            // key 1: 6 >= 5 ✓, key 2: 4 < 10 ✗
            let results = output.collect_sorted::<Vec<_>>().await;
            assert_eq!(results, vec![(1, 5)]);
        });

        assert!(count > 0);
    }

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

        let values = node
            .source_iter(q!(vec![(1, 6usize), (2, 4usize)]))
            .into_keyed()
            .first();
        let threshold = node.singleton(q!(5usize));

        let output = values
            .threshold_greater_or_equal_uniform(threshold)
            .entries()
            .sim_output();

        let count = flow.sim().exhaustive(async || {
            // key 1: 6 >= 5 ✓, key 2: 4 < 5 ✗
            let results = output.collect_sorted::<Vec<_>>().await;
            assert_eq!(results, vec![(1, 5)]);
        });

        assert!(count > 0);
    }
}