dfir_lang 0.17.0-alpha.1

use quote::quote_spanned;

use super::{
    OperatorCategory, OperatorConstraints, OperatorWriteOutput, Persistence, RANGE_0, RANGE_1,
    WriteContextArgs,
};

/// > 1 input stream of type `T`, 1 output stream of type `(usize, T)`
///
/// For each item passed in, enumerate it with its index: `(0, x_0)`, `(1, x_1)`, etc.
///
/// `enumerate` can also be provided with one generic lifetime persistence argument, either
/// `'tick` or `'static`, to specify if indexing resets. If `'tick` (the default) is specified, indexing will
/// restart at zero at the start of each tick. Otherwise `'static` will never reset
/// and count monotonically upwards.
///
/// ```dfir
/// source_iter(vec!["hello", "world"])
///     -> enumerate()
///     -> assert_eq([(0, "hello"), (1, "world")]);
/// ```
pub const ENUMERATE: OperatorConstraints = OperatorConstraints {
    name: "enumerate",
    categories: &[OperatorCategory::Map],
    hard_range_inn: RANGE_1,
    soft_range_inn: RANGE_1,
    hard_range_out: RANGE_1,
    soft_range_out: RANGE_1,
    num_args: 0,
    persistence_args: &(0..=1),
    type_args: RANGE_0,
    is_external_input: false,
    flo_type: None,
    ports_inn: None,
    ports_out: None,
    input_delaytype_fn: |_| None,
    write_fn: |wc @ &WriteContextArgs {
                   root,
                   op_span,
                   ident,
                   inputs,
                   outputs,
                   is_pull,
                   ..
               },
               diagnostics| {
        let [persistence] = wc.persistence_args_disallow_mutable(diagnostics);

        let input = &inputs[0];
        let output = &outputs[0];

        let counter_ident = wc.make_ident("counterdata");

        let write_prologue = quote_spanned! {op_span=>
            let mut #counter_ident = 0..;
        };
        let write_tick_end = match persistence {
            Persistence::None | Persistence::Tick => quote_spanned! {op_span=>
                #counter_ident = 0..;
            },
            _ => Default::default(),
        };

        let map_fn = quote_spanned! {op_span=>
            |item| {
                (::std::iter::Iterator::next(&mut #counter_ident).unwrap(), item)
            }
        };
        let write_iterator = if is_pull {
            quote_spanned! {op_span=>
                let #ident = #root::dfir_pipes::pull::Pull::map(#input, #map_fn);
            }
        } else {
            quote_spanned! {op_span=>
                let #ident = #root::dfir_pipes::push::map(#map_fn, #output);
            }
        };

        Ok(OperatorWriteOutput {
            write_prologue,
            write_iterator,
            write_tick_end,
            ..Default::default()
        })
    },
};