Skip to main content

hydro_lang/live_collections/sliced/
mod.rs

1//! Utilities for transforming live collections via slicing.
2
3pub mod style;
4
5use super::boundedness::{Bounded, Unbounded};
6use super::stream::{Ordering, Retries};
7use crate::location::{Location, Tick};
8
9#[doc(hidden)]
10#[macro_export]
11macro_rules! __sliced_parse_uses__ {
12    // Parse immutable use statements with style: let name = use::style(args...);
13    (
14        @uses [$($uses:tt)*]
15        @states [$($states:tt)*]
16        let $name:ident = use:: $invocation:expr; $($rest:tt)*
17    ) => {
18        $crate::__sliced_parse_uses__!(
19            @uses [$($uses)* { $name, $invocation, $invocation }]
20            @states [$($states)*]
21            $($rest)*
22        )
23    };
24
25    // Parse immutable use statements without style: let name = use(args...);
26    (
27        @uses [$($uses:tt)*]
28        @states [$($states:tt)*]
29        let $name:ident = use($($args:expr),* $(,)?); $($rest:tt)*
30    ) => {
31        $crate::__sliced_parse_uses__!(
32            @uses [$($uses)* { $name, $crate::macro_support::copy_span::copy_span!($($args,)* default)($($args),*), $($args),* }]
33            @states [$($states)*]
34            $($rest)*
35        )
36    };
37
38    // Parse mutable state statements: let mut name = use::style::<Type>(args);
39    (
40        @uses [$($uses:tt)*]
41        @states [$($states:tt)*]
42        let mut $name:ident = use:: $style:ident $(::<$ty:ty>)? ($($args:expr)?); $($rest:tt)*
43    ) => {
44        $crate::__sliced_parse_uses__!(
45            @uses [$($uses)*]
46            @states [$($states)* { $name, $style, (($($ty)?), ($($args)?)) }]
47            $($rest)*
48        )
49    };
50
51    // Terminal case: no uses, only states
52    (
53        @uses []
54        @states [$({ $state_name:ident, $state_style:ident, $state_arg:tt })+]
55        $($body:tt)*
56    ) => {
57        {
58            // We need at least one use to get a tick, so panic if there are none
59            compile_error!("sliced! requires at least one `let name = use(...)` statement to determine the tick")
60        }
61    };
62
63    // Terminal case: uses with optional states
64    (
65        @uses [$({ $use_name:ident, $invocation:expr, $($invocation_spans:expr),* })+]
66        @states [$({ $state_name:ident, $state_style:ident, (($($state_ty:ty)?), ($($state_arg:expr)?)) })*]
67        $($body:tt)*
68    ) => {
69        {
70            use $crate::live_collections::sliced::style::*;
71            let __styled = (
72                $($invocation,)+
73            );
74
75            let __tick = $crate::live_collections::sliced::Slicable::create_tick(&__styled.0);
76            let __backtraces = {
77                use $crate::compile::ir::backtrace::__macro_get_backtrace;
78                (
79                    $($crate::macro_support::copy_span::copy_span!($($invocation_spans,)* {
80                        __macro_get_backtrace(1)
81                    }),)+
82                )
83            };
84            let __sliced = $crate::live_collections::sliced::Slicable::slice(__styled, &__tick, __backtraces);
85            let (
86                $($use_name,)+
87            ) = __sliced;
88
89            // Create all cycles and pack handles/values into tuples.
90            //
91            // The `copy_span!` wrapper re-spans the macro-generated tokens (the style function
92            // path and the `build` call) to the user's tokens, so that errors arising from the
93            // state creation (e.g. an initializer that is not general enough over lifetimes) are
94            // attributed to the originating `use` statement instead of the entire `sliced!`
95            // invocation. The `$state_ty` and `$state_arg` fragments are passed through untouched
96            // (as interpolated fragments), preserving the precise spans of errors within the
97            // user-provided argument.
98            //
99            // Each state borrows its own clone of the tick (created inside the `copy_span!`
100            // target and spanned to the style name, e.g. `state`), so lifetime errors caused
101            // by a bad initializer point at the originating `use` statement rather than the
102            // shared `__tick` local, whose span covers the entire macro invocation.
103            let (__handles, __states) = $crate::live_collections::sliced::unzip_cycles((
104                $($crate::macro_support::copy_span::copy_span!($state_style, {
105                    $crate::live_collections::sliced::style::$state_style$(::<$state_ty, _>)?(& __tick.clone()).build($($state_arg)?)
106                }),)*
107            ));
108
109            // Unpack mutable state values
110            let (
111                $(mut $state_name,)*
112            ) = __states;
113
114            // Execute the body
115            let __body_result = {
116                $($body)*
117            };
118
119            // Re-pack the final state values and complete cycles
120            let __final_states = (
121                $($state_name,)*
122            );
123            $crate::live_collections::sliced::complete_cycles(__handles, __final_states);
124
125            // Unslice the result
126            $crate::live_collections::sliced::Unslicable::unslice(__body_result)
127        }
128    };
129}
130
131#[macro_export]
132/// Transforms a live collection with a computation relying on a slice of another live collection.
133/// This is useful for reading a snapshot of an asynchronously updated collection while processing another
134/// collection, such as joining a stream with the latest values from a singleton.
135///
136/// # Syntax
137/// The `sliced!` macro takes in a closure-like syntax specifying the live collections to be sliced
138/// and the body of the transformation. Each `use` statement indicates a live collection to be sliced,
139/// along with a non-determinism explanation. Optionally, a style can be specified to control how the
140/// live collection is sliced (e.g., atomically). All `use` statements must appear before the body.
141///
142/// ```rust,ignore
143/// let stream = sliced! {
144///     let name1 = use(collection1, nondet!(/** explanation */));
145///     let name2 = use::atomic(collection2, nondet!(/** explanation */));
146///
147///     // arbitrary statements can follow
148///     let intermediate = name1.map(...);
149///     intermediate.cross_singleton(name2)
150/// };
151/// ```
152///
153/// # Stateful Computations
154/// The `sliced!` macro also supports stateful computations across iterations using `let mut` bindings
155/// with `use::state` or `use::state_null`. These create cycles that persist values between iterations.
156///
157/// - `use::state(|l| initial)`: Creates a cycle with an initial value. The closure receives
158///   the slice location and returns the initial state for the first iteration.
159/// - `use::state_null::<Type>()`: Creates a cycle that starts as null/empty on the first iteration.
160///
161/// The mutable binding can be reassigned in the body, and the final value will be passed to the
162/// next iteration.
163///
164/// ```rust,ignore
165/// let counter_stream = sliced! {
166///     let batch = use(input_stream, nondet!(/** explanation */));
167///     let mut counter = use::state(|l| l.singleton(q!(0)));
168///
169///     // Increment counter by the number of items in this batch
170///     let new_count = counter.clone().zip(batch.count())
171///         .map(q!(|(old, add)| old + add));
172///     counter = new_count.clone();
173///     new_count.into_stream()
174/// };
175/// ```
176macro_rules! __sliced__ {
177    ($($tt:tt)*) => {
178        $crate::__sliced_parse_uses__!(
179            @uses []
180            @states []
181            $($tt)*
182        )
183    };
184}
185
186pub use crate::__sliced__ as sliced;
187
188/// Marks this live collection as atomically-yielded, which means that the output outside
189/// `sliced` will be at an atomic location that is synchronous with respect to the body
190/// of the slice.
191pub fn yield_atomic<T>(t: T) -> style::Atomic<T> {
192    style::Atomic {
193        collection: t,
194        // yield_atomic doesn't need a nondet since it's for output, not input
195        nondet: crate::nondet::NonDet,
196    }
197}
198
199/// A trait for live collections which can be sliced into bounded versions at a tick.
200pub trait Slicable<'a, L: Location<'a>> {
201    /// The sliced version of this live collection.
202    type Slice;
203
204    /// The type of backtrace associated with this slice.
205    type Backtrace;
206
207    /// Gets the location associated with this live collection.
208    fn get_location(&self) -> L;
209
210    /// Creates a tick that is appropriate for the collection's location.
211    fn create_tick(&self) -> Tick<L> {
212        self.get_location().try_tick().unwrap()
213    }
214
215    /// Slices this live collection at the given tick.
216    ///
217    /// # Non-Determinism
218    /// Slicing a live collection may involve non-determinism, such as choosing which messages
219    /// to include in a batch.
220    fn slice(self, tick: &Tick<L>, backtrace: Self::Backtrace) -> Self::Slice;
221}
222
223/// A trait for live collections which can be yielded out of a slice back into their original form.
224pub trait Unslicable {
225    /// The unsliced version of this live collection.
226    type Unsliced;
227
228    /// Unslices a sliced live collection back into its original form.
229    fn unslice(self) -> Self::Unsliced;
230}
231
232/// A trait for unzipping a tuple of (handle, state) pairs into separate tuples.
233#[doc(hidden)]
234pub trait UnzipCycles {
235    /// The tuple of cycle handles.
236    type Handles;
237    /// The tuple of state values.
238    type States;
239
240    /// Unzips the cycles into handles and states.
241    fn unzip(self) -> (Self::Handles, Self::States);
242}
243
244/// Unzips a tuple of cycles into handles and states.
245#[doc(hidden)]
246pub fn unzip_cycles<T: UnzipCycles>(cycles: T) -> (T::Handles, T::States) {
247    cycles.unzip()
248}
249
250/// A trait for completing a tuple of cycle handles with their final state values.
251#[doc(hidden)]
252pub trait CompleteCycles<States> {
253    /// Completes all cycles with the provided state values.
254    fn complete(self, states: States);
255}
256
257/// Completes a tuple of cycle handles with their final state values.
258#[doc(hidden)]
259pub fn complete_cycles<H: CompleteCycles<S>, S>(handles: H, states: S) {
260    handles.complete(states);
261}
262
263impl<'a, L: Location<'a>> Slicable<'a, L> for () {
264    type Slice = ();
265    type Backtrace = ();
266
267    fn get_location(&self) -> L {
268        unreachable!()
269    }
270
271    fn slice(self, _tick: &Tick<L>, _backtrace: Self::Backtrace) -> Self::Slice {}
272}
273
274impl Unslicable for () {
275    type Unsliced = ();
276
277    fn unslice(self) -> Self::Unsliced {}
278}
279
280macro_rules! impl_slicable_for_tuple {
281    ($($T:ident, $T_bt:ident, $idx:tt),+) => {
282        impl<'a, L: Location<'a>, $($T: Slicable<'a, L>),+> Slicable<'a, L> for ($($T,)+) {
283            type Slice = ($($T::Slice,)+);
284            type Backtrace = ($($T::Backtrace,)+);
285
286            fn get_location(&self) -> L {
287                self.0.get_location()
288            }
289
290            #[expect(non_snake_case, reason = "macro codegen")]
291            fn slice(self, tick: &Tick<L>, backtrace: Self::Backtrace) -> Self::Slice {
292                let ($($T,)+) = self;
293                let ($($T_bt,)+) = backtrace;
294                ($($T.slice(tick, $T_bt),)+)
295            }
296        }
297
298        impl<$($T: Unslicable),+> Unslicable for ($($T,)+) {
299            type Unsliced = ($($T::Unsliced,)+);
300
301            #[expect(non_snake_case, reason = "macro codegen")]
302            fn unslice(self) -> Self::Unsliced {
303                let ($($T,)+) = self;
304                ($($T.unslice(),)+)
305            }
306        }
307    };
308}
309
310#[cfg(stageleft_runtime)]
311impl_slicable_for_tuple!(S1, S1_bt, 0);
312#[cfg(stageleft_runtime)]
313impl_slicable_for_tuple!(S1, S1_bt, 0, S2, S2_bt, 1);
314#[cfg(stageleft_runtime)]
315impl_slicable_for_tuple!(S1, S1_bt, 0, S2, S2_bt, 1, S3, S3_bt, 2);
316#[cfg(stageleft_runtime)]
317impl_slicable_for_tuple!(S1, S1_bt, 0, S2, S2_bt, 1, S3, S3_bt, 2, S4, S4_bt, 3);
318#[cfg(stageleft_runtime)]
319impl_slicable_for_tuple!(
320    S1, S1_bt, 0, S2, S2_bt, 1, S3, S3_bt, 2, S4, S4_bt, 3, S5, S5_bt, 4
321);
322#[cfg(stageleft_runtime)]
323impl_slicable_for_tuple!(
324    S1, S1_bt, 0, S2, S2_bt, 1, S3, S3_bt, 2, S4, S4_bt, 3, S5, S5_bt, 4, S6, S6_bt, 5
325);
326#[cfg(stageleft_runtime)]
327impl_slicable_for_tuple!(
328    S1, S1_bt, 0, S2, S2_bt, 1, S3, S3_bt, 2, S4, S4_bt, 3, S5, S5_bt, 4, S6, S6_bt, 5, S7, S7_bt,
329    6
330);
331#[cfg(stageleft_runtime)]
332impl_slicable_for_tuple!(
333    S1, S1_bt, 0, S2, S2_bt, 1, S3, S3_bt, 2, S4, S4_bt, 3, S5, S5_bt, 4, S6, S6_bt, 5, S7, S7_bt,
334    6, S8, S8_bt, 7
335);
336#[cfg(stageleft_runtime)]
337impl_slicable_for_tuple!(
338    S1, S1_bt, 0, S2, S2_bt, 1, S3, S3_bt, 2, S4, S4_bt, 3, S5, S5_bt, 4, S6, S6_bt, 5, S7, S7_bt,
339    6, S8, S8_bt, 7, S9, S9_bt, 8
340);
341#[cfg(stageleft_runtime)]
342impl_slicable_for_tuple!(
343    S1, S1_bt, 0, S2, S2_bt, 1, S3, S3_bt, 2, S4, S4_bt, 3, S5, S5_bt, 4, S6, S6_bt, 5, S7, S7_bt,
344    6, S8, S8_bt, 7, S9, S9_bt, 8, S10, S10_bt, 9
345);
346#[cfg(stageleft_runtime)]
347impl_slicable_for_tuple!(
348    S1, S1_bt, 0, S2, S2_bt, 1, S3, S3_bt, 2, S4, S4_bt, 3, S5, S5_bt, 4, S6, S6_bt, 5, S7, S7_bt,
349    6, S8, S8_bt, 7, S9, S9_bt, 8, S10, S10_bt, 9, S11, S11_bt, 10
350);
351#[cfg(stageleft_runtime)]
352impl_slicable_for_tuple!(
353    S1, S1_bt, 0, S2, S2_bt, 1, S3, S3_bt, 2, S4, S4_bt, 3, S5, S5_bt, 4, S6, S6_bt, 5, S7, S7_bt,
354    6, S8, S8_bt, 7, S9, S9_bt, 8, S10, S10_bt, 9, S11, S11_bt, 10, S12, S12_bt, 11
355);
356
357macro_rules! impl_cycles_for_tuple {
358    ($($H:ident, $S:ident, $idx:tt),*) => {
359        impl<$($H, $S),*> UnzipCycles for ($(($H, $S),)*) {
360            type Handles = ($($H,)*);
361            type States = ($($S,)*);
362
363            #[expect(clippy::allow_attributes, reason = "macro codegen")]
364            #[allow(non_snake_case, reason = "macro codegen")]
365            fn unzip(self) -> (Self::Handles, Self::States) {
366                let ($($H,)*) = self;
367                (
368                    ($($H.0,)*),
369                    ($($H.1,)*),
370                )
371            }
372        }
373
374        impl<$($H: crate::forward_handle::CompleteCycle<$S>, $S),*> CompleteCycles<($($S,)*)> for ($($H,)*) {
375            #[expect(clippy::allow_attributes, reason = "macro codegen")]
376            #[allow(non_snake_case, reason = "macro codegen")]
377            fn complete(self, states: ($($S,)*)) {
378                let ($($H,)*) = self;
379                let ($($S,)*) = states;
380                $($H.complete_next_tick($S);)*
381            }
382        }
383    };
384}
385
386#[cfg(stageleft_runtime)]
387impl_cycles_for_tuple!();
388#[cfg(stageleft_runtime)]
389impl_cycles_for_tuple!(H1, S1, 0);
390#[cfg(stageleft_runtime)]
391impl_cycles_for_tuple!(H1, S1, 0, H2, S2, 1);
392#[cfg(stageleft_runtime)]
393impl_cycles_for_tuple!(H1, S1, 0, H2, S2, 1, H3, S3, 2);
394#[cfg(stageleft_runtime)]
395impl_cycles_for_tuple!(H1, S1, 0, H2, S2, 1, H3, S3, 2, H4, S4, 3);
396#[cfg(stageleft_runtime)]
397impl_cycles_for_tuple!(H1, S1, 0, H2, S2, 1, H3, S3, 2, H4, S4, 3, H5, S5, 4);
398#[cfg(stageleft_runtime)]
399impl_cycles_for_tuple!(
400    H1, S1, 0, H2, S2, 1, H3, S3, 2, H4, S4, 3, H5, S5, 4, H6, S6, 5
401);
402#[cfg(stageleft_runtime)]
403impl_cycles_for_tuple!(
404    H1, S1, 0, H2, S2, 1, H3, S3, 2, H4, S4, 3, H5, S5, 4, H6, S6, 5, H7, S7, 6
405);
406#[cfg(stageleft_runtime)]
407impl_cycles_for_tuple!(
408    H1, S1, 0, H2, S2, 1, H3, S3, 2, H4, S4, 3, H5, S5, 4, H6, S6, 5, H7, S7, 6, H8, S8, 7
409);
410#[cfg(stageleft_runtime)]
411impl_cycles_for_tuple!(
412    H1, S1, 0, H2, S2, 1, H3, S3, 2, H4, S4, 3, H5, S5, 4, H6, S6, 5, H7, S7, 6, H8, S8, 7, H9, S9,
413    8
414);
415#[cfg(stageleft_runtime)]
416impl_cycles_for_tuple!(
417    H1, S1, 0, H2, S2, 1, H3, S3, 2, H4, S4, 3, H5, S5, 4, H6, S6, 5, H7, S7, 6, H8, S8, 7, H9, S9,
418    8, H10, S10, 9
419);
420#[cfg(stageleft_runtime)]
421impl_cycles_for_tuple!(
422    H1, S1, 0, H2, S2, 1, H3, S3, 2, H4, S4, 3, H5, S5, 4, H6, S6, 5, H7, S7, 6, H8, S8, 7, H9, S9,
423    8, H10, S10, 9, H11, S11, 10
424);
425#[cfg(stageleft_runtime)]
426impl_cycles_for_tuple!(
427    H1, S1, 0, H2, S2, 1, H3, S3, 2, H4, S4, 3, H5, S5, 4, H6, S6, 5, H7, S7, 6, H8, S8, 7, H9, S9,
428    8, H10, S10, 9, H11, S11, 10, H12, S12, 11
429);
430
431// Unslicable implementations for plain collections (used when returning from sliced! body)
432impl<'a, T, L: Location<'a>, O: Ordering, R: Retries> Unslicable
433    for super::Stream<T, Tick<L>, Bounded, O, R>
434{
435    type Unsliced = super::Stream<T, L, Unbounded, O, R>;
436
437    fn unslice(self) -> Self::Unsliced {
438        self.all_ticks()
439    }
440}
441
442impl<'a, T, L: Location<'a>> Unslicable for super::Singleton<T, Tick<L>, Bounded> {
443    type Unsliced = super::Singleton<T, L, Unbounded>;
444
445    fn unslice(self) -> Self::Unsliced {
446        self.latest()
447    }
448}
449
450impl<'a, T, L: Location<'a>> Unslicable for super::Optional<T, Tick<L>, Bounded> {
451    type Unsliced = super::Optional<T, L, Unbounded>;
452
453    fn unslice(self) -> Self::Unsliced {
454        self.latest()
455    }
456}
457
458impl<'a, K, V, L: Location<'a>, O: Ordering, R: Retries> Unslicable
459    for super::KeyedStream<K, V, Tick<L>, Bounded, O, R>
460{
461    type Unsliced = super::KeyedStream<K, V, L, Unbounded, O, R>;
462
463    fn unslice(self) -> Self::Unsliced {
464        self.all_ticks()
465    }
466}
467
468// Unslicable implementations for Atomic-wrapped bounded collections
469impl<'a, T, L: Location<'a>, O: Ordering, R: Retries> Unslicable
470    for style::Atomic<super::Stream<T, Tick<L>, Bounded, O, R>>
471{
472    type Unsliced = super::Stream<T, crate::location::Atomic<L>, Unbounded, O, R>;
473
474    fn unslice(self) -> Self::Unsliced {
475        self.collection.all_ticks_atomic()
476    }
477}
478
479impl<'a, T, L: Location<'a>> Unslicable for style::Atomic<super::Singleton<T, Tick<L>, Bounded>> {
480    type Unsliced = super::Singleton<T, crate::location::Atomic<L>, Unbounded>;
481
482    fn unslice(self) -> Self::Unsliced {
483        self.collection.latest_atomic()
484    }
485}
486
487impl<'a, T, L: Location<'a>> Unslicable for style::Atomic<super::Optional<T, Tick<L>, Bounded>> {
488    type Unsliced = super::Optional<T, crate::location::Atomic<L>, Unbounded>;
489
490    fn unslice(self) -> Self::Unsliced {
491        self.collection.latest_atomic()
492    }
493}
494
495impl<'a, K, V, L: Location<'a>, O: Ordering, R: Retries> Unslicable
496    for style::Atomic<super::KeyedStream<K, V, Tick<L>, Bounded, O, R>>
497{
498    type Unsliced = super::KeyedStream<K, V, crate::location::Atomic<L>, Unbounded, O, R>;
499
500    fn unslice(self) -> Self::Unsliced {
501        self.collection.all_ticks_atomic()
502    }
503}
504
505#[cfg(feature = "sim")]
506#[cfg(test)]
507mod tests {
508    use stageleft::q;
509
510    use super::sliced;
511    use crate::location::Location;
512    use crate::nondet::nondet;
513    use crate::prelude::FlowBuilder;
514
515    /// Test a counter using `use::state` with an initial singleton value.
516    /// Each input increments the counter, and we verify the output after each tick.
517    #[test]
518    fn sim_state_counter() {
519        let mut flow = FlowBuilder::new();
520        let node = flow.process::<()>();
521
522        let (input_send, input) = node.sim_input::<i32, _, _>();
523
524        let out_recv = sliced! {
525            let batch = use(input, nondet!(/** test */));
526            let mut counter = use::state(|l| l.singleton(q!(0)));
527
528            let new_count = counter.clone().zip(batch.count())
529                .map(q!(|(old, add)| old + add));
530            counter = new_count.clone();
531            new_count.into_stream()
532        }
533        .sim_output();
534
535        flow.sim().exhaustive(async || {
536            input_send.send(1);
537            assert_eq!(out_recv.next().await.unwrap(), 1);
538
539            input_send.send(1);
540            assert_eq!(out_recv.next().await.unwrap(), 2);
541
542            input_send.send(1);
543            assert_eq!(out_recv.next().await.unwrap(), 3);
544        });
545    }
546
547    /// Test `use::state_null` with an Optional that starts as None.
548    #[cfg(feature = "sim")]
549    #[test]
550    fn sim_state_null_optional() {
551        use crate::live_collections::Optional;
552        use crate::live_collections::boundedness::Bounded;
553        use crate::location::{Location, Tick};
554
555        let mut flow = FlowBuilder::new();
556        let node = flow.process::<()>();
557
558        let (input_send, input) = node.sim_input::<i32, _, _>();
559
560        let out_recv = sliced! {
561            let batch = use(input, nondet!(/** test */));
562            let mut prev = use::state_null::<Optional<i32, Tick<_>, Bounded>>();
563
564            // Output the previous value (or -1 if none)
565            let output = prev.clone().unwrap_or(prev.location().singleton(q!(-1)));
566            // Store the current batch's first value for next tick
567            prev = batch.first();
568            output.into_stream()
569        }
570        .sim_output();
571
572        flow.sim().exhaustive(async || {
573            input_send.send(10);
574            // First tick: prev is None, so output is -1
575            assert_eq!(out_recv.next().await.unwrap(), -1);
576
577            input_send.send(20);
578            // Second tick: prev is Some(10), so output is 10
579            assert_eq!(out_recv.next().await.unwrap(), 10);
580
581            input_send.send(30);
582            // Third tick: prev is Some(20), so output is 20
583            assert_eq!(out_recv.next().await.unwrap(), 20);
584        });
585    }
586
587    /// Test `use::state` with `source_iter` to initialize a stream state.
588    /// On the first tick, the state is the initial `[10, 20]` from `source_iter`.
589    /// On subsequent ticks, the state is the batch from the previous tick.
590    #[test]
591    fn sim_state_source_iter() {
592        let mut flow = FlowBuilder::new();
593        let node = flow.process::<()>();
594
595        let (input_send, input) = node.sim_input::<i32, _, _>();
596
597        let out_recv = sliced! {
598            let batch = use(input, nondet!(/** test */));
599            let mut items = use::state(|l| l.source_iter(q!([10, 20])));
600
601            // Output the current state, then replace it with the batch
602            let output = items.clone();
603            items = batch;
604            output
605        }
606        .sim_output();
607
608        flow.sim().exhaustive(async || {
609            input_send.send(3);
610            // First tick: items = initial [10, 20], output = [10, 20]
611            let mut results = vec![];
612            results.push(out_recv.next().await.unwrap());
613            results.push(out_recv.next().await.unwrap());
614            results.sort();
615            assert_eq!(results, vec![10, 20]);
616
617            input_send.send(4);
618            // Second tick: items = [3] (from previous batch), output = [3]
619            assert_eq!(out_recv.next().await.unwrap(), 3);
620
621            input_send.send(5);
622            // Third tick: items = [4] (from previous batch), output = [4]
623            assert_eq!(out_recv.next().await.unwrap(), 4);
624        });
625    }
626
627    /// Test atomic slicing with keyed streams.
628    #[test]
629    fn sim_sliced_atomic_keyed_stream() {
630        let mut flow = FlowBuilder::new();
631        let node = flow.process::<()>();
632
633        let (input_send, input) = node.sim_input::<(i32, i32), _, _>();
634        let atomic_keyed_input = input.into_keyed().atomic();
635        let accumulated_inputs = atomic_keyed_input
636            .clone()
637            .assume_ordering(nondet!(/** Test */))
638            .fold(
639                q!(|| 0),
640                q!(|curr, new| {
641                    *curr += new;
642                }),
643            );
644
645        let out_recv = sliced! {
646            let atomic_keyed_input = use::atomic(atomic_keyed_input, nondet!(/** test */));
647            let accumulated_inputs = use::atomic(accumulated_inputs, nondet!(/** test */));
648            accumulated_inputs.join_keyed_stream(atomic_keyed_input)
649                .map(q!(|(sum, _input)| sum))
650                .entries()
651        }
652        .assume_ordering_trusted(nondet!(/** test */))
653        .sim_output();
654
655        flow.sim().exhaustive(async || {
656            input_send.send((1, 1));
657            assert_eq!(out_recv.next().await.unwrap(), (1, 1));
658
659            input_send.send((1, 2));
660            assert_eq!(out_recv.next().await.unwrap(), (1, 3));
661
662            input_send.send((2, 1));
663            assert_eq!(out_recv.next().await.unwrap(), (2, 1));
664
665            input_send.send((1, 3));
666            assert_eq!(out_recv.next().await.unwrap(), (1, 6));
667        });
668    }
669}