dfir_lang/graph/ops/
flatten.rs

1use quote::quote_spanned;
2
3use super::{
4    OperatorCategory, OperatorConstraints, OperatorWriteOutput, WriteContextArgs,
5    RANGE_0, RANGE_1,
6};
7
8/// > 1 input stream, 1 output stream
9///
10/// For each item `i` passed in, treat `i` as an iterator and produce its items one by one.
11/// The type of the input items must be iterable.
12///
13/// ```dfir
14/// source_iter(vec![[1, 2], [3, 4], [5, 6]])
15///     -> flatten()
16///     -> assert_eq([1, 2, 3, 4, 5, 6]);
17/// ```
18pub const FLATTEN: OperatorConstraints = OperatorConstraints {
19    name: "flatten",
20    categories: &[OperatorCategory::Flatten],
21    hard_range_inn: RANGE_1,
22    soft_range_inn: RANGE_1,
23    hard_range_out: RANGE_1,
24    soft_range_out: RANGE_1,
25    num_args: 0,
26    persistence_args: RANGE_0,
27    type_args: RANGE_0,
28    is_external_input: false,
29    has_singleton_output: false,
30    flo_type: None,
31    ports_inn: None,
32    ports_out: None,
33    input_delaytype_fn: |_| None,
34    write_fn: |&WriteContextArgs {
35                   root,
36                   op_span,
37                   ident,
38                   inputs,
39                   outputs,
40                   is_pull,
41                   ..
42               },
43               _| {
44        let write_iterator = if is_pull {
45            let input = &inputs[0];
46            quote_spanned! {op_span=>
47                let #ident = #input.flatten();
48            }
49        } else {
50            let output = &outputs[0];
51            quote_spanned! {op_span=>
52                let #ident = #root::pusherator::flatten::Flatten::new(#output);
53            }
54        };
55        Ok(OperatorWriteOutput {
56            write_iterator,
57            ..Default::default()
58        })
59    },
60};