From 6a391b1ef5278042424089f339a2e8d69c3dac9c Mon Sep 17 00:00:00 2001 From: Owais Jamil Date: Thu, 18 Jun 2026 22:10:06 -0500 Subject: [PATCH] feat: export arena helpers for managed Wasmtime values --- crates/cli/src/commands.rs | 100 +++++++++++++++++- crates/cli/src/commands/runner.rs | 16 ++- crates/cli/tests/build.rs | 24 +++++ crates/core/src/wasm/codegen.rs | 28 +++++ .../core/src/wasm/fragments/allocation.wat.rs | 4 +- docs/website/development/runtime-memory.md | 5 + 6 files changed, 169 insertions(+), 8 deletions(-) diff --git a/crates/cli/src/commands.rs b/crates/cli/src/commands.rs index 93fb527..a9e6702 100644 --- a/crates/cli/src/commands.rs +++ b/crates/cli/src/commands.rs @@ -9,6 +9,7 @@ use std::path::{Path, PathBuf}; use std::process::ExitCode; use compiler_core::source::SourceFile; +use compiler_core::types::Type; use compiler_core::{diagnostic::Diagnostics, target::CompileTarget}; use super::args::{Command, DebugCommand}; @@ -107,7 +108,7 @@ fn list(input: &Path) -> ExitCode { } } -fn run_wasm_export(bytes: &[u8], function: &str, args: &[String]) -> Result<(), String> { +fn run_wasm_export(bytes: &[u8], function: &str, args: &[String], return_type: Option<&Type>) -> Result<(), String> { use wasmtime::{Caller, Engine, Linker, Module, Store}; let engine = Engine::default(); @@ -156,15 +157,38 @@ fn run_wasm_export(bytes: &[u8], function: &str, args: &[String]) -> Result<(), .map(|(arg, type_)| parse_wasm_arg(arg, type_)) .collect::, _>>()?; let mut result_values = results.iter().map(default_wasm_value).collect::, _>>()?; - func.call(&mut store, &values, &mut result_values) - .map_err(|error| error.to_string())?; + let arena_mark = instance.get_typed_func::<(), i32>(&mut store, "__arena_mark").ok(); + let arena_reset = instance.get_typed_func::(&mut store, "__arena_reset").ok(); + let mark = match &arena_mark { + Some(mark) => Some(mark.call(&mut store, ()).map_err(|error| error.to_string())?), + None => None, + }; + let call_result = func + .call(&mut store, &values, &mut result_values) + .map_err(|error| error.to_string()) + .and_then(|_| format_wasm_results(&instance, &mut store, &result_values, return_type)); + let reset_result = reset_arena(&mut store, arena_reset.as_ref(), mark); + let formatted = match (call_result, reset_result) { + (Ok(formatted), Ok(())) => formatted, + (Err(error), Ok(())) => return Err(error), + (Ok(_), Err(error)) | (Err(_), Err(error)) => return Err(error), + }; - for result in result_values { - println!("{}", format_wasm_value(&result)); + for result in formatted { + println!("{result}"); } Ok(()) } +fn reset_arena( + store: &mut wasmtime::Store<()>, arena_reset: Option<&wasmtime::TypedFunc>, mark: Option, +) -> Result<(), String> { + let (Some(reset), Some(mark)) = (arena_reset, mark) else { + return Ok(()); + }; + reset.call(store, mark).map_err(|error| error.to_string()) +} + fn parse_wasm_arg(arg: &str, type_: &wasmtime::ValType) -> Result { use wasmtime::{Val, ValType}; @@ -203,6 +227,72 @@ fn format_wasm_value(value: &wasmtime::Val) -> String { } } +fn format_wasm_results( + instance: &wasmtime::Instance, store: &mut wasmtime::Store<()>, values: &[wasmtime::Val], + return_type: Option<&Type>, +) -> Result, String> { + if values.len() == 1 + && let Some(type_) = return_type + && let Some(value) = values.first() + && let Some(formatted) = format_typed_wasm_value(instance, store, value, type_)? + { + return Ok(vec![formatted]); + } + Ok(values.iter().map(format_wasm_value).collect()) +} + +fn format_typed_wasm_value( + instance: &wasmtime::Instance, store: &mut wasmtime::Store<()>, value: &wasmtime::Val, type_: &Type, +) -> Result, String> { + match (type_, value) { + (Type::String, wasmtime::Val::I32(ptr)) => read_memory_string(instance, store, *ptr).map(Some), + ( + Type::BitArray + | Type::Tuple(_) + | Type::List(_) + | Type::Record { .. } + | Type::Custom { .. } + | Type::Opaque { .. } + | Type::Function { .. } + | Type::Generic(_), + wasmtime::Val::I32(ptr), + ) => read_memory_debug(instance, store, *ptr).map(Some), + _ => Ok(None), + } +} + +fn read_memory_string( + instance: &wasmtime::Instance, store: &mut wasmtime::Store<()>, ptr: i32, +) -> Result { + let memory = instance + .get_memory(&mut *store, "memory") + .ok_or_else(|| "missing memory export".to_string())?; + let ptr = ptr as usize; + let mut header = [0; 8]; + memory + .read(&mut *store, ptr, &mut header) + .map_err(|error| error.to_string())?; + if u32::from_le_bytes(header[0..4].try_into().expect("string tag header")) != 1 { + return Err("managed return value is not a string".into()); + } + let len = u32::from_le_bytes(header[4..8].try_into().expect("string length header")) as usize; + let mut bytes = vec![0; len]; + memory + .read(&mut *store, ptr + 8, &mut bytes) + .map_err(|error| error.to_string())?; + String::from_utf8(bytes).map_err(|error| error.to_string()) +} + +fn read_memory_debug( + instance: &wasmtime::Instance, store: &mut wasmtime::Store<()>, ptr: i32, +) -> Result { + let memory = instance + .get_memory(&mut *store, "memory") + .ok_or_else(|| "missing memory export".to_string())?; + let data = memory.data(&mut *store); + compiler_core::runtime::debug_render(data, ptr as u32).ok_or_else(|| "could not decode managed return value".into()) +} + fn read_host_string(caller: &mut wasmtime::Caller<'_, ()>, ptr: i32) -> String { let Some(memory) = caller.get_export("memory").and_then(|export| export.into_memory()) else { return "".into(); diff --git a/crates/cli/src/commands/runner.rs b/crates/cli/src/commands/runner.rs index 3d9bb93..a855122 100644 --- a/crates/cli/src/commands/runner.rs +++ b/crates/cli/src/commands/runner.rs @@ -33,7 +33,21 @@ impl Runner<'_> { Err(diagnostics) => return echo::fail_with_diagnostics("compile", self.input.display(), &diagnostics), }; - match super::run_wasm_export(&compiled.wasm.bytes, self.function, self.args) { + let return_type = compiled + .ir + .exports + .iter() + .find(|export| export.name == self.function) + .and_then(|export| { + compiled + .ir + .functions + .iter() + .find(|function| function.name == export.backend_name()) + }) + .map(|function| &function.return_type); + + match super::run_wasm_export(&compiled.wasm.bytes, self.function, self.args, return_type) { Ok(()) => ExitCode::SUCCESS, Err(message) => echo::fail("run", self.function, message), } diff --git a/crates/cli/tests/build.rs b/crates/cli/tests/build.rs index 7e513fc..de6c945 100644 --- a/crates/cli/tests/build.rs +++ b/crates/cli/tests/build.rs @@ -70,6 +70,30 @@ diagnostic ProjectError: duplicate module `app` in examples/diagnostics/duplicat let _ = fs::remove_dir_all(out_dir); } +#[test] +fn run_decodes_managed_string_return_before_arena_reset() { + let temp = unique_temp_dir("regulus_cli_run_string"); + fs::create_dir_all(&temp).expect("create temp dir"); + let input = temp.join("app.gleam"); + fs::write(&input, r#"pub fn main() -> String { "Ada" <> " Lovelace" }"#).expect("write Gleam input"); + + let output = Command::new(env!("CARGO_BIN_EXE_reggie")) + .arg("run") + .arg(&input) + .output() + .expect("run reggie run"); + + assert!( + output.status.success(), + "run failed\nstdout:\n{}\nstderr:\n{}", + String::from_utf8_lossy(&output.stdout), + String::from_utf8_lossy(&output.stderr) + ); + assert_eq!(String::from_utf8_lossy(&output.stdout), "Ada Lovelace\n"); + + let _ = fs::remove_dir_all(temp); +} + #[test] fn build_emit_writes_wat_and_debug_artifacts_without_wasm() { let fixture = workspace_root().join("fixtures/projects/generated_names/dependency_module_overlap"); diff --git a/crates/core/src/wasm/codegen.rs b/crates/core/src/wasm/codegen.rs index 2dd0ca8..96a34cd 100644 --- a/crates/core/src/wasm/codegen.rs +++ b/crates/core/src/wasm/codegen.rs @@ -268,6 +268,10 @@ impl<'a> StructuredEmitter<'a> { if self.options.target.is_js_host() { self.emit_js_host_abi_helpers(); } + if self.options.target == WasmTarget::Wasmtime && module_exports_arena_scoped_values(source) { + self.runtime_helper_roots.insert("__arena_mark".into()); + self.runtime_helper_roots.insert("__arena_reset".into()); + } if !self.runtime_helper_roots.is_empty() { self.ensure_memory(); @@ -3402,6 +3406,30 @@ fn is_js_host_opaque_handle(module: &ir::Module, type_: &Type) -> bool { } } +fn module_exports_arena_scoped_values(module: &ir::Module) -> bool { + module.functions.iter().any(|function| { + matches!(function.abi.boundary, ir::CallBoundary::ModuleExport) + && needs_allocation(function) + && (is_heap_managed_type(&function.return_type) + || function.params.iter().any(|param| is_heap_managed_type(¶m.type_))) + }) +} + +fn is_heap_managed_type(type_: &Type) -> bool { + matches!( + type_, + Type::String + | Type::BitArray + | Type::Tuple(_) + | Type::List(_) + | Type::Record { .. } + | Type::Custom { .. } + | Type::Opaque { .. } + | Type::Function { .. } + | Type::Generic(_) + ) +} + fn value_type(type_: &Type, span: Span) -> StructuredResult { maybe_value_type(type_).ok_or_else(|| { StructuredError::Diagnostics(vec![ diff --git a/crates/core/src/wasm/fragments/allocation.wat.rs b/crates/core/src/wasm/fragments/allocation.wat.rs index b2cd875..d75c72f 100644 --- a/crates/core/src/wasm/fragments/allocation.wat.rs +++ b/crates/core/src/wasm/fragments/allocation.wat.rs @@ -4,10 +4,10 @@ pub const ALLOC_HELPER: &str = r#" (func $__last_panic (export "__last_panic") (result i32) global.get $__last_panic_payload ) - (func $__arena_mark (result i32) + (func $__arena_mark (export "__arena_mark") (result i32) global.get $__heap ) - (func $__arena_reset (param $mark i32) + (func $__arena_reset (export "__arena_reset") (param $mark i32) local.get $mark i32.const {heap_start} i32.lt_u diff --git a/docs/website/development/runtime-memory.md b/docs/website/development/runtime-memory.md index 0ad0b4b..d1db46d 100644 --- a/docs/website/development/runtime-memory.md +++ b/docs/website/development/runtime-memory.md @@ -129,6 +129,11 @@ return into JS-owned data, and resets in a `finally` block. Raw Wasm and Wasmtime exports are not automatically reset because those callers may inspect borrowed managed pointers after the call. +The CLI `run` command is ABI-aware for managed returns. When arena helpers are +available, it marks before calling an export, decodes the result for display, +and resets before printing. String returns are printed as text. Other managed +returns use the runtime debug renderer. + Compiler-generated code must not return or retain pointers allocated after a mark that will be reset. General internal reset scopes still require escape analysis or region tracking. -- 2.51.2