mumu-array 0.1.2

// array/src/simple.rs

use mumu::{
    parser::interpreter::Interpreter,
    parser::types::Value,
    parser::types::FunctionValue,
};
use mumu::parser::interpreter::deep_flatten_upcast_array;

pub fn array_filter(
    interp: &mut Interpreter,
    args: Vec<Value>,
) -> Result<Value, String> {
    match args.len() {
        0 => Ok(make_two_arg_partial(None, None)),
        1 => {
            let a = &args[0];
            if matches!(a, Value::Placeholder) {
                Ok(make_two_arg_partial(None, None))
            } else {
                Ok(make_two_arg_partial(Some(a.clone()), None))
            }
        }
        2 => {
            let a1 = &args[0];
            let a2 = &args[1];
            let a1_is_pl = matches!(a1, Value::Placeholder);
            let a2_is_pl = matches!(a2, Value::Placeholder);
            if a1_is_pl || a2_is_pl {
                Ok(make_two_arg_partial(
                    if a1_is_pl { None } else { Some(a1.clone()) },
                    if a2_is_pl { None } else { Some(a2.clone()) },
                ))
            } else {
                // Both are real => perform filtering
                do_filter(interp, a1.clone(), a2.clone())
            }
        }
        n => Err(format!("array:filter expects up to 2 arguments: function, array (got {})", n)),
    }
}

fn do_filter(
    interp: &mut Interpreter,
    func: Value,
    arr: Value,
) -> Result<Value, String> {
    match arr {
        Value::IntArray(xs) => {
            let mut out = Vec::new();
            for x in xs {
                let res = match &func {
                    Value::Function(fb) => mumu::parser::interpreter::apply_one_function_value(interp, fb.clone(), Value::Int(x))?,
                    _ => return Err("array:filter: first arg must be function".to_string()),
                };
                match res {
                    Value::Bool(true) => out.push(x),
                    Value::Bool(false) => {}
                    _ => return Err("array:filter: predicate did not return bool".to_string()),
                }
            }
            Ok(Value::IntArray(out))
        }
        Value::FloatArray(xs) => {
            let mut out = Vec::new();
            for x in xs {
                let res = match &func {
                    Value::Function(fb) => mumu::parser::interpreter::apply_one_function_value(interp, fb.clone(), Value::Float(x))?,
                    _ => return Err("array:filter: first arg must be function".to_string()),
                };
                match res {
                    Value::Bool(true) => out.push(x),
                    Value::Bool(false) => {}
                    _ => return Err("array:filter: predicate did not return bool".to_string()),
                }
            }
            Ok(Value::FloatArray(out))
        }
        Value::StrArray(xs) => {
            let mut out = Vec::new();
            for x in xs {
                let res = match &func {
                    Value::Function(fb) => mumu::parser::interpreter::apply_one_function_value(interp, fb.clone(), Value::SingleString(x.clone()))?,
                    _ => return Err("array:filter: first arg must be function".to_string()),
                };
                match res {
                    Value::Bool(true) => out.push(x),
                    Value::Bool(false) => {}
                    _ => return Err("array:filter: predicate did not return bool".to_string()),
                }
            }
            Ok(Value::StrArray(out))
        }
        Value::BoolArray(xs) => {
            let mut out = Vec::new();
            for x in xs {
                let res = match &func {
                    Value::Function(fb) => mumu::parser::interpreter::apply_one_function_value(interp, fb.clone(), Value::Bool(x))?,
                    _ => return Err("array:filter: first arg must be function".to_string()),
                };
                match res {
                    Value::Bool(true) => out.push(x),
                    Value::Bool(false) => {}
                    _ => return Err("array:filter: predicate did not return bool".to_string()),
                }
            }
            Ok(Value::BoolArray(out))
        }
        _ => Err("array:filter: only IntArray, FloatArray, StrArray, and BoolArray supported".to_string()),
    }
}

// ==== array:zip with partial and placeholder support ====

fn is_placeholder(v: &Value) -> bool {
    matches!(v, Value::Placeholder)
        || matches!(v, Value::SingleString(s) if s == "_")
        || matches!(v, Value::StrArray(ss) if ss.len() == 1 && ss[0] == "_")
}

pub fn array_zip(
    _interp: &mut Interpreter,
    args: Vec<Value>,
) -> Result<Value, String> {
    match args.len() {
        0 => Ok(make_zip_partial(None, None)),
        1 => {
            let a = &args[0];
            if is_placeholder(a) {
                Ok(make_zip_partial(None, None))
            } else {
                Ok(make_zip_partial(Some(a.clone()), None))
            }
        }
        2 => {
            let a1 = &args[0];
            let a2 = &args[1];
            let a1_is_pl = is_placeholder(a1);
            let a2_is_pl = is_placeholder(a2);
            if a1_is_pl || a2_is_pl {
                Ok(make_zip_partial(
                    if a1_is_pl { None } else { Some(a1.clone()) },
                    if a2_is_pl { None } else { Some(a2.clone()) },
                ))
            } else {
                do_zip(a1.clone(), a2.clone())
            }
        }
        n => Err(format!("array:zip expects up to 2 arguments, got {}", n)),
    }
}

fn do_zip(a: Value, b: Value) -> Result<Value, String> {
    match (a, b) {
        (Value::IntArray(xs), Value::IntArray(ys)) => {
            let n = xs.len().min(ys.len());
            let mut out = Vec::with_capacity(n);
            for i in 0..n {
                out.push(vec![xs[i], ys[i]]);
            }
            Ok(Value::Int2DArray(out))
        }
        (Value::StrArray(xs), Value::StrArray(ys)) => {
            let n = xs.len().min(ys.len());
            let mut out = Vec::with_capacity(n);
            for i in 0..n {
                out.push(vec![xs[i].clone(), ys[i].clone()]);
            }
            // There is NO Str2DArray; wrap each as StrArray for MixedArray
            Ok(Value::MixedArray(out.into_iter().map(Value::StrArray).collect()))
        }
        _ => Err("array:zip: only IntArray/IntArray or StrArray/StrArray".to_string()),
    }
}

fn make_zip_partial(a_opt: Option<Value>, b_opt: Option<Value>) -> Value {
    use mumu::parser::types::FunctionValue::RustClosure;
    use std::sync::{Arc, Mutex};
    let closure = move |_interp: &mut Interpreter, new_args: Vec<Value>| {
        let mut a = a_opt.clone();
        let mut b = b_opt.clone();

        for arg in new_args {
            if a.is_none() {
                if is_placeholder(&arg) {
                    // remain None
                } else {
                    a = Some(arg);
                }
                continue;
            }
            if b.is_none() {
                if is_placeholder(&arg) {
                    // remain None
                } else {
                    b = Some(arg);
                }
                continue;
            }
            return Err("array:zip partial: too many arguments".to_string());
        }

        if a.is_some() && b.is_some() {
            do_zip(a.unwrap(), b.unwrap())
        } else {
            Ok(make_zip_partial(a, b))
        }
    };

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

// === rest of file unchanged ===

pub fn cmp_value(a: &Value, b: &Value) -> std::cmp::Ordering {
    match (a, b) {
        (Value::Int(a), Value::Int(b)) => a.cmp(b),
        (Value::Float(a), Value::Float(b)) => a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal),
        (Value::SingleString(a), Value::SingleString(b)) => a.cmp(b),
        (Value::StrArray(a), Value::StrArray(b)) => a.cmp(b),
        _ => std::cmp::Ordering::Equal,
    }
}

pub fn array_sort(interp: &mut Interpreter, mut args: Vec<Value>) -> Result<Value, String> {
    if args.len() != 1 && args.len() != 2 {
        return Err("array:sort expects array and optional function".to_string());
    }
    let arr = args.remove(0);
    if args.is_empty() {
        match arr {
            Value::IntArray(mut xs) => {
                xs.sort();
                Ok(Value::IntArray(xs))
            }
            Value::StrArray(mut xs) => {
                xs.sort();
                Ok(Value::StrArray(xs))
            }
            Value::FloatArray(mut xs) => {
                xs.sort_by(|a, b| a.partial_cmp(b).unwrap());
                Ok(Value::FloatArray(xs))
            }
            Value::MixedArray(items) => {
                if items.iter().all(|v| matches!(v, Value::Int(_))) {
                    let mut vals: Vec<i32> = items.iter().map(|v| match v { Value::Int(i) => *i, _ => 0 }).collect();
                    vals.sort();
                    Ok(Value::IntArray(vals))
                } else if items.iter().all(|v| matches!(v, Value::Float(_))) {
                    let mut vals: Vec<f64> = items.iter().map(|v| match v { Value::Float(f) => *f, _ => 0.0 }).collect();
                    vals.sort_by(|a, b| a.partial_cmp(b).unwrap());
                    Ok(Value::FloatArray(vals))
                } else if items.iter().all(|v| matches!(v, Value::SingleString(_))) {
                    let mut vals: Vec<String> = items.iter().map(|v| match v { Value::SingleString(s) => s.clone(), _ => "".to_string() }).collect();
                    vals.sort();
                    Ok(Value::StrArray(vals))
                } else {
                    Ok(Value::MixedArray(items))
                }
            }
            _ => Err("array:sort: only IntArray, FloatArray, StrArray, or MixedArray".to_string()),
        }
    } else {
        let func = args.remove(0);
        match arr {
            Value::IntArray(xs) => {
                let mut decorated: Vec<(i32, Value)> = Vec::new();
                for x in &xs {
                    let key = match &func {
                        Value::Function(fb) => mumu::parser::interpreter::apply_one_function_value(interp, fb.clone(), Value::Int(*x))?,
                        _ => return Err("array:sort: 2nd arg must be function".to_string()),
                    };
                    decorated.push((*x, key));
                }
                decorated.sort_by(|a, b| cmp_value(&a.1, &b.1));
                Ok(Value::IntArray(decorated.into_iter().map(|(x,_)| x).collect()))
            }
            Value::StrArray(xs) => {
                let mut decorated: Vec<(String, Value)> = Vec::new();
                for x in &xs {
                    let key = match &func {
                        Value::Function(fb) => mumu::parser::interpreter::apply_one_function_value(interp, fb.clone(), Value::SingleString(x.clone()))?,
                        _ => return Err("array:sort: 2nd arg must be function".to_string()),
                    };
                    decorated.push((x.clone(), key));
                }
                decorated.sort_by(|a, b| cmp_value(&a.1, &b.1));
                Ok(Value::StrArray(decorated.into_iter().map(|(x,_)| x).collect()))
            }
            Value::FloatArray(xs) => {
                let mut decorated: Vec<(f64, Value)> = Vec::new();
                for x in &xs {
                    let key = match &func {
                        Value::Function(fb) => mumu::parser::interpreter::apply_one_function_value(interp, fb.clone(), Value::Float(*x))?,
                        _ => return Err("array:sort: 2nd arg must be function".to_string()),
                    };
                    decorated.push((*x, key));
                }
                decorated.sort_by(|a, b| cmp_value(&a.1, &b.1));
                Ok(Value::FloatArray(decorated.into_iter().map(|(x,_)| x).collect()))
            }
            Value::MixedArray(items) => {
                // Use key function, but preserve result type if possible
                let mut decorated: Vec<(Value, Value)> = Vec::new();
                for item in &items {
                    let key = match &func {
                        Value::Function(fb) => mumu::parser::interpreter::apply_one_function_value(interp, fb.clone(), item.clone())?,
                        _ => return Err("array:sort: 2nd arg must be function".to_string()),
                    };
                    decorated.push((item.clone(), key));
                }
                decorated.sort_by(|a, b| cmp_value(&a.1, &b.1));
                let values: Vec<Value> = decorated.into_iter().map(|(v, _)| v).collect();
                // Coerce if all elements of the same type
                if values.iter().all(|v| matches!(v, Value::Int(_))) {
                    Ok(Value::IntArray(values.into_iter().map(|v| match v { Value::Int(i) => i, _ => 0 }).collect()))
                } else if values.iter().all(|v| matches!(v, Value::Float(_))) {
                    Ok(Value::FloatArray(values.into_iter().map(|v| match v { Value::Float(f) => f, _ => 0.0 }).collect()))
                } else if values.iter().all(|v| matches!(v, Value::SingleString(_))) {
                    Ok(Value::StrArray(values.into_iter().map(|v| match v { Value::SingleString(s) => s, _ => "".to_string() }).collect()))
                } else {
                    Ok(Value::MixedArray(values))
                }
            }
            _ => Err("array:sort: only IntArray, FloatArray, StrArray, MixedArray".to_string()),
        }
    }
}

// SHALLOW flatten: flattens one level only.
pub fn array_flatten(_interp: &mut Interpreter, mut args: Vec<Value>) -> Result<Value, String> {
    if args.len() != 1 {
        return Err("array:flatten expects one array".to_string());
    }
    let arr = args.remove(0);
    match arr {
        Value::IntArray(xs) => Ok(Value::IntArray(xs)),
        Value::FloatArray(xs) => Ok(Value::FloatArray(xs)),
        Value::StrArray(xs) => Ok(Value::StrArray(xs)),
        Value::BoolArray(xs) => Ok(Value::BoolArray(xs)),
        Value::MixedArray(items) => {
            let mut flat = Vec::new();
            for v in items {
                match v {
                    Value::IntArray(x) => for i in x { flat.push(Value::Int(i)); }
                    Value::FloatArray(f) => for ff in f { flat.push(Value::Float(ff)); }
                    Value::StrArray(s) => for ss in s { flat.push(Value::SingleString(ss)); }
                    Value::BoolArray(b) => for bb in b { flat.push(Value::Bool(bb)); }
                    Value::MixedArray(nested) => flat.push(Value::MixedArray(nested)),
                    other => flat.push(other),
                }
            }
            Ok(Value::MixedArray(flat))
        }
        _ => Err("array:flatten: unsupported type".to_string()),
    }
}

pub fn array_join(_interp: &mut Interpreter, mut args: Vec<Value>) -> Result<Value, String> {
    if args.len() != 2 {
        return Err("array:join expects 2 arguments".to_string());
    }
    let arr = args.remove(0);
    let delim = args.remove(0);
    let delim = match delim {
        Value::SingleString(s) => s,
        Value::StrArray(xs) if xs.len() == 1 => xs[0].clone(),
        _ => return Err("array:join => second argument must be StrArray or string".to_string()),
    };
    match arr {
        Value::StrArray(xs) => Ok(Value::SingleString(xs.join(&delim))),
        Value::IntArray(xs) => Ok(Value::SingleString(xs.iter().map(|i| i.to_string()).collect::<Vec<_>>().join(&delim))),
        Value::FloatArray(xs) => Ok(Value::SingleString(xs.iter().map(|f| f.to_string()).collect::<Vec<_>>().join(&delim))),
        Value::BoolArray(xs) => Ok(Value::SingleString(xs.iter().map(|b| b.to_string()).collect::<Vec<_>>().join(&delim))),
        _ => Err("array:join => first argument must be StrArray, IntArray, FloatArray, or BoolArray".to_string()),
    }
}

pub fn array_zipwith(interp: &mut Interpreter, mut args: Vec<Value>) -> Result<Value, String> {
    if args.len() != 3 {
        return Err("array:zipWith expects function and two arrays".to_string());
    }
    let func = args.remove(0);
    let a = args.remove(0);
    let b = args.remove(0);
    match (a, b) {
        (Value::IntArray(xs), Value::IntArray(ys)) => {
            let n = xs.len().min(ys.len());
            let mut out = Vec::with_capacity(n);
            for i in 0..n {
                let res = match &func {
                    Value::Function(fb) => mumu::parser::interpreter::apply_n_ary_function_value(
                        interp, fb.clone(), vec![Value::Int(xs[i]), Value::Int(ys[i])]
                    )?,
                    _ => return Err("array:zipWith: first arg must be function".to_string()),
                };
                match res {
                    Value::Int(z) => out.push(z),
                    _ => return Err("array:zipWith: function must return int".to_string()),
                }
            }
            Ok(Value::IntArray(out))
        }
        _ => Err("array:zipWith: only IntArray/IntArray supported".to_string()),
    }
}

pub fn array_prop(_interp: &mut Interpreter, args: Vec<Value>) -> Result<Value, String> {
    if args.len() != 2 {
        return Err("array:prop expects 2 args".to_string());
    }
    let key = match &args[0] {
        Value::SingleString(s) => s,
        Value::StrArray(ss) if ss.len() == 1 => &ss[0],
        _ => return Err("array:prop expects a string key as the first argument".to_string()),
    };
    let arr = &args[1];
    match arr {
        Value::KeyedArray(map) => map.get(key).cloned().ok_or_else(|| format!("Key '{}' not found", key)),
        _ => Err("array:prop expects keyed array as the second argument".to_string()),
    }
}

fn make_two_arg_partial(
    fn_opt: Option<Value>,
    arr_opt: Option<Value>,
) -> Value {
    use mumu::parser::types::FunctionValue::RustClosure;
    use std::sync::{Arc, Mutex};
    let closure = move |interp: &mut Interpreter, new_args: Vec<Value>| {
        let mut f_cur = fn_opt.clone();
        let mut d_cur = arr_opt.clone();

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

        if f_cur.is_some() && d_cur.is_some() {
            do_filter(
                interp,
                f_cur.as_ref().unwrap().clone(),
                d_cur.as_ref().unwrap().clone()
            )
        } else {
            Ok(make_two_arg_partial(f_cur, d_cur))
        }
    };

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