dfir_lang/graph/ops/
flat_map.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/// > Arguments: A Rust closure that handles an iterator
11///
12/// For each item `i` passed in, treat `i` as an iterator and map the closure to that
13/// iterator to produce items one by one. The type of the input items must be iterable.
14///
15/// > Note: The closure has access to the [`context` object](surface_flows.mdx#the-context-object).
16///
17/// ```dfir
18/// // should print out each character of each word on a separate line
19/// source_iter(vec!["hello", "world"])
20///     -> flat_map(|x| x.chars())
21///     -> assert_eq(['h', 'e', 'l', 'l', 'o', 'w', 'o', 'r', 'l', 'd']);
22/// ```
23pub const FLAT_MAP: OperatorConstraints = OperatorConstraints {
24    name: "flat_map",
25    categories: &[OperatorCategory::Flatten],
26    hard_range_inn: RANGE_1,
27    soft_range_inn: RANGE_1,
28    hard_range_out: RANGE_1,
29    soft_range_out: RANGE_1,
30    num_args: 1,
31    persistence_args: RANGE_0,
32    type_args: RANGE_0,
33    is_external_input: false,
34    has_singleton_output: false,
35    flo_type: None,
36    ports_inn: None,
37    ports_out: None,
38    input_delaytype_fn: |_| None,
39    write_fn: |&WriteContextArgs {
40                   root,
41                   op_span,
42                   ident,
43                   inputs,
44                   outputs,
45                   is_pull,
46                   arguments,
47                   ..
48               },
49               _| {
50        let write_iterator = if is_pull {
51            let input = &inputs[0];
52            quote_spanned! {op_span=>
53                let #ident = #input.flat_map(#arguments);
54            }
55        } else {
56            let output = &outputs[0];
57            quote_spanned! {op_span=>
58                let #ident = #root::pusherator::map::Map::new(
59                    #arguments,
60                    #root::pusherator::flatten::Flatten::new(#output)
61                );
62            }
63        };
64        Ok(OperatorWriteOutput {
65            write_iterator,
66            ..Default::default()
67        })
68    },
69};