feat(hydro_lang): add API for safely broadcasting a stream across several keys

Author: shadaj

hydro_lang/src/live_collections/stream/mod.rs

@@ -15,6 +15,7 @@ use super::boundedness::{Bounded, Boundedness, Unbounded};
 use super::keyed_stream::KeyedStream;
 use super::optional::Optional;
 use super::singleton::Singleton;
+use super::keyed_singleton::KeyedSingleton;
 use crate::compile::ir::{HydroIrOpMetadata, HydroNode, HydroRoot, TeeNode};
 #[cfg(stageleft_runtime)]
 use crate::forward_handle::{CycleCollection, ReceiverComplete};
@@ -1783,6 +1784,54 @@ where
             },
         )
     }
+
+    /// Creates a [`KeyedStream`] with the same set of keys as `keys`, but with the elements in
+    /// `self` used as the values for *each* key.
+    ///
+    /// This is helpful when "broadcasting" a set of values so that all the keys have the same
+    /// values. For example, it can be used to send the same set of elements to several cluster
+    /// members, if the membership information is available as a [`KeyedSingleton`].
+    ///
+    /// # Example
+    /// ```rust
+    /// # 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: () }
+    /// # process
+    /// #     .source_iter(q!(vec![(1, ()), (2, ())]))
+    /// #     .into_keyed()
+    /// #     .batch(&tick, nondet!(/** test */))
+    /// #     .first();
+    /// let stream = // [ "a", "b" ]
+    /// # process
+    /// #     .source_iter(q!(vec!["a".to_string(), "b".to_string()]))
+    /// #     .batch(&tick, nondet!(/** test */));
+    /// stream.repeat_with_keys(keyed_singleton)
+    /// # .entries().all_ticks()
+    /// # }, |mut stream| async move {
+    /// // { 1: ["a", "b" ], 2: ["a", "b"] }
+    /// # let mut results = Vec::new();
+    /// # for _ in 0..4 {
+    /// #     results.push(stream.next().await.unwrap());
+    /// # }
+    /// # results.sort();
+    /// # assert_eq!(results, vec![(1, "a".to_string()), (1, "b".to_string()), (2, "a".to_string()), (2, "b".to_string())]);
+    /// # }));
+    /// ```
+    pub fn repeat_with_keys<K, V2>(
+        self,
+        keys: KeyedSingleton<K, V2, L, Bounded>,
+    ) -> KeyedStream<K, T, L, Bounded, O, R>
+    where
+        K: Clone,
+        T: Clone,
+    {
+        keys.keys().weaken_retries().cross_product_nested_loop(self).into_keyed().assume_ordering(
+            nondet!(/** keyed stream does not depend on ordering of keys, cross_product_nested_loop preserves order of values */)
+        )
+    }
 }
 
 impl<'a, K, V1, L, B: Boundedness, O: Ordering, R: Retries> Stream<(K, V1), L, B, O, R>

hydro_lang/src/live_collections/stream/networking.rs

@@ -18,7 +18,7 @@ use crate::location::dynamic::DynLocation;
 use crate::location::external_process::ExternalBincodeStream;
 use crate::location::tick::NoAtomic;
 use crate::location::{Cluster, External, Location, MemberId, MembershipEvent, NoTick, Process};
-use crate::nondet::{NonDet, nondet};
+use crate::nondet::NonDet;
 use crate::staging_util::get_this_crate;
 
 // same as the one in `hydro_std`, but internal use only
@@ -185,15 +185,10 @@ impl<'a, T, L, B: Boundedness, O: Ordering, R: Retries> Stream<T, Process<'a, L>
     {
         let ids = track_membership(self.location.source_cluster_members(other));
         let join_tick = self.location.tick();
-        let current_members = ids.snapshot(&join_tick, nondet_membership).keys();
+        let current_members = ids.snapshot(&join_tick, nondet_membership);
 
-        current_members
-            .weaker_retries()
-            .assume_ordering::<TotalOrder>(
-                nondet!(/** we send to each member independently, order does not matter */),
-            )
-            .cross_product_nested_loop(self.batch(&join_tick, nondet_membership))
-            .weaken_ordering::<O>()
+        self.batch(&join_tick, nondet_membership)
+            .repeat_with_keys(current_members)
             .all_ticks()
             .demux_bincode(other)
     }
@@ -511,14 +506,10 @@ impl<'a, T, L, B: Boundedness, O: Ordering, R: Retries> Stream<T, Cluster<'a, L>
     {
         let ids = track_membership(self.location.source_cluster_members(other));
         let join_tick = self.location.tick();
-        let current_members = ids.snapshot(&join_tick, nondet_membership).keys();
+        let current_members = ids.snapshot(&join_tick, nondet_membership);
 
-        current_members
-            .weaker_retries()
-            .assume_ordering::<TotalOrder>(
-                nondet!(/** we send to each member independently, order does not matter */),
-            )
-            .cross_product_nested_loop(self.batch(&join_tick, nondet_membership))
+        self.batch(&join_tick, nondet_membership)
+            .repeat_with_keys(current_members)
             .all_ticks()
             .demux_bincode(other)
     }