dfir_lang/graph/ops/
lattice_reduce.rs

1use quote::quote_spanned;
2use syn::parse_quote_spanned;
3
4use super::{
5    DelayType, OperatorCategory, OperatorConstraints, OperatorWriteOutput,
6    WriteContextArgs, RANGE_0, RANGE_1,
7};
8
9/// > 1 input stream, 1 output stream
10///
11/// A specialized operator for merging lattices together into an accumulated value. Like [`reduce()`](#reduce)
12/// but specialized for lattice types. `lattice_reduce()` is equivalent to `reduce(dfir_rs::lattices::Merge::merge)`.
13///
14/// `lattice_reduce` can also be provided with one generic lifetime persistence argument, either
15/// `'tick` or `'static`, to specify how data persists. With `'tick`, values will only be collected
16/// within the same tick. With `'static`, values will be remembered across ticks and will be
17/// aggregated with pairs arriving in later ticks. When not explicitly specified persistence
18/// defaults to `'tick`.
19///
20/// `lattice_reduce` is differentiated from `lattice_fold` in that `lattice_reduce` does not require the accumulating type to have a sensible default value.
21/// But it also means that the accumulating function inputs and the accumulating type must be the same.
22///
23/// ```dfir
24/// use dfir_rs::lattices::Max;
25///
26/// source_iter([1, 2, 3, 4, 5])
27///     -> map(Max::new)
28///     -> lattice_reduce()
29///     -> assert_eq([Max::new(5)]);
30/// ```
31pub const LATTICE_REDUCE: OperatorConstraints = OperatorConstraints {
32    name: "lattice_reduce",
33    categories: &[OperatorCategory::LatticeFold],
34    hard_range_inn: RANGE_1,
35    soft_range_inn: RANGE_1,
36    hard_range_out: RANGE_1,
37    soft_range_out: RANGE_1,
38    num_args: 0,
39    persistence_args: &(0..=1),
40    type_args: RANGE_0,
41    is_external_input: false,
42    has_singleton_output: false,
43    flo_type: None,
44    ports_inn: None,
45    ports_out: None,
46    input_delaytype_fn: |_| Some(DelayType::MonotoneAccum),
47    write_fn: |wc @ &WriteContextArgs {
48                   root,
49                   inputs,
50                   op_span,
51                   is_pull,
52                   ..
53               },
54               diagnostics| {
55        assert!(is_pull);
56
57        let arguments = &parse_quote_spanned! {op_span=>
58            |acc, item| { #root::lattices::Merge::<_>::merge(acc, item); }
59        };
60        let wc = WriteContextArgs {
61            arguments,
62            ..wc.clone()
63        };
64
65        let OperatorWriteOutput {
66            write_prologue,
67            write_iterator,
68            write_iterator_after,
69        } = (super::reduce::REDUCE.write_fn)(&wc, diagnostics)?;
70
71        assert_eq!(1, inputs.len());
72        let input = &inputs[0];
73
74        let write_iterator = quote_spanned! {op_span=>
75            let #input = {
76                #[inline(always)]
77                fn check_inputs<Lat>(
78                    input: impl ::std::iter::Iterator<Item = Lat>
79                ) -> impl ::std::iter::Iterator<Item = Lat>
80                where
81                    Lat: #root::lattices::Merge<Lat>,
82                {
83                    input
84                }
85                check_inputs::<_>(#input)
86            };
87            #write_iterator
88        };
89        Ok(OperatorWriteOutput {
90            write_prologue,
91            write_iterator,
92            write_iterator_after,
93        })
94    },
95};