array-mumu 0.2.0-rc.5

// FILE: array/src/util.rs

use crate::apply::apply_n_ary_function_value;
use mumu::parser::types::{Value, FunctionValue};
use mumu::parser::interpreter::Interpreter;
use std::collections::HashSet;

// -------- flatten (shallow) --------
pub fn array_flatten_bridge(
    _interp: &mut Interpreter,
    args: Vec<Value>,
) -> Result<Value, String> {
    if args.len() != 1 {
        return Err("array:flatten expects exactly one argument".to_string());
    }
    let v = &args[0];
    match v {
        Value::MixedArray(xs) => {
            let mut out = Vec::new();
            for x in xs {
                match x {
                    Value::MixedArray(inner) => out.extend(inner.clone()),
                    _ => out.push(x.clone()),
                }
            }
            Ok(Value::MixedArray(out))
        }
        Value::IntArray(_) | Value::FloatArray(_) | Value::StrArray(_) | Value::BoolArray(_) => Ok(v.clone()),
        _ => Err("array:flatten expects a MixedArray or array".to_string()),
    }
}

// -------- flatten_deep (recursive) --------
pub fn array_flatten_deep_bridge(
    _interp: &mut Interpreter,
    args: Vec<Value>,
) -> Result<Value, String> {
    if args.len() != 1 {
        return Err("array:flatten_deep expects exactly one argument".to_string());
    }
    Ok(flatten_deep(&args[0]))
}

fn flatten_deep(v: &Value) -> Value {
    match v {
        Value::MixedArray(xs) => {
            let mut out = Vec::new();
            for x in xs {
                match x {
                    Value::MixedArray(inner) => {
                        let deep = flatten_deep(&Value::MixedArray(inner.clone()));
                        match deep {
                            Value::MixedArray(inner_deep) => out.extend(inner_deep),
                            _ => out.push(deep),
                        }
                    }
                    _ => out.push(x.clone()),
                }
            }
            Value::MixedArray(out)
        }
        _ => v.clone(),
    }
}

// -------- uniq --------
pub fn array_uniq_bridge(
    _interp: &mut Interpreter,
    args: Vec<Value>,
) -> Result<Value, String> {
    if args.len() != 1 {
        return Err("array:uniq expects exactly one argument".to_string());
    }
    let v = &args[0];
    match v {
        Value::IntArray(xs) => {
            let mut seen = HashSet::new();
            let mut out = Vec::new();
            for &x in xs {
                if seen.insert(x) {
                    out.push(x);
                }
            }
            Ok(Value::IntArray(out))
        }
        Value::BoolArray(xs) => {
            let mut seen = HashSet::new();
            let mut out = Vec::new();
            for &x in xs {
                if seen.insert(x) {
                    out.push(x);
                }
            }
            Ok(Value::BoolArray(out))
        }
        Value::StrArray(xs) => {
            let mut seen = HashSet::new();
            let mut out = Vec::new();
            for x in xs {
                if seen.insert(x) {
                    out.push(x.clone());
                }
            }
            Ok(Value::StrArray(out))
        }
        Value::MixedArray(xs) => {
            let mut seen = Vec::new();
            let mut out = Vec::new();
            for x in xs {
                if !seen.iter().any(|y| y == x) {
                    seen.push(x.clone());
                    out.push(x.clone());
                }
            }
            Ok(Value::MixedArray(out))
        }
        Value::FloatArray(xs) => {
            let mut seen: Vec<f64> = Vec::new();
            let mut out = Vec::new();
            for &x in xs {
                if !seen.iter().any(|y: &f64| (*y - x).abs() < std::f64::EPSILON) {
                    seen.push(x);
                    out.push(x);
                }
            }
            Ok(Value::FloatArray(out))
        }
        _ => Err("array:uniq => unsupported type".to_string()),
    }
}

// -------- compact --------
pub fn array_compact_bridge(
    _interp: &mut Interpreter,
    args: Vec<Value>,
) -> Result<Value, String> {
    if args.len() != 1 {
        return Err("array:compact expects exactly one argument".to_string());
    }
    let v = &args[0];
    match v {
        Value::IntArray(xs) => Ok(Value::IntArray(xs.iter().filter(|&&x| x != 0).copied().collect())),
        Value::FloatArray(xs) => Ok(Value::FloatArray(xs.iter().filter(|&&x| x != 0.0).copied().collect())),
        Value::BoolArray(xs) => Ok(Value::BoolArray(xs.iter().filter(|&&x| x).copied().collect())),
        Value::StrArray(xs) => Ok(Value::StrArray(xs.iter().filter(|x| !x.is_empty()).cloned().collect())),
        Value::MixedArray(xs) => Ok(Value::MixedArray(xs.iter().filter(|x| !is_falsy(x)).cloned().collect())),
        _ => Err("array:compact => unsupported type".to_string()),
    }
}

fn is_falsy(v: &Value) -> bool {
    match v {
        Value::Bool(false) => true,
        Value::Int(0) => true,
        Value::Float(x) if *x == 0.0 => true,
        Value::SingleString(s) if s.is_empty() => true,
        Value::Placeholder => true,
        _ => false,
    }
}

// -------- reverse --------
pub fn array_reverse_bridge(
    _interp: &mut Interpreter,
    args: Vec<Value>,
) -> Result<Value, String> {
    if args.len() != 1 {
        return Err("array:reverse expects exactly one argument".to_string());
    }
    let v = &args[0];
    match v {
        Value::IntArray(xs) => {
            let mut out = xs.clone();
            out.reverse();
            Ok(Value::IntArray(out))
        }
        Value::FloatArray(xs) => {
            let mut out = xs.clone();
            out.reverse();
            Ok(Value::FloatArray(out))
        }
        Value::BoolArray(xs) => {
            let mut out = xs.clone();
            out.reverse();
            Ok(Value::BoolArray(out))
        }
        Value::StrArray(xs) => {
            let mut out = xs.clone();
            out.reverse();
            Ok(Value::StrArray(out))
        }
        Value::MixedArray(xs) => {
            let mut out = xs.clone();
            out.reverse();
            Ok(Value::MixedArray(out))
        }
        _ => Err("array:reverse => unsupported type".to_string()),
    }
}

// -------- shuffle --------
pub fn array_shuffle_bridge(
    _interp: &mut Interpreter,
    args: Vec<Value>,
) -> Result<Value, String> {
    use rand::seq::SliceRandom;
    use rand::thread_rng;

    if args.len() != 1 {
        return Err("array:shuffle expects exactly one argument".to_string());
    }
    let v = &args[0];
    match v {
        Value::IntArray(xs) => {
            let mut out = xs.clone();
            out.shuffle(&mut thread_rng());
            Ok(Value::IntArray(out))
        }
        Value::FloatArray(xs) => {
            let mut out = xs.clone();
            out.shuffle(&mut thread_rng());
            Ok(Value::FloatArray(out))
        }
        Value::BoolArray(xs) => {
            let mut out = xs.clone();
            out.shuffle(&mut thread_rng());
            Ok(Value::BoolArray(out))
        }
        Value::StrArray(xs) => {
            let mut out = xs.clone();
            out.shuffle(&mut thread_rng());
            Ok(Value::StrArray(out))
        }
        Value::MixedArray(xs) => {
            let mut out = xs.clone();
            out.shuffle(&mut thread_rng());
            Ok(Value::MixedArray(out))
        }
        _ => Err("array:shuffle => unsupported type".to_string()),
    }
}

// -------- partition --------
pub fn array_partition_bridge(
    interp: &mut Interpreter,
    args: Vec<Value>,
) -> Result<Value, String> {
    if args.len() != 2 {
        return Err("array:partition => expects (fn, array)".to_string());
    }
    let pred_fn = match &args[0] {
        Value::Function(f) => f,
        other => return Err(format!("array:partition => first argument must be a function, got {:?}", other)),
    };
    match &args[1] {
        Value::IntArray(xs) => {
            let (yes, no) = partition_values(interp, &**pred_fn, xs.iter().map(|&x| Value::Int(x)).collect());
            Ok(Value::MixedArray(vec![Value::MixedArray(yes), Value::MixedArray(no)]))
        }
        Value::FloatArray(xs) => {
            let (yes, no) = partition_values(interp, &**pred_fn, xs.iter().map(|&x| Value::Float(x)).collect());
            Ok(Value::MixedArray(vec![Value::MixedArray(yes), Value::MixedArray(no)]))
        }
        Value::StrArray(xs) => {
            let (yes, no) = partition_values(interp, &**pred_fn, xs.iter().map(|x| Value::SingleString(x.clone())).collect());
            Ok(Value::MixedArray(vec![Value::MixedArray(yes), Value::MixedArray(no)]))
        }
        Value::BoolArray(xs) => {
            let (yes, no) = partition_values(interp, &**pred_fn, xs.iter().map(|&x| Value::Bool(x)).collect());
            Ok(Value::MixedArray(vec![Value::MixedArray(yes), Value::MixedArray(no)]))
        }
        Value::MixedArray(xs) => {
            let (yes, no) = partition_values(interp, &**pred_fn, xs.clone());
            Ok(Value::MixedArray(vec![Value::MixedArray(yes), Value::MixedArray(no)]))
        }
        other => Err(format!("array:partition => unsupported array type: {:?}", other)),
    }
}

fn partition_values(
    interp: &mut Interpreter,
    pred: &FunctionValue,
    items: Vec<Value>,
) -> (Vec<Value>, Vec<Value>) {
    let mut yes = Vec::new();
    let mut no = Vec::new();
    for v in items {
        let res = apply_n_ary_function_value(interp, Box::new(pred.clone()), vec![v.clone()]);
        match res {
            Ok(Value::Bool(true)) => yes.push(v),
            Ok(Value::Bool(false)) => no.push(v),
            _ => no.push(v),
        }
    }
    (yes, no)
}

// --- Also keep the original head, last, etc. here ---

pub fn array_head_bridge(
    _interp: &mut Interpreter,
    args: Vec<Value>,
) -> Result<Value, String> {
    if args.len() != 1 {
        return Err("array:head => expects 1 argument".to_string());
    }
    let v = &args[0];
    match v {
        Value::IntArray(xs) => xs.get(0).map(|&x| Value::Int(x)).ok_or("_".to_string()),
        Value::FloatArray(xs) => xs.get(0).map(|&x| Value::Float(x)).ok_or("_".to_string()),
        Value::BoolArray(xs) => xs.get(0).map(|&x| Value::Bool(x)).ok_or("_".to_string()),
        Value::StrArray(xs) => xs.get(0).map(|x| Value::SingleString(x.clone())).ok_or("_".to_string()),
        Value::MixedArray(xs) => xs.get(0).cloned().ok_or("_".to_string()),
        Value::SingleString(s) => s.chars().next().map(|c| Value::SingleString(c.to_string())).ok_or("_".to_string()),
        _ => Err("array:head => unsupported type".to_string()),
    }
}

pub fn array_last_bridge(
    _interp: &mut Interpreter,
    args: Vec<Value>,
) -> Result<Value, String> {
    if args.len() != 1 {
        return Err("array:last => expects 1 argument".to_string());
    }
    let v = &args[0];
    match v {
        Value::IntArray(xs) => xs.last().map(|&x| Value::Int(x)).ok_or("_".to_string()),
        Value::FloatArray(xs) => xs.last().map(|&x| Value::Float(x)).ok_or("_".to_string()),
        Value::BoolArray(xs) => xs.last().map(|&x| Value::Bool(x)).ok_or("_".to_string()),
        Value::StrArray(xs) => xs.last().map(|x| Value::SingleString(x.clone())).ok_or("_".to_string()),
        Value::MixedArray(xs) => xs.last().cloned().ok_or("_".to_string()),
        Value::SingleString(s) => s.chars().last().map(|c| Value::SingleString(c.to_string())).ok_or("_".to_string()),
        _ => Err("array:last => unsupported type".to_string()),
    }
}