harn-vm 0.8.4

use std::collections::{BTreeMap, HashSet};
use std::rc::Rc;

use crate::value::{value_structural_hash_key, VmError, VmValue};
use crate::vm::Vm;

fn current_async_vm(builtin: &str) -> Result<Vm, VmError> {
    crate::vm::clone_async_builtin_child_vm().ok_or_else(|| {
        VmError::Runtime(format!("{builtin}: builtin requires VM execution context"))
    })
}

fn list_arg<'a>(args: &'a [VmValue], builtin: &str) -> Result<&'a Rc<Vec<VmValue>>, VmError> {
    match args.first() {
        Some(VmValue::List(items)) => Ok(items),
        Some(other) => Err(VmError::TypeError(format!(
            "{builtin}: first argument must be a list, got {}",
            other.type_name()
        ))),
        None => Err(VmError::Runtime(format!(
            "{builtin}: first argument must be a list"
        ))),
    }
}

fn positive_usize_arg(args: &[VmValue], index: usize, default: usize, _builtin: &str) -> usize {
    args.get(index)
        .and_then(VmValue::as_int)
        .unwrap_or(default as i64)
        .max(1) as usize
}

pub(crate) fn register_collection_builtins(vm: &mut Vm) {
    vm.register_async_builtin("chunk", |args| async move {
        let items = list_arg(&args, "chunk")?;
        let size = positive_usize_arg(&args, 1, 1, "chunk");
        Ok(VmValue::List(Rc::new(
            items
                .chunks(size)
                .map(|chunk| VmValue::List(Rc::new(chunk.to_vec())))
                .collect(),
        )))
    });

    vm.register_async_builtin("window", |args| async move {
        let items = list_arg(&args, "window")?;
        let size = positive_usize_arg(&args, 1, 2, "window");
        let step = positive_usize_arg(&args, 2, 1, "window");
        if size > items.len() {
            return Ok(VmValue::List(Rc::new(Vec::new())));
        }
        let mut windows = Vec::new();
        let mut start = 0;
        while start + size <= items.len() {
            windows.push(VmValue::List(Rc::new(items[start..start + size].to_vec())));
            start += step;
        }
        Ok(VmValue::List(Rc::new(windows)))
    });

    vm.register_async_builtin("group_by", |args| async move {
        let items = list_arg(&args, "group_by")?;
        let callable = args
            .get(1)
            .ok_or_else(|| VmError::Runtime("group_by: callback is required".to_string()))?;
        if !Vm::is_callable_value(callable) {
            return Err(VmError::TypeError(format!(
                "group_by: callback must be callable, got {}",
                callable.type_name()
            )));
        }
        let mut vm = current_async_vm("group_by")?;
        let mut groups: BTreeMap<String, Vec<VmValue>> = BTreeMap::new();
        for item in items.iter() {
            let key = vm.call_callable_value(callable, &[item.clone()]).await?;
            groups.entry(key.display()).or_default().push(item.clone());
        }
        Ok(VmValue::Dict(Rc::new(
            groups
                .into_iter()
                .map(|(key, values)| (key, VmValue::List(Rc::new(values))))
                .collect(),
        )))
    });

    vm.register_async_builtin("partition", |args| async move {
        let items = list_arg(&args, "partition")?;
        let callable = args
            .get(1)
            .ok_or_else(|| VmError::Runtime("partition: callback is required".to_string()))?;
        if !Vm::is_callable_value(callable) {
            return Err(VmError::TypeError(format!(
                "partition: callback must be callable, got {}",
                callable.type_name()
            )));
        }
        let mut vm = current_async_vm("partition")?;
        let mut matched = Vec::new();
        let mut no_match = Vec::new();
        for item in items.iter() {
            let result = vm.call_callable_value(callable, &[item.clone()]).await?;
            if result.is_truthy() {
                matched.push(item.clone());
            } else {
                no_match.push(item.clone());
            }
        }
        Ok(VmValue::Dict(Rc::new(BTreeMap::from([
            ("match".to_string(), VmValue::List(Rc::new(matched))),
            ("no_match".to_string(), VmValue::List(Rc::new(no_match))),
        ]))))
    });

    vm.register_async_builtin("dedup_by", |args| async move {
        let items = list_arg(&args, "dedup_by")?;
        let callable = args
            .get(1)
            .ok_or_else(|| VmError::Runtime("dedup_by: callback is required".to_string()))?;
        if !Vm::is_callable_value(callable) {
            return Err(VmError::TypeError(format!(
                "dedup_by: callback must be callable, got {}",
                callable.type_name()
            )));
        }
        let mut vm = current_async_vm("dedup_by")?;
        let mut seen = HashSet::new();
        let mut out = Vec::new();
        for item in items.iter() {
            let key = vm.call_callable_value(callable, &[item.clone()]).await?;
            if seen.insert(value_structural_hash_key(&key)) {
                out.push(item.clone());
            }
        }
        Ok(VmValue::List(Rc::new(out)))
    });

    vm.register_async_builtin("flat_map", |args| async move {
        let items = list_arg(&args, "flat_map")?;
        let callable = args
            .get(1)
            .ok_or_else(|| VmError::Runtime("flat_map: callback is required".to_string()))?;
        if !Vm::is_callable_value(callable) {
            return Err(VmError::TypeError(format!(
                "flat_map: callback must be callable, got {}",
                callable.type_name()
            )));
        }
        let mut vm = current_async_vm("flat_map")?;
        let mut out = Vec::new();
        for item in items.iter() {
            match vm.call_callable_value(callable, &[item.clone()]).await? {
                VmValue::List(inner) => out.extend(inner.iter().cloned()),
                other => out.push(other),
            }
        }
        Ok(VmValue::List(Rc::new(out)))
    });
}