nms/

interpreter.rs

use scripting::host::{WasmHost, define_host_functions};
use scripting::nomiscript::{
    Compiler, Expr, Fraction, GIT_REVISION, Program, Reader, Symbol, SymbolKind, SymbolTable, Value,
};
use scripting::runtime::{EngineOpts, ProfilerStrategy, build_engine};
use scripting_format::{
    ContextType, DEBUG_VALUE_DATA_SIZE, DebugValueData, ENTITY_HEADER_SIZE, EntityType,
    GlobalHeader, OUTPUT_HEADER_SIZE, OutputHeader, ValueType,
};
use thiserror::Error;
use tracing::debug;
use tracing_subscriber::EnvFilter;
use tracing_subscriber::prelude::*;
use tracing_subscriber::reload;
use wasmtime::{Linker, Module, Store};

#[derive(Error, Debug)]
pub enum Error {
    #[error("{0}")]
    Script(#[from] scripting::nomiscript::Error),

    #[error("{0}")]
    Runtime(String),
}

impl Error {
    #[must_use]
    pub fn render(&self, use_color: bool) -> String {
        match self {
            Error::Script(e) => e.render(use_color),
            Error::Runtime(msg) => {
                if use_color {
                    format!("\x1b[31merror:\x1b[0m {msg}")
                } else {
                    format!("error: {msg}")
                }
            }
        }
    }
}

pub type Result<T> = std::result::Result<T, Error>;

pub struct Interpreter {
    host: WasmHost,
    compiler: Compiler,
    reload_handle: reload::Handle<EnvFilter, tracing_subscriber::Registry>,
}

const DEFAULT_OUTPUT_SIZE: u32 = 64 * 1024;
const WASM_PAGE_SIZE: u32 = 65536;

/// `build_engine` turns on epoch interruption, so every store needs a
/// deadline or it traps at the default tick 0. The interactive
/// interpreter has no wall-clock budget to enforce (unlike the rpc
/// `Session`, which spawns an epoch bumper) — it just runs the user's
/// REPL input to completion. A deadline that's never advanced past
/// gives the old no-epoch-limit behaviour back.
fn set_interactive_epoch_deadline<T>(store: &mut Store<T>) {
    store.set_epoch_deadline(u64::MAX);
}

impl Interpreter {
    pub fn new(debug_mode: bool) -> anyhow::Result<Self> {
        Self::with_profiler(debug_mode, ProfilerStrategy::None)
    }

    pub fn with_profiler(debug_mode: bool, profiler: ProfilerStrategy) -> anyhow::Result<Self> {
        // Route through `build_engine` so the wasm feature set (GC,
        // function-references, exceptions) stays single-sourced — a flag
        // added there can't silently miss the nms interpreter path.
        let engine = build_engine(EngineOpts::baseline().with_profiler(profiler))?;

        let mut symbols = SymbolTable::with_builtins();
        symbols.define(
            Symbol::new("REVISION", SymbolKind::Variable)
                .with_value(Expr::String(GIT_REVISION.to_string())),
        );

        let default_filter = if debug_mode { "debug" } else { "warn" };
        let filter =
            EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new(default_filter));
        let (filter_layer, reload_handle) = reload::Layer::new(filter);
        tracing_subscriber::registry()
            .with(filter_layer)
            .with(tracing_subscriber::fmt::layer().with_writer(std::io::stderr))
            .try_init()
            .ok();

        let mut interp = Self {
            host: WasmHost::new(engine, symbols),
            compiler: Compiler::new(),
            reload_handle,
        };
        interp.load_stdlib()?;
        Ok(interp)
    }

    fn load_stdlib(&mut self) -> Result<()> {
        const STDLIB: &str = include_str!("stdlib.lisp");
        let program = Reader::parse(STDLIB)?;
        let mut symbols = self
            .host
            .symbol_table()
            .write()
            .map_err(|e| Error::Runtime(format!("failed to write symbol table: {e}")))?;
        self.compiler.compile(&program, &mut symbols)?;
        Ok(())
    }

    pub fn eval(&mut self, input: &str) -> Result<Vec<Value>> {
        let program = Reader::parse(input)?;
        let debug_mode = program.annotations.iter().any(|a| a.name == "debug");
        if debug_mode {
            self.reload_handle
                .modify(|f| *f = EnvFilter::new("debug"))
                .ok();
        }
        let result = self.eval_program(&program);
        if debug_mode {
            self.reload_handle
                .modify(|f| {
                    *f = EnvFilter::try_from_default_env()
                        .unwrap_or_else(|_| EnvFilter::new("warn"));
                })
                .ok();
        }
        result
    }

    fn eval_program(&mut self, program: &Program) -> Result<Vec<Value>> {
        if program.exprs.is_empty() {
            return Ok(vec![]);
        }

        let wasm = {
            let mut symbols = self
                .host
                .symbol_table()
                .write()
                .map_err(|e| Error::Runtime(format!("failed to write symbol table: {e}")))?;
            self.compiler.compile(program, &mut symbols)?
        };

        let value = self.run_wasm(&wasm)?;
        Ok(vec![value])
    }

    #[must_use]
    pub fn struct_fields(&self, name: &str) -> Option<Vec<String>> {
        self.host
            .symbol_table()
            .read()
            .ok()
            .and_then(|st| st.struct_fields(name).map(<[std::string::String]>::to_vec))
    }

    pub fn compile_to_wasm(&mut self, input: &str) -> Result<Vec<u8>> {
        let program = Reader::parse(input)?;
        if program.exprs.is_empty() {
            return Err(Error::Runtime("nothing to compile".to_string()));
        }
        let mut symbols = self
            .host
            .symbol_table()
            .write()
            .map_err(|e| Error::Runtime(format!("failed to write symbol table: {e}")))?;
        Ok(self.compiler.compile(&program, &mut symbols)?)
    }

    pub fn run_wasm(&self, wasm: &[u8]) -> Result<Value> {
        let input = build_minimal_input(DEFAULT_OUTPUT_SIZE);
        self.run_wasm_with_input(wasm, &input)
    }

    pub fn run_wasm_with_input(&self, wasm: &[u8], input: &[u8]) -> Result<Value> {
        debug!(wasm_size = wasm.len(), "creating WASM module");
        let module =
            Module::new(self.host.engine(), wasm).map_err(|e| Error::Runtime(e.to_string()))?;

        let output_size = DEFAULT_OUTPUT_SIZE;
        let input_offset = scripting_format::BASE_OFFSET;
        let output_offset = input_offset + input.len() as u32;
        let strings_offset = {
            let header = GlobalHeader::from_bytes(input).expect("minimal input must be valid");
            header.strings_pool_offset
        };
        debug!(
            input_offset,
            output_offset, strings_offset, "memory layout offsets"
        );

        let exec_state = self
            .host
            .execution_state(input_offset, output_offset, strings_offset);
        let mut store = Store::new(self.host.engine(), exec_state);
        set_interactive_epoch_deadline(&mut store);

        let mut linker = Linker::new(self.host.engine());
        define_host_functions(&mut linker).map_err(|e| Error::Runtime(e.to_string()))?;

        let instance = linker
            .instantiate(&mut store, &module)
            .map_err(|e| Error::Runtime(e.to_string()))?;

        let memory = instance
            .get_memory(&mut store, "memory")
            .ok_or_else(|| Error::Runtime("no memory export".to_string()))?;

        store.data_mut().memory = Some(memory);

        let total_size = input.len() + output_size as usize;
        let required_pages = (input_offset as usize + total_size).div_ceil(WASM_PAGE_SIZE as usize);
        let current_pages = memory.size(&store) as usize;

        if required_pages > current_pages {
            debug!(current_pages, required_pages, "growing memory");
            memory
                .grow(&mut store, (required_pages - current_pages) as u64)
                .map_err(|e| Error::Runtime(e.to_string()))?;
        }

        let mem_data = memory.data_mut(&mut store);
        let input_start = input_offset as usize;
        mem_data[input_start..input_start + input.len()].copy_from_slice(input);

        let output_start = output_offset as usize;
        let output_header = OutputHeader::new(0);
        mem_data[output_start..output_start + OUTPUT_HEADER_SIZE]
            .copy_from_slice(&output_header.to_bytes());

        let should_apply = instance
            .get_typed_func::<(), i32>(&mut store, "should_apply")
            .map_err(|e| Error::Runtime(e.to_string()))?;
        debug!("calling should_apply");
        let apply = should_apply
            .call(&mut store, ())
            .map_err(|e| Error::Runtime(e.to_string()))?;
        debug!(result = apply, "should_apply returned");
        if apply != 1 {
            return Err(Error::Runtime(format!(
                "should_apply returned {apply}, expected 1"
            )));
        }

        let process = instance
            .get_typed_func::<(), ()>(&mut store, "process")
            .map_err(|e| Error::Runtime(e.to_string()))?;
        debug!("calling process");
        process
            .call(&mut store, ())
            .map_err(|e| Error::Runtime(e.to_string()))?;

        let mem_data = memory.data(&store);
        let output_data = &mem_data[output_start..];

        let result = decode_result(output_data);
        debug!("result decoded");
        result
    }

    pub fn run_wasm_with_input_raw(&self, wasm: &[u8], input: &[u8]) -> Result<Vec<u8>> {
        let module =
            Module::new(self.host.engine(), wasm).map_err(|e| Error::Runtime(e.to_string()))?;

        let output_size = DEFAULT_OUTPUT_SIZE;
        let input_offset = scripting_format::BASE_OFFSET;
        let output_offset = input_offset + input.len() as u32;
        let strings_offset = {
            let header = GlobalHeader::from_bytes(input).expect("minimal input must be valid");
            header.strings_pool_offset
        };

        let exec_state = self
            .host
            .execution_state(input_offset, output_offset, strings_offset);
        let mut store = Store::new(self.host.engine(), exec_state);
        set_interactive_epoch_deadline(&mut store);

        let mut linker = Linker::new(self.host.engine());
        define_host_functions(&mut linker).map_err(|e| Error::Runtime(e.to_string()))?;

        let instance = linker
            .instantiate(&mut store, &module)
            .map_err(|e| Error::Runtime(e.to_string()))?;

        let memory = instance
            .get_memory(&mut store, "memory")
            .ok_or_else(|| Error::Runtime("no memory export".to_string()))?;

        store.data_mut().memory = Some(memory);

        let total_size = input.len() + output_size as usize;
        let required_pages = (input_offset as usize + total_size).div_ceil(WASM_PAGE_SIZE as usize);
        let current_pages = memory.size(&store) as usize;

        if required_pages > current_pages {
            memory
                .grow(&mut store, (required_pages - current_pages) as u64)
                .map_err(|e| Error::Runtime(e.to_string()))?;
        }

        let mem_data = memory.data_mut(&mut store);
        let input_start = input_offset as usize;
        mem_data[input_start..input_start + input.len()].copy_from_slice(input);

        let output_start = output_offset as usize;
        let output_header = OutputHeader::new(0);
        mem_data[output_start..output_start + OUTPUT_HEADER_SIZE]
            .copy_from_slice(&output_header.to_bytes());

        let should_apply = instance
            .get_typed_func::<(), i32>(&mut store, "should_apply")
            .map_err(|e| Error::Runtime(e.to_string()))?;
        let apply = should_apply
            .call(&mut store, ())
            .map_err(|e| Error::Runtime(e.to_string()))?;
        if apply != 1 {
            return Err(Error::Runtime(format!(
                "should_apply returned {apply}, expected 1"
            )));
        }

        let process = instance
            .get_typed_func::<(), ()>(&mut store, "process")
            .map_err(|e| Error::Runtime(e.to_string()))?;
        process
            .call(&mut store, ())
            .map_err(|e| Error::Runtime(e.to_string()))?;

        let mem_data = memory.data(&store);
        Ok(mem_data[output_start..output_start + output_size as usize].to_vec())
    }
}

impl Default for Interpreter {
    fn default() -> Self {
        Self::new(false).expect("failed to create Interpreter")
    }
}

fn build_minimal_input(output_size: u32) -> Vec<u8> {
    use scripting::MemorySerializer;
    let mut ser = MemorySerializer::new();
    ser.set_context(ContextType::BatchProcess, EntityType::Transaction);
    ser.finalize(output_size)
}

fn decode_result(data: &[u8]) -> Result<Value> {
    let output_header = OutputHeader::from_bytes(data)
        .ok_or_else(|| Error::Runtime("invalid output header".to_string()))?;

    if output_header.output_entity_count == 0 {
        return Err(Error::Runtime("no output entities".to_string()));
    }

    // The eval result is the program's tail value — a `DebugValue` entity. Other
    // entities (a side-effecting `create-tag`) may precede it in the same output
    // stream, so walk to the DebugValue rather than assuming entity 0. Its data
    // is at `data_offset`; any string payload is inline right after the fixed
    // `DebugValueData` (the writer counts those bytes in `data_size`), so it is
    // resolved per-entity, never via the shared `strings_offset` the entity
    // parser uses for tag strings.
    let mut offset = OUTPUT_HEADER_SIZE;
    let mut found: Option<usize> = None;
    for _ in 0..output_header.output_entity_count {
        let header = scripting_format::EntityHeader::from_bytes(&data[offset..])
            .ok_or_else(|| Error::Runtime("invalid entity header".to_string()))?;
        if header.entity_type == EntityType::DebugValue as u8 {
            found = Some(header.data_offset as usize);
        }
        offset += ENTITY_HEADER_SIZE + header.data_size as usize;
    }
    let data_offset =
        found.ok_or_else(|| Error::Runtime("no debug value in output".to_string()))?;

    let value_data = DebugValueData::from_bytes(&data[data_offset..])
        .ok_or_else(|| Error::Runtime("invalid debug value data".to_string()))?;

    let value_type = ValueType::try_from(value_data.value_type)
        .map_err(|()| Error::Runtime("unknown value type".to_string()))?;

    match value_type {
        ValueType::Nil => Ok(Value::Nil),
        ValueType::Bool => Ok(Value::Bool(value_data.data1 != 0)),
        ValueType::Number => Ok(Value::Number(Fraction::new(
            value_data.data1,
            value_data.data2,
        ))),
        ValueType::String | ValueType::Symbol => {
            // String bytes are inline right after DebugValueData, at
            // `data_offset + DEBUG_VALUE_DATA_SIZE + data1` (data1 is 0 today).
            let start = data_offset + DEBUG_VALUE_DATA_SIZE + value_data.data1 as usize;
            let len = value_data.data2 as usize;
            let end = start + len;
            if end > data.len() {
                return Err(Error::Runtime("string data out of bounds".to_string()));
            }
            let s = std::str::from_utf8(&data[start..end])
                .map_err(|_| Error::Runtime("invalid UTF-8 in string".to_string()))?;
            if value_type == ValueType::Symbol {
                Ok(Value::Symbol(s.to_string()))
            } else {
                Ok(Value::String(s.to_string()))
            }
        }
    }
}