dfir_lang/graph/ops/
_lattice_fold_batch.rs

Help
1use quote::{quote_spanned, ToTokens};
2use syn::parse_quote;
3
4use super::{
5    DelayType, OpInstGenerics, OperatorCategory, OperatorConstraints, OperatorInstance,
6    OperatorWriteOutput, PortListSpec, WriteContextArgs, RANGE_0, RANGE_1,
7};
8
9/// > 2 input streams, 1 output stream, no arguments.
10///
11/// Batches streaming input and releases it downstream when a signal is delivered. This allows for buffering data and delivering it later while also folding it into a single lattice data structure.
12/// This operator is similar to `defer_signal` in that it batches input and releases it when a signal is given. It is also similar to `lattice_fold` in that it folds the input into a single lattice.
13/// So, `_lattice_fold_batch` is a combination of both `defer_signal` and `lattice_fold`. This operator is useful when trying to combine a sequence of `defer_signal` and `lattice_fold` operators without unnecessary memory consumption.
14///
15/// There are two inputs to `_lattice_fold_batch`, they are `input` and `signal`.
16/// `input` is the input data flow. Data that is delivered on this input is collected in order inside of the `_lattice_fold_batch` operator.
17/// When anything is sent to `signal` the collected data is released downstream. The entire `signal` input is consumed each tick, so sending 5 things on `signal` will not release inputs on the next 5 consecutive ticks.
18///
19/// ```dfir
20/// use lattices::set_union::SetUnionHashSet;
21/// use lattices::set_union::SetUnionSingletonSet;
22///
23/// source_iter([1, 2, 3])
24///     -> map(SetUnionSingletonSet::new_from)
25///     -> [input]batcher;
26///
27/// source_iter([()])
28///     -> [signal]batcher;
29///
30/// batcher = _lattice_fold_batch::<SetUnionHashSet<usize>>()
31///     -> assert_eq([SetUnionHashSet::new_from([1, 2, 3])]);
32/// ```
33pub const _LATTICE_FOLD_BATCH: OperatorConstraints = OperatorConstraints {
34    name: "_lattice_fold_batch",
35    categories: &[OperatorCategory::CompilerFusionOperator],
36    persistence_args: RANGE_0,
37    type_args: &(0..=1),
38    hard_range_inn: &(2..=2),
39    soft_range_inn: &(2..=2),
40    hard_range_out: RANGE_1,
41    soft_range_out: RANGE_1,
42    num_args: 0,
43    is_external_input: false,
44    has_singleton_output: false,
45    flo_type: None,
46    ports_inn: Some(|| PortListSpec::Fixed(parse_quote! { input, signal })),
47    ports_out: None,
48    input_delaytype_fn: |_| Some(DelayType::MonotoneAccum),
49    write_fn: |wc @ &WriteContextArgs {
50                   context,
51                   df_ident,
52                   ident,
53                   op_span,
54                   root,
55                   inputs,
56                   is_pull,
57                   op_inst:
58                       OperatorInstance {
59                           generics: OpInstGenerics { type_args, .. },
60                           ..
61                       },
62                   ..
63               },
64               _| {
65        assert!(is_pull);
66
67        let lattice_type = type_args
68            .first()
69            .map(ToTokens::to_token_stream)
70            .unwrap_or(quote_spanned!(op_span=> _));
71
72        let lattice_ident = wc.make_ident("lattice");
73
74        let write_prologue = quote_spanned! {op_span=>
75            let #lattice_ident = #df_ident.add_state(::std::cell::RefCell::new(<#lattice_type as ::std::default::Default>::default()));
76        };
77
78        let input = &inputs[0];
79        let signal = &inputs[1];
80
81        let write_iterator = {
82            quote_spanned! {op_span=>
83
84                {
85                    let mut __lattice = unsafe {
86                        // SAFETY: handle from `#df_ident.add_state(..)`.
87                        #context.state_ref_unchecked(#lattice_ident)
88                    }.borrow_mut();
89
90                    for __item in #input {
91                        #root::lattices::Merge::merge(&mut *__lattice, __item);
92                    }
93                }
94
95                let #ident = if #signal.count() > 0 {
96                    ::std::option::Option::Some(unsafe {
97                        // SAFETY: handle from `#df_ident.add_state(..)`.
98                        #context.state_ref_unchecked(#lattice_ident)
99                    }.take())
100                } else {
101                    ::std::option::Option::None
102                }.into_iter();
103            }
104        };
105
106        Ok(OperatorWriteOutput {
107            write_prologue,
108            write_iterator,
109            write_iterator_after: Default::default(),
110            ..Default::default()
111        })
112    },
113};
dfir_lang/graph/ops/_lattice_fold_batch.rs

dfir_lang/graph/ops/
_lattice_fold_batch.rs
