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 stageleft::{IntoQuotedMut, QuotedWithContext, q};

use super::boundedness::{Bounded, Boundedness, IsBounded, Unbounded};
use super::keyed_stream::KeyedStream;
use super::optional::Optional;
use super::singleton::Singleton;
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, NoTick, Tick, 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 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 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 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 struct BoundedValue;

impl KeyedSingletonBound for BoundedValue {
    type UnderlyingBound = Unbounded;
    type ValueBound = Bounded;
    type WithBoundedValue = BoundedValue;
    type KeyedStreamToMonotone = 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 struct MonotonicValue;

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

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

/// 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: 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: HydroNode::Tee {
                    inner: SharedNode(inner.0.clone()),
                    metadata: metadata.clone(),
                }
                .into(),
                _phantom: PhantomData,
            }
        } else {
            unreachable!()
        }
    }
}

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

    fn create_source(cycle_id: CycleId, location: L) -> Self {
        KeyedSingleton {
            flow_state: location.flow_state().clone(),
            location: location.clone(),
            ir_node: RefCell::new(HydroNode::CycleSource {
                cycle_id,
                metadata: location.new_node_metadata(Self::collection_kind()),
            }),
            _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> + NoTick,
{
    fn complete(self, cycle_id: CycleId, expected_location: LocationId) {
        assert_eq!(
            Location::id(&self.location),
            expected_location,
            "locations do not match"
        );
        self.location
            .flow_state()
            .borrow_mut()
            .push_root(HydroRoot::CycleSink {
                cycle_id,
                input: Box::new(self.ir_node.replace(HydroNode::Placeholder)),
                op_metadata: HydroIrOpMetadata::new(),
            });
    }
}

impl<'a, K, 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();
        KeyedSingleton {
            location,
            flow_state,
            ir_node: RefCell::new(ir_node),
            _phantom: PhantomData,
        }
    }

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

#[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, L> + Copy,
    ) -> KeyedSingleton<K, U, L, B::EraseMonotonic>
    where
        F: Fn(V) -> U + 'a,
    {
        let f: ManualExpr<F, _> = ManualExpr::new(move |ctx: &L| f.splice_fn1_ctx(ctx));
        let map_f = q!({
            let orig = f;
            move |(k, v)| (k, orig(v))
        })
        .splice_fn1_ctx::<(K, V), (K, U)>(&self.location)
        .into();

        KeyedSingleton::new(
            self.location.clone(),
            HydroNode::Map {
                f: map_f,
                singleton_refs: Vec::new(),
                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, L> + Copy,
    ) -> KeyedSingleton<K, U, L, B::EraseMonotonic>
    where
        F: Fn((K, V)) -> U + 'a,
        K: Clone,
    {
        let f: ManualExpr<F, _> = ManualExpr::new(move |ctx: &L| f.splice_fn1_ctx(ctx));
        let map_f = q!({
            let orig = f;
            move |(k, v)| {
                let out = orig((Clone::clone(&k), v));
                (k, out)
            }
        })
        .splice_fn1_ctx::<(K, V), (K, U)>(&self.location)
        .into();

        KeyedSingleton::new(
            self.location.clone(),
            HydroNode::Map {
                f: map_f,
                singleton_refs: Vec::new(),
                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: RefCell::new(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
        {
            let me: KeyedSingleton<K, V, L, Unbounded> = KeyedSingleton::new(
                self.location.clone(),
                self.ir_node.replace(HydroNode::Placeholder),
            );

            let out =
                key_count_inside_tick(me.snapshot(&tick, nondet!(/** eventually stabilizes */)))
                    .latest();
            Singleton::new(
                out.location.clone(),
                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,
    {
        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: RefCell::new(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
        {
            let me: KeyedSingleton<K, V, L, Unbounded> = KeyedSingleton::new(
                self.location.clone(),
                self.ir_node.replace(HydroNode::Placeholder),
            );

            let out = into_singleton_inside_tick(
                me.snapshot(&tick, nondet!(/** eventually stabilizes */)),
            )
            .latest();
            Singleton::new(
                out.location.clone(),
                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)
    }
}

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>(&self.location)
            .into();

        Stream::new(
            self.location.clone(),
            HydroNode::Map {
                f: map_f,
                singleton_refs: Vec::new(),
                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, L> + Copy) -> Self
    where
        F: Fn(&V) + 'a,
    {
        let f: ManualExpr<F, _> = ManualExpr::new(move |ctx: &L| f.splice_fn1_borrow_ctx(ctx));
        let inspect_f = q!({
            let orig = f;
            move |t: &(_, _)| orig(&t.1)
        })
        .splice_fn1_borrow_ctx::<(K, V), ()>(&self.location)
        .into();

        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, L>) -> Self
    where
        F: Fn(&(K, V)) + 'a,
    {
        let inspect_f = f.splice_fn1_borrow_ctx::<(K, V), ()>(&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>
    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>,
{
    /// Shifts this 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> {
        let id = self.location.flow_state().borrow_mut().next_clock_id();
        let out_location = Atomic {
            tick: Tick {
                id,
                l: self.location.clone(),
            },
        };
        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> + NoTick,
{
    /// 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(
        self,
        tick: &Tick<L>,
        _nondet: NonDet,
    ) -> KeyedSingleton<K, V, Tick<L>, Bounded> {
        assert_eq!(Location::id(tick.outer()), Location::id(&self.location));
        KeyedSingleton::new(
            tick.clone(),
            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> + NoTick,
{
    /// 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(
        self,
        tick: &Tick<L>,
        _nondet: NonDet,
    ) -> KeyedSingleton<K, V, Tick<L>, Bounded> {
        KeyedSingleton::new(
            tick.clone(),
            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, L> + Copy) -> KeyedSingleton<K, V, L, B>
    where
        F: Fn(&V) -> bool + 'a,
    {
        let f: ManualExpr<F, _> = ManualExpr::new(move |ctx: &L| f.splice_fn1_borrow_ctx(ctx));
        let filter_f = q!({
            let orig = f;
            move |t: &(_, _)| orig(&t.1)
        })
        .splice_fn1_borrow_ctx::<(K, V), bool>(&self.location)
        .into();

        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, L> + Copy,
    ) -> KeyedSingleton<K, U, L, B::EraseMonotonic>
    where
        F: Fn(V) -> Option<U> + 'a,
    {
        let f: ManualExpr<F, _> = ManualExpr::new(move |ctx: &L| f.splice_fn1_ctx(ctx));
        let filter_map_f = q!({
            let orig = f;
            move |(k, v)| orig(v).map(|o| (k, o))
        })
        .splice_fn1_ctx::<(K, V), Option<(K, U)>>(&self.location)
        .into();

        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(self, tick: &Tick<L>, _nondet: NonDet) -> KeyedSingleton<K, V, Tick<L>, Bounded>
    where
        L: NoTick,
    {
        assert_eq!(Location::id(tick.outer()), Location::id(&self.location));
        KeyedSingleton::new(
            tick.clone(),
            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> + NoTick,
{
    /// 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(
        self,
        tick: &Tick<L>,
        nondet: NonDet,
    ) -> KeyedSingleton<K, V, Tick<L>, Bounded> {
        let _ = nondet;
        KeyedSingleton::new(
            tick.clone(),
            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))]);
    }
}