refactor(dfir_lang): remove stratum, add push codegen, test reduce, reduce_keyed, reduce_no_replay

PR: #2966

Handle outputs.is_empty() in reduce_keyed push path

Restructured the write_iterator logic in reduce_keyed.rs to follow the
standard pattern: `if is_pull { ... } else if outputs.is_empty() { ... } else { ... }`.

Previously, the `!is_pull` branch unconditionally accessed `outputs[0]`,
which would panic if reduce_keyed was used as a terminal push operator
with no downstream outputs (e.g., as a singleton reference target).

The new `outputs.is_empty()` branch generates a `for_each`-style sink
that accumulates into the hashtable without forwarding downstream,
matching the pattern used by reduce, fold, reduce_no_replay, and
fold_no_replay.

Co-authored-by: Infinity 🤖 <infinity@hydro.run>
PR: #2966

Author: MingweiSamuel

dfir_lang/src/graph/ops/reduce.rs

@@ -1,7 +1,7 @@
 use quote::quote_spanned;
 
 use super::{
-    DelayType, OperatorCategory, OperatorConstraints, OperatorWriteOutput, Persistence, RANGE_0,
+    OperatorCategory, OperatorConstraints, OperatorWriteOutput, Persistence, RANGE_0,
     RANGE_1, WriteContextArgs,
 };
 
@@ -44,14 +44,15 @@ pub const REDUCE: OperatorConstraints = OperatorConstraints {
     flo_type: None,
     ports_inn: None,
     ports_out: None,
-    input_delaytype_fn: |_| Some(DelayType::Stratum),
+    input_delaytype_fn: |_| None,
     write_fn: |wc @ &WriteContextArgs {
                    root,
                    op_span,
                    work_fn,
                    work_fn_async,
                    ident,
                    inputs,
+                   outputs,
                    is_pull,
                    arguments,
                    ..
@@ -117,15 +118,30 @@ pub const REDUCE: OperatorConstraints = OperatorConstraints {
                     )
                 );
             }
-        } else {
-            // Is only push when used as a singleton, so no need to push to `outputs[0]`.
+        } else if outputs.is_empty() {
+            // Terminal push: reduce is a singleton reference target with no downstream.
             quote_spanned! {op_span=>
                 let #ident = #root::dfir_pipes::push::for_each(|#item_ident| {
                     #assign_accum_ident
 
                     #foreach_body
                 });
             }
+        } else {
+            let output = &outputs[0];
+            quote_spanned! {op_span=>
+                let #ident = #root::dfir_pipes::push::reduce_ref(
+                    &mut #singleton_output_ident,
+                    |#accumulator_ident: &mut _, #item_ident| {
+                        #[allow(clippy::redundant_closure_call)]
+                        (#func)(#accumulator_ident, #item_ident);
+                    },
+                    #root::dfir_pipes::push::map(
+                        |__val: &mut _| ::std::clone::Clone::clone(&*__val),
+                        #output,
+                    ),
+                );
+            }
         };
 
         Ok(OperatorWriteOutput {

dfir_lang/src/graph/ops/reduce_keyed.rs

@@ -1,7 +1,7 @@
 use quote::{ToTokens, quote_spanned};
 
 use super::{
-    DelayType, OpInstGenerics, OperatorCategory, OperatorConstraints, OperatorInstance,
+    OpInstGenerics, OperatorCategory, OperatorConstraints, OperatorInstance,
     OperatorWriteOutput, Persistence, RANGE_1, WriteContextArgs,
 };
 
@@ -68,15 +68,15 @@ pub const REDUCE_KEYED: OperatorConstraints = OperatorConstraints {
     flo_type: None,
     ports_inn: None,
     ports_out: None,
-    input_delaytype_fn: |_| Some(DelayType::Stratum),
+    input_delaytype_fn: |_| None,
     write_fn: |wc @ &WriteContextArgs {
                    op_span,
                    ident,
                    inputs,
+                   outputs,
                    is_pull,
                    work_fn_async,
                    root,
-                   op_name,
                    op_inst:
                        OperatorInstance {
                            generics: OpInstGenerics { type_args, .. },
@@ -86,8 +86,6 @@ pub const REDUCE_KEYED: OperatorConstraints = OperatorConstraints {
                    ..
                },
                diagnostics| {
-        assert!(is_pull, "TODO(mingwei): `{}` only supports pull.", op_name);
-
         let [persistence] = wc.persistence_args_disallow_mutable(diagnostics);
 
         let generic_type_args = [
@@ -118,7 +116,7 @@ pub const REDUCE_KEYED: OperatorConstraints = OperatorConstraints {
             _ => Default::default(),
         };
 
-        let write_iterator = {
+        let write_iterator = if is_pull {
             let iter_expr = match persistence {
                 Persistence::None | Persistence::Tick => quote_spanned! {op_span=>
                     #hashtable_ident.drain()
@@ -184,6 +182,33 @@ pub const REDUCE_KEYED: OperatorConstraints = OperatorConstraints {
                 let #ident = #iter_expr;
                 let #ident = #root::dfir_pipes::pull::iter(#ident);
             }
+        } else if outputs.is_empty() {
+            // Terminal push: reduce_keyed is a singleton reference target with no downstream.
+            quote_spanned! {op_span=>
+                let #ident = #root::dfir_pipes::push::for_each(|kv: (#( #generic_type_args ),*)| {
+                    match #singleton_output_ident.entry(kv.0) {
+                        ::std::collections::hash_map::Entry::Vacant(vacant) => {
+                            vacant.insert(kv.1);
+                        }
+                        ::std::collections::hash_map::Entry::Occupied(mut occupied) => {
+                            #[inline(always)]
+                            fn call_comb_type<A>(acc: &mut A, item: A, f: impl Fn(&mut A, A)) {
+                                let () = (f)(acc, item);
+                            }
+                            call_comb_type(occupied.get_mut(), kv.1, #aggfn);
+                        }
+                    }
+                });
+            }
+        } else {
+            let output = &outputs[0];
+            quote_spanned! {op_span=>
+                let #ident = #root::dfir_pipes::push::ReduceKeyed::new(
+                    &mut #singleton_output_ident,
+                    #aggfn,
+                    #output,
+                );
+            }
         };
 
         Ok(OperatorWriteOutput {

dfir_lang/src/graph/ops/reduce_no_replay.rs

@@ -1,8 +1,8 @@
 use quote::quote_spanned;
 
 use super::{
-    DelayType, OperatorCategory, OperatorConstraints, OperatorWriteOutput, Persistence, RANGE_0,
-    RANGE_1, WriteContextArgs,
+    OperatorCategory, OperatorConstraints, OperatorWriteOutput, Persistence, RANGE_0, RANGE_1,
+    WriteContextArgs,
 };
 
 /// > 1 input stream, 1 output stream
@@ -27,7 +27,7 @@ pub const REDUCE_NO_REPLAY: OperatorConstraints = OperatorConstraints {
     flo_type: None,
     ports_inn: None,
     ports_out: None,
-    input_delaytype_fn: |_| Some(DelayType::Stratum),
+    input_delaytype_fn: |_| None,
     write_fn: |wc @ &WriteContextArgs {
                    root,
                    context,
@@ -36,6 +36,7 @@ pub const REDUCE_NO_REPLAY: OperatorConstraints = OperatorConstraints {
                    work_fn_async,
                    ident,
                    inputs,
+                   outputs,
                    is_pull,
                    arguments,
                    ..
@@ -108,15 +109,40 @@ pub const REDUCE_NO_REPLAY: OperatorConstraints = OperatorConstraints {
                     )
                 };
             }
-        } else {
-            // Is only push when used as a singleton, so no need to push to `outputs[0]`.
+        } else if outputs.is_empty() {
+            // Terminal push: reduce_no_replay is a singleton reference target with no downstream.
             quote_spanned! {op_span=>
                 let #ident = #root::dfir_pipes::push::for_each(|#item_ident| {
                     #assign_accum_ident
 
                     #foreach_body
                 });
             }
+        } else {
+            let output = &outputs[0];
+            let was_updated_ident = wc.make_ident("was_updated");
+            quote_spanned! {op_span=>
+                let #was_updated_ident = ::std::cell::Cell::new(false);
+                let #ident = #root::dfir_pipes::push::reduce_ref(
+                    &mut #singleton_output_ident,
+                    |#accumulator_ident: &mut _, #item_ident| {
+                        #was_updated_ident.set(true);
+                        #[allow(clippy::redundant_closure_call)]
+                        (#func)(#accumulator_ident, #item_ident);
+                    },
+                    #root::dfir_pipes::push::filter(
+                        {
+                            let __was_updated = &#was_updated_ident;
+                            let __context: &_ = #context;
+                            move |_| __was_updated.get() || __context.current_tick().0 == 0
+                        },
+                        #root::dfir_pipes::push::map(
+                            |__val: &mut _| ::std::clone::Clone::clone(&*__val),
+                            #output,
+                        ),
+                    ),
+                );
+            }
         };
 
         Ok(OperatorWriteOutput {

dfir_rs/tests/compile-fail/stable/surface_fold_keyed_badtype_int.stderr

@@ -1,4 +1,4 @@
-error[E0271]: type mismatch resolving `<Iter<Drain<'_, '_, {integer}>> as Pull>::Item == (_, _)`
+error[E0271]: type mismatch resolving `<impl Pull<Item = {integer}, Meta = (), CanPend = <Iter<&mut impl Iterator<Item = {integer}>> as Pull>::CanPend, CanEnd = <Iter<&mut impl Iterator<Item = {integer}>> as Pull>::CanEnd> as Pull>::Item == (_, _)`
  --> tests/compile-fail/stable/surface_fold_keyed_badtype_int.rs:3:9
   |
 3 |         source_iter(0..1)
@@ -14,7 +14,7 @@ note: required by a bound in `check_input`
 4 |             -> reduce_keyed(|old: &mut u32, val: u32| { *old += val; })
   |                ^^^^^^^^^^^^ required by this bound in `check_input`
 
-error[E0271]: type mismatch resolving `<Iter<Drain<'_, '_, {integer}>> as Pull>::Item == (_, _)`
+error[E0271]: type mismatch resolving `<impl Pull<Item = {integer}, Meta = (), CanPend = <Iter<&mut impl Iterator<Item = {integer}>> as Pull>::CanPend, CanEnd = <Iter<&mut impl Iterator<Item = {integer}>> as Pull>::CanEnd> as Pull>::Item == (_, _)`
  --> tests/compile-fail/stable/surface_fold_keyed_badtype_int.rs:4:16
   |
 4 |             -> reduce_keyed(|old: &mut u32, val: u32| { *old += val; })

dfir_rs/tests/compile-fail/stable/surface_fold_keyed_badtype_option.stderr

@@ -1,4 +1,4 @@
-error[E0271]: type mismatch resolving `<Iter<Drain<'_, '_, Option<{integer}>>> as Pull>::Item == (_, _)`
+error[E0271]: type mismatch resolving `<impl Pull<Item = Option<{integer}>, Meta = (), CanPend = <Iter<&mut impl Iterator<Item = Option<{integer}>>> as Pull>::CanPend, CanEnd = <Iter<&mut impl Iterator<Item = Option<{integer}>>> as Pull>::CanEnd> as Pull>::Item == (_, _)`
  --> tests/compile-fail/stable/surface_fold_keyed_badtype_option.rs:3:9
   |
 3 |         source_iter([ Some(5), None, Some(12) ])
@@ -14,7 +14,7 @@ note: required by a bound in `check_input`
 4 |             -> reduce_keyed(|old: &mut u32, val: u32| { *old += val; })
   |                ^^^^^^^^^^^^ required by this bound in `check_input`
 
-error[E0271]: type mismatch resolving `<Iter<Drain<'_, '_, Option<{integer}>>> as Pull>::Item == (_, _)`
+error[E0271]: type mismatch resolving `<impl Pull<Item = Option<{integer}>, Meta = (), CanPend = <Iter<&mut impl Iterator<Item = Option<{integer}>>> as Pull>::CanPend, CanEnd = <Iter<&mut impl Iterator<Item = Option<{integer}>>> as Pull>::CanEnd> as Pull>::Item == (_, _)`
  --> tests/compile-fail/stable/surface_fold_keyed_badtype_option.rs:4:16
   |
 4 |             -> reduce_keyed(|old: &mut u32, val: u32| { *old += val; })

dfir_rs/tests/snapshots/surface_reduce__reduce_static@graphvis_dot.snap

@@ -8,22 +8,14 @@ digraph {
     n1v1 [label="(n1v1) source_stream(items_recv)", shape=invhouse, fillcolor="#88aaff"]
     n2v1 [label="(n2v1) reduce::<'static>(|acc: &mut u32, next: u32| *acc += next)", shape=invhouse, fillcolor="#88aaff"]
     n3v1 [label="(n3v1) for_each(|v| result_send.send(v).unwrap())", shape=house, fillcolor="#ffff88"]
-    n4v1 [label="(n4v1) handoff", shape=parallelogram, fillcolor="#ddddff"]
     n2v1 -> n3v1
-    n1v1 -> n4v1
-    n4v1 -> n2v1 [color=red]
+    n1v1 -> n2v1
     subgraph sg_1v1 {
         cluster=true
         fillcolor="#dddddd"
         style=filled
         label = "sg_1v1"
         n1v1
-    }
-    subgraph sg_2v1 {
-        cluster=true
-        fillcolor="#dddddd"
-        style=filled
-        label = "sg_2v1"
         n2v1
         n3v1
     }

dfir_rs/tests/snapshots/surface_reduce__reduce_static@graphvis_mermaid.snap

@@ -11,14 +11,10 @@ linkStyle default stroke:#aaa
 1v1[\"(1v1) <code>source_stream(items_recv)</code>"/]:::pullClass
 2v1[\"(2v1) <code>reduce::&lt;'static&gt;(|acc: &amp;mut u32, next: u32| *acc += next)</code>"/]:::pullClass
 3v1[/"(3v1) <code>for_each(|v| result_send.send(v).unwrap())</code>"\]:::pushClass
-4v1["(4v1) <code>handoff</code>"]:::otherClass
 2v1-->3v1
-1v1-->4v1
-4v1--x2v1; linkStyle 2 stroke:red
+1v1-->2v1
 subgraph sg_1v1 ["sg_1v1"]
     1v1
-end
-subgraph sg_2v1 ["sg_2v1"]
     2v1
     3v1
 end

dfir_rs/tests/snapshots/surface_reduce__reduce_sum@graphvis_dot.snap

@@ -8,22 +8,14 @@ digraph {
     n1v1 [label="(n1v1) source_stream(items_recv)", shape=invhouse, fillcolor="#88aaff"]
     n2v1 [label="(n2v1) reduce(|a: &mut _, b| *a += b)", shape=invhouse, fillcolor="#88aaff"]
     n3v1 [label="(n3v1) for_each(|v| print!(\"{:?}\", v))", shape=house, fillcolor="#ffff88"]
-    n4v1 [label="(n4v1) handoff", shape=parallelogram, fillcolor="#ddddff"]
     n2v1 -> n3v1
-    n1v1 -> n4v1
-    n4v1 -> n2v1 [color=red]
+    n1v1 -> n2v1
     subgraph sg_1v1 {
         cluster=true
         fillcolor="#dddddd"
         style=filled
         label = "sg_1v1"
         n1v1
-    }
-    subgraph sg_2v1 {
-        cluster=true
-        fillcolor="#dddddd"
-        style=filled
-        label = "sg_2v1"
         n2v1
         n3v1
     }

dfir_rs/tests/snapshots/surface_reduce__reduce_sum@graphvis_mermaid.snap

@@ -11,14 +11,10 @@ linkStyle default stroke:#aaa
 1v1[\"(1v1) <code>source_stream(items_recv)</code>"/]:::pullClass
 2v1[\"(2v1) <code>reduce(|a: &amp;mut _, b| *a += b)</code>"/]:::pullClass
 3v1[/"(3v1) <code>for_each(|v| print!(&quot;{:?}&quot;, v))</code>"\]:::pushClass
-4v1["(4v1) <code>handoff</code>"]:::otherClass
 2v1-->3v1
-1v1-->4v1
-4v1--x2v1; linkStyle 2 stroke:red
+1v1-->2v1
 subgraph sg_1v1 ["sg_1v1"]
     1v1
-end
-subgraph sg_2v1 ["sg_2v1"]
     2v1
     3v1
 end

dfir_rs/tests/snapshots/surface_reduce__reduce_tick@graphvis_dot.snap

@@ -8,22 +8,14 @@ digraph {
     n1v1 [label="(n1v1) source_stream(items_recv)", shape=invhouse, fillcolor="#88aaff"]
     n2v1 [label="(n2v1) reduce::<'tick>(|acc: &mut u32, next: u32| *acc += next)", shape=invhouse, fillcolor="#88aaff"]
     n3v1 [label="(n3v1) for_each(|v| result_send.send(v).unwrap())", shape=house, fillcolor="#ffff88"]
-    n4v1 [label="(n4v1) handoff", shape=parallelogram, fillcolor="#ddddff"]
     n2v1 -> n3v1
-    n1v1 -> n4v1
-    n4v1 -> n2v1 [color=red]
+    n1v1 -> n2v1
     subgraph sg_1v1 {
         cluster=true
         fillcolor="#dddddd"
         style=filled
         label = "sg_1v1"
         n1v1
-    }
-    subgraph sg_2v1 {
-        cluster=true
-        fillcolor="#dddddd"
-        style=filled
-        label = "sg_2v1"
         n2v1
         n3v1
     }