mumu-array 0.1.2

// array/src/multi.rs

use mumu::parser::types::{FunctionValue, Value};
use mumu::parser::interpreter::Interpreter;
use std::collections::HashSet;
use std::sync::{Arc, Mutex};

fn is_placeholder(val: &Value) -> bool {
    match val {
        Value::Placeholder => true,
        Value::SingleString(s) if s == "_" => true,
        Value::StrArray(arr) if arr.len() == 1 && arr[0] == "_" => true,
        _ => false,
    }
}

fn make_two_arg_partial(
    finalize_fn: fn(Value, Value) -> Result<Value, String>,
    a_opt: Option<Value>,
    b_opt: Option<Value>
) -> Value {
    use FunctionValue::RustClosure;
    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) { a = Some(arg); }
                continue;
            }
            if b.is_none() {
                if !is_placeholder(&arg) { b = Some(arg); }
                continue;
            }
            return Err("Too many arguments for partial function".to_string());
        }
        if let (Some(aa), Some(bb)) = (a.clone(), b.clone()) {
            finalize_fn(aa, bb)
        } else {
            Ok(make_two_arg_partial(finalize_fn, a, b))
        }
    };
    Value::Function(Box::new(RustClosure(
        "array-2arg-partial".to_string(),
        Arc::new(Mutex::new(closure)),
        0,
    )))
}

fn make_three_arg_partial(
    finalize: fn(&mut Interpreter, Value, Value, Value) -> Result<Value, String>,
    a_opt: Option<Value>,
    b_opt: Option<Value>,
    c_opt: Option<Value>,
) -> Value {
    use FunctionValue::RustClosure;
    let closure = move |interp: &mut Interpreter, new_args: Vec<Value>| {
        let mut a = a_opt.clone();
        let mut b = b_opt.clone();
        let mut c = c_opt.clone();
        for arg in new_args {
            if a.is_none() {
                if !is_placeholder(&arg) {
                    a = Some(arg);
                }
                continue;
            }
            if b.is_none() {
                if !is_placeholder(&arg) {
                    b = Some(arg);
                }
                continue;
            }
            if c.is_none() {
                if !is_placeholder(&arg) {
                    c = Some(arg);
                }
                continue;
            }
            return Err("Too many arguments for partial function".to_string());
        }
        if let (Some(aa), Some(bb), Some(cc)) = (a.clone(), b.clone(), c.clone()) {
            finalize(interp, aa, bb, cc)
        } else {
            Ok(make_three_arg_partial(finalize, a, b, c))
        }
    };
    Value::Function(Box::new(RustClosure(
        "array-3arg-partial".to_string(),
        Arc::new(Mutex::new(closure)),
        0,
    )))
}

// ========================== array:concat ===============================
pub fn array_concat_bridge(_interp: &mut Interpreter, args: Vec<Value>) -> Result<Value, String> {
    match args.len() {
        0 => Ok(make_two_arg_partial(do_concat, None, None)),
        1 => {
            if is_placeholder(&args[0]) {
                Ok(make_two_arg_partial(do_concat, None, None))
            } else {
                Ok(make_two_arg_partial(do_concat, Some(args[0].clone()), None))
            }
        }
        2 => {
            let a1_pl = is_placeholder(&args[0]);
            let a2_pl = is_placeholder(&args[1]);
            if !a1_pl && !a2_pl {
                do_concat(args[0].clone(), args[1].clone())
            } else {
                Ok(make_two_arg_partial(
                    do_concat,
                    if a1_pl { None } else { Some(args[0].clone()) },
                    if a2_pl { None } else { Some(args[1].clone()) },
                ))
            }
        }
        n => Err(format!("array:concat => expected up to 2 arguments, got {}", n)),
    }
}
fn do_concat(a: Value, b: Value) -> Result<Value, String> {
    match (a, b) {
        (Value::IntArray(mut xs), Value::IntArray(ys)) => { xs.extend(ys); Ok(Value::IntArray(xs)) }
        (Value::StrArray(mut xs), Value::StrArray(ys)) => { xs.extend(ys); Ok(Value::StrArray(xs)) }
        (Value::FloatArray(mut xs), Value::FloatArray(ys)) => { xs.extend(ys); Ok(Value::FloatArray(xs)) }
        (Value::BoolArray(mut xs), Value::BoolArray(ys)) => { xs.extend(ys); Ok(Value::BoolArray(xs)) }
        (Value::MixedArray(mut xs), Value::MixedArray(ys)) => { xs.extend(ys); Ok(Value::MixedArray(xs)) }
        (Value::MixedArray(mut xs), v) => { xs.push(v); Ok(Value::MixedArray(xs)) }
        (v, Value::MixedArray(mut xs)) => { xs.insert(0, v); Ok(Value::MixedArray(xs)) }
        (a, b) => Ok(Value::MixedArray(vec![a, b])),
    }
}

// ========================== array:difference ===============================
pub fn array_difference_bridge(_interp: &mut Interpreter, args: Vec<Value>) -> Result<Value, String> {
    match args.len() {
        0 => Ok(make_two_arg_partial(do_difference, None, None)),
        1 => {
            if is_placeholder(&args[0]) {
                Ok(make_two_arg_partial(do_difference, None, None))
            } else {
                Ok(make_two_arg_partial(do_difference, Some(args[0].clone()), None))
            }
        }
        2 => {
            let a1_pl = is_placeholder(&args[0]);
            let a2_pl = is_placeholder(&args[1]);
            if !a1_pl && !a2_pl {
                do_difference(args[0].clone(), args[1].clone())
            } else {
                Ok(make_two_arg_partial(
                    do_difference,
                    if a1_pl { None } else { Some(args[0].clone()) },
                    if a2_pl { None } else { Some(args[1].clone()) },
                ))
            }
        }
        n => Err(format!("array:difference => expected up to 2 arguments, got {}", n)),
    }
}
fn do_difference(a: Value, b: Value) -> Result<Value, String> {
    match (&a, &b) {
        (Value::IntArray(xs), Value::IntArray(ys)) => {
            let ys_set: HashSet<_> = ys.iter().copied().collect();
            Ok(Value::IntArray(xs.iter().filter(|x| !ys_set.contains(x)).copied().collect()))
        }
        (Value::StrArray(xs), Value::StrArray(ys)) => {
            let ys_set: HashSet<_> = ys.iter().cloned().collect();
            Ok(Value::StrArray(xs.iter().filter(|x| !ys_set.contains(*x)).cloned().collect()))
        }
        (Value::MixedArray(xs), Value::MixedArray(ys)) => {
            Ok(Value::MixedArray(xs.iter().filter(|x| !ys.iter().any(|y| x == &y)).cloned().collect()))
        }
        _ => Err("array:difference => both arguments must be arrays of the same type".to_string()),
    }
}

// ========================== array:intersection ===============================
pub fn array_intersection_bridge(_interp: &mut Interpreter, args: Vec<Value>) -> Result<Value, String> {
    match args.len() {
        0 => Ok(make_two_arg_partial(do_intersection, None, None)),
        1 => {
            if is_placeholder(&args[0]) {
                Ok(make_two_arg_partial(do_intersection, None, None))
            } else {
                Ok(make_two_arg_partial(do_intersection, Some(args[0].clone()), None))
            }
        }
        2 => {
            let a1_pl = is_placeholder(&args[0]);
            let a2_pl = is_placeholder(&args[1]);
            if !a1_pl && !a2_pl {
                do_intersection(args[0].clone(), args[1].clone())
            } else {
                Ok(make_two_arg_partial(
                    do_intersection,
                    if a1_pl { None } else { Some(args[0].clone()) },
                    if a2_pl { None } else { Some(args[1].clone()) },
                ))
            }
        }
        n => Err(format!("array:intersection => expected up to 2 arguments, got {}", n)),
    }
}
fn do_intersection(a: Value, b: Value) -> Result<Value, String> {
    match (&a, &b) {
        (Value::IntArray(xs), Value::IntArray(ys)) => {
            let ys_set: HashSet<_> = ys.iter().copied().collect();
            Ok(Value::IntArray(xs.iter().filter(|x| ys_set.contains(x)).copied().collect()))
        }
        (Value::StrArray(xs), Value::StrArray(ys)) => {
            let ys_set: HashSet<_> = ys.iter().cloned().collect();
            Ok(Value::StrArray(xs.iter().filter(|x| ys_set.contains(*x)).cloned().collect()))
        }
        (Value::MixedArray(xs), Value::MixedArray(ys)) => {
            Ok(Value::MixedArray(xs.iter().filter(|x| ys.iter().any(|y| x == &y)).cloned().collect()))
        }
        _ => Err("array:intersection => both arguments must be arrays of the same type".to_string()),
    }
}

// ========================== array:union ===============================
pub fn array_union_bridge(_interp: &mut Interpreter, args: Vec<Value>) -> Result<Value, String> {
    match args.len() {
        0 => Ok(make_two_arg_partial(do_union, None, None)),
        1 => {
            if is_placeholder(&args[0]) {
                Ok(make_two_arg_partial(do_union, None, None))
            } else {
                Ok(make_two_arg_partial(do_union, Some(args[0].clone()), None))
            }
        }
        2 => {
            let a1_pl = is_placeholder(&args[0]);
            let a2_pl = is_placeholder(&args[1]);
            if !a1_pl && !a2_pl {
                do_union(args[0].clone(), args[1].clone())
            } else {
                Ok(make_two_arg_partial(
                    do_union,
                    if a1_pl { None } else { Some(args[0].clone()) },
                    if a2_pl { None } else { Some(args[1].clone()) },
                ))
            }
        }
        n => Err(format!("array:union => expected up to 2 arguments, got {}", n)),
    }
}
fn do_union(a: Value, b: Value) -> Result<Value, String> {
    match (a, b) {
        (Value::IntArray(mut xs), Value::IntArray(ys)) => {
            for y in ys {
                if !xs.contains(&y) { xs.push(y); }
            }
            Ok(Value::IntArray(xs))
        }
        (Value::StrArray(mut xs), Value::StrArray(ys)) => {
            for y in ys {
                if !xs.contains(&y) { xs.push(y); }
            }
            Ok(Value::StrArray(xs))
        }
        (Value::MixedArray(mut xs), Value::MixedArray(ys)) => {
            for y in ys {
                if !xs.iter().any(|x| x == &y) { xs.push(y); }
            }
            Ok(Value::MixedArray(xs))
        }
        _ => Err("array:union => both arguments must be arrays of the same type".to_string()),
    }
}

// ========================== array:zip ===============================
pub fn array_zip_bridge(_interp: &mut Interpreter, args: Vec<Value>) -> Result<Value, String> {
    match args.len() {
        0 => Ok(make_two_arg_partial(do_zip, None, None)),
        1 => {
            if is_placeholder(&args[0]) {
                Ok(make_two_arg_partial(do_zip, None, None))
            } else {
                Ok(make_two_arg_partial(do_zip, Some(args[0].clone()), None))
            }
        }
        2 => {
            let a1_pl = is_placeholder(&args[0]);
            let a2_pl = is_placeholder(&args[1]);
            if !a1_pl && !a2_pl {
                do_zip(args[0].clone(), args[1].clone())
            } else {
                Ok(make_two_arg_partial(
                    do_zip,
                    if a1_pl { None } else { Some(args[0].clone()) },
                    if a2_pl { None } else { Some(args[1].clone()) },
                ))
            }
        }
        n => Err(format!("array:zip => expected up to 2 arguments, got {}", n)),
    }
}
fn do_zip(a: Value, b: Value) -> Result<Value, String> {
    let len = match (&a, &b) {
        (Value::IntArray(xs), Value::IntArray(ys)) => xs.len().min(ys.len()),
        (Value::StrArray(xs), Value::StrArray(ys)) => xs.len().min(ys.len()),
        (Value::FloatArray(xs), Value::FloatArray(ys)) => xs.len().min(ys.len()),
        (Value::BoolArray(xs), Value::BoolArray(ys)) => xs.len().min(ys.len()),
        (Value::MixedArray(xs), Value::MixedArray(ys)) => xs.len().min(ys.len()),
        _ => return Err("array:zip => both args must be arrays of same type".to_string()),
    };
    let out = match (a, b) {
        (Value::IntArray(xs), Value::IntArray(ys)) =>
            xs.into_iter().zip(ys.into_iter()).take(len).map(|(x, y)| Value::MixedArray(vec![Value::Int(x), Value::Int(y)])).collect(),
        (Value::StrArray(xs), Value::StrArray(ys)) =>
            xs.into_iter().zip(ys.into_iter()).take(len).map(|(x, y)| Value::MixedArray(vec![Value::SingleString(x), Value::SingleString(y)])).collect(),
        (Value::FloatArray(xs), Value::FloatArray(ys)) =>
            xs.into_iter().zip(ys.into_iter()).take(len).map(|(x, y)| Value::MixedArray(vec![Value::Float(x), Value::Float(y)])).collect(),
        (Value::BoolArray(xs), Value::BoolArray(ys)) =>
            xs.into_iter().zip(ys.into_iter()).take(len).map(|(x, y)| Value::MixedArray(vec![Value::Bool(x), Value::Bool(y)])).collect(),
        (Value::MixedArray(xs), Value::MixedArray(ys)) =>
            xs.into_iter().zip(ys.into_iter()).take(len).map(|(x, y)| Value::MixedArray(vec![x, y])).collect(),
        _ => unreachable!(),
    };
    Ok(Value::MixedArray(out))
}

// ========================== array:zipWith ===============================
pub fn array_zip_with_bridge(interp: &mut Interpreter, args: Vec<Value>) -> Result<Value, String> {
    match args.len() {
        3 => {
            let f = &args[0];
            let a = &args[1];
            let b = &args[2];
            let f_pl = is_placeholder(f);
            let a_pl = is_placeholder(a);
            let b_pl = is_placeholder(b);
            if !f_pl && !a_pl && !b_pl {
                do_zip_with(interp, f.clone(), a.clone(), b.clone())
            } else {
                Ok(make_three_arg_partial(
                    do_zip_with,
                    if f_pl { None } else { Some(f.clone()) },
                    if a_pl { None } else { Some(a.clone()) },
                    if b_pl { None } else { Some(b.clone()) },
                ))
            }
        }
        2 => {
            Ok(make_three_arg_partial(
                do_zip_with,
                Some(args[0].clone()),
                Some(args[1].clone()),
                None,
            ))
        }
        1 => Ok(make_three_arg_partial(do_zip_with, Some(args[0].clone()), None, None)),
        0 => Ok(make_three_arg_partial(do_zip_with, None, None, None)),
        n => Err(format!("array:zipWith expects 3 arguments: fn, arr1, arr2, got {}", n)),
    }
}

fn do_zip_with(interp: &mut Interpreter, f: Value, a: Value, b: Value) -> Result<Value, String> {
    let f = match f {
        Value::Function(fb) => fb,
        _ => return Err("array:zipWith => first argument must be function".to_string()),
    };
    let (len, a_iter, b_iter) = match (&a, &b) {
        (Value::IntArray(xs), Value::IntArray(ys)) => (xs.len().min(ys.len()), xs.iter().map(|&x| Value::Int(x)).collect::<Vec<_>>(), ys.iter().map(|&x| Value::Int(x)).collect::<Vec<_>>()),
        (Value::StrArray(xs), Value::StrArray(ys)) => (xs.len().min(ys.len()), xs.iter().map(|x| Value::SingleString(x.clone())).collect::<Vec<_>>(), ys.iter().map(|x| Value::SingleString(x.clone())).collect::<Vec<_>>()),
        (Value::FloatArray(xs), Value::FloatArray(ys)) => (xs.len().min(ys.len()), xs.iter().map(|&x| Value::Float(x)).collect::<Vec<_>>(), ys.iter().map(|&x| Value::Float(x)).collect::<Vec<_>>()),
        (Value::BoolArray(xs), Value::BoolArray(ys)) => (xs.len().min(ys.len()), xs.iter().map(|&x| Value::Bool(x)).collect::<Vec<_>>(), ys.iter().map(|&x| Value::Bool(x)).collect::<Vec<_>>()),
        _ => return Err("array:zipWith => both args must be arrays of same type".to_string()),
    };
    let mut out = Vec::with_capacity(len);
    for i in 0..len {
        let res = mumu::parser::interpreter::apply_n_ary_function_value(
            interp, f.clone(), vec![a_iter[i].clone(), b_iter[i].clone()]
        ).unwrap_or(Value::Placeholder);
        out.push(res);
    }
    Ok(Value::MixedArray(out))
}