array-mumu 0.2.0-rc.5

use mumu::parser::types::{FunctionValue, Value};
use mumu::parser::interpreter::Interpreter;
use crate::apply::apply_n_ary_function_value;
use indexmap::IndexMap;
use super::common::{is_placeholder, make_two_arg_partial};

pub fn array_group_by_bridge(interp: &mut Interpreter, args: Vec<Value>) -> Result<Value, String> {
    match args.len() {
        0 => Ok(make_two_arg_partial(do_group_by, None, None, "array-transform-group_by-partial")),
        1 => {
            if is_placeholder(&args[0]) {
                Ok(make_two_arg_partial(do_group_by, None, None, "array-transform-group_by-partial"))
            } else {
                Ok(make_two_arg_partial(do_group_by, Some(args[0].clone()), None, "array-transform-group_by-partial"))
            }
        }
        2 => {
            let a1_pl = is_placeholder(&args[0]);
            let a2_pl = is_placeholder(&args[1]);
            if !a1_pl && !a2_pl {
                do_group_by(interp, args[0].clone(), args[1].clone())
            } else {
                Ok(make_two_arg_partial(
                    do_group_by,
                    if a1_pl { None } else { Some(args[0].clone()) },
                    if a2_pl { None } else { Some(args[1].clone()) },
                    "array-transform-group_by-partial"
                ))
            }
        }
        _ => Err("array:group_by expects 2 arguments (fn, array)".to_string()),
    }
}

fn to_mixed_items(v: Value) -> Result<Vec<Value>, String> {
    Ok(match v {
        Value::MixedArray(xs) => xs,
        Value::IntArray(xs) => xs.into_iter().map(Value::Int).collect(),
        Value::FloatArray(xs) => xs.into_iter().map(Value::Float).collect(),
        Value::StrArray(xs) => xs.into_iter().map(Value::SingleString).collect(),
        Value::BoolArray(xs) => xs.into_iter().map(Value::Bool).collect(),
        other => return Err(format!("array:group_by => second argument must be an array, got {:?}", other)),
    })
}

fn do_group_by(interp: &mut Interpreter, func: Value, arr: Value) -> Result<Value, String> {
    let func = match func {
        Value::Function(fb) => fb,
        _ => return Err("array:group_by => first argument must be a function".to_string()),
    };

    let items = to_mixed_items(arr)?;

    let mut groups: IndexMap<String, Vec<Value>> = IndexMap::new();
    for item in items {
        let key_val = apply_n_ary_function_value(interp, func.clone(), vec![item.clone()])?;
        let key_str = match key_val {
            Value::SingleString(s) => s,
            Value::StrArray(a) if a.len() == 1 => a[0].clone(),
            Value::Int(i) => i.to_string(),
            Value::Float(f) => f.to_string(),
            Value::Bool(b) => b.to_string(),
            _ => "_".to_string(),
        };
        groups.entry(key_str).or_default().push(item);
    }
    let mut result = IndexMap::new();
    for (k, v) in groups {
        result.insert(k, Value::MixedArray(v));
    }
    Ok(Value::KeyedArray(result))
}