wasm_mumu/

lib.rs

// src/lib.rs
//
// WASM wrapper for the MuMu/Lava interpreter.
// - Compiles against `core-mumu` (browser-safe set of modules).
// - Bakes in array (`array-mumu`) and math (`math-mumu`) plugins.
// - **Intentionally does NOT pull std-mumu’s host-only parts**.
//   We register only the std functions that make sense on the web:
//     • std:put(x)   — print without newline (also drains core iterators)
//     • std:log(x)   — print with newline   (also drains core iterators)
//     • std:deep(x)  — deep-printed value (newline), returns x
//
// Extras exported to JS:
//   - `symbols()` and `complete(prefix)` for REPL-style completion.
//   - `poll()` to drive non-blocking iterator printers one tick.
//   - `reset()` to re-initialize the VM in-place.
//   - `drain_prints()` to fetch only the *new* prints since last drain.
//
// Returned values are mirrored via a small JSON view so the UI can render
// them without poking at the internal `Value` representation.
//

use wasm_bindgen::prelude::*;
use serde::Serialize;
use serde_wasm_bindgen as swb;

use std::collections::BTreeSet;
use std::sync::{Arc, Mutex};

use core_mumu::{Interpreter, Value};
use core_mumu::parser::core::driver::parse_tokens;
use core_mumu::parser::lexer::tokenize;
use core_mumu::parser::types::{FunctionValue, IteratorKind};

/// Better error messages in the browser console
#[wasm_bindgen(start)]
pub fn wasm_start() {
    console_error_panic_hook::set_once();
}

/// A minimal JSON-serializable mirror of `Value` for JS interop.
#[derive(Serialize)]
#[serde(tag = "kind", rename_all = "snake_case")]
enum JsValueView {
    Int { value: i32 },
    Long { value: i64 },
    Float { value: f64 },
    Bool { value: bool },
    String { value: String },
    IntArray { value: Vec<i32> },
    FloatArray { value: Vec<f64> },
    BoolArray { value: Vec<bool> },
    StrArray { value: Vec<String> },
    Int2DArray { value: Vec<Vec<i32>> },
    Float2DArray { value: Vec<Vec<f64>> },
    MixedArray { value: Vec<JsValueView> },
    KeyedArray { value: Vec<(String, JsValueView)> },
    Function,
    Iterator,
    Stream { label: String },
    Tensor,
    Regex { pattern: String, flags: String },
    Placeholder,
    Undefined,
}

/// Minimal conversion for interop (skip Ref; deref by reading inner).
fn to_js_view(v: &Value) -> JsValueView {
    use Value::*;
    match v {
        Int(n) => JsValueView::Int { value: *n },
        Long(l) => JsValueView::Long { value: *l },
        Float(f) => JsValueView::Float { value: *f },
        Bool(b) => JsValueView::Bool { value: *b },
        SingleString(s) => JsValueView::String { value: s.clone() },
        IntArray(xs) => JsValueView::IntArray { value: xs.clone() },
        FloatArray(xs) => JsValueView::FloatArray { value: xs.clone() },
        BoolArray(xs) => JsValueView::BoolArray { value: xs.clone() },
        StrArray(xs) => JsValueView::StrArray { value: xs.clone() },
        Int2DArray(rows) => JsValueView::Int2DArray { value: rows.clone() },
        Float2DArray(rows) => JsValueView::Float2DArray { value: rows.clone() },
        MixedArray(items) => JsValueView::MixedArray {
            value: items.iter().map(to_js_view).collect(),
        },
        KeyedArray(map) => JsValueView::KeyedArray {
            value: map
                .iter()
                .map(|(k, v)| (k.clone(), to_js_view(v)))
                .collect(),
        },
        Function(_) => JsValueView::Function,
        Iterator(_) => JsValueView::Iterator,
        Stream(h) => JsValueView::Stream {
            label: h.label.clone(),
        },
        Tensor(_) => JsValueView::Tensor,
        Regex(rx) => JsValueView::Regex {
            pattern: rx.pattern.clone(),
            flags: rx.flags.clone(),
        },
        Ref(cell) => to_js_view(&cell.lock().unwrap()),
        Placeholder => JsValueView::Placeholder,
        Undefined => JsValueView::Undefined,
    }
}

/// Return a short type label (aligned with core `type(...)`).
fn short_type_name(v: &Value) -> &'static str {
    use Value::*;
    match v {
        Int(_) => "int",
        IntArray(_) => "int_array",
        Int2DArray(_) => "int2d_array",
        Float(_) => "float",
        FloatArray(_) => "float_array",
        Float2DArray(_) => "float2d_array",
        Long(_) => "long",
        Bool(_) => "bool",
        BoolArray(_) => "bool_array",
        Undefined => "undefined",
        Placeholder => "placeholder",
        SingleString(_) => "string",
        StrArray(_) => "str_array",
        KeyedArray(_) => "keyed_array",
        Function(_) => "function",
        Stream(_) => "stream",
        Iterator(_) => "iterator",
        Tensor(_) => "tensor",
        MixedArray(_) => "mixed_array",
        Ref(_) => "ref",
        Regex(_) => "regex",
    }
}

#[derive(Serialize)]
struct ExecOk {
    ok: bool,
    value: JsValueView,
    value_string: String,
    value_type: String,
    prints: Vec<String>,
}

#[derive(Serialize)]
struct ExecErr {
    ok: bool,
    error: String,
}

/// Public WASM wrapper around the interpreter
#[wasm_bindgen]
pub struct WasmMumu {
    interp: Interpreter,
    verbose: bool,
    drained_cursor: usize, // index into print buffer for incremental drains
}

fn bootstrap_interpreter(verbose: bool) -> Interpreter {
    let mut i = Interpreter::new();
    i.set_verbose(verbose);

    // Register **web-safe** core modules (no libloading / filesystem here)
    core_mumu::modules::compose::register_compose_and_pipe(&mut i);
    core_mumu::modules::sput::register_sput(&mut i);
    core_mumu::modules::slog::register_slog(&mut i);
    core_mumu::modules::step::register_step(&mut i);
    core_mumu::modules::r#type::register_type(&mut i);
    core_mumu::modules::check::register_check(&mut i);

    // Disallow extend/include in browser
    i.disallow_extend();

    // Bake-in plugins
    mumuarray::register_all(&mut i);
    mumumath::register_all(&mut i);

    // Register only the **usable** subset of std (no input/input_iter in web)
    register_web_std(&mut i);

    i
}

/// ---- Minimal, web-safe std:* registration (put/log/deep only) ----
fn register_web_std(interp: &mut Interpreter) {
    macro_rules! reg {
        ($name:expr, $f:expr) => {{
            let func = Arc::new(Mutex::new($f));
            interp.register_dynamic_function($name, func);
            interp.set_variable(
                $name,
                Value::Function(Box::new(FunctionValue::Named($name.into()))),
            );
        }};
    }

    fn to_printable(v: &Value) -> String {
        match v {
            Value::SingleString(s) => s.clone(),
            Value::Int(i) => i.to_string(),
            Value::Long(l) => l.to_string(),
            Value::Float(f) => f.to_string(),
            Value::Bool(b) => b.to_string(),
            _ => format!("{:?}", v),
        }
    }

    // Drain core/plugin iterators synchronously (like std:put/log)
    fn drain_iterator(
        interp: &mut Interpreter,
        it: &core_mumu::parser::types::IteratorHandle,
        newline: bool,
    ) -> Result<(), String> {
        match &it.kind {
            IteratorKind::Core(state_arc) => {
                let mut guard = state_arc
                    .lock()
                    .map_err(|_| "std (wasm): iterator lock error".to_string())?;
                while !guard.done && guard.current < guard.end {
                    let item = guard.current;
                    guard.current += 1;
                    if guard.current >= guard.end {
                        guard.done = true;
                    }
                    drop(guard);
                    let s = item.to_string();
                    if newline {
                        interp.push_print(&(s.clone() + "\n"));
                    } else {
                        interp.push_print(&s);
                    }
                    guard = state_arc
                        .lock()
                        .map_err(|_| "std (wasm): iterator lock error".to_string())?;
                }
                Ok(())
            }
            IteratorKind::Plugin(plugin_arc) => {
                let mut plugin = plugin_arc
                    .lock()
                    .map_err(|_| "std (wasm): plugin iterator lock error".to_string())?;
                loop {
                    match plugin.next_value() {
                        Ok(val) => {
                            let s = to_printable(&val);
                            if newline {
                                interp.push_print(&(s.clone() + "\n"));
                            } else {
                                interp.push_print(&s);
                            }
                        }
                        Err(e) if e == "NO_MORE_DATA" => break,
                        Err(e) => return Err(e),
                    }
                }
                Ok(())
            }
        }
    }

    // std:put(x)
    fn std_put_bridge(interp: &mut Interpreter, mut args: Vec<Value>) -> Result<Value, String> {
        if args.len() != 1 {
            return Err(format!(
                "std:put => expected exactly 1 argument, got {}",
                args.len()
            ));
        }
        let val = args.remove(0);
        match &val {
            Value::Iterator(handle) => {
                drain_iterator(interp, handle, false)?;
                Ok(val)
            }
            other => {
                let s = to_printable(other);
                interp.push_print(&s);
                Ok(val)
            }
        }
    }

    // std:log(x)
    fn std_log_bridge(interp: &mut Interpreter, mut args: Vec<Value>) -> Result<Value, String> {
        if args.len() != 1 {
            return Err(format!(
                "std:log => expected exactly 1 argument, got {}",
                args.len()
            ));
        }
        let val = args.remove(0);
        match &val {
            Value::Iterator(handle) => {
                drain_iterator(interp, handle, true)?;
                Ok(val)
            }
            other => {
                let s = to_printable(other);
                interp.push_print(&(s + "\n"));
                Ok(val)
            }
        }
    }

    // Deep string (slightly condensed; uses core value_to_string for stability)
    fn std_deep_bridge(interp: &mut Interpreter, mut args: Vec<Value>) -> Result<Value, String> {
        if args.len() != 1 {
            return Err(format!(
                "std:deep => expected exactly 1 argument, got {}",
                args.len()
            ));
        }
        let val = args.remove(0);
        let s = core_mumu::value_to_string(&val);
        interp.push_print(&(s.clone() + "\n"));
        Ok(val)
    }

    reg!("std:put", std_put_bridge);
    reg!("std:log", std_log_bridge);
    reg!("std:deep", std_deep_bridge);
}

impl WasmMumu {
    /// Gather all known symbol names (dynamic functions + top-level vars).
    fn collect_symbols(&self) -> Vec<String> {
        let mut set: BTreeSet<String> = BTreeSet::new();

        // Dynamic functions (bridges)
        for k in self.interp.get_dynamic_functions_for_clone().keys() {
            set.insert(k.clone());
        }
        // Variables currently bound in the scopes
        for k in self.interp.get_variables().keys() {
            set.insert(k.clone());
        }

        set.into_iter().collect()
    }
}

#[wasm_bindgen]
impl WasmMumu {
    #[wasm_bindgen(constructor)]
    pub fn new(verbose: bool) -> WasmMumu {
        WasmMumu {
            interp: bootstrap_interpreter(verbose),
            verbose,
            drained_cursor: 0,
        }
    }

    /// Recreate the interpreter in-place (alternative to constructing a new WasmMumu on the JS side).
    #[wasm_bindgen]
    pub fn reset(&mut self) {
        self.interp = bootstrap_interpreter(self.verbose);
        self.drained_cursor = 0;
    }

    /// Execute a full MuMu/Lava source string and return a structured result.
    /// JS receives an object: { ok, value, value_string, value_type, prints }.
    #[wasm_bindgen]
    pub fn exec(&mut self, source: &str) -> JsValue {
        let res = (|| -> Result<ExecOk, ExecErr> {
            let tokens = tokenize(source, self.interp.is_verbose())
                .map_err(|e| ExecErr {
                    ok: false,
                    error: e.to_string(),
                })?;
            let ast = parse_tokens(&tokens, self.interp.is_verbose())
                .map_err(|e| ExecErr {
                    ok: false,
                    error: e.to_string(),
                })?;

            let mut last = Value::Bool(true);
            for stmt in &ast {
                last = self
                    .interp
                    .exec_statement(stmt)
                    .map_err(|e| ExecErr { ok: false, error: e })?;
            }

            let view = to_js_view(&last);
            let value_string = core_mumu::value_to_string(&last);
            let value_type = short_type_name(&last).to_string();
            let prints = self.interp.get_prints().clone();

            Ok(ExecOk {
                ok: true,
                value: view,
                value_string,
                value_type,
                prints,
            })
        })();

        match res {
            Ok(ok) => swb::to_value(&ok).unwrap(),
            Err(err) => swb::to_value(&err).unwrap(),
        }
    }

    /// Drive any non-blocking iterator printers once (slog/sput with Iterator).
    /// Returns number of items processed this tick.
    #[wasm_bindgen]
    pub fn poll(&mut self) -> u32 {
        self.interp.poll_all() as u32
    }

    /// Return **all** known symbols (functions + variables) for client-side completion menus.
    #[wasm_bindgen]
    pub fn symbols(&self) -> JsValue {
        let v = self.collect_symbols();
        swb::to_value(&v).unwrap()
    }

    /// Return completions for a case-insensitive `prefix`, sorted and deduped (REPL-like).
    #[wasm_bindgen]
    pub fn complete(&self, prefix: &str) -> JsValue {
        let lower = prefix.to_ascii_lowercase();
        let mut out: Vec<String> = self
            .collect_symbols()
            .into_iter()
            .filter(|s| s.to_ascii_lowercase().starts_with(&lower))
            .collect();
        out.sort();
        out.dedup();
        swb::to_value(&out).unwrap()
    }

    /// Drain and return only the *new* lines printed since the previous drain.
    /// (If called for the first time, returns everything printed so far.)
    #[wasm_bindgen]
    pub fn drain_prints(&mut self) -> JsValue {
        let all = self.interp.get_prints().clone();
        let start = if self.drained_cursor <= all.len() {
            self.drained_cursor
        } else {
            0
        };
        let new: Vec<String> = all[start..].to_vec();
        self.drained_cursor = all.len();
        swb::to_value(&new).unwrap()
    }

    /// Enable/disable verbose parser/interpreter tracing.
    #[wasm_bindgen]
    pub fn set_verbose(&mut self, v: bool) {
        self.verbose = v;
        self.interp.set_verbose(v);
    }

    /// Return a simple version string from the core crate (if any) or a static label.
    #[wasm_bindgen]
    pub fn version(&self) -> String {
        // If core_mumu exposes a version at runtime in future, plumb it here.
        "wasm-mumu/0.1".to_string()
    }
}