hydro_lang/stream.rs
1use std::cell::RefCell;
2use std::future::Future;
3use std::hash::Hash;
4use std::marker::PhantomData;
5use std::ops::Deref;
6use std::rc::Rc;
7
8use bytes::Bytes;
9use serde::Serialize;
10use serde::de::DeserializeOwned;
11use stageleft::{IntoQuotedMut, QuotedWithContext, q};
12use syn::parse_quote;
13use tokio::time::Instant;
14
15use crate::builder::FLOW_USED_MESSAGE;
16use crate::cycle::{CycleCollection, CycleComplete, DeferTick, ForwardRefMarker, TickCycleMarker};
17use crate::ir::{DebugInstantiate, HydroLeaf, HydroNode, TeeNode};
18use crate::location::external_process::{ExternalBincodeStream, ExternalBytesPort};
19use crate::location::tick::{Atomic, NoAtomic};
20use crate::location::{
21 CanSend, External, Location, LocationId, NoTick, Tick, check_matching_location,
22};
23use crate::staging_util::get_this_crate;
24use crate::{Bounded, Cluster, ClusterId, Optional, Singleton, Unbounded};
25
26/// Marks the stream as being totally ordered, which means that there are
27/// no sources of non-determinism (other than intentional ones) that will
28/// affect the order of elements.
29pub enum TotalOrder {}
30
31/// Marks the stream as having no order, which means that the order of
32/// elements may be affected by non-determinism.
33///
34/// This restricts certain operators, such as `fold` and `reduce`, to only
35/// be used with commutative aggregation functions.
36pub enum NoOrder {}
37
38/// Helper trait for determining the weakest of two orderings.
39#[sealed::sealed]
40pub trait MinOrder<Other> {
41 /// The weaker of the two orderings.
42 type Min;
43}
44
45#[sealed::sealed]
46impl<T> MinOrder<T> for T {
47 type Min = T;
48}
49
50#[sealed::sealed]
51impl MinOrder<NoOrder> for TotalOrder {
52 type Min = NoOrder;
53}
54
55#[sealed::sealed]
56impl MinOrder<TotalOrder> for NoOrder {
57 type Min = NoOrder;
58}
59
60/// Marks the stream as having deterministic message cardinality, with no
61/// possibility of duplicates.
62pub enum ExactlyOnce {}
63
64/// Marks the stream as having non-deterministic message cardinality, which
65/// means that duplicates may occur, but messages will not be dropped.
66pub enum AtLeastOnce {}
67
68/// Helper trait for determining the weakest of two retry guarantees.
69#[sealed::sealed]
70pub trait MinRetries<Other> {
71 /// The weaker of the two retry guarantees.
72 type Min;
73}
74
75#[sealed::sealed]
76impl<T> MinRetries<T> for T {
77 type Min = T;
78}
79
80#[sealed::sealed]
81impl MinRetries<ExactlyOnce> for AtLeastOnce {
82 type Min = AtLeastOnce;
83}
84
85#[sealed::sealed]
86impl MinRetries<AtLeastOnce> for ExactlyOnce {
87 type Min = ExactlyOnce;
88}
89
90/// An ordered sequence stream of elements of type `T`.
91///
92/// Type Parameters:
93/// - `Type`: the type of elements in the stream
94/// - `Loc`: the location where the stream is being materialized
95/// - `Bound`: the boundedness of the stream, which is either [`Bounded`]
96/// or [`Unbounded`]
97/// - `Order`: the ordering of the stream, which is either [`TotalOrder`]
98/// or [`NoOrder`] (default is [`TotalOrder`])
99pub struct Stream<Type, Loc, Bound, Order = TotalOrder, Retries = ExactlyOnce> {
100 location: Loc,
101 pub(crate) ir_node: RefCell<HydroNode>,
102
103 _phantom: PhantomData<(Type, Loc, Bound, Order, Retries)>,
104}
105
106impl<'a, T, L, O, R> From<Stream<T, L, Bounded, O, R>> for Stream<T, L, Unbounded, O, R>
107where
108 L: Location<'a>,
109{
110 fn from(stream: Stream<T, L, Bounded, O, R>) -> Stream<T, L, Unbounded, O, R> {
111 Stream {
112 location: stream.location,
113 ir_node: stream.ir_node,
114 _phantom: PhantomData,
115 }
116 }
117}
118
119impl<'a, T, L, B, R> From<Stream<T, L, B, TotalOrder, R>> for Stream<T, L, B, NoOrder, R>
120where
121 L: Location<'a>,
122{
123 fn from(stream: Stream<T, L, B, TotalOrder, R>) -> Stream<T, L, B, NoOrder, R> {
124 Stream {
125 location: stream.location,
126 ir_node: stream.ir_node,
127 _phantom: PhantomData,
128 }
129 }
130}
131
132impl<'a, T, L, B, O> From<Stream<T, L, B, O, ExactlyOnce>> for Stream<T, L, B, O, AtLeastOnce>
133where
134 L: Location<'a>,
135{
136 fn from(stream: Stream<T, L, B, O, ExactlyOnce>) -> Stream<T, L, B, O, AtLeastOnce> {
137 Stream {
138 location: stream.location,
139 ir_node: stream.ir_node,
140 _phantom: PhantomData,
141 }
142 }
143}
144
145impl<'a, T, L, O, R> DeferTick for Stream<T, Tick<L>, Bounded, O, R>
146where
147 L: Location<'a>,
148{
149 fn defer_tick(self) -> Self {
150 Stream::defer_tick(self)
151 }
152}
153
154impl<'a, T, L, O, R> CycleCollection<'a, TickCycleMarker> for Stream<T, Tick<L>, Bounded, O, R>
155where
156 L: Location<'a>,
157{
158 type Location = Tick<L>;
159
160 fn create_source(ident: syn::Ident, location: Tick<L>) -> Self {
161 Stream::new(
162 location.clone(),
163 HydroNode::CycleSource {
164 ident,
165 metadata: location.new_node_metadata::<T>(),
166 },
167 )
168 }
169}
170
171impl<'a, T, L, O, R> CycleComplete<'a, TickCycleMarker> for Stream<T, Tick<L>, Bounded, O, R>
172where
173 L: Location<'a>,
174{
175 fn complete(self, ident: syn::Ident, expected_location: LocationId) {
176 assert_eq!(
177 self.location.id(),
178 expected_location,
179 "locations do not match"
180 );
181 self.location
182 .flow_state()
183 .borrow_mut()
184 .leaves
185 .as_mut()
186 .expect(FLOW_USED_MESSAGE)
187 .push(HydroLeaf::CycleSink {
188 ident,
189 input: Box::new(self.ir_node.into_inner()),
190 metadata: self.location.new_node_metadata::<T>(),
191 });
192 }
193}
194
195impl<'a, T, L, B, O, R> CycleCollection<'a, ForwardRefMarker> for Stream<T, L, B, O, R>
196where
197 L: Location<'a> + NoTick,
198{
199 type Location = L;
200
201 fn create_source(ident: syn::Ident, location: L) -> Self {
202 Stream::new(
203 location.clone(),
204 HydroNode::Persist {
205 inner: Box::new(HydroNode::CycleSource {
206 ident,
207 metadata: location.new_node_metadata::<T>(),
208 }),
209 metadata: location.new_node_metadata::<T>(),
210 },
211 )
212 }
213}
214
215impl<'a, T, L, B, O, R> CycleComplete<'a, ForwardRefMarker> for Stream<T, L, B, O, R>
216where
217 L: Location<'a> + NoTick,
218{
219 fn complete(self, ident: syn::Ident, expected_location: LocationId) {
220 assert_eq!(
221 self.location.id(),
222 expected_location,
223 "locations do not match"
224 );
225 let metadata = self.location.new_node_metadata::<T>();
226 self.location
227 .flow_state()
228 .borrow_mut()
229 .leaves
230 .as_mut()
231 .expect(FLOW_USED_MESSAGE)
232 .push(HydroLeaf::CycleSink {
233 ident,
234 input: Box::new(HydroNode::Unpersist {
235 inner: Box::new(self.ir_node.into_inner()),
236 metadata: metadata.clone(),
237 }),
238 metadata,
239 });
240 }
241}
242
243impl<'a, T, L, B, O, R> Stream<T, L, B, O, R>
244where
245 L: Location<'a>,
246{
247 pub(crate) fn new(location: L, ir_node: HydroNode) -> Self {
248 Stream {
249 location,
250 ir_node: RefCell::new(ir_node),
251 _phantom: PhantomData,
252 }
253 }
254}
255
256impl<'a, T, L, B, O, R> Clone for Stream<T, L, B, O, R>
257where
258 T: Clone,
259 L: Location<'a>,
260{
261 fn clone(&self) -> Self {
262 if !matches!(self.ir_node.borrow().deref(), HydroNode::Tee { .. }) {
263 let orig_ir_node = self.ir_node.replace(HydroNode::Placeholder);
264 *self.ir_node.borrow_mut() = HydroNode::Tee {
265 inner: TeeNode(Rc::new(RefCell::new(orig_ir_node))),
266 metadata: self.location.new_node_metadata::<T>(),
267 };
268 }
269
270 if let HydroNode::Tee { inner, metadata } = self.ir_node.borrow().deref() {
271 Stream {
272 location: self.location.clone(),
273 ir_node: HydroNode::Tee {
274 inner: TeeNode(inner.0.clone()),
275 metadata: metadata.clone(),
276 }
277 .into(),
278 _phantom: PhantomData,
279 }
280 } else {
281 unreachable!()
282 }
283 }
284}
285
286impl<'a, T, L, B, O, R> Stream<T, L, B, O, R>
287where
288 L: Location<'a>,
289{
290 /// Produces a stream based on invoking `f` on each element in order.
291 /// If you do not want to modify the stream and instead only want to view
292 /// each item use [`Stream::inspect`] instead.
293 ///
294 /// # Example
295 /// ```rust
296 /// # use hydro_lang::*;
297 /// # use futures::StreamExt;
298 /// # tokio_test::block_on(test_util::stream_transform_test(|process| {
299 /// let words = process.source_iter(q!(vec!["hello", "world"]));
300 /// words.map(q!(|x| x.to_uppercase()))
301 /// # }, |mut stream| async move {
302 /// # for w in vec!["HELLO", "WORLD"] {
303 /// # assert_eq!(stream.next().await.unwrap(), w);
304 /// # }
305 /// # }));
306 /// ```
307 pub fn map<U, F>(self, f: impl IntoQuotedMut<'a, F, L>) -> Stream<U, L, B, O, R>
308 where
309 F: Fn(T) -> U + 'a,
310 {
311 let f = f.splice_fn1_ctx(&self.location).into();
312 Stream::new(
313 self.location.clone(),
314 HydroNode::Map {
315 f,
316 input: Box::new(self.ir_node.into_inner()),
317 metadata: self.location.new_node_metadata::<U>(),
318 },
319 )
320 }
321
322 /// For each item `i` in the input stream, transform `i` using `f` and then treat the
323 /// result as an [`Iterator`] to produce items one by one. The implementation for [`Iterator`]
324 /// for the output type `U` must produce items in a **deterministic** order.
325 ///
326 /// For example, `U` could be a `Vec`, but not a `HashSet`. If the order of the items in `U` is
327 /// not deterministic, use [`Stream::flat_map_unordered`] instead.
328 ///
329 /// # Example
330 /// ```rust
331 /// # use hydro_lang::*;
332 /// # use futures::StreamExt;
333 /// # tokio_test::block_on(test_util::stream_transform_test(|process| {
334 /// process
335 /// .source_iter(q!(vec![vec![1, 2], vec![3, 4]]))
336 /// .flat_map_ordered(q!(|x| x))
337 /// # }, |mut stream| async move {
338 /// // 1, 2, 3, 4
339 /// # for w in (1..5) {
340 /// # assert_eq!(stream.next().await.unwrap(), w);
341 /// # }
342 /// # }));
343 /// ```
344 pub fn flat_map_ordered<U, I, F>(self, f: impl IntoQuotedMut<'a, F, L>) -> Stream<U, L, B, O, R>
345 where
346 I: IntoIterator<Item = U>,
347 F: Fn(T) -> I + 'a,
348 {
349 let f = f.splice_fn1_ctx(&self.location).into();
350 Stream::new(
351 self.location.clone(),
352 HydroNode::FlatMap {
353 f,
354 input: Box::new(self.ir_node.into_inner()),
355 metadata: self.location.new_node_metadata::<U>(),
356 },
357 )
358 }
359
360 /// Like [`Stream::flat_map_ordered`], but allows the implementation of [`Iterator`]
361 /// for the output type `U` to produce items in any order.
362 ///
363 /// # Example
364 /// ```rust
365 /// # use hydro_lang::{*, stream::ExactlyOnce};
366 /// # use futures::StreamExt;
367 /// # tokio_test::block_on(test_util::stream_transform_test::<_, _, NoOrder, ExactlyOnce>(|process| {
368 /// process
369 /// .source_iter(q!(vec![
370 /// std::collections::HashSet::<i32>::from_iter(vec![1, 2]),
371 /// std::collections::HashSet::from_iter(vec![3, 4]),
372 /// ]))
373 /// .flat_map_unordered(q!(|x| x))
374 /// # }, |mut stream| async move {
375 /// // 1, 2, 3, 4, but in no particular order
376 /// # let mut results = Vec::new();
377 /// # for w in (1..5) {
378 /// # results.push(stream.next().await.unwrap());
379 /// # }
380 /// # results.sort();
381 /// # assert_eq!(results, vec![1, 2, 3, 4]);
382 /// # }));
383 /// ```
384 pub fn flat_map_unordered<U, I, F>(
385 self,
386 f: impl IntoQuotedMut<'a, F, L>,
387 ) -> Stream<U, L, B, NoOrder, R>
388 where
389 I: IntoIterator<Item = U>,
390 F: Fn(T) -> I + 'a,
391 {
392 let f = f.splice_fn1_ctx(&self.location).into();
393 Stream::new(
394 self.location.clone(),
395 HydroNode::FlatMap {
396 f,
397 input: Box::new(self.ir_node.into_inner()),
398 metadata: self.location.new_node_metadata::<U>(),
399 },
400 )
401 }
402
403 /// For each item `i` in the input stream, treat `i` as an [`Iterator`] and produce its items one by one.
404 /// The implementation for [`Iterator`] for the element type `T` must produce items in a **deterministic** order.
405 ///
406 /// For example, `T` could be a `Vec`, but not a `HashSet`. If the order of the items in `T` is
407 /// not deterministic, use [`Stream::flatten_unordered`] instead.
408 ///
409 /// ```rust
410 /// # use hydro_lang::*;
411 /// # use futures::StreamExt;
412 /// # tokio_test::block_on(test_util::stream_transform_test(|process| {
413 /// process
414 /// .source_iter(q!(vec![vec![1, 2], vec![3, 4]]))
415 /// .flatten_ordered()
416 /// # }, |mut stream| async move {
417 /// // 1, 2, 3, 4
418 /// # for w in (1..5) {
419 /// # assert_eq!(stream.next().await.unwrap(), w);
420 /// # }
421 /// # }));
422 /// ```
423 pub fn flatten_ordered<U>(self) -> Stream<U, L, B, O, R>
424 where
425 T: IntoIterator<Item = U>,
426 {
427 self.flat_map_ordered(q!(|d| d))
428 }
429
430 /// Like [`Stream::flatten_ordered`], but allows the implementation of [`Iterator`]
431 /// for the element type `T` to produce items in any order.
432 ///
433 /// # Example
434 /// ```rust
435 /// # use hydro_lang::{*, stream::ExactlyOnce};
436 /// # use futures::StreamExt;
437 /// # tokio_test::block_on(test_util::stream_transform_test::<_, _, NoOrder, ExactlyOnce>(|process| {
438 /// process
439 /// .source_iter(q!(vec![
440 /// std::collections::HashSet::<i32>::from_iter(vec![1, 2]),
441 /// std::collections::HashSet::from_iter(vec![3, 4]),
442 /// ]))
443 /// .flatten_unordered()
444 /// # }, |mut stream| async move {
445 /// // 1, 2, 3, 4, but in no particular order
446 /// # let mut results = Vec::new();
447 /// # for w in (1..5) {
448 /// # results.push(stream.next().await.unwrap());
449 /// # }
450 /// # results.sort();
451 /// # assert_eq!(results, vec![1, 2, 3, 4]);
452 /// # }));
453 pub fn flatten_unordered<U>(self) -> Stream<U, L, B, NoOrder, R>
454 where
455 T: IntoIterator<Item = U>,
456 {
457 self.flat_map_unordered(q!(|d| d))
458 }
459
460 /// Creates a stream containing only the elements of the input stream that satisfy a predicate
461 /// `f`, preserving the order of the elements.
462 ///
463 /// The closure `f` receives a reference `&T` rather than an owned value `T` because filtering does
464 /// not modify or take ownership of the values. If you need to modify the values while filtering
465 /// use [`Stream::filter_map`] instead.
466 ///
467 /// # Example
468 /// ```rust
469 /// # use hydro_lang::*;
470 /// # use futures::StreamExt;
471 /// # tokio_test::block_on(test_util::stream_transform_test(|process| {
472 /// process
473 /// .source_iter(q!(vec![1, 2, 3, 4]))
474 /// .filter(q!(|&x| x > 2))
475 /// # }, |mut stream| async move {
476 /// // 3, 4
477 /// # for w in (3..5) {
478 /// # assert_eq!(stream.next().await.unwrap(), w);
479 /// # }
480 /// # }));
481 /// ```
482 pub fn filter<F>(self, f: impl IntoQuotedMut<'a, F, L>) -> Stream<T, L, B, O, R>
483 where
484 F: Fn(&T) -> bool + 'a,
485 {
486 let f = f.splice_fn1_borrow_ctx(&self.location).into();
487 Stream::new(
488 self.location.clone(),
489 HydroNode::Filter {
490 f,
491 input: Box::new(self.ir_node.into_inner()),
492 metadata: self.location.new_node_metadata::<T>(),
493 },
494 )
495 }
496
497 /// An operator that both filters and maps. It yields only the items for which the supplied closure `f` returns `Some(value)`.
498 ///
499 /// # Example
500 /// ```rust
501 /// # use hydro_lang::*;
502 /// # use futures::StreamExt;
503 /// # tokio_test::block_on(test_util::stream_transform_test(|process| {
504 /// process
505 /// .source_iter(q!(vec!["1", "hello", "world", "2"]))
506 /// .filter_map(q!(|s| s.parse::<usize>().ok()))
507 /// # }, |mut stream| async move {
508 /// // 1, 2
509 /// # for w in (1..3) {
510 /// # assert_eq!(stream.next().await.unwrap(), w);
511 /// # }
512 /// # }));
513 pub fn filter_map<U, F>(self, f: impl IntoQuotedMut<'a, F, L>) -> Stream<U, L, B, O, R>
514 where
515 F: Fn(T) -> Option<U> + 'a,
516 {
517 let f = f.splice_fn1_ctx(&self.location).into();
518 Stream::new(
519 self.location.clone(),
520 HydroNode::FilterMap {
521 f,
522 input: Box::new(self.ir_node.into_inner()),
523 metadata: self.location.new_node_metadata::<U>(),
524 },
525 )
526 }
527
528 /// Generates a stream that maps each input element `i` to a tuple `(i, x)`,
529 /// where `x` is the final value of `other`, a bounded [`Singleton`].
530 ///
531 /// # Example
532 /// ```rust
533 /// # use hydro_lang::*;
534 /// # use futures::StreamExt;
535 /// # tokio_test::block_on(test_util::stream_transform_test(|process| {
536 /// let tick = process.tick();
537 /// let batch = unsafe {
538 /// process
539 /// .source_iter(q!(vec![1, 2, 3, 4]))
540 /// .tick_batch(&tick)
541 /// };
542 /// let count = batch.clone().count(); // `count()` returns a singleton
543 /// batch.cross_singleton(count).all_ticks()
544 /// # }, |mut stream| async move {
545 /// // (1, 4), (2, 4), (3, 4), (4, 4)
546 /// # for w in vec![(1, 4), (2, 4), (3, 4), (4, 4)] {
547 /// # assert_eq!(stream.next().await.unwrap(), w);
548 /// # }
549 /// # }));
550 pub fn cross_singleton<O2>(
551 self,
552 other: impl Into<Optional<O2, L, Bounded>>,
553 ) -> Stream<(T, O2), L, B, O, R>
554 where
555 O2: Clone,
556 {
557 let other: Optional<O2, L, Bounded> = other.into();
558 check_matching_location(&self.location, &other.location);
559
560 Stream::new(
561 self.location.clone(),
562 HydroNode::CrossSingleton {
563 left: Box::new(self.ir_node.into_inner()),
564 right: Box::new(other.ir_node.into_inner()),
565 metadata: self.location.new_node_metadata::<(T, O2)>(),
566 },
567 )
568 }
569
570 /// Allow this stream through if the argument (a Bounded Optional) is non-empty, otherwise the output is empty.
571 pub fn continue_if<U>(self, signal: Optional<U, L, Bounded>) -> Stream<T, L, B, O, R> {
572 self.cross_singleton(signal.map(q!(|_u| ())))
573 .map(q!(|(d, _signal)| d))
574 }
575
576 /// Allow this stream through if the argument (a Bounded Optional) is empty, otherwise the output is empty.
577 pub fn continue_unless<U>(self, other: Optional<U, L, Bounded>) -> Stream<T, L, B, O, R> {
578 self.continue_if(other.into_stream().count().filter(q!(|c| *c == 0)))
579 }
580
581 /// Forms the cross-product (Cartesian product, cross-join) of the items in the 2 input streams, returning all
582 /// tupled pairs in a non-deterministic order.
583 ///
584 /// # Example
585 /// ```rust
586 /// # use hydro_lang::*;
587 /// # use std::collections::HashSet;
588 /// # use futures::StreamExt;
589 /// # tokio_test::block_on(test_util::stream_transform_test(|process| {
590 /// let tick = process.tick();
591 /// let stream1 = process.source_iter(q!(vec!['a', 'b', 'c']));
592 /// let stream2 = process.source_iter(q!(vec![1, 2, 3]));
593 /// stream1.cross_product(stream2)
594 /// # }, |mut stream| async move {
595 /// # let expected = HashSet::from([('a', 1), ('b', 1), ('c', 1), ('a', 2), ('b', 2), ('c', 2), ('a', 3), ('b', 3), ('c', 3)]);
596 /// # stream.map(|i| assert!(expected.contains(&i)));
597 /// # }));
598 pub fn cross_product<O2>(
599 self,
600 other: Stream<O2, L, B, O, R>,
601 ) -> Stream<(T, O2), L, B, NoOrder, R>
602 where
603 T: Clone,
604 O2: Clone,
605 {
606 check_matching_location(&self.location, &other.location);
607
608 Stream::new(
609 self.location.clone(),
610 HydroNode::CrossProduct {
611 left: Box::new(self.ir_node.into_inner()),
612 right: Box::new(other.ir_node.into_inner()),
613 metadata: self.location.new_node_metadata::<(T, O2)>(),
614 },
615 )
616 }
617
618 /// Takes one stream as input and filters out any duplicate occurrences. The output
619 /// contains all unique values from the input.
620 ///
621 /// # Example
622 /// ```rust
623 /// # use hydro_lang::*;
624 /// # use futures::StreamExt;
625 /// # tokio_test::block_on(test_util::stream_transform_test(|process| {
626 /// let tick = process.tick();
627 /// process.source_iter(q!(vec![1, 2, 3, 2, 1, 4])).unique()
628 /// # }, |mut stream| async move {
629 /// # for w in vec![1, 2, 3, 4] {
630 /// # assert_eq!(stream.next().await.unwrap(), w);
631 /// # }
632 /// # }));
633 pub fn unique(self) -> Stream<T, L, B, O, ExactlyOnce>
634 where
635 T: Eq + Hash,
636 {
637 Stream::new(
638 self.location.clone(),
639 HydroNode::Unique {
640 input: Box::new(self.ir_node.into_inner()),
641 metadata: self.location.new_node_metadata::<T>(),
642 },
643 )
644 }
645
646 /// Outputs everything in this stream that is *not* contained in the `other` stream.
647 ///
648 /// The `other` stream must be [`Bounded`], since this function will wait until
649 /// all its elements are available before producing any output.
650 /// # Example
651 /// ```rust
652 /// # use hydro_lang::*;
653 /// # use futures::StreamExt;
654 /// # tokio_test::block_on(test_util::stream_transform_test(|process| {
655 /// let tick = process.tick();
656 /// let stream = unsafe {
657 /// process
658 /// .source_iter(q!(vec![ 1, 2, 3, 4 ]))
659 /// .tick_batch(&tick)
660 /// };
661 /// let batch = unsafe {
662 /// process
663 /// .source_iter(q!(vec![1, 2]))
664 /// .tick_batch(&tick)
665 /// };
666 /// stream.filter_not_in(batch).all_ticks()
667 /// # }, |mut stream| async move {
668 /// # for w in vec![3, 4] {
669 /// # assert_eq!(stream.next().await.unwrap(), w);
670 /// # }
671 /// # }));
672 pub fn filter_not_in<O2>(
673 self,
674 other: Stream<T, L, Bounded, O2, R>,
675 ) -> Stream<T, L, Bounded, O, R>
676 where
677 T: Eq + Hash,
678 {
679 check_matching_location(&self.location, &other.location);
680
681 Stream::new(
682 self.location.clone(),
683 HydroNode::Difference {
684 pos: Box::new(self.ir_node.into_inner()),
685 neg: Box::new(other.ir_node.into_inner()),
686 metadata: self.location.new_node_metadata::<T>(),
687 },
688 )
689 }
690
691 /// An operator which allows you to "inspect" each element of a stream without
692 /// modifying it. The closure `f` is called on a reference to each item. This is
693 /// mainly useful for debugging, and should not be used to generate side-effects.
694 ///
695 /// # Example
696 /// ```rust
697 /// # use hydro_lang::*;
698 /// # use futures::StreamExt;
699 /// # tokio_test::block_on(test_util::stream_transform_test(|process| {
700 /// let nums = process.source_iter(q!(vec![1, 2]));
701 /// // prints "1 * 10 = 10" and "2 * 10 = 20"
702 /// nums.inspect(q!(|x| println!("{} * 10 = {}", x, x * 10)))
703 /// # }, |mut stream| async move {
704 /// # for w in vec![1, 2] {
705 /// # assert_eq!(stream.next().await.unwrap(), w);
706 /// # }
707 /// # }));
708 /// ```
709 pub fn inspect<F>(self, f: impl IntoQuotedMut<'a, F, L>) -> Stream<T, L, B, O, R>
710 where
711 F: Fn(&T) + 'a,
712 {
713 let f = f.splice_fn1_borrow_ctx(&self.location).into();
714
715 if L::is_top_level() {
716 Stream::new(
717 self.location.clone(),
718 HydroNode::Persist {
719 inner: Box::new(HydroNode::Inspect {
720 f,
721 input: Box::new(HydroNode::Unpersist {
722 inner: Box::new(self.ir_node.into_inner()),
723 metadata: self.location.new_node_metadata::<T>(),
724 }),
725 metadata: self.location.new_node_metadata::<T>(),
726 }),
727 metadata: self.location.new_node_metadata::<T>(),
728 },
729 )
730 } else {
731 Stream::new(
732 self.location.clone(),
733 HydroNode::Inspect {
734 f,
735 input: Box::new(self.ir_node.into_inner()),
736 metadata: self.location.new_node_metadata::<T>(),
737 },
738 )
739 }
740 }
741
742 /// Explicitly "casts" the stream to a type with a different ordering
743 /// guarantee. Useful in unsafe code where the ordering cannot be proven
744 /// by the type-system.
745 ///
746 /// # Safety
747 /// This function is used as an escape hatch, and any mistakes in the
748 /// provided ordering guarantee will propagate into the guarantees
749 /// for the rest of the program.
750 ///
751 /// # Example
752 /// # TODO: more sensible code after Shadaj merges
753 /// ```rust
754 /// # use hydro_lang::*;
755 /// # use std::collections::HashSet;
756 /// # use futures::StreamExt;
757 /// # tokio_test::block_on(test_util::stream_transform_test(|process| {
758 /// let nums = process.source_iter(q!({
759 /// let now = std::time::SystemTime::now();
760 /// match now.elapsed().unwrap().as_secs() % 2 {
761 /// 0 => vec![5, 4, 3, 2, 1],
762 /// _ => vec![1, 2, 3, 4, 5],
763 /// }
764 /// .into_iter()
765 /// }));
766 /// // despite being generated by `source_iter`, the order of `nums` across runs is non-deterministic
767 /// let stream = unsafe { nums.assume_ordering::<NoOrder>() };
768 /// stream
769 /// # }, |mut stream| async move {
770 /// # for w in vec![1, 2, 3, 4, 5] {
771 /// # assert!((1..=5).contains(&stream.next().await.unwrap()));
772 /// # }
773 /// # }));
774 /// ```
775 pub unsafe fn assume_ordering<O2>(self) -> Stream<T, L, B, O2, R> {
776 Stream::new(self.location, self.ir_node.into_inner())
777 }
778
779 /// Explicitly "casts" the stream to a type with a different retries
780 /// guarantee. Useful in unsafe code where the lack of retries cannot
781 /// be proven by the type-system.
782 ///
783 /// # Safety
784 /// This function is used as an escape hatch, and any mistakes in the
785 /// provided retries guarantee will propagate into the guarantees
786 /// for the rest of the program.
787 pub unsafe fn assume_retries<R2>(self) -> Stream<T, L, B, O, R2> {
788 Stream::new(self.location, self.ir_node.into_inner())
789 }
790
791 pub fn weakest_retries(self) -> Stream<T, L, B, O, AtLeastOnce> {
792 unsafe {
793 // SAFETY: this is a weaker retry guarantee, so it is safe to assume
794 self.assume_retries::<AtLeastOnce>()
795 }
796 }
797}
798
799impl<'a, T, L, B, O, R> Stream<&T, L, B, O, R>
800where
801 L: Location<'a>,
802{
803 /// Clone each element of the stream; akin to `map(q!(|d| d.clone()))`.
804 ///
805 /// # Example
806 /// ```rust
807 /// # use hydro_lang::*;
808 /// # use futures::StreamExt;
809 /// # tokio_test::block_on(test_util::stream_transform_test(|process| {
810 /// process.source_iter(q!(&[1, 2, 3])).cloned()
811 /// # }, |mut stream| async move {
812 /// // 1, 2, 3
813 /// # for w in vec![1, 2, 3] {
814 /// # assert_eq!(stream.next().await.unwrap(), w);
815 /// # }
816 /// # }));
817 /// ```
818 pub fn cloned(self) -> Stream<T, L, B, O, R>
819 where
820 T: Clone,
821 {
822 self.map(q!(|d| d.clone()))
823 }
824}
825
826impl<'a, T, L, B, O, R> Stream<T, L, B, O, R>
827where
828 L: Location<'a>,
829{
830 /// Combines elements of the stream into a [`Singleton`], by starting with an initial value,
831 /// generated by the `init` closure, and then applying the `comb` closure to each element in the stream.
832 /// Unlike iterators, `comb` takes the accumulator by `&mut` reference, so that it can be modified in place.
833 ///
834 /// The `comb` closure must be **commutative** AND **idempotent**, as the order of input items is not guaranteed
835 /// and there may be duplicates.
836 ///
837 /// # Example
838 /// ```rust
839 /// # use hydro_lang::*;
840 /// # use futures::StreamExt;
841 /// # tokio_test::block_on(test_util::stream_transform_test(|process| {
842 /// let tick = process.tick();
843 /// let bools = process.source_iter(q!(vec![false, true, false]));
844 /// let batch = unsafe { bools.tick_batch(&tick) };
845 /// batch
846 /// .fold_commutative_idempotent(q!(|| false), q!(|acc, x| *acc |= x))
847 /// .all_ticks()
848 /// # }, |mut stream| async move {
849 /// // true
850 /// # assert_eq!(stream.next().await.unwrap(), true);
851 /// # }));
852 /// ```
853 pub fn fold_commutative_idempotent<A, I, F>(
854 self,
855 init: impl IntoQuotedMut<'a, I, L>,
856 comb: impl IntoQuotedMut<'a, F, L>,
857 ) -> Singleton<A, L, B>
858 where
859 I: Fn() -> A + 'a,
860 F: Fn(&mut A, T),
861 {
862 unsafe {
863 // SAFETY: the combinator function is commutative and idempotent
864 self.assume_ordering().assume_retries()
865 }
866 .fold(init, comb)
867 }
868
869 /// Combines elements of the stream into an [`Optional`], by starting with the first element in the stream,
870 /// and then applying the `comb` closure to each element in the stream. The [`Optional`] will be empty
871 /// until the first element in the input arrives. Unlike iterators, `comb` takes the accumulator by `&mut`
872 /// reference, so that it can be modified in place.
873 ///
874 /// The `comb` closure must be **commutative** AND **idempotent**, as the order of input items is not guaranteed
875 /// and there may be duplicates.
876 ///
877 /// # Example
878 /// ```rust
879 /// # use hydro_lang::*;
880 /// # use futures::StreamExt;
881 /// # tokio_test::block_on(test_util::stream_transform_test(|process| {
882 /// let tick = process.tick();
883 /// let bools = process.source_iter(q!(vec![false, true, false]));
884 /// let batch = unsafe { bools.tick_batch(&tick) };
885 /// batch
886 /// .reduce_commutative_idempotent(q!(|acc, x| *acc |= x))
887 /// .all_ticks()
888 /// # }, |mut stream| async move {
889 /// // true
890 /// # assert_eq!(stream.next().await.unwrap(), true);
891 /// # }));
892 /// ```
893 pub fn reduce_commutative_idempotent<F>(
894 self,
895 comb: impl IntoQuotedMut<'a, F, L>,
896 ) -> Optional<T, L, B>
897 where
898 F: Fn(&mut T, T) + 'a,
899 {
900 unsafe {
901 // SAFETY: the combinator function is commutative and idempotent
902 self.assume_ordering().assume_retries()
903 }
904 .reduce(comb)
905 }
906
907 /// Computes the maximum element in the stream as an [`Optional`], which
908 /// will be empty until the first element in the input arrives.
909 ///
910 /// # Example
911 /// ```rust
912 /// # use hydro_lang::*;
913 /// # use futures::StreamExt;
914 /// # tokio_test::block_on(test_util::stream_transform_test(|process| {
915 /// let tick = process.tick();
916 /// let numbers = process.source_iter(q!(vec![1, 2, 3, 4]));
917 /// let batch = unsafe { numbers.tick_batch(&tick) };
918 /// batch.max().all_ticks()
919 /// # }, |mut stream| async move {
920 /// // 4
921 /// # assert_eq!(stream.next().await.unwrap(), 4);
922 /// # }));
923 /// ```
924 pub fn max(self) -> Optional<T, L, B>
925 where
926 T: Ord,
927 {
928 self.reduce_commutative_idempotent(q!(|curr, new| {
929 if new > *curr {
930 *curr = new;
931 }
932 }))
933 }
934
935 /// Computes the maximum element in the stream as an [`Optional`], where the
936 /// maximum is determined according to the `key` function. The [`Optional`] will
937 /// be empty until the first element in the input arrives.
938 ///
939 /// # Example
940 /// ```rust
941 /// # use hydro_lang::*;
942 /// # use futures::StreamExt;
943 /// # tokio_test::block_on(test_util::stream_transform_test(|process| {
944 /// let tick = process.tick();
945 /// let numbers = process.source_iter(q!(vec![1, 2, 3, 4]));
946 /// let batch = unsafe { numbers.tick_batch(&tick) };
947 /// batch.max_by_key(q!(|x| -x)).all_ticks()
948 /// # }, |mut stream| async move {
949 /// // 1
950 /// # assert_eq!(stream.next().await.unwrap(), 1);
951 /// # }));
952 /// ```
953 pub fn max_by_key<K, F>(self, key: impl IntoQuotedMut<'a, F, L> + Copy) -> Optional<T, L, B>
954 where
955 K: Ord,
956 F: Fn(&T) -> K + 'a,
957 {
958 let f = key.splice_fn1_borrow_ctx(&self.location);
959
960 let wrapped: syn::Expr = parse_quote!({
961 let key_fn = #f;
962 move |curr, new| {
963 if key_fn(&new) > key_fn(&*curr) {
964 *curr = new;
965 }
966 }
967 });
968
969 let mut core = HydroNode::Reduce {
970 f: wrapped.into(),
971 input: Box::new(self.ir_node.into_inner()),
972 metadata: self.location.new_node_metadata::<T>(),
973 };
974
975 if L::is_top_level() {
976 core = HydroNode::Persist {
977 inner: Box::new(core),
978 metadata: self.location.new_node_metadata::<T>(),
979 };
980 }
981
982 Optional::new(self.location, core)
983 }
984
985 /// Computes the minimum element in the stream as an [`Optional`], which
986 /// will be empty until the first element in the input arrives.
987 ///
988 /// # Example
989 /// ```rust
990 /// # use hydro_lang::*;
991 /// # use futures::StreamExt;
992 /// # tokio_test::block_on(test_util::stream_transform_test(|process| {
993 /// let tick = process.tick();
994 /// let numbers = process.source_iter(q!(vec![1, 2, 3, 4]));
995 /// let batch = unsafe { numbers.tick_batch(&tick) };
996 /// batch.min().all_ticks()
997 /// # }, |mut stream| async move {
998 /// // 1
999 /// # assert_eq!(stream.next().await.unwrap(), 1);
1000 /// # }));
1001 /// ```
1002 pub fn min(self) -> Optional<T, L, B>
1003 where
1004 T: Ord,
1005 {
1006 self.reduce_commutative_idempotent(q!(|curr, new| {
1007 if new < *curr {
1008 *curr = new;
1009 }
1010 }))
1011 }
1012}
1013
1014impl<'a, T, L, B, O> Stream<T, L, B, O, ExactlyOnce>
1015where
1016 L: Location<'a>,
1017{
1018 /// Combines elements of the stream into a [`Singleton`], by starting with an initial value,
1019 /// generated by the `init` closure, and then applying the `comb` closure to each element in the stream.
1020 /// Unlike iterators, `comb` takes the accumulator by `&mut` reference, so that it can be modified in place.
1021 ///
1022 /// The `comb` closure must be **commutative**, as the order of input items is not guaranteed.
1023 ///
1024 /// # Example
1025 /// ```rust
1026 /// # use hydro_lang::*;
1027 /// # use futures::StreamExt;
1028 /// # tokio_test::block_on(test_util::stream_transform_test(|process| {
1029 /// let tick = process.tick();
1030 /// let numbers = process.source_iter(q!(vec![1, 2, 3, 4]));
1031 /// let batch = unsafe { numbers.tick_batch(&tick) };
1032 /// batch
1033 /// .fold_commutative(q!(|| 0), q!(|acc, x| *acc += x))
1034 /// .all_ticks()
1035 /// # }, |mut stream| async move {
1036 /// // 10
1037 /// # assert_eq!(stream.next().await.unwrap(), 10);
1038 /// # }));
1039 /// ```
1040 pub fn fold_commutative<A, I, F>(
1041 self,
1042 init: impl IntoQuotedMut<'a, I, L>,
1043 comb: impl IntoQuotedMut<'a, F, L>,
1044 ) -> Singleton<A, L, B>
1045 where
1046 I: Fn() -> A + 'a,
1047 F: Fn(&mut A, T),
1048 {
1049 unsafe {
1050 // SAFETY: the combinator function is commutative
1051 self.assume_ordering()
1052 }
1053 .fold(init, comb)
1054 }
1055
1056 /// Combines elements of the stream into a [`Optional`], by starting with the first element in the stream,
1057 /// and then applying the `comb` closure to each element in the stream. The [`Optional`] will be empty
1058 /// until the first element in the input arrives. Unlike iterators, `comb` takes the accumulator by `&mut`
1059 /// reference, so that it can be modified in place.
1060 ///
1061 /// The `comb` closure must be **commutative**, as the order of input items is not guaranteed.
1062 ///
1063 /// # Example
1064 /// ```rust
1065 /// # use hydro_lang::*;
1066 /// # use futures::StreamExt;
1067 /// # tokio_test::block_on(test_util::stream_transform_test(|process| {
1068 /// let tick = process.tick();
1069 /// let numbers = process.source_iter(q!(vec![1, 2, 3, 4]));
1070 /// let batch = unsafe { numbers.tick_batch(&tick) };
1071 /// batch
1072 /// .reduce_commutative(q!(|curr, new| *curr += new))
1073 /// .all_ticks()
1074 /// # }, |mut stream| async move {
1075 /// // 10
1076 /// # assert_eq!(stream.next().await.unwrap(), 10);
1077 /// # }));
1078 /// ```
1079 pub fn reduce_commutative<F>(self, comb: impl IntoQuotedMut<'a, F, L>) -> Optional<T, L, B>
1080 where
1081 F: Fn(&mut T, T) + 'a,
1082 {
1083 unsafe {
1084 // SAFETY: the combinator function is commutative
1085 self.assume_ordering()
1086 }
1087 .reduce(comb)
1088 }
1089
1090 /// Computes the number of elements in the stream as a [`Singleton`].
1091 ///
1092 /// # Example
1093 /// ```rust
1094 /// # use hydro_lang::*;
1095 /// # use futures::StreamExt;
1096 /// # tokio_test::block_on(test_util::stream_transform_test(|process| {
1097 /// let tick = process.tick();
1098 /// let numbers = process.source_iter(q!(vec![1, 2, 3, 4]));
1099 /// let batch = unsafe { numbers.tick_batch(&tick) };
1100 /// batch.count().all_ticks()
1101 /// # }, |mut stream| async move {
1102 /// // 4
1103 /// # assert_eq!(stream.next().await.unwrap(), 4);
1104 /// # }));
1105 /// ```
1106 pub fn count(self) -> Singleton<usize, L, B> {
1107 self.fold_commutative(q!(|| 0usize), q!(|count, _| *count += 1))
1108 }
1109}
1110
1111impl<'a, T, L, B, R> Stream<T, L, B, TotalOrder, R>
1112where
1113 L: Location<'a>,
1114{
1115 /// Combines elements of the stream into a [`Singleton`], by starting with an initial value,
1116 /// generated by the `init` closure, and then applying the `comb` closure to each element in the stream.
1117 /// Unlike iterators, `comb` takes the accumulator by `&mut` reference, so that it can be modified in place.
1118 ///
1119 /// The `comb` closure must be **idempotent**, as there may be non-deterministic duplicates.
1120 ///
1121 /// # Example
1122 /// ```rust
1123 /// # use hydro_lang::*;
1124 /// # use futures::StreamExt;
1125 /// # tokio_test::block_on(test_util::stream_transform_test(|process| {
1126 /// let tick = process.tick();
1127 /// let bools = process.source_iter(q!(vec![false, true, false]));
1128 /// let batch = unsafe { bools.tick_batch(&tick) };
1129 /// batch
1130 /// .fold_idempotent(q!(|| false), q!(|acc, x| *acc |= x))
1131 /// .all_ticks()
1132 /// # }, |mut stream| async move {
1133 /// // true
1134 /// # assert_eq!(stream.next().await.unwrap(), true);
1135 /// # }));
1136 /// ```
1137 pub fn fold_idempotent<A, I, F>(
1138 self,
1139 init: impl IntoQuotedMut<'a, I, L>,
1140 comb: impl IntoQuotedMut<'a, F, L>,
1141 ) -> Singleton<A, L, B>
1142 where
1143 I: Fn() -> A + 'a,
1144 F: Fn(&mut A, T),
1145 {
1146 unsafe {
1147 // SAFETY: the combinator function is idempotent
1148 self.assume_retries()
1149 }
1150 .fold(init, comb)
1151 }
1152
1153 /// Combines elements of the stream into an [`Optional`], by starting with the first element in the stream,
1154 /// and then applying the `comb` closure to each element in the stream. The [`Optional`] will be empty
1155 /// until the first element in the input arrives. Unlike iterators, `comb` takes the accumulator by `&mut`
1156 /// reference, so that it can be modified in place.
1157 ///
1158 /// The `comb` closure must be **idempotent**, as there may be non-deterministic duplicates.
1159 ///
1160 /// # Example
1161 /// ```rust
1162 /// # use hydro_lang::*;
1163 /// # use futures::StreamExt;
1164 /// # tokio_test::block_on(test_util::stream_transform_test(|process| {
1165 /// let tick = process.tick();
1166 /// let bools = process.source_iter(q!(vec![false, true, false]));
1167 /// let batch = unsafe { bools.tick_batch(&tick) };
1168 /// batch.reduce_idempotent(q!(|acc, x| *acc |= x)).all_ticks()
1169 /// # }, |mut stream| async move {
1170 /// // true
1171 /// # assert_eq!(stream.next().await.unwrap(), true);
1172 /// # }));
1173 /// ```
1174 pub fn reduce_idempotent<F>(self, comb: impl IntoQuotedMut<'a, F, L>) -> Optional<T, L, B>
1175 where
1176 F: Fn(&mut T, T) + 'a,
1177 {
1178 unsafe {
1179 // SAFETY: the combinator function is idempotent
1180 self.assume_retries()
1181 }
1182 .reduce(comb)
1183 }
1184}
1185
1186impl<'a, T, L, B> Stream<T, L, B, TotalOrder, ExactlyOnce>
1187where
1188 L: Location<'a>,
1189{
1190 /// Returns a stream with the current count tupled with each element in the input stream.
1191 ///
1192 /// # Example
1193 /// ```rust
1194 /// # use hydro_lang::{*, stream::ExactlyOnce};
1195 /// # use futures::StreamExt;
1196 /// # tokio_test::block_on(test_util::stream_transform_test::<_, _, TotalOrder, ExactlyOnce>(|process| {
1197 /// let tick = process.tick();
1198 /// let numbers = process.source_iter(q!(vec![1, 2, 3, 4]));
1199 /// numbers.enumerate()
1200 /// # }, |mut stream| async move {
1201 /// // (0, 1), (1, 2), (2, 3), (3, 4)
1202 /// # for w in vec![(0, 1), (1, 2), (2, 3), (3, 4)] {
1203 /// # assert_eq!(stream.next().await.unwrap(), w);
1204 /// # }
1205 /// # }));
1206 /// ```
1207 pub fn enumerate(self) -> Stream<(usize, T), L, B, TotalOrder, ExactlyOnce> {
1208 if L::is_top_level() {
1209 Stream::new(
1210 self.location.clone(),
1211 HydroNode::Persist {
1212 inner: Box::new(HydroNode::Enumerate {
1213 is_static: true,
1214 input: Box::new(HydroNode::Unpersist {
1215 inner: Box::new(self.ir_node.into_inner()),
1216 metadata: self.location.new_node_metadata::<T>(),
1217 }),
1218 metadata: self.location.new_node_metadata::<(usize, T)>(),
1219 }),
1220 metadata: self.location.new_node_metadata::<(usize, T)>(),
1221 },
1222 )
1223 } else {
1224 Stream::new(
1225 self.location.clone(),
1226 HydroNode::Enumerate {
1227 is_static: false,
1228 input: Box::new(self.ir_node.into_inner()),
1229 metadata: self.location.new_node_metadata::<(usize, T)>(),
1230 },
1231 )
1232 }
1233 }
1234
1235 /// Computes the first element in the stream as an [`Optional`], which
1236 /// will be empty until the first element in the input arrives.
1237 ///
1238 /// This requires the stream to have a [`TotalOrder`] guarantee, otherwise
1239 /// re-ordering of elements may cause the first element to change.
1240 ///
1241 /// # Example
1242 /// ```rust
1243 /// # use hydro_lang::*;
1244 /// # use futures::StreamExt;
1245 /// # tokio_test::block_on(test_util::stream_transform_test(|process| {
1246 /// let tick = process.tick();
1247 /// let numbers = process.source_iter(q!(vec![1, 2, 3, 4]));
1248 /// let batch = unsafe { numbers.tick_batch(&tick) };
1249 /// batch.first().all_ticks()
1250 /// # }, |mut stream| async move {
1251 /// // 1
1252 /// # assert_eq!(stream.next().await.unwrap(), 1);
1253 /// # }));
1254 /// ```
1255 pub fn first(self) -> Optional<T, L, B> {
1256 self.reduce(q!(|_, _| {}))
1257 }
1258
1259 /// Computes the last element in the stream as an [`Optional`], which
1260 /// will be empty until an element in the input arrives.
1261 ///
1262 /// This requires the stream to have a [`TotalOrder`] guarantee, otherwise
1263 /// re-ordering of elements may cause the last element to change.
1264 ///
1265 /// # Example
1266 /// ```rust
1267 /// # use hydro_lang::*;
1268 /// # use futures::StreamExt;
1269 /// # tokio_test::block_on(test_util::stream_transform_test(|process| {
1270 /// let tick = process.tick();
1271 /// let numbers = process.source_iter(q!(vec![1, 2, 3, 4]));
1272 /// let batch = unsafe { numbers.tick_batch(&tick) };
1273 /// batch.last().all_ticks()
1274 /// # }, |mut stream| async move {
1275 /// // 4
1276 /// # assert_eq!(stream.next().await.unwrap(), 4);
1277 /// # }));
1278 /// ```
1279 pub fn last(self) -> Optional<T, L, B> {
1280 self.reduce(q!(|curr, new| *curr = new))
1281 }
1282
1283 /// Combines elements of the stream into a [`Singleton`], by starting with an intitial value,
1284 /// generated by the `init` closure, and then applying the `comb` closure to each element in the stream.
1285 /// Unlike iterators, `comb` takes the accumulator by `&mut` reference, so that it can be modified in place.
1286 ///
1287 /// The input stream must have a [`TotalOrder`] guarantee, which means that the `comb` closure is allowed
1288 /// to depend on the order of elements in the stream.
1289 ///
1290 /// # Example
1291 /// ```rust
1292 /// # use hydro_lang::*;
1293 /// # use futures::StreamExt;
1294 /// # tokio_test::block_on(test_util::stream_transform_test(|process| {
1295 /// let tick = process.tick();
1296 /// let words = process.source_iter(q!(vec!["HELLO", "WORLD"]));
1297 /// let batch = unsafe { words.tick_batch(&tick) };
1298 /// batch
1299 /// .fold(q!(|| String::new()), q!(|acc, x| acc.push_str(x)))
1300 /// .all_ticks()
1301 /// # }, |mut stream| async move {
1302 /// // "HELLOWORLD"
1303 /// # assert_eq!(stream.next().await.unwrap(), "HELLOWORLD");
1304 /// # }));
1305 /// ```
1306 pub fn fold<A, I: Fn() -> A + 'a, F: Fn(&mut A, T)>(
1307 self,
1308 init: impl IntoQuotedMut<'a, I, L>,
1309 comb: impl IntoQuotedMut<'a, F, L>,
1310 ) -> Singleton<A, L, B> {
1311 let init = init.splice_fn0_ctx(&self.location).into();
1312 let comb = comb.splice_fn2_borrow_mut_ctx(&self.location).into();
1313
1314 let mut core = HydroNode::Fold {
1315 init,
1316 acc: comb,
1317 input: Box::new(self.ir_node.into_inner()),
1318 metadata: self.location.new_node_metadata::<A>(),
1319 };
1320
1321 if L::is_top_level() {
1322 // top-level (possibly unbounded) singletons are represented as
1323 // a stream which produces all values from all ticks every tick,
1324 // so Unpersist will always give the lastest aggregation
1325 core = HydroNode::Persist {
1326 inner: Box::new(core),
1327 metadata: self.location.new_node_metadata::<A>(),
1328 };
1329 }
1330
1331 Singleton::new(self.location, core)
1332 }
1333
1334 /// Applies a function to each element of the stream, maintaining an internal state (accumulator)
1335 /// and emitting each intermediate result.
1336 ///
1337 /// Unlike `fold` which only returns the final accumulated value, `scan` produces a new stream
1338 /// containing all intermediate accumulated values. The scan operation can also terminate early
1339 /// by returning `None`.
1340 ///
1341 /// The function takes a mutable reference to the accumulator and the current element, and returns
1342 /// an `Option<U>`. If the function returns `Some(value)`, `value` is emitted to the output stream.
1343 /// If the function returns `None`, the stream is terminated and no more elements are processed.
1344 ///
1345 /// # Examples
1346 ///
1347 /// Basic usage - running sum:
1348 /// ```rust
1349 /// # use hydro_lang::*;
1350 /// # use futures::StreamExt;
1351 /// # tokio_test::block_on(test_util::stream_transform_test(|process| {
1352 /// process.source_iter(q!(vec![1, 2, 3, 4])).scan(
1353 /// q!(|| 0),
1354 /// q!(|acc, x| {
1355 /// *acc += x;
1356 /// Some(*acc)
1357 /// }),
1358 /// )
1359 /// # }, |mut stream| async move {
1360 /// // Output: 1, 3, 6, 10
1361 /// # for w in vec![1, 3, 6, 10] {
1362 /// # assert_eq!(stream.next().await.unwrap(), w);
1363 /// # }
1364 /// # }));
1365 /// ```
1366 ///
1367 /// Early termination example:
1368 /// ```rust
1369 /// # use hydro_lang::*;
1370 /// # use futures::StreamExt;
1371 /// # tokio_test::block_on(test_util::stream_transform_test(|process| {
1372 /// process.source_iter(q!(vec![1, 2, 3, 4])).scan(
1373 /// q!(|| 1),
1374 /// q!(|state, x| {
1375 /// *state = *state * x;
1376 /// if *state > 6 {
1377 /// None // Terminate the stream
1378 /// } else {
1379 /// Some(-*state)
1380 /// }
1381 /// }),
1382 /// )
1383 /// # }, |mut stream| async move {
1384 /// // Output: -1, -2, -6
1385 /// # for w in vec![-1, -2, -6] {
1386 /// # assert_eq!(stream.next().await.unwrap(), w);
1387 /// # }
1388 /// # }));
1389 /// ```
1390 pub fn scan<A, U, I, F>(
1391 self,
1392 init: impl IntoQuotedMut<'a, I, L>,
1393 f: impl IntoQuotedMut<'a, F, L>,
1394 ) -> Stream<U, L, B, TotalOrder, ExactlyOnce>
1395 where
1396 I: Fn() -> A + 'a,
1397 F: Fn(&mut A, T) -> Option<U> + 'a,
1398 {
1399 let init = init.splice_fn0_ctx(&self.location).into();
1400 let f = f.splice_fn2_borrow_mut_ctx(&self.location).into();
1401
1402 if L::is_top_level() {
1403 Stream::new(
1404 self.location.clone(),
1405 HydroNode::Persist {
1406 inner: Box::new(HydroNode::Scan {
1407 init,
1408 acc: f,
1409 input: Box::new(HydroNode::Unpersist {
1410 inner: Box::new(self.ir_node.into_inner()),
1411 metadata: self.location.new_node_metadata::<U>(),
1412 }),
1413 metadata: self.location.new_node_metadata::<U>(),
1414 }),
1415 metadata: self.location.new_node_metadata::<U>(),
1416 },
1417 )
1418 } else {
1419 Stream::new(
1420 self.location.clone(),
1421 HydroNode::Scan {
1422 init,
1423 acc: f,
1424 input: Box::new(self.ir_node.into_inner()),
1425 metadata: self.location.new_node_metadata::<U>(),
1426 },
1427 )
1428 }
1429 }
1430
1431 /// Combines elements of the stream into an [`Optional`], by starting with the first element in the stream,
1432 /// and then applying the `comb` closure to each element in the stream. The [`Optional`] will be empty
1433 /// until the first element in the input arrives.
1434 ///
1435 /// The input stream must have a [`TotalOrder`] guarantee, which means that the `comb` closure is allowed
1436 /// to depend on the order of elements in the stream.
1437 ///
1438 /// # Example
1439 /// ```rust
1440 /// # use hydro_lang::*;
1441 /// # use futures::StreamExt;
1442 /// # tokio_test::block_on(test_util::stream_transform_test(|process| {
1443 /// let tick = process.tick();
1444 /// let words = process.source_iter(q!(vec!["HELLO", "WORLD"]));
1445 /// let batch = unsafe { words.tick_batch(&tick) };
1446 /// batch
1447 /// .map(q!(|x| x.to_string()))
1448 /// .reduce(q!(|curr, new| curr.push_str(&new)))
1449 /// .all_ticks()
1450 /// # }, |mut stream| async move {
1451 /// // "HELLOWORLD"
1452 /// # assert_eq!(stream.next().await.unwrap(), "HELLOWORLD");
1453 /// # }));
1454 /// ```
1455 pub fn reduce<F: Fn(&mut T, T) + 'a>(
1456 self,
1457 comb: impl IntoQuotedMut<'a, F, L>,
1458 ) -> Optional<T, L, B> {
1459 let f = comb.splice_fn2_borrow_mut_ctx(&self.location).into();
1460 let mut core = HydroNode::Reduce {
1461 f,
1462 input: Box::new(self.ir_node.into_inner()),
1463 metadata: self.location.new_node_metadata::<T>(),
1464 };
1465
1466 if L::is_top_level() {
1467 core = HydroNode::Persist {
1468 inner: Box::new(core),
1469 metadata: self.location.new_node_metadata::<T>(),
1470 };
1471 }
1472
1473 Optional::new(self.location, core)
1474 }
1475}
1476
1477impl<'a, T, L: Location<'a> + NoTick + NoAtomic, O, R> Stream<T, L, Unbounded, O, R> {
1478 /// Produces a new stream that interleaves the elements of the two input streams.
1479 /// The result has [`NoOrder`] because the order of interleaving is not guaranteed.
1480 ///
1481 /// Currently, both input streams must be [`Unbounded`]. When the streams are
1482 /// [`Bounded`], you can use [`Stream::chain`] instead.
1483 ///
1484 /// # Example
1485 /// ```rust
1486 /// # use hydro_lang::*;
1487 /// # use futures::StreamExt;
1488 /// # tokio_test::block_on(test_util::stream_transform_test(|process| {
1489 /// let numbers = process.source_iter(q!(vec![1, 2, 3, 4]));
1490 /// numbers.clone().map(q!(|x| x + 1)).union(numbers)
1491 /// # }, |mut stream| async move {
1492 /// // 2, 3, 4, 5, and 1, 2, 3, 4 interleaved in unknown order
1493 /// # for w in vec![2, 3, 4, 5, 1, 2, 3, 4] {
1494 /// # assert_eq!(stream.next().await.unwrap(), w);
1495 /// # }
1496 /// # }));
1497 /// ```
1498 pub fn union<O2, R2: MinRetries<R>>(
1499 self,
1500 other: Stream<T, L, Unbounded, O2, R2>,
1501 ) -> Stream<T, L, Unbounded, NoOrder, R2::Min> {
1502 let tick = self.location.tick();
1503 unsafe {
1504 // SAFETY: Because the outputs are unordered,
1505 // we can interleave batches from both streams.
1506 self.tick_batch(&tick)
1507 .assume_ordering::<NoOrder>()
1508 .assume_retries::<R2::Min>()
1509 .chain(
1510 other
1511 .tick_batch(&tick)
1512 .assume_ordering::<NoOrder>()
1513 .assume_retries::<R2::Min>(),
1514 )
1515 .all_ticks()
1516 .assume_ordering()
1517 }
1518 }
1519}
1520
1521impl<'a, T, L, O, R> Stream<T, L, Bounded, O, R>
1522where
1523 L: Location<'a>,
1524{
1525 /// Produces a new stream that emits the input elements in sorted order.
1526 ///
1527 /// The input stream can have any ordering guarantee, but the output stream
1528 /// will have a [`TotalOrder`] guarantee. This operator will block until all
1529 /// elements in the input stream are available, so it requires the input stream
1530 /// to be [`Bounded`].
1531 ///
1532 /// # Example
1533 /// ```rust
1534 /// # use hydro_lang::*;
1535 /// # use futures::StreamExt;
1536 /// # tokio_test::block_on(test_util::stream_transform_test(|process| {
1537 /// let tick = process.tick();
1538 /// let numbers = process.source_iter(q!(vec![4, 2, 3, 1]));
1539 /// let batch = unsafe { numbers.tick_batch(&tick) };
1540 /// batch.sort().all_ticks()
1541 /// # }, |mut stream| async move {
1542 /// // 1, 2, 3, 4
1543 /// # for w in (1..5) {
1544 /// # assert_eq!(stream.next().await.unwrap(), w);
1545 /// # }
1546 /// # }));
1547 /// ```
1548 pub fn sort(self) -> Stream<T, L, Bounded, TotalOrder, R>
1549 where
1550 T: Ord,
1551 {
1552 Stream::new(
1553 self.location.clone(),
1554 HydroNode::Sort {
1555 input: Box::new(self.ir_node.into_inner()),
1556 metadata: self.location.new_node_metadata::<T>(),
1557 },
1558 )
1559 }
1560
1561 /// Produces a new stream that first emits the elements of the `self` stream,
1562 /// and then emits the elements of the `other` stream. The output stream has
1563 /// a [`TotalOrder`] guarantee if and only if both input streams have a
1564 /// [`TotalOrder`] guarantee.
1565 ///
1566 /// Currently, both input streams must be [`Bounded`]. This operator will block
1567 /// on the first stream until all its elements are available. In a future version,
1568 /// we will relax the requirement on the `other` stream.
1569 ///
1570 /// # Example
1571 /// ```rust
1572 /// # use hydro_lang::*;
1573 /// # use futures::StreamExt;
1574 /// # tokio_test::block_on(test_util::stream_transform_test(|process| {
1575 /// let tick = process.tick();
1576 /// let numbers = process.source_iter(q!(vec![1, 2, 3, 4]));
1577 /// let batch = unsafe { numbers.tick_batch(&tick) };
1578 /// batch.clone().map(q!(|x| x + 1)).chain(batch).all_ticks()
1579 /// # }, |mut stream| async move {
1580 /// // 2, 3, 4, 5, 1, 2, 3, 4
1581 /// # for w in vec![2, 3, 4, 5, 1, 2, 3, 4] {
1582 /// # assert_eq!(stream.next().await.unwrap(), w);
1583 /// # }
1584 /// # }));
1585 /// ```
1586 pub fn chain<O2>(self, other: Stream<T, L, Bounded, O2, R>) -> Stream<T, L, Bounded, O::Min, R>
1587 where
1588 O: MinOrder<O2>,
1589 {
1590 check_matching_location(&self.location, &other.location);
1591
1592 Stream::new(
1593 self.location.clone(),
1594 HydroNode::Chain {
1595 first: Box::new(self.ir_node.into_inner()),
1596 second: Box::new(other.ir_node.into_inner()),
1597 metadata: self.location.new_node_metadata::<T>(),
1598 },
1599 )
1600 }
1601}
1602
1603impl<'a, K, V1, L, B, O, R> Stream<(K, V1), L, B, O, R>
1604where
1605 L: Location<'a>,
1606{
1607 /// Given two streams of pairs `(K, V1)` and `(K, V2)`, produces a new stream of nested pairs `(K, (V1, V2))`
1608 /// by equi-joining the two streams on the key attribute `K`.
1609 ///
1610 /// # Example
1611 /// ```rust
1612 /// # use hydro_lang::*;
1613 /// # use std::collections::HashSet;
1614 /// # use futures::StreamExt;
1615 /// # tokio_test::block_on(test_util::stream_transform_test(|process| {
1616 /// let tick = process.tick();
1617 /// let stream1 = process.source_iter(q!(vec![(1, 'a'), (2, 'b')]));
1618 /// let stream2 = process.source_iter(q!(vec![(1, 'x'), (2, 'y')]));
1619 /// stream1.join(stream2)
1620 /// # }, |mut stream| async move {
1621 /// // (1, ('a', 'x')), (2, ('b', 'y'))
1622 /// # let expected = HashSet::from([(1, ('a', 'x')), (2, ('b', 'y'))]);
1623 /// # stream.map(|i| assert!(expected.contains(&i)));
1624 /// # }));
1625 pub fn join<V2, O2>(
1626 self,
1627 n: Stream<(K, V2), L, B, O2, R>,
1628 ) -> Stream<(K, (V1, V2)), L, B, NoOrder, R>
1629 where
1630 K: Eq + Hash,
1631 {
1632 check_matching_location(&self.location, &n.location);
1633
1634 Stream::new(
1635 self.location.clone(),
1636 HydroNode::Join {
1637 left: Box::new(self.ir_node.into_inner()),
1638 right: Box::new(n.ir_node.into_inner()),
1639 metadata: self.location.new_node_metadata::<(K, (V1, V2))>(),
1640 },
1641 )
1642 }
1643
1644 /// Given a stream of pairs `(K, V1)` and a bounded stream of keys `K`,
1645 /// computes the anti-join of the items in the input -- i.e. returns
1646 /// unique items in the first input that do not have a matching key
1647 /// in the second input.
1648 ///
1649 /// # Example
1650 /// ```rust
1651 /// # use hydro_lang::*;
1652 /// # use futures::StreamExt;
1653 /// # tokio_test::block_on(test_util::stream_transform_test(|process| {
1654 /// let tick = process.tick();
1655 /// let stream = unsafe {
1656 /// process
1657 /// .source_iter(q!(vec![ (1, 'a'), (2, 'b'), (3, 'c'), (4, 'd') ]))
1658 /// .tick_batch(&tick)
1659 /// };
1660 /// let batch = unsafe {
1661 /// process
1662 /// .source_iter(q!(vec![1, 2]))
1663 /// .tick_batch(&tick)
1664 /// };
1665 /// stream.anti_join(batch).all_ticks()
1666 /// # }, |mut stream| async move {
1667 /// # for w in vec![(3, 'c'), (4, 'd')] {
1668 /// # assert_eq!(stream.next().await.unwrap(), w);
1669 /// # }
1670 /// # }));
1671 pub fn anti_join<O2>(self, n: Stream<K, L, Bounded, O2, R>) -> Stream<(K, V1), L, B, O, R>
1672 where
1673 K: Eq + Hash,
1674 {
1675 check_matching_location(&self.location, &n.location);
1676
1677 Stream::new(
1678 self.location.clone(),
1679 HydroNode::AntiJoin {
1680 pos: Box::new(self.ir_node.into_inner()),
1681 neg: Box::new(n.ir_node.into_inner()),
1682 metadata: self.location.new_node_metadata::<(K, V1)>(),
1683 },
1684 )
1685 }
1686}
1687
1688impl<'a, K, V, L> Stream<(K, V), Tick<L>, Bounded, TotalOrder, ExactlyOnce>
1689where
1690 K: Eq + Hash,
1691 L: Location<'a>,
1692{
1693 /// A special case of [`Stream::fold`], in the spirit of SQL's GROUP BY and aggregation constructs. The input
1694 /// tuples are partitioned into groups by the first element ("keys"), and for each group the values
1695 /// in the second element are accumulated via the `comb` closure.
1696 ///
1697 /// The input stream must have a [`TotalOrder`] guarantee, which means that the `comb` closure is allowed
1698 /// to depend on the order of elements in the stream.
1699 ///
1700 /// If the input and output value types are the same and do not require initialization then use
1701 /// [`Stream::reduce_keyed`].
1702 ///
1703 /// # Example
1704 /// ```rust
1705 /// # use hydro_lang::*;
1706 /// # use futures::StreamExt;
1707 /// # tokio_test::block_on(test_util::stream_transform_test(|process| {
1708 /// let tick = process.tick();
1709 /// let numbers = process.source_iter(q!(vec![(1, 2), (2, 3), (1, 3), (2, 4)]));
1710 /// let batch = unsafe { numbers.tick_batch(&tick) };
1711 /// batch
1712 /// .fold_keyed(q!(|| 0), q!(|acc, x| *acc += x))
1713 /// .all_ticks()
1714 /// # }, |mut stream| async move {
1715 /// // (1, 5), (2, 7)
1716 /// # assert_eq!(stream.next().await.unwrap(), (1, 5));
1717 /// # assert_eq!(stream.next().await.unwrap(), (2, 7));
1718 /// # }));
1719 /// ```
1720 pub fn fold_keyed<A, I, F>(
1721 self,
1722 init: impl IntoQuotedMut<'a, I, Tick<L>>,
1723 comb: impl IntoQuotedMut<'a, F, Tick<L>>,
1724 ) -> Stream<(K, A), Tick<L>, Bounded, NoOrder, ExactlyOnce>
1725 where
1726 I: Fn() -> A + 'a,
1727 F: Fn(&mut A, V) + 'a,
1728 {
1729 let init = init.splice_fn0_ctx(&self.location).into();
1730 let comb = comb.splice_fn2_borrow_mut_ctx(&self.location).into();
1731
1732 Stream::new(
1733 self.location.clone(),
1734 HydroNode::FoldKeyed {
1735 init,
1736 acc: comb,
1737 input: Box::new(self.ir_node.into_inner()),
1738 metadata: self.location.new_node_metadata::<(K, A)>(),
1739 },
1740 )
1741 }
1742
1743 /// A special case of [`Stream::reduce`], in the spirit of SQL's GROUP BY and aggregation constructs. The input
1744 /// tuples are partitioned into groups by the first element ("keys"), and for each group the values
1745 /// in the second element are accumulated via the `comb` closure.
1746 ///
1747 /// The input stream must have a [`TotalOrder`] guarantee, which means that the `comb` closure is allowed
1748 /// to depend on the order of elements in the stream.
1749 ///
1750 /// If you need the accumulated value to have a different type than the input, use [`Stream::fold_keyed`].
1751 ///
1752 /// # Example
1753 /// ```rust
1754 /// # use hydro_lang::*;
1755 /// # use futures::StreamExt;
1756 /// # tokio_test::block_on(test_util::stream_transform_test(|process| {
1757 /// let tick = process.tick();
1758 /// let numbers = process.source_iter(q!(vec![(1, 2), (2, 3), (1, 3), (2, 4)]));
1759 /// let batch = unsafe { numbers.tick_batch(&tick) };
1760 /// batch.reduce_keyed(q!(|acc, x| *acc += x)).all_ticks()
1761 /// # }, |mut stream| async move {
1762 /// // (1, 5), (2, 7)
1763 /// # assert_eq!(stream.next().await.unwrap(), (1, 5));
1764 /// # assert_eq!(stream.next().await.unwrap(), (2, 7));
1765 /// # }));
1766 /// ```
1767 pub fn reduce_keyed<F>(
1768 self,
1769 comb: impl IntoQuotedMut<'a, F, Tick<L>>,
1770 ) -> Stream<(K, V), Tick<L>, Bounded, NoOrder, ExactlyOnce>
1771 where
1772 F: Fn(&mut V, V) + 'a,
1773 {
1774 let f = comb.splice_fn2_borrow_mut_ctx(&self.location).into();
1775
1776 Stream::new(
1777 self.location.clone(),
1778 HydroNode::ReduceKeyed {
1779 f,
1780 input: Box::new(self.ir_node.into_inner()),
1781 metadata: self.location.new_node_metadata::<(K, V)>(),
1782 },
1783 )
1784 }
1785}
1786
1787impl<'a, K, V, L, O, R> Stream<(K, V), Tick<L>, Bounded, O, R>
1788where
1789 K: Eq + Hash,
1790 L: Location<'a>,
1791{
1792 /// A special case of [`Stream::fold_commutative_idempotent`], in the spirit of SQL's GROUP BY and aggregation constructs.
1793 /// The input tuples are partitioned into groups by the first element ("keys"), and for each group the values
1794 /// in the second element are accumulated via the `comb` closure.
1795 ///
1796 /// The `comb` closure must be **commutative**, as the order of input items is not guaranteed, and **idempotent**,
1797 /// as there may be non-deterministic duplicates.
1798 ///
1799 /// If the input and output value types are the same and do not require initialization then use
1800 /// [`Stream::reduce_keyed_commutative_idempotent`].
1801 ///
1802 /// # Example
1803 /// ```rust
1804 /// # use hydro_lang::*;
1805 /// # use futures::StreamExt;
1806 /// # tokio_test::block_on(test_util::stream_transform_test(|process| {
1807 /// let tick = process.tick();
1808 /// let numbers = process.source_iter(q!(vec![(1, false), (2, true), (1, false), (2, false)]));
1809 /// let batch = unsafe { numbers.tick_batch(&tick) };
1810 /// batch
1811 /// .fold_keyed_commutative_idempotent(q!(|| false), q!(|acc, x| *acc |= x))
1812 /// .all_ticks()
1813 /// # }, |mut stream| async move {
1814 /// // (1, false), (2, true)
1815 /// # assert_eq!(stream.next().await.unwrap(), (1, false));
1816 /// # assert_eq!(stream.next().await.unwrap(), (2, true));
1817 /// # }));
1818 /// ```
1819 pub fn fold_keyed_commutative_idempotent<A, I, F>(
1820 self,
1821 init: impl IntoQuotedMut<'a, I, Tick<L>>,
1822 comb: impl IntoQuotedMut<'a, F, Tick<L>>,
1823 ) -> Stream<(K, A), Tick<L>, Bounded, NoOrder, ExactlyOnce>
1824 where
1825 I: Fn() -> A + 'a,
1826 F: Fn(&mut A, V) + 'a,
1827 {
1828 unsafe {
1829 // SAFETY: aggregation is commutative and idempotent
1830 self.assume_ordering().assume_retries()
1831 }
1832 .fold_keyed(init, comb)
1833 }
1834
1835 /// Given a stream of pairs `(K, V)`, produces a new stream of unique keys `K`.
1836 /// # Example
1837 /// ```rust
1838 /// # use hydro_lang::*;
1839 /// # use futures::StreamExt;
1840 /// # tokio_test::block_on(test_util::stream_transform_test(|process| {
1841 /// let tick = process.tick();
1842 /// let numbers = process.source_iter(q!(vec![(1, 2), (2, 3), (1, 3), (2, 4)]));
1843 /// let batch = unsafe { numbers.tick_batch(&tick) };
1844 /// batch.keys().all_ticks()
1845 /// # }, |mut stream| async move {
1846 /// // 1, 2
1847 /// # assert_eq!(stream.next().await.unwrap(), 1);
1848 /// # assert_eq!(stream.next().await.unwrap(), 2);
1849 /// # }));
1850 /// ```
1851 pub fn keys(self) -> Stream<K, Tick<L>, Bounded, NoOrder, ExactlyOnce> {
1852 self.fold_keyed_commutative_idempotent(q!(|| ()), q!(|_, _| {}))
1853 .map(q!(|(k, _)| k))
1854 }
1855
1856 /// A special case of [`Stream::reduce_commutative_idempotent`], in the spirit of SQL's GROUP BY and aggregation constructs.
1857 /// The input tuples are partitioned into groups by the first element ("keys"), and for each group the values
1858 /// in the second element are accumulated via the `comb` closure.
1859 ///
1860 /// The `comb` closure must be **commutative**, as the order of input items is not guaranteed, and **idempotent**,
1861 /// as there may be non-deterministic duplicates.
1862 ///
1863 /// If you need the accumulated value to have a different type than the input, use [`Stream::fold_keyed_commutative_idempotent`].
1864 ///
1865 /// # Example
1866 /// ```rust
1867 /// # use hydro_lang::*;
1868 /// # use futures::StreamExt;
1869 /// # tokio_test::block_on(test_util::stream_transform_test(|process| {
1870 /// let tick = process.tick();
1871 /// let numbers = process.source_iter(q!(vec![(1, false), (2, true), (1, false), (2, false)]));
1872 /// let batch = unsafe { numbers.tick_batch(&tick) };
1873 /// batch
1874 /// .reduce_keyed_commutative_idempotent(q!(|acc, x| *acc |= x))
1875 /// .all_ticks()
1876 /// # }, |mut stream| async move {
1877 /// // (1, false), (2, true)
1878 /// # assert_eq!(stream.next().await.unwrap(), (1, false));
1879 /// # assert_eq!(stream.next().await.unwrap(), (2, true));
1880 /// # }));
1881 /// ```
1882 pub fn reduce_keyed_commutative_idempotent<F>(
1883 self,
1884 comb: impl IntoQuotedMut<'a, F, Tick<L>>,
1885 ) -> Stream<(K, V), Tick<L>, Bounded, NoOrder, ExactlyOnce>
1886 where
1887 F: Fn(&mut V, V) + 'a,
1888 {
1889 unsafe {
1890 // SAFETY: aggregation is commutative and idempotent
1891 self.assume_ordering().assume_retries()
1892 }
1893 .reduce_keyed(comb)
1894 }
1895}
1896
1897impl<'a, K, V, L, O> Stream<(K, V), Tick<L>, Bounded, O, ExactlyOnce>
1898where
1899 K: Eq + Hash,
1900 L: Location<'a>,
1901{
1902 /// A special case of [`Stream::fold_commutative`], in the spirit of SQL's GROUP BY and aggregation constructs. The input
1903 /// tuples are partitioned into groups by the first element ("keys"), and for each group the values
1904 /// in the second element are accumulated via the `comb` closure.
1905 ///
1906 /// The `comb` closure must be **commutative**, as the order of input items is not guaranteed.
1907 ///
1908 /// If the input and output value types are the same and do not require initialization then use
1909 /// [`Stream::reduce_keyed_commutative`].
1910 ///
1911 /// # Example
1912 /// ```rust
1913 /// # use hydro_lang::*;
1914 /// # use futures::StreamExt;
1915 /// # tokio_test::block_on(test_util::stream_transform_test(|process| {
1916 /// let tick = process.tick();
1917 /// let numbers = process.source_iter(q!(vec![(1, 2), (2, 3), (1, 3), (2, 4)]));
1918 /// let batch = unsafe { numbers.tick_batch(&tick) };
1919 /// batch
1920 /// .fold_keyed_commutative(q!(|| 0), q!(|acc, x| *acc += x))
1921 /// .all_ticks()
1922 /// # }, |mut stream| async move {
1923 /// // (1, 5), (2, 7)
1924 /// # assert_eq!(stream.next().await.unwrap(), (1, 5));
1925 /// # assert_eq!(stream.next().await.unwrap(), (2, 7));
1926 /// # }));
1927 /// ```
1928 pub fn fold_keyed_commutative<A, I, F>(
1929 self,
1930 init: impl IntoQuotedMut<'a, I, Tick<L>>,
1931 comb: impl IntoQuotedMut<'a, F, Tick<L>>,
1932 ) -> Stream<(K, A), Tick<L>, Bounded, NoOrder, ExactlyOnce>
1933 where
1934 I: Fn() -> A + 'a,
1935 F: Fn(&mut A, V) + 'a,
1936 {
1937 unsafe {
1938 // SAFETY: aggregation is commutative
1939 self.assume_ordering()
1940 }
1941 .fold_keyed(init, comb)
1942 }
1943
1944 /// A special case of [`Stream::reduce_commutative`], in the spirit of SQL's GROUP BY and aggregation constructs. The input
1945 /// tuples are partitioned into groups by the first element ("keys"), and for each group the values
1946 /// in the second element are accumulated via the `comb` closure.
1947 ///
1948 /// The `comb` closure must be **commutative**, as the order of input items is not guaranteed.
1949 ///
1950 /// If you need the accumulated value to have a different type than the input, use [`Stream::fold_keyed_commutative`].
1951 ///
1952 /// # Example
1953 /// ```rust
1954 /// # use hydro_lang::*;
1955 /// # use futures::StreamExt;
1956 /// # tokio_test::block_on(test_util::stream_transform_test(|process| {
1957 /// let tick = process.tick();
1958 /// let numbers = process.source_iter(q!(vec![(1, 2), (2, 3), (1, 3), (2, 4)]));
1959 /// let batch = unsafe { numbers.tick_batch(&tick) };
1960 /// batch
1961 /// .reduce_keyed_commutative(q!(|acc, x| *acc += x))
1962 /// .all_ticks()
1963 /// # }, |mut stream| async move {
1964 /// // (1, 5), (2, 7)
1965 /// # assert_eq!(stream.next().await.unwrap(), (1, 5));
1966 /// # assert_eq!(stream.next().await.unwrap(), (2, 7));
1967 /// # }));
1968 /// ```
1969 pub fn reduce_keyed_commutative<F>(
1970 self,
1971 comb: impl IntoQuotedMut<'a, F, Tick<L>>,
1972 ) -> Stream<(K, V), Tick<L>, Bounded, NoOrder, ExactlyOnce>
1973 where
1974 F: Fn(&mut V, V) + 'a,
1975 {
1976 unsafe {
1977 // SAFETY: aggregation is commutative
1978 self.assume_ordering()
1979 }
1980 .reduce_keyed(comb)
1981 }
1982}
1983
1984impl<'a, K, V, L, R> Stream<(K, V), Tick<L>, Bounded, TotalOrder, R>
1985where
1986 K: Eq + Hash,
1987 L: Location<'a>,
1988{
1989 /// A special case of [`Stream::fold_idempotent`], in the spirit of SQL's GROUP BY and aggregation constructs.
1990 /// The input tuples are partitioned into groups by the first element ("keys"), and for each group the values
1991 /// in the second element are accumulated via the `comb` closure.
1992 ///
1993 /// The `comb` closure must be **idempotent** as there may be non-deterministic duplicates.
1994 ///
1995 /// If the input and output value types are the same and do not require initialization then use
1996 /// [`Stream::reduce_keyed_idempotent`].
1997 ///
1998 /// # Example
1999 /// ```rust
2000 /// # use hydro_lang::*;
2001 /// # use futures::StreamExt;
2002 /// # tokio_test::block_on(test_util::stream_transform_test(|process| {
2003 /// let tick = process.tick();
2004 /// let numbers = process.source_iter(q!(vec![(1, false), (2, true), (1, false), (2, false)]));
2005 /// let batch = unsafe { numbers.tick_batch(&tick) };
2006 /// batch
2007 /// .fold_keyed_idempotent(q!(|| false), q!(|acc, x| *acc |= x))
2008 /// .all_ticks()
2009 /// # }, |mut stream| async move {
2010 /// // (1, false), (2, true)
2011 /// # assert_eq!(stream.next().await.unwrap(), (1, false));
2012 /// # assert_eq!(stream.next().await.unwrap(), (2, true));
2013 /// # }));
2014 /// ```
2015 pub fn fold_keyed_idempotent<A, I, F>(
2016 self,
2017 init: impl IntoQuotedMut<'a, I, Tick<L>>,
2018 comb: impl IntoQuotedMut<'a, F, Tick<L>>,
2019 ) -> Stream<(K, A), Tick<L>, Bounded, NoOrder, ExactlyOnce>
2020 where
2021 I: Fn() -> A + 'a,
2022 F: Fn(&mut A, V) + 'a,
2023 {
2024 unsafe {
2025 // SAFETY: aggregation is idempotent
2026 self.assume_retries()
2027 }
2028 .fold_keyed(init, comb)
2029 }
2030
2031 /// A special case of [`Stream::reduce_idempotent`], in the spirit of SQL's GROUP BY and aggregation constructs.
2032 /// The input tuples are partitioned into groups by the first element ("keys"), and for each group the values
2033 /// in the second element are accumulated via the `comb` closure.
2034 ///
2035 /// The `comb` closure must be **idempotent**, as there may be non-deterministic duplicates.
2036 ///
2037 /// If you need the accumulated value to have a different type than the input, use [`Stream::fold_keyed_idempotent`].
2038 ///
2039 /// # Example
2040 /// ```rust
2041 /// # use hydro_lang::*;
2042 /// # use futures::StreamExt;
2043 /// # tokio_test::block_on(test_util::stream_transform_test(|process| {
2044 /// let tick = process.tick();
2045 /// let numbers = process.source_iter(q!(vec![(1, false), (2, true), (1, false), (2, false)]));
2046 /// let batch = unsafe { numbers.tick_batch(&tick) };
2047 /// batch
2048 /// .reduce_keyed_idempotent(q!(|acc, x| *acc |= x))
2049 /// .all_ticks()
2050 /// # }, |mut stream| async move {
2051 /// // (1, false), (2, true)
2052 /// # assert_eq!(stream.next().await.unwrap(), (1, false));
2053 /// # assert_eq!(stream.next().await.unwrap(), (2, true));
2054 /// # }));
2055 /// ```
2056 pub fn reduce_keyed_idempotent<F>(
2057 self,
2058 comb: impl IntoQuotedMut<'a, F, Tick<L>>,
2059 ) -> Stream<(K, V), Tick<L>, Bounded, NoOrder, ExactlyOnce>
2060 where
2061 F: Fn(&mut V, V) + 'a,
2062 {
2063 unsafe {
2064 // SAFETY: aggregation is idempotent
2065 self.assume_retries()
2066 }
2067 .reduce_keyed(comb)
2068 }
2069}
2070
2071impl<'a, T, L, B, O, R> Stream<T, Atomic<L>, B, O, R>
2072where
2073 L: Location<'a> + NoTick,
2074{
2075 /// Returns a stream corresponding to the latest batch of elements being atomically
2076 /// processed. These batches are guaranteed to be contiguous across ticks and preserve
2077 /// the order of the input.
2078 ///
2079 /// # Safety
2080 /// The batch boundaries are non-deterministic and may change across executions.
2081 pub unsafe fn tick_batch(self) -> Stream<T, Tick<L>, Bounded, O, R> {
2082 Stream::new(
2083 self.location.clone().tick,
2084 HydroNode::Unpersist {
2085 inner: Box::new(self.ir_node.into_inner()),
2086 metadata: self.location.new_node_metadata::<T>(),
2087 },
2088 )
2089 }
2090
2091 pub fn end_atomic(self) -> Stream<T, L, B, O, R> {
2092 Stream::new(self.location.tick.l, self.ir_node.into_inner())
2093 }
2094
2095 pub fn atomic_source(&self) -> Tick<L> {
2096 self.location.tick.clone()
2097 }
2098}
2099
2100impl<'a, T, L, B, O, R> Stream<T, L, B, O, R>
2101where
2102 L: Location<'a> + NoTick + NoAtomic,
2103{
2104 pub fn atomic(self, tick: &Tick<L>) -> Stream<T, Atomic<L>, B, O, R> {
2105 Stream::new(Atomic { tick: tick.clone() }, self.ir_node.into_inner())
2106 }
2107
2108 /// Consumes a stream of `Future<T>`, produces a new stream of the resulting `T` outputs.
2109 /// Future outputs are produced as available, regardless of input arrival order.
2110 ///
2111 /// # Example
2112 /// ```rust
2113 /// # use std::collections::HashSet;
2114 /// # use futures::StreamExt;
2115 /// # use hydro_lang::*;
2116 /// # tokio_test::block_on(test_util::stream_transform_test(|process| {
2117 /// process.source_iter(q!([2, 3, 1, 9, 6, 5, 4, 7, 8]))
2118 /// .map(q!(|x| async move {
2119 /// tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
2120 /// x
2121 /// }))
2122 /// .resolve_futures()
2123 /// # },
2124 /// # |mut stream| async move {
2125 /// // 1, 2, 3, 4, 5, 6, 7, 8, 9 (in any order)
2126 /// # let mut output = HashSet::new();
2127 /// # for _ in 1..10 {
2128 /// # output.insert(stream.next().await.unwrap());
2129 /// # }
2130 /// # assert_eq!(
2131 /// # output,
2132 /// # HashSet::<i32>::from_iter(1..10)
2133 /// # );
2134 /// # },
2135 /// # ));
2136 pub fn resolve_futures<T2>(self) -> Stream<T2, L, B, NoOrder, R>
2137 where
2138 T: Future<Output = T2>,
2139 {
2140 Stream::new(
2141 self.location.clone(),
2142 HydroNode::ResolveFutures {
2143 input: Box::new(self.ir_node.into_inner()),
2144 metadata: self.location.new_node_metadata::<T2>(),
2145 },
2146 )
2147 }
2148
2149 /// Given a tick, returns a stream corresponding to a batch of elements segmented by
2150 /// that tick. These batches are guaranteed to be contiguous across ticks and preserve
2151 /// the order of the input.
2152 ///
2153 /// # Safety
2154 /// The batch boundaries are non-deterministic and may change across executions.
2155 pub unsafe fn tick_batch(self, tick: &Tick<L>) -> Stream<T, Tick<L>, Bounded, O, R> {
2156 unsafe { self.atomic(tick).tick_batch() }
2157 }
2158
2159 /// Given a time interval, returns a stream corresponding to samples taken from the
2160 /// stream roughly at that interval. The output will have elements in the same order
2161 /// as the input, but with arbitrary elements skipped between samples. There is also
2162 /// no guarantee on the exact timing of the samples.
2163 ///
2164 /// # Safety
2165 /// The output stream is non-deterministic in which elements are sampled, since this
2166 /// is controlled by a clock.
2167 pub unsafe fn sample_every(
2168 self,
2169 interval: impl QuotedWithContext<'a, std::time::Duration, L> + Copy + 'a,
2170 ) -> Stream<T, L, Unbounded, O, AtLeastOnce> {
2171 let samples = unsafe {
2172 // SAFETY: source of intentional non-determinism
2173 self.location.source_interval(interval)
2174 };
2175
2176 let tick = self.location.tick();
2177 unsafe {
2178 // SAFETY: source of intentional non-determinism
2179 self.tick_batch(&tick)
2180 .continue_if(samples.tick_batch(&tick).first())
2181 .all_ticks()
2182 .weakest_retries()
2183 }
2184 }
2185
2186 /// Given a timeout duration, returns an [`Optional`] which will have a value if the
2187 /// stream has not emitted a value since that duration.
2188 ///
2189 /// # Safety
2190 /// Timeout relies on non-deterministic sampling of the stream, so depending on when
2191 /// samples take place, timeouts may be non-deterministically generated or missed,
2192 /// and the notification of the timeout may be delayed as well. There is also no
2193 /// guarantee on how long the [`Optional`] will have a value after the timeout is
2194 /// detected based on when the next sample is taken.
2195 pub unsafe fn timeout(
2196 self,
2197 duration: impl QuotedWithContext<'a, std::time::Duration, Tick<L>> + Copy + 'a,
2198 ) -> Optional<(), L, Unbounded> {
2199 let tick = self.location.tick();
2200
2201 let latest_received = unsafe { self.assume_retries() }.fold_commutative(
2202 q!(|| None),
2203 q!(|latest, _| {
2204 *latest = Some(Instant::now());
2205 }),
2206 );
2207
2208 unsafe {
2209 // SAFETY: Non-deterministic delay in detecting a timeout is expected.
2210 latest_received.latest_tick(&tick)
2211 }
2212 .filter_map(q!(move |latest_received| {
2213 if let Some(latest_received) = latest_received {
2214 if Instant::now().duration_since(latest_received) > duration {
2215 Some(())
2216 } else {
2217 None
2218 }
2219 } else {
2220 Some(())
2221 }
2222 }))
2223 .latest()
2224 }
2225}
2226
2227impl<'a, F, T, L, B, O, R> Stream<F, L, B, O, R>
2228where
2229 L: Location<'a> + NoTick + NoAtomic,
2230 F: Future<Output = T>,
2231{
2232 /// Consumes a stream of `Future<T>`, produces a new stream of the resulting `T` outputs.
2233 /// Future outputs are produced in the same order as the input stream.
2234 ///
2235 /// # Example
2236 /// ```rust
2237 /// # use std::collections::HashSet;
2238 /// # use futures::StreamExt;
2239 /// # use hydro_lang::*;
2240 /// # tokio_test::block_on(test_util::stream_transform_test(|process| {
2241 /// process.source_iter(q!([2, 3, 1, 9, 6, 5, 4, 7, 8]))
2242 /// .map(q!(|x| async move {
2243 /// tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
2244 /// x
2245 /// }))
2246 /// .resolve_futures_ordered()
2247 /// # },
2248 /// # |mut stream| async move {
2249 /// // 2, 3, 1, 9, 6, 5, 4, 7, 8
2250 /// # let mut output = Vec::new();
2251 /// # for _ in 1..10 {
2252 /// # output.push(stream.next().await.unwrap());
2253 /// # }
2254 /// # assert_eq!(
2255 /// # output,
2256 /// # vec![2, 3, 1, 9, 6, 5, 4, 7, 8]
2257 /// # );
2258 /// # },
2259 /// # ));
2260 pub fn resolve_futures_ordered(self) -> Stream<T, L, B, O, R> {
2261 Stream::new(
2262 self.location.clone(),
2263 HydroNode::ResolveFuturesOrdered {
2264 input: Box::new(self.ir_node.into_inner()),
2265 metadata: self.location.new_node_metadata::<T>(),
2266 },
2267 )
2268 }
2269}
2270
2271impl<'a, T, L, B, O, R> Stream<T, L, B, O, R>
2272where
2273 L: Location<'a> + NoTick,
2274{
2275 pub fn for_each<F: Fn(T) + 'a>(self, f: impl IntoQuotedMut<'a, F, L>) {
2276 let f = f.splice_fn1_ctx(&self.location).into();
2277 let metadata = self.location.new_node_metadata::<T>();
2278 self.location
2279 .flow_state()
2280 .borrow_mut()
2281 .leaves
2282 .as_mut()
2283 .expect(FLOW_USED_MESSAGE)
2284 .push(HydroLeaf::ForEach {
2285 input: Box::new(HydroNode::Unpersist {
2286 inner: Box::new(self.ir_node.into_inner()),
2287 metadata: metadata.clone(),
2288 }),
2289 f,
2290 metadata,
2291 });
2292 }
2293
2294 pub fn dest_sink<S>(self, sink: impl QuotedWithContext<'a, S, L>)
2295 where
2296 S: 'a + futures::Sink<T> + Unpin,
2297 {
2298 self.location
2299 .flow_state()
2300 .borrow_mut()
2301 .leaves
2302 .as_mut()
2303 .expect(FLOW_USED_MESSAGE)
2304 .push(HydroLeaf::DestSink {
2305 sink: sink.splice_typed_ctx(&self.location).into(),
2306 input: Box::new(self.ir_node.into_inner()),
2307 metadata: self.location.new_node_metadata::<T>(),
2308 });
2309 }
2310}
2311
2312impl<'a, T, L, O, R> Stream<T, Tick<L>, Bounded, O, R>
2313where
2314 L: Location<'a>,
2315{
2316 pub fn all_ticks(self) -> Stream<T, L, Unbounded, O, R> {
2317 Stream::new(
2318 self.location.outer().clone(),
2319 HydroNode::Persist {
2320 inner: Box::new(self.ir_node.into_inner()),
2321 metadata: self.location.new_node_metadata::<T>(),
2322 },
2323 )
2324 }
2325
2326 pub fn all_ticks_atomic(self) -> Stream<T, Atomic<L>, Unbounded, O, R> {
2327 Stream::new(
2328 Atomic {
2329 tick: self.location.clone(),
2330 },
2331 HydroNode::Persist {
2332 inner: Box::new(self.ir_node.into_inner()),
2333 metadata: self.location.new_node_metadata::<T>(),
2334 },
2335 )
2336 }
2337
2338 pub fn persist(self) -> Stream<T, Tick<L>, Bounded, O, R>
2339 where
2340 T: Clone,
2341 {
2342 Stream::new(
2343 self.location.clone(),
2344 HydroNode::Persist {
2345 inner: Box::new(self.ir_node.into_inner()),
2346 metadata: self.location.new_node_metadata::<T>(),
2347 },
2348 )
2349 }
2350
2351 pub fn defer_tick(self) -> Stream<T, Tick<L>, Bounded, O, R> {
2352 Stream::new(
2353 self.location.clone(),
2354 HydroNode::DeferTick {
2355 input: Box::new(self.ir_node.into_inner()),
2356 metadata: self.location.new_node_metadata::<T>(),
2357 },
2358 )
2359 }
2360
2361 pub fn delta(self) -> Stream<T, Tick<L>, Bounded, O, R> {
2362 Stream::new(
2363 self.location.clone(),
2364 HydroNode::Delta {
2365 inner: Box::new(self.ir_node.into_inner()),
2366 metadata: self.location.new_node_metadata::<T>(),
2367 },
2368 )
2369 }
2370}
2371
2372pub fn serialize_bincode_with_type(is_demux: bool, t_type: &syn::Type) -> syn::Expr {
2373 let root = get_this_crate();
2374
2375 if is_demux {
2376 parse_quote! {
2377 ::#root::runtime_support::stageleft::runtime_support::fn1_type_hint::<(#root::ClusterId<_>, #t_type), _>(
2378 |(id, data)| {
2379 (id.raw_id, #root::runtime_support::bincode::serialize(&data).unwrap().into())
2380 }
2381 )
2382 }
2383 } else {
2384 parse_quote! {
2385 ::#root::runtime_support::stageleft::runtime_support::fn1_type_hint::<#t_type, _>(
2386 |data| {
2387 #root::runtime_support::bincode::serialize(&data).unwrap().into()
2388 }
2389 )
2390 }
2391 }
2392}
2393
2394fn serialize_bincode<T: Serialize>(is_demux: bool) -> syn::Expr {
2395 serialize_bincode_with_type(is_demux, &stageleft::quote_type::<T>())
2396}
2397
2398pub fn deserialize_bincode_with_type(tagged: Option<&syn::Type>, t_type: &syn::Type) -> syn::Expr {
2399 let root = get_this_crate();
2400
2401 if let Some(c_type) = tagged {
2402 parse_quote! {
2403 |res| {
2404 let (id, b) = res.unwrap();
2405 (#root::ClusterId::<#c_type>::from_raw(id), #root::runtime_support::bincode::deserialize::<#t_type>(&b).unwrap())
2406 }
2407 }
2408 } else {
2409 parse_quote! {
2410 |res| {
2411 #root::runtime_support::bincode::deserialize::<#t_type>(&res.unwrap()).unwrap()
2412 }
2413 }
2414 }
2415}
2416
2417pub(super) fn deserialize_bincode<T: DeserializeOwned>(tagged: Option<&syn::Type>) -> syn::Expr {
2418 deserialize_bincode_with_type(tagged, &stageleft::quote_type::<T>())
2419}
2420
2421impl<'a, T, L, B, O, R> Stream<T, L, B, O, R>
2422where
2423 L: Location<'a> + NoTick,
2424{
2425 #[expect(
2426 clippy::type_complexity,
2427 reason = "Complex signatures for CanSend trait"
2428 )]
2429 pub fn send_bincode<L2, CoreType>(
2430 self,
2431 other: &L2,
2432 ) -> Stream<<L::Root as CanSend<'a, L2>>::Out<CoreType>, L2, Unbounded, O::Min, R>
2433 where
2434 L::Root: CanSend<'a, L2, In<CoreType> = T>,
2435 L2: Location<'a>,
2436 CoreType: Serialize + DeserializeOwned,
2437 O: MinOrder<<L::Root as CanSend<'a, L2>>::OutStrongestOrder<O>>,
2438 {
2439 let serialize_pipeline = Some(serialize_bincode::<CoreType>(L::Root::is_demux()));
2440
2441 let deserialize_pipeline = Some(deserialize_bincode::<CoreType>(
2442 L::Root::tagged_type().as_ref(),
2443 ));
2444
2445 Stream::new(
2446 other.clone(),
2447 HydroNode::Network {
2448 serialize_fn: serialize_pipeline.map(|e| e.into()),
2449 instantiate_fn: DebugInstantiate::Building,
2450 deserialize_fn: deserialize_pipeline.map(|e| e.into()),
2451 input: Box::new(self.ir_node.into_inner()),
2452 metadata: other.new_node_metadata::<CoreType>(),
2453 },
2454 )
2455 }
2456
2457 pub fn send_bincode_external<L2, CoreType>(
2458 self,
2459 other: &External<L2>,
2460 ) -> ExternalBincodeStream<L::Out<CoreType>>
2461 where
2462 L: CanSend<'a, External<'a, L2>, In<CoreType> = T, Out<CoreType> = CoreType>,
2463 L2: 'a,
2464 CoreType: Serialize + DeserializeOwned,
2465 // for now, we restirct Out<CoreType> to be CoreType, which means no tagged cluster -> external
2466 {
2467 let serialize_pipeline = Some(serialize_bincode::<CoreType>(L::is_demux()));
2468
2469 let mut flow_state_borrow = self.location.flow_state().borrow_mut();
2470
2471 let external_key = flow_state_borrow.next_external_out;
2472 flow_state_borrow.next_external_out += 1;
2473
2474 let leaves = flow_state_borrow.leaves.as_mut().expect("Attempted to add a leaf to a flow that has already been finalized. No leaves can be added after the flow has been compiled()");
2475
2476 leaves.push(HydroLeaf::SendExternal {
2477 to_external_id: other.id,
2478 to_key: external_key,
2479 serialize_fn: serialize_pipeline.map(|e| e.into()),
2480 instantiate_fn: DebugInstantiate::Building,
2481 input: Box::new(HydroNode::Unpersist {
2482 inner: Box::new(self.ir_node.into_inner()),
2483 metadata: self.location.new_node_metadata::<T>(),
2484 }),
2485 });
2486
2487 ExternalBincodeStream {
2488 process_id: other.id,
2489 port_id: external_key,
2490 _phantom: PhantomData,
2491 }
2492 }
2493
2494 #[expect(
2495 clippy::type_complexity,
2496 reason = "Complex signatures for CanSend trait"
2497 )]
2498 pub fn send_bytes<L2>(
2499 self,
2500 other: &L2,
2501 ) -> Stream<<L::Root as CanSend<'a, L2>>::Out<Bytes>, L2, Unbounded, O::Min, R>
2502 where
2503 L2: Location<'a>,
2504 L::Root: CanSend<'a, L2, In<Bytes> = T>,
2505 O: MinOrder<<L::Root as CanSend<'a, L2>>::OutStrongestOrder<O>>,
2506 {
2507 let root = get_this_crate();
2508 Stream::new(
2509 other.clone(),
2510 HydroNode::Network {
2511 serialize_fn: None,
2512 instantiate_fn: DebugInstantiate::Building,
2513 deserialize_fn: if let Some(c_type) = L::Root::tagged_type() {
2514 let expr: syn::Expr = parse_quote!(|(id, b)| (#root::ClusterId<#c_type>::from_raw(id), b.unwrap().freeze()));
2515 Some(expr.into())
2516 } else {
2517 let expr: syn::Expr = parse_quote!(|b| b.unwrap().freeze());
2518 Some(expr.into())
2519 },
2520 input: Box::new(self.ir_node.into_inner()),
2521 metadata: other.new_node_metadata::<Bytes>(),
2522 },
2523 )
2524 }
2525
2526 pub fn send_bytes_external<L2>(self, other: &External<L2>) -> ExternalBytesPort
2527 where
2528 L2: 'a,
2529 L::Root: CanSend<'a, External<'a, L2>, In<Bytes> = T, Out<Bytes> = Bytes>,
2530 {
2531 let mut flow_state_borrow = self.location.flow_state().borrow_mut();
2532 let external_key = flow_state_borrow.next_external_out;
2533 flow_state_borrow.next_external_out += 1;
2534
2535 let leaves = flow_state_borrow.leaves.as_mut().expect("Attempted to add a leaf to a flow that has already been finalized. No leaves can be added after the flow has been compiled()");
2536
2537 leaves.push(HydroLeaf::SendExternal {
2538 to_external_id: other.id,
2539 to_key: external_key,
2540 serialize_fn: None,
2541 instantiate_fn: DebugInstantiate::Building,
2542 input: Box::new(HydroNode::Unpersist {
2543 inner: Box::new(self.ir_node.into_inner()),
2544 metadata: self.location.new_node_metadata::<T>(),
2545 }),
2546 });
2547
2548 ExternalBytesPort {
2549 process_id: other.id,
2550 port_id: external_key,
2551 }
2552 }
2553
2554 pub fn send_bincode_anonymous<L2, Tag, CoreType>(
2555 self,
2556 other: &L2,
2557 ) -> Stream<CoreType, L2, Unbounded, O::Min, R>
2558 where
2559 L2: Location<'a>,
2560 L::Root: CanSend<'a, L2, In<CoreType> = T, Out<CoreType> = (Tag, CoreType)>,
2561 CoreType: Serialize + DeserializeOwned,
2562 O: MinOrder<<L::Root as CanSend<'a, L2>>::OutStrongestOrder<O>>,
2563 {
2564 self.send_bincode::<L2, CoreType>(other).map(q!(|(_, b)| b))
2565 }
2566
2567 pub fn send_bytes_anonymous<L2, Tag>(
2568 self,
2569 other: &L2,
2570 ) -> Stream<Bytes, L2, Unbounded, O::Min, R>
2571 where
2572 L2: Location<'a>,
2573 L::Root: CanSend<'a, L2, In<Bytes> = T, Out<Bytes> = (Tag, Bytes)>,
2574 O: MinOrder<<L::Root as CanSend<'a, L2>>::OutStrongestOrder<O>>,
2575 {
2576 self.send_bytes::<L2>(other).map(q!(|(_, b)| b))
2577 }
2578
2579 #[expect(clippy::type_complexity, reason = "ordering semantics for broadcast")]
2580 pub fn broadcast_bincode<C2>(
2581 self,
2582 other: &Cluster<'a, C2>,
2583 ) -> Stream<
2584 <L::Root as CanSend<'a, Cluster<'a, C2>>>::Out<T>,
2585 Cluster<'a, C2>,
2586 Unbounded,
2587 O::Min,
2588 R,
2589 >
2590 where
2591 C2: 'a,
2592 L::Root: CanSend<'a, Cluster<'a, C2>, In<T> = (ClusterId<C2>, T)>,
2593 T: Clone + Serialize + DeserializeOwned,
2594 O: MinOrder<<L::Root as CanSend<'a, Cluster<'a, C2>>>::OutStrongestOrder<O>>,
2595 {
2596 let ids = other.members();
2597 self.flat_map_ordered(q!(|v| { ids.iter().map(move |id| (*id, v.clone())) }))
2598 .send_bincode(other)
2599 }
2600
2601 pub fn broadcast_bincode_anonymous<C2, Tag>(
2602 self,
2603 other: &Cluster<'a, C2>,
2604 ) -> Stream<T, Cluster<'a, C2>, Unbounded, O::Min, R>
2605 where
2606 C2: 'a,
2607 L::Root: CanSend<'a, Cluster<'a, C2>, In<T> = (ClusterId<C2>, T), Out<T> = (Tag, T)>,
2608 T: Clone + Serialize + DeserializeOwned,
2609 O: MinOrder<<L::Root as CanSend<'a, Cluster<'a, C2>>>::OutStrongestOrder<O>>,
2610 {
2611 self.broadcast_bincode(other).map(q!(|(_, b)| b))
2612 }
2613
2614 #[expect(clippy::type_complexity, reason = "ordering semantics for broadcast")]
2615 pub fn broadcast_bytes<C2>(
2616 self,
2617 other: &Cluster<'a, C2>,
2618 ) -> Stream<
2619 <L::Root as CanSend<'a, Cluster<'a, C2>>>::Out<Bytes>,
2620 Cluster<'a, C2>,
2621 Unbounded,
2622 O::Min,
2623 R,
2624 >
2625 where
2626 C2: 'a,
2627 L::Root: CanSend<'a, Cluster<'a, C2>, In<Bytes> = (ClusterId<C2>, T)>,
2628 T: Clone,
2629 O: MinOrder<<L::Root as CanSend<'a, Cluster<'a, C2>>>::OutStrongestOrder<O>>,
2630 {
2631 let ids = other.members();
2632
2633 self.flat_map_ordered(q!(|b| ids.iter().map(move |id| (
2634 ::std::clone::Clone::clone(id),
2635 ::std::clone::Clone::clone(&b)
2636 ))))
2637 .send_bytes(other)
2638 }
2639
2640 pub fn broadcast_bytes_anonymous<C2, Tag>(
2641 self,
2642 other: &Cluster<'a, C2>,
2643 ) -> Stream<Bytes, Cluster<'a, C2>, Unbounded, O::Min, R>
2644 where
2645 C2: 'a,
2646 L::Root: CanSend<'a, Cluster<'a, C2>, In<Bytes> = (ClusterId<C2>, T), Out<Bytes> = (Tag, Bytes)>
2647 + 'a,
2648 T: Clone,
2649 O: MinOrder<<L::Root as CanSend<'a, Cluster<'a, C2>>>::OutStrongestOrder<O>>,
2650 {
2651 self.broadcast_bytes(other).map(q!(|(_, b)| b))
2652 }
2653}
2654
2655#[expect(clippy::type_complexity, reason = "ordering semantics for round-robin")]
2656impl<'a, T, L, B> Stream<T, L, B, TotalOrder, ExactlyOnce>
2657where
2658 L: Location<'a> + NoTick,
2659{
2660 pub fn round_robin_bincode<C2>(
2661 self,
2662 other: &Cluster<'a, C2>,
2663 ) -> Stream<
2664 <L::Root as CanSend<'a, Cluster<'a, C2>>>::Out<T>,
2665 Cluster<'a, C2>,
2666 Unbounded,
2667 <TotalOrder as MinOrder<
2668 <L::Root as CanSend<'a, Cluster<'a, C2>>>::OutStrongestOrder<TotalOrder>,
2669 >>::Min,
2670 ExactlyOnce,
2671 >
2672 where
2673 C2: 'a,
2674 L::Root: CanSend<'a, Cluster<'a, C2>, In<T> = (ClusterId<C2>, T)>,
2675 T: Clone + Serialize + DeserializeOwned,
2676 TotalOrder:
2677 MinOrder<<L::Root as CanSend<'a, Cluster<'a, C2>>>::OutStrongestOrder<TotalOrder>>,
2678 {
2679 let ids = other.members();
2680
2681 self.enumerate()
2682 .map(q!(|(i, w)| (ids[i % ids.len()], w)))
2683 .send_bincode(other)
2684 }
2685
2686 pub fn round_robin_bincode_anonymous<C2, Tag>(
2687 self,
2688 other: &Cluster<'a, C2>,
2689 ) -> Stream<
2690 T,
2691 Cluster<'a, C2>,
2692 Unbounded,
2693 <TotalOrder as MinOrder<
2694 <L::Root as CanSend<'a, Cluster<'a, C2>>>::OutStrongestOrder<TotalOrder>,
2695 >>::Min,
2696 ExactlyOnce,
2697 >
2698 where
2699 C2: 'a,
2700 L::Root: CanSend<'a, Cluster<'a, C2>, In<T> = (ClusterId<C2>, T), Out<T> = (Tag, T)> + 'a,
2701 T: Clone + Serialize + DeserializeOwned,
2702 TotalOrder:
2703 MinOrder<<L::Root as CanSend<'a, Cluster<'a, C2>>>::OutStrongestOrder<TotalOrder>>,
2704 {
2705 self.round_robin_bincode(other).map(q!(|(_, b)| b))
2706 }
2707
2708 pub fn round_robin_bytes<C2>(
2709 self,
2710 other: &Cluster<'a, C2>,
2711 ) -> Stream<
2712 <L::Root as CanSend<'a, Cluster<'a, C2>>>::Out<Bytes>,
2713 Cluster<'a, C2>,
2714 Unbounded,
2715 <TotalOrder as MinOrder<
2716 <L::Root as CanSend<'a, Cluster<'a, C2>>>::OutStrongestOrder<TotalOrder>,
2717 >>::Min,
2718 ExactlyOnce,
2719 >
2720 where
2721 C2: 'a,
2722 L::Root: CanSend<'a, Cluster<'a, C2>, In<Bytes> = (ClusterId<C2>, T)> + 'a,
2723 T: Clone,
2724 TotalOrder:
2725 MinOrder<<L::Root as CanSend<'a, Cluster<'a, C2>>>::OutStrongestOrder<TotalOrder>>,
2726 {
2727 let ids = other.members();
2728
2729 self.enumerate()
2730 .map(q!(|(i, w)| (ids[i % ids.len()], w)))
2731 .send_bytes(other)
2732 }
2733
2734 pub fn round_robin_bytes_anonymous<C2, Tag>(
2735 self,
2736 other: &Cluster<'a, C2>,
2737 ) -> Stream<
2738 Bytes,
2739 Cluster<'a, C2>,
2740 Unbounded,
2741 <TotalOrder as MinOrder<
2742 <L::Root as CanSend<'a, Cluster<'a, C2>>>::OutStrongestOrder<TotalOrder>,
2743 >>::Min,
2744 ExactlyOnce,
2745 >
2746 where
2747 C2: 'a,
2748 L::Root: CanSend<'a, Cluster<'a, C2>, In<Bytes> = (ClusterId<C2>, T), Out<Bytes> = (Tag, Bytes)>
2749 + 'a,
2750 T: Clone,
2751 TotalOrder:
2752 MinOrder<<L::Root as CanSend<'a, Cluster<'a, C2>>>::OutStrongestOrder<TotalOrder>>,
2753 {
2754 self.round_robin_bytes(other).map(q!(|(_, b)| b))
2755 }
2756}
2757
2758#[cfg(test)]
2759mod tests {
2760 use futures::StreamExt;
2761 use hydro_deploy::Deployment;
2762 use serde::{Deserialize, Serialize};
2763 use stageleft::q;
2764
2765 use crate::FlowBuilder;
2766 use crate::location::Location;
2767
2768 struct P1 {}
2769 struct P2 {}
2770
2771 #[derive(Serialize, Deserialize, Debug)]
2772 struct SendOverNetwork {
2773 n: u32,
2774 }
2775
2776 #[tokio::test]
2777 async fn first_ten_distributed() {
2778 let mut deployment = Deployment::new();
2779
2780 let flow = FlowBuilder::new();
2781 let first_node = flow.process::<P1>();
2782 let second_node = flow.process::<P2>();
2783 let external = flow.external::<P2>();
2784
2785 let numbers = first_node.source_iter(q!(0..10));
2786 let out_port = numbers
2787 .map(q!(|n| SendOverNetwork { n }))
2788 .send_bincode(&second_node)
2789 .send_bincode_external(&external);
2790
2791 let nodes = flow
2792 .with_process(&first_node, deployment.Localhost())
2793 .with_process(&second_node, deployment.Localhost())
2794 .with_external(&external, deployment.Localhost())
2795 .deploy(&mut deployment);
2796
2797 deployment.deploy().await.unwrap();
2798
2799 let mut external_out = nodes.connect_source_bincode(out_port).await;
2800
2801 deployment.start().await.unwrap();
2802
2803 for i in 0..10 {
2804 assert_eq!(external_out.next().await.unwrap().n, i);
2805 }
2806 }
2807
2808 #[tokio::test]
2809 async fn first_cardinality() {
2810 let mut deployment = Deployment::new();
2811
2812 let flow = FlowBuilder::new();
2813 let node = flow.process::<()>();
2814 let external = flow.external::<()>();
2815
2816 let node_tick = node.tick();
2817 let count = node_tick
2818 .singleton(q!([1, 2, 3]))
2819 .into_stream()
2820 .flatten_ordered()
2821 .first()
2822 .into_stream()
2823 .count()
2824 .all_ticks()
2825 .send_bincode_external(&external);
2826
2827 let nodes = flow
2828 .with_process(&node, deployment.Localhost())
2829 .with_external(&external, deployment.Localhost())
2830 .deploy(&mut deployment);
2831
2832 deployment.deploy().await.unwrap();
2833
2834 let mut external_out = nodes.connect_source_bincode(count).await;
2835
2836 deployment.start().await.unwrap();
2837
2838 assert_eq!(external_out.next().await.unwrap(), 1);
2839 }
2840}