array-mumu 0.2.0-rc.5

// array/src/multi.rs

use crate::apply::apply_n_ary_function_value;
use mumu::parser::interpreter::Interpreter;
use mumu::parser::types::{FunctionValue, Value};
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::FloatArray(xs), Value::FloatArray(ys)) => {
            Ok(Value::FloatArray(
                xs.iter()
                    .cloned()
                    .filter(|x| !ys.iter().any(|y| x == y))
                    .collect(),
            ))
        }
        (Value::BoolArray(xs), Value::BoolArray(ys)) => {
            Ok(Value::BoolArray(
                xs.iter()
                    .cloned()
                    .filter(|x| !ys.iter().any(|y| x == y))
                    .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::FloatArray(xs), Value::FloatArray(ys)) => Ok(Value::FloatArray(
            xs.iter()
                .cloned()
                .filter(|x| ys.iter().any(|y| x == y))
                .collect(),
        )),
        (Value::BoolArray(xs), Value::BoolArray(ys)) => Ok(Value::BoolArray(
            xs.iter()
                .cloned()
                .filter(|x| ys.iter().any(|y| x == y))
                .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::FloatArray(mut xs), Value::FloatArray(ys)) => {
            for y in ys {
                if !xs.iter().any(|x| x == &y) {
                    xs.push(y);
                }
            }
            Ok(Value::FloatArray(xs))
        }
        (Value::BoolArray(mut xs), Value::BoolArray(ys)) => {
            for y in ys {
                if !xs.iter().any(|x| x == &y) {
                    xs.push(y);
                }
            }
            Ok(Value::BoolArray(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:zip_with ===============================
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 => {
            // Respect placeholders for both arguments provided.
            let f = &args[0];
            let a = &args[1];
            Ok(make_three_arg_partial(
                do_zip_with,
                if is_placeholder(f) { None } else { Some(f.clone()) },
                if is_placeholder(a) { None } else { Some(a.clone()) },
                None,
            ))
        }
        1 => {
            let f = &args[0];
            Ok(make_three_arg_partial(
                do_zip_with,
                if is_placeholder(f) { None } else { Some(f.clone()) },
                None,
                None,
            ))
        }
        0 => Ok(make_three_arg_partial(do_zip_with, None, None, None)),
        n => Err(format!(
            "array:zip_with 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:zip_with => 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:zip_with => both args must be arrays of same type".to_string())
        }
    };
    let mut out = Vec::with_capacity(len);
    for i in 0..len {
        let res = 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))
}