feat(hydro_lang)!: start adding APIs for various forms of keyed streams "joins", add KVS counter server (#1996)

Also fixes incorrect output metadata for `KeyedStream::fold`, and
removes `KeyedSingleton::all_ticks` for now until we a) actually need it
and b) can think carefully about its semantics.

Author: shadaj

hydro_deploy/hydro_deploy_integration/src/multi_connection.rs

@@ -249,6 +249,10 @@ impl<
             conn.is_some()
         });
 
+        if me.poll_cursor == me.active_connections.len() {
+            me.poll_cursor = 0;
+        }
+
         out
     }
 }

hydro_lang/src/keyed_singleton.rs

@@ -1,3 +1,5 @@
+use std::hash::Hash;
+
 use stageleft::{IntoQuotedMut, QuotedWithContext, q};
 
 use crate::cycle::{CycleCollection, CycleComplete, ForwardRefMarker};
@@ -6,7 +8,10 @@ use crate::location::{LocationId, NoTick};
 use crate::manual_expr::ManualExpr;
 use crate::stream::ExactlyOnce;
 use crate::unsafety::NonDet;
-use crate::{Atomic, Bounded, Location, NoOrder, Singleton, Stream, Tick, Unbounded};
+use crate::{
+    Atomic, Bounded, KeyedStream, Location, NoOrder, Optional, Singleton, Stream, Tick, TotalOrder,
+    nondet,
+};
 
 pub struct KeyedSingleton<K, V, Loc, Bound> {
     pub(crate) underlying: Stream<(K, V), Loc, Bound, NoOrder, ExactlyOnce>,
@@ -104,6 +109,77 @@ impl<'a, K, V, L: Location<'a>, B> KeyedSingleton<K, V, L, B> {
     }
 }
 
+impl<'a, K: Hash + Eq, V, L: Location<'a>> KeyedSingleton<K, V, Tick<L>, Bounded> {
+    /// Gets the value associated with a specific key from the keyed singleton.
+    ///
+    /// # Example
+    /// ```rust
+    /// # use hydro_lang::*;
+    /// # use futures::StreamExt;
+    /// # tokio_test::block_on(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 */))
+    ///     .fold(q!(|| 0), q!(|acc, x| *acc = x));
+    /// 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: Singleton<K, Tick<L>, Bounded>) -> Optional<V, Tick<L>, Bounded> {
+        self.entries()
+            .join(key.into_stream().map(q!(|k| (k, ()))))
+            .map(q!(|(_, (v, _))| v))
+            .assume_ordering::<TotalOrder>(nondet!(/** only a single key, so totally ordered */))
+            .first()
+    }
+
+    /// Given a keyed stream of lookup requests, where the key is the lookup and the value
+    /// is some additional metadata, emits a keyed stream of lookup results where the key
+    /// is the same as before, but the value is a tuple of the lookup result and the metadata
+    /// of the request.
+    ///
+    /// # Example
+    /// ```rust
+    /// # use hydro_lang::*;
+    /// # use futures::StreamExt;
+    /// # tokio_test::block_on(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 */))
+    ///     .fold(q!(|| 0), q!(|acc, x| *acc = x));
+    /// let other_data = process
+    ///     .source_iter(q!(vec![(1, 100), (2, 200), (1, 101)]))
+    ///     .into_keyed()
+    ///     .batch(&tick, nondet!(/** test */));
+    /// keyed_data.get_many(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 get_many<O2, R2, V2>(
+        self,
+        with: KeyedStream<K, V2, Tick<L>, Bounded, O2, R2>,
+    ) -> KeyedStream<K, (V, V2), Tick<L>, Bounded, NoOrder, R2> {
+        self.entries()
+            .weaker_retries()
+            .join(with.entries())
+            .into_keyed()
+    }
+}
+
 impl<'a, K, V, L, B> KeyedSingleton<K, V, L, B>
 where
     L: Location<'a> + NoTick + NoAtomic,
@@ -151,14 +227,3 @@ where
         }
     }
 }
-
-impl<'a, K, V, L> KeyedSingleton<K, V, Tick<L>, Bounded>
-where
-    L: Location<'a>,
-{
-    pub fn all_ticks(self) -> KeyedSingleton<K, V, L, Unbounded> {
-        KeyedSingleton {
-            underlying: self.underlying.all_ticks(),
-        }
-    }
-}

hydro_lang/src/keyed_stream.rs

@@ -9,7 +9,7 @@ use crate::keyed_singleton::KeyedSingleton;
 use crate::location::tick::NoAtomic;
 use crate::location::{LocationId, NoTick, check_matching_location};
 use crate::manual_expr::ManualExpr;
-use crate::stream::ExactlyOnce;
+use crate::stream::{ExactlyOnce, MinRetries};
 use crate::unsafety::NonDet;
 use crate::*;
 
@@ -188,10 +188,54 @@ impl<'a, K, V, L: Location<'a>, B, O, R> KeyedStream<K, V, L, B, O, R> {
         }
     }
 
+    /// Transforms each value by invoking `f` on each key-value pair. The resulting values are **not**
+    /// re-grouped even they are tuples; instead they will be grouped under the original key.
+    ///
+    /// If you do not want to modify the stream and instead only want to view
+    /// each item use [`KeyedStream::inspect_with_key`] instead.
+    ///
+    /// # Example
+    /// ```rust
+    /// # use hydro_lang::*;
+    /// # use futures::StreamExt;
+    /// # tokio_test::block_on(test_util::stream_transform_test(|process| {
+    /// process
+    ///     .source_iter(q!(vec![(1, 2), (1, 3), (2, 4)]))
+    ///     .into_keyed()
+    ///     .map_with_key(q!(|(k, v)| k + v))
+    /// #   .entries()
+    /// # }, |mut stream| async move {
+    /// // { 1: [3, 4], 2: [6] }
+    /// # for w in vec![(1, 3), (1, 4), (2, 6)] {
+    /// #     assert_eq!(stream.next().await.unwrap(), w);
+    /// # }
+    /// # }));
+    /// ```
+    pub fn map_with_key<U, F>(
+        self,
+        f: impl IntoQuotedMut<'a, F, L> + Copy,
+    ) -> KeyedStream<K, U, L, B, O, R>
+    where
+        F: Fn((K, V)) -> U + 'a,
+        K: Clone,
+    {
+        let f: ManualExpr<F, _> = ManualExpr::new(move |ctx: &L| f.splice_fn1_ctx(ctx));
+        KeyedStream {
+            underlying: self.underlying.map(q!({
+                let orig = f;
+                move |(k, v)| {
+                    let out = orig((k.clone(), v));
+                    (k, out)
+                }
+            })),
+            _phantom_order: Default::default(),
+        }
+    }
+
     /// Creates a stream containing only the elements of each group stream that satisfy a predicate
     /// `f`, preserving the order of the elements within the group.
     ///
-    /// The closure `f` receives a reference `&T` rather than an owned value `T` because filtering does
+    /// The closure `f` receives a reference `&V` rather than an owned value `v` because filtering does
     /// not modify or take ownership of the values. If you need to modify the values while filtering
     /// use [`KeyedStream::filter_map`] instead.
     ///
@@ -226,11 +270,12 @@ impl<'a, K, V, L: Location<'a>, B, O, R> KeyedStream<K, V, L, B, O, R> {
         }
     }
 
-    /// Transforms each value by invoking `f` on each key-value pair. The resulting values are **not**
-    /// re-grouped even they are tuples; instead they will be grouped under the original key.
+    /// Creates a stream containing only the elements of each group stream that satisfy a predicate
+    /// `f` (which receives the key-value tuple), preserving the order of the elements within the group.
     ///
-    /// If you do not want to modify the stream and instead only want to view
-    /// each item use [`KeyedStream::inspect_with_key`] instead.
+    /// The closure `f` receives a reference `&(K, V)` rather than an owned value `(K, V)` because filtering does
+    /// not modify or take ownership of the values. If you need to modify the values while filtering
+    /// use [`KeyedStream::filter_map_with_key`] instead.
     ///
     /// # Example
     /// ```rust
@@ -240,32 +285,24 @@ impl<'a, K, V, L: Location<'a>, B, O, R> KeyedStream<K, V, L, B, O, R> {
     /// process
     ///     .source_iter(q!(vec![(1, 2), (1, 3), (2, 4)]))
     ///     .into_keyed()
-    ///     .map_with_key(q!(|(k, v)| k + v))
+    ///     .filter_with_key(q!(|&(k, v)| v - k == 2))
     /// #   .entries()
     /// # }, |mut stream| async move {
-    /// // { 1: [3, 4], 2: [6] }
-    /// # for w in vec![(1, 3), (1, 4), (2, 6)] {
+    /// // { 1: [3], 2: [4] }
+    /// # for w in vec![(1, 3), (2, 4)] {
     /// #     assert_eq!(stream.next().await.unwrap(), w);
     /// # }
     /// # }));
     /// ```
-    pub fn map_with_key<U, F>(
+    pub fn filter_with_key<F>(
         self,
         f: impl IntoQuotedMut<'a, F, L> + Copy,
-    ) -> KeyedStream<K, U, L, B, O, R>
+    ) -> KeyedStream<K, V, L, B, O, R>
     where
-        F: Fn((K, V)) -> U + 'a,
-        K: Clone,
+        F: Fn(&(K, V)) -> bool + 'a,
     {
-        let f: ManualExpr<F, _> = ManualExpr::new(move |ctx: &L| f.splice_fn1_ctx(ctx));
         KeyedStream {
-            underlying: self.underlying.map(q!({
-                let orig = f;
-                move |(k, v)| {
-                    let out = orig((k.clone(), v));
-                    (k, out)
-                }
-            })),
+            underlying: self.underlying.filter(f),
             _phantom_order: Default::default(),
         }
     }
@@ -418,6 +455,41 @@ impl<'a, K, V, L: Location<'a>, B, O, R> KeyedStream<K, V, L, B, O, R> {
     }
 }
 
+impl<'a, K, V, L: Location<'a> + NoTick + NoAtomic, O, R> KeyedStream<K, V, L, Unbounded, O, R> {
+    /// Produces a new keyed stream that "merges" the inputs by interleaving the elements
+    /// of any overlapping groups. The result has [`NoOrder`] on each group because the
+    /// order of interleaving is not guaranteed. If the keys across both inputs do not overlap,
+    /// the ordering will be deterministic and you can safely use [`Self::assume_ordering`].
+    ///
+    /// Currently, both input streams must be [`Unbounded`].
+    ///
+    /// # Example
+    /// ```rust
+    /// # use hydro_lang::*;
+    /// # use futures::StreamExt;
+    /// # tokio_test::block_on(test_util::stream_transform_test(|process| {
+    /// let numbers1 = process.source_iter(q!(vec![(1, 2), (3, 4)])).into_keyed();
+    /// let numbers2 = process.source_iter(q!(vec![(1, 3), (3, 5)])).into_keyed();
+    /// numbers1.interleave(numbers2)
+    /// #   .entries()
+    /// # }, |mut stream| async move {
+    /// // { 1: [2, 3], 3: [4, 5] } with each group in unknown order
+    /// # for w in vec![(1, 2), (3, 4), (1, 3), (3, 5)] {
+    /// #     assert_eq!(stream.next().await.unwrap(), w);
+    /// # }
+    /// # }));
+    /// ```
+    pub fn interleave<O2, R2: MinRetries<R>>(
+        self,
+        other: KeyedStream<K, V, L, Unbounded, O2, R2>,
+    ) -> KeyedStream<K, V, L, Unbounded, NoOrder, R::Min>
+    where
+        R: MinRetries<R2, Min = R2::Min>,
+    {
+        self.entries().interleave(other.entries()).into_keyed()
+    }
+}
+
 impl<'a, K, V, L, B> KeyedStream<K, V, L, B, TotalOrder, ExactlyOnce>
 where
     K: Eq + Hash,
@@ -571,7 +643,7 @@ where
             init,
             acc: comb,
             input: Box::new(self.underlying.ir_node.into_inner()),
-            metadata: self.underlying.location.new_node_metadata::<(K, V)>(),
+            metadata: self.underlying.location.new_node_metadata::<(K, A)>(),
         };
 
         KeyedSingleton {
@@ -1019,6 +1091,37 @@ where
             .assume_retries::<ExactlyOnce>(nondet!(/** the combinator function is idempotent */))
             .reduce_watermark(other, comb)
     }
+
+    /// Given a bounded stream of keys `K`, returns a new keyed stream containing only the groups
+    /// whose keys are not in the bounded stream.
+    ///
+    /// # Example
+    /// ```rust
+    /// # use hydro_lang::*;
+    /// # use futures::StreamExt;
+    /// # tokio_test::block_on(test_util::stream_transform_test(|process| {
+    /// let tick = process.tick();
+    /// let keyed_stream = process
+    ///     .source_iter(q!(vec![ (1, 'a'), (2, 'b'), (3, 'c'), (4, 'd') ]))
+    ///     .batch(&tick, nondet!(/** test */))
+    ///     .into_keyed();
+    /// let keys_to_remove = process
+    ///     .source_iter(q!(vec![1, 2]))
+    ///     .batch(&tick, nondet!(/** test */));
+    /// keyed_stream.filter_key_not_in(keys_to_remove).all_ticks()
+    /// #   .entries()
+    /// # }, |mut stream| async move {
+    /// // { 3: ['c'], 4: ['d'] }
+    /// # for w in vec![(3, 'c'), (4, 'd')] {
+    /// #     assert_eq!(stream.next().await.unwrap(), w);
+    /// # }
+    /// # }));
+    pub fn filter_key_not_in<O2, R2>(self, other: Stream<K, L, Bounded, O2, R2>) -> Self {
+        KeyedStream {
+            underlying: self.entries().anti_join(other),
+            _phantom_order: Default::default(),
+        }
+    }
 }
 
 impl<'a, K, V, L, B, O, R> KeyedStream<K, V, L, B, O, R>

hydro_lang/src/rewrites/snapshots/hydro_lang__rewrites__properties__tests__property_optimized.snap

@@ -40,7 +40,7 @@ expression: built.ir()
                     ),
                 ),
                 output_type: Some(
-                    (std :: string :: String , ()),
+                    (std :: string :: String , i32),
                 ),
             },
         },

hydro_lang/src/stream.rs

@@ -782,6 +782,19 @@ where
     }
 }
 
+impl<'a, T, L, B, O> Stream<T, L, B, O, ExactlyOnce>
+where
+    L: Location<'a>,
+{
+    /// Given a stream with [`ExactlyOnce`] retry guarantees, weakens it to an arbitrary guarantee
+    /// `R2`, which is safe because all guarantees are equal to or weaker than [`ExactlyOnce`]
+    pub fn weaker_retries<R2>(self) -> Stream<T, L, B, O, R2> {
+        self.assume_retries(
+            nondet!(/** any retry ordering is the same or weaker than ExactlyOnce */),
+        )
+    }
+}
+
 impl<'a, T, L, B, O, R> Stream<&T, L, B, O, R>
 where
     L: Location<'a>,
@@ -1636,7 +1649,7 @@ where
     /// #     assert_eq!(stream.next().await.unwrap(), w);
     /// # }
     /// # }));
-    pub fn anti_join<O2>(self, n: Stream<K, L, Bounded, O2, R>) -> Stream<(K, V1), L, B, O, R>
+    pub fn anti_join<O2, R2>(self, n: Stream<K, L, Bounded, O2, R2>) -> Stream<(K, V1), L, B, O, R>
     where
         K: Eq + Hash,
     {