mumu-array 0.1.2

// FILE: array/src/reduce.rs

use mumu::{
    parser::interpreter::apply_n_ary_function_value, // For calling the bridging function
    parser::types::{Value, InkIteratorKind}, // Only keep what's used
    parser::interpreter::Interpreter,
};
use std::sync::{Arc, Mutex};

/// The bridging function: `array:reduce(fn, init, data)`.
/// We allow placeholders => partial usage. Once all 3 are known, we do the final reduce.
/// Now we also handle InkIterator data.
pub fn array_reduce_bridge(
    interp: &mut Interpreter,
    args: Vec<Value>
) -> Result<Value, String> {
    match args.len() {
        0 => {
            Ok(make_three_arg_partial(None, None, None))
        }
        1 => {
            let v = &args[0];
            if matches!(v, Value::Placeholder) {
                Ok(make_three_arg_partial(None, None, None))
            } else {
                Ok(make_three_arg_partial(Some(v.clone()), None, None))
            }
        }
        2 => {
            let a = &args[0];
            let b = &args[1];
            let a_opt = if matches!(a, Value::Placeholder) { None } else { Some(a.clone()) };
            let b_opt = if matches!(b, Value::Placeholder) { None } else { Some(b.clone()) };
            Ok(make_three_arg_partial(a_opt, b_opt, None))
        }
        3 => {
            let a_is_pl = matches!(&args[0], Value::Placeholder);
            let b_is_pl = matches!(&args[1], Value::Placeholder);
            let c_is_pl = matches!(&args[2], Value::Placeholder);

            if a_is_pl || b_is_pl || c_is_pl {
                let a_opt = if a_is_pl { None } else { Some(args[0].clone()) };
                let b_opt = if b_is_pl { None } else { Some(args[1].clone()) };
                let c_opt = if c_is_pl { None } else { Some(args[2].clone()) };
                Ok(make_three_arg_partial(a_opt, b_opt, c_opt))
            } else {
                do_reduce(interp, &args[0], &args[1], &args[2])
            }
        }
        n => {
            Err(format!("array:reduce => expected up to 3 arguments, got {}", n))
        }
    }
}

fn do_reduce(
    interp: &mut Interpreter,
    func_val: &Value,
    init_val: &Value,
    data_val: &Value
) -> Result<Value, String> {
    let f_boxed = match func_val {
        Value::Function(fb) => fb,  
        other => return Err(format!("array:reduce => first arg must be Function(...), got {:?}", other)),
    };

    match data_val {
        Value::IntArray(xs) => {
            let mut acc = init_val.clone();
            for &x in xs {
                let partial_val = apply_n_ary_function_value(interp, f_boxed.clone(), vec![acc])?;
                let second_func = match partial_val {
                    Value::Function(f2) => f2,
                    other => {
                        return Err(format!("reduce => expected function after first arg, got {:?}", other));
                    }
                };
                let new_acc = apply_n_ary_function_value(interp, second_func, vec![Value::Int(x)])?;
                acc = new_acc;
            }
            Ok(acc)
        }

        Value::StrArray(ss) => {
            let mut acc = init_val.clone();
            for s in ss {
                let partial_val = apply_n_ary_function_value(interp, f_boxed.clone(), vec![acc])?;
                let second_func = match partial_val {
                    Value::Function(f2) => f2,
                    other => {
                        return Err(format!("reduce => expected function after first arg, got {:?}", other));
                    }
                };
                let new_acc = apply_n_ary_function_value(
                    interp,
                    second_func,
                    vec![Value::SingleString(s.clone())]
                )?;
                acc = new_acc;
            }
            Ok(acc)
        }

        Value::InkIterator(ink_handle) => {
            match &ink_handle.kind {
                InkIteratorKind::Core(state_arc) => {
                    let mut acc = init_val.clone();
                    let mut state = state_arc.lock().map_err(|_| {
                        "array:reduce => cannot lock InkIteratorState".to_string()
                    })?;
                    while !state.done {
                        if state.current >= state.end {
                            state.done = true;
                            break;
                        }
                        let partial_val = apply_n_ary_function_value(interp, f_boxed.clone(), vec![acc])?;
                        let second_func = match partial_val {
                            Value::Function(f2) => f2,
                            other => {
                                return Err(format!("reduce => expected function after first arg, got {:?}", other));
                            }
                        };
                        let new_acc = apply_n_ary_function_value(
                            interp,
                            second_func,
                            vec![Value::Int(state.current)]
                        )?;
                        acc = new_acc;
                        state.current += 1;
                    }
                    Ok(acc)
                }
                InkIteratorKind::Plugin(plugin_arc) => {
                    let mut acc = init_val.clone();
                    let mut plugin = plugin_arc.lock().map_err(|_| {
                        "array:reduce => cannot lock plugin InkIterator".to_string()
                    })?;
                    loop {
                        match plugin.next_value() {
                            Ok(val) => {
                                let partial_val = apply_n_ary_function_value(interp, f_boxed.clone(), vec![acc])?;
                                let second_func = match partial_val {
                                    Value::Function(f2) => f2,
                                    other => {
                                        return Err(format!("reduce => expected function after first arg, got {:?}", other));
                                    }
                                };
                                let new_acc = apply_n_ary_function_value(
                                    interp,
                                    second_func,
                                    vec![val]
                                )?;
                                acc = new_acc;
                            }
                            Err(e) if e == "NO_MORE_DATA" => break,
                            Err(e) => return Err(e),
                        }
                    }
                    Ok(acc)
                }
            }
        }

        other => {
            Err(format!(
                "array:reduce => data must be IntArray, StrArray, or InkIterator, got {:?}",
                other
            ))
        }
    }
}

fn make_three_arg_partial(
    fn_opt: Option<Value>,
    init_opt: Option<Value>,
    data_opt: Option<Value>,
) -> Value {
    use mumu::parser::types::FunctionValue::RustClosure;

    let closure = move |interp: &mut Interpreter, new_args: Vec<Value>| {
        let mut f_cur = fn_opt.clone();
        let mut i_cur = init_opt.clone();
        let mut d_cur = data_opt.clone();

        for arg in new_args {
            if f_cur.is_none() {
                if matches!(arg, Value::Placeholder) {
                } else {
                    f_cur = Some(arg);
                }
                continue;
            }
            if i_cur.is_none() {
                if matches!(arg, Value::Placeholder) {
                } else {
                    i_cur = Some(arg);
                }
                continue;
            }
            if d_cur.is_none() {
                if matches!(arg, Value::Placeholder) {
                } else {
                    d_cur = Some(arg);
                }
                continue;
            }
            return Err("array:reduce => partial => too many arguments".to_string());
        }

        if f_cur.is_some() && i_cur.is_some() && d_cur.is_some() {
            do_reduce(
                interp,
                f_cur.as_ref().unwrap(),
                i_cur.as_ref().unwrap(),
                d_cur.as_ref().unwrap()
            )
        } else {
            Ok(make_three_arg_partial(f_cur, i_cur, d_cur))
        }
    };

    Value::Function(Box::new(RustClosure(
        "array:reduce-partial".to_string(),
        Arc::new(Mutex::new(closure)),
        0,
    )))
}