Make Param be Copy and mark exceptions raised from Ruby (#158)

Author: ianks

.github/workflows/memcheck.yml

@@ -7,13 +7,18 @@ on:
       ruby-version:
         description: "Ruby version to memcheck"
         required: true
-        default: "3.1"
+        default: "3.2"
         type: choice
         options:
           - "head"
           - "3.2"
           - "3.1"
           - "3.0"
+      debug:
+        description: "Enable debug mode"
+        required: false
+        default: "false"
+        type: boolean
   push:
     branches: ["*"]
     tags-ignore: ["v*"] # Skip Memcheck for releases
@@ -49,6 +54,8 @@ jobs:
           RSPEC_FORMATTER: "progress"
           RSPEC_FAILURE_EXIT_CODE: "0"
           GC_AT_EXIT: "1"
+          DEBUG: ${{ inputs.debug || 'false' }}
+          RB_SYS_CARGO_PROFILE: ${{ inputs.debug == 'true' && 'dev' || 'release' }}
           WASMTIME_TARGET: "x86_64-unknown-linux-gnu" # use generic target for memcheck
         run: |
           if ! bundle exec rake mem:check; then

Cargo.lock

@@ -2,19 +2,24 @@
 
 Bench.ips do |x|
   engine = Wasmtime::Engine.new
-  mod = Wasmtime::Module.new(engine, <<~WAT)
-    (module
-      (import "host" "succ" (func (param i32) (result i32)))
-      (export "run" (func 0)))
-  WAT
-  linker = Wasmtime::Linker.new(engine)
-  linker.func_new("host", "succ", [:i32], [:i32]) do |_caller, arg1|
-    arg1.succ
-  end
+  [4, 16, 64, 128, 256].each do |n|
+    result_type_wat = Array.new(n) { |_| :i32 }.join(" ")
+    mod = Wasmtime::Module.new(engine, <<~WAT)
+      (module
+        (import "host" "succ" (func (param i32) (result #{result_type_wat})))
+        (export "run" (func 0)))
+    WAT
+    linker = Wasmtime::Linker.new(engine)
+    results = Array.new(n) { |_| :i32 }
+    result_array = Array.new(n) { |i| i }
+    linker.func_new("host", "succ", [:i32], results) do |_caller, arg1|
+      result_array
+    end
 
-  x.report("Call host func") do
-    store = Wasmtime::Store.new(engine)
-    linker.instantiate(store, mod).invoke("run", 101)
+    x.report("Call host func (#{n} args)") do
+      store = Wasmtime::Store.new(engine)
+      linker.instantiate(store, mod).invoke("run", 101)
+    end
   end
 
   x.compare!

bench/host_call.rb

@@ -27,6 +27,7 @@ anyhow = "*" # Use whatever Wasmtime uses
 wat = "1.0.59"
 tokio = { version = "1.25.0", features = ["rt", "rt-multi-thread", "time", "net"], optional = true }
 async-timer = { version = "1.0.0-beta.8", features = ["tokio1"], optional = true }
+static_assertions = "1.1.0"
 
 [build-dependencies]
 rb-sys-env = "0.1.2"

ext/Cargo.toml

@@ -59,11 +59,11 @@ impl ToRubyValue for Val {
     }
 }
 pub trait ToWasmVal {
-    fn to_wasm_val(&self, ty: &ValType) -> Result<Val, Error>;
+    fn to_wasm_val(&self, ty: ValType) -> Result<Val, Error>;
 }
 
 impl ToWasmVal for Value {
-    fn to_wasm_val(&self, ty: &ValType) -> Result<Val, Error> {
+    fn to_wasm_val(&self, ty: ValType) -> Result<Val, Error> {
         match ty {
             ValType::I32 => Ok(i32::try_convert(*self)?.into()),
             ValType::I64 => Ok(i64::try_convert(*self)?.into()),

ext/src/ruby_api/convert.rs

@@ -208,6 +208,21 @@ impl From<&Func<'_>> for wasmtime::Extern {
     }
 }
 
+macro_rules! caller_error {
+    ($store:expr, $caller:expr, $error:expr) => {{
+        $store.set_last_error($error);
+        $caller.expire();
+        Err(anyhow::anyhow!(""))
+    }};
+}
+
+macro_rules! result_error {
+    ($store:expr, $caller:expr, $msg:expr) => {{
+        let error = Error::new(result_error(), $msg);
+        caller_error!($store, $caller, error)
+    }};
+}
+
 pub fn make_func_closure(
     ty: &wasmtime::FuncType,
     callable: Proc,
@@ -216,6 +231,11 @@ pub fn make_func_closure(
     let ty = ty.to_owned();
     let callable = ShareableProc(callable);
 
+    // The error handling here is a bit tricky. We want to return a Ruby exception,
+    // but doing so directly can easily cause an early Ruby GC and segfault. So to
+    // be safe, we store all Ruby errors on the store context so it can be marked.
+    // We then return a generic error here. The caller will check for a stored error
+    // and raise it if it exists.
     move |caller_impl: CallerImpl<'_, StoreData>, params: &[Val], results: &mut [Val]| {
         let wrapped_caller = Obj::wrap(Caller::new(caller_impl));
         let store_context = StoreContextValue::from(wrapped_caller);
@@ -232,49 +252,59 @@ pub fn make_func_closure(
 
         let callable = callable.0;
 
-        let result = callable
-            .call(unsafe { rparams.as_slice() })
-            .and_then(|proc_result| {
-                match results.len() {
-                    0 => Ok(()), // Ignore return value
-                    n => {
-                        // For len=1, accept both `val` and `[val]`
-                        let proc_result = RArray::to_ary(proc_result)?;
-                        if proc_result.len() != n {
-                            return Err(Error::new(
-                                result_error(),
-                                format!(
-                                    "wrong number of results (given {}, expected {}) in {}",
-                                    proc_result.len(),
-                                    n,
-                                    callable,
-                                ),
-                            ));
-                        }
-                        for (i, ((rb_val, wasm_val), ty)) in unsafe { proc_result.as_slice() }
-                            .iter()
-                            .zip(results.iter_mut())
-                            .zip(ty.results())
-                            .enumerate()
-                        {
-                            *wasm_val = rb_val.to_wasm_val(&ty).map_err(|e| {
-                                Error::new(
-                                    result_error(),
-                                    format!("{e} (result index {i} in {callable})"),
-                                )
-                            })?;
+        match (callable.call(unsafe { rparams.as_slice() }), results.len()) {
+            (Ok(_proc_result), 0) => {
+                wrapped_caller.get().expire();
+                Ok(())
+            }
+            (Ok(proc_result), n) => {
+                // For len=1, accept both `val` and `[val]`
+                let Ok(proc_result) = RArray::to_ary(proc_result) else {
+                    return result_error!(
+                        store_context,
+                        wrapped_caller.get(),
+                        format!("could not convert {} to results array", callable)
+                    );
+                };
+
+                if proc_result.len() != results.len() {
+                    return result_error!(
+                        store_context,
+                        wrapped_caller.get(),
+                        format!(
+                            "wrong number of results (given {}, expected {}) in {}",
+                            proc_result.len(),
+                            n,
+                            callable
+                        )
+                    );
+                }
+
+                for (i, ((rb_val, wasm_val), ty)) in unsafe { proc_result.as_slice() }
+                    .iter()
+                    .zip(results.iter_mut())
+                    .zip(ty.results())
+                    .enumerate()
+                {
+                    match rb_val.to_wasm_val(ty) {
+                        Ok(val) => *wasm_val = val,
+                        Err(e) => {
+                            return result_error!(
+                                store_context,
+                                wrapped_caller.get(),
+                                format!("invalid result at index {i}: {e} in {callable}")
+                            );
                         }
-                        Ok(())
                     }
                 }
-            })
-            .map_err(|e| anyhow::anyhow!(e));
-
-        // Drop the wasmtime::Caller so it does not outlive the Func call, if e.g. the user
-        // assigned the Ruby Wasmtime::Caller instance to a global.
-        wrapped_caller.get().expire();
 
-        result
+                wrapped_caller.get().expire();
+                Ok(())
+            }
+            (Err(e), _) => {
+                caller_error!(store_context, wrapped_caller.get(), e)
+            }
+        }
     }
 }
 

ext/src/ruby_api/func.rs

@@ -55,7 +55,7 @@ impl<'a> Global<'a> {
         mutability: Mutability,
     ) -> Result<Self, Error> {
         let wasm_type = value_type.to_val_type()?;
-        let wasm_default = default.to_wasm_val(&wasm_type)?;
+        let wasm_default = default.to_wasm_val(wasm_type.clone())?;
         let store = s.get();
         let inner = GlobalImpl::new(
             store.context_mut(),
@@ -116,7 +116,7 @@ impl<'a> Global<'a> {
         self.inner
             .set(
                 self.store.context_mut()?,
-                value.to_wasm_val(&self.value_type()?)?,
+                value.to_wasm_val(self.value_type()?)?,
             )
             .map_err(|e| error!("{}", e))
             .and_then(|result| {

ext/src/ruby_api/global.rs

@@ -1,30 +1,41 @@
 use super::convert::ToWasmVal;
 use magnus::{exception::arg_error, Error, ExceptionClass, Value};
+use static_assertions::assert_eq_size;
 use wasmtime::{FuncType, ValType};
 
-#[derive(Debug)]
-struct Param<'a> {
-    index: usize,
-    ty: ValType,
-    val: &'a Value,
+#[derive(Debug, Copy, Clone)]
+#[repr(C)]
+struct Param {
+    val: Value,
+    index: u32,
+    ty: ValTypeCopy,
 }
+// Keep `Param` small so copying it to the stack is cheap, typically anything
+// less than 3usize is good
+assert_eq_size!(Param, [u64; 2]);
 
-impl<'a> Param<'a> {
-    pub fn new(index: usize, ty: ValType, val: &'a Value) -> Self {
-        Self { index, ty, val }
+impl Param {
+    pub fn new(index: u32, ty: ValType, val: Value) -> Self {
+        Self {
+            index,
+            ty: ty.into(),
+            val,
+        }
     }
 
-    fn to_wasmtime_val(&self) -> Result<wasmtime::Val, Error> {
-        self.val.to_wasm_val(&self.ty).map_err(|error| match error {
-            Error::Error(class, msg) => {
-                Error::new(class, format!("{} (param index {}) ", msg, self.index))
-            }
-            Error::Exception(exception) => Error::new(
-                ExceptionClass::from_value(exception.class().into()).unwrap_or_else(arg_error),
-                format!("{} (param index {}) ", exception, self.index),
-            ),
-            _ => error,
-        })
+    fn to_wasmtime_val(self) -> Result<wasmtime::Val, Error> {
+        self.val
+            .to_wasm_val(self.ty.into())
+            .map_err(|error| match error {
+                Error::Error(class, msg) => {
+                    Error::new(class, format!("{} (param at index {})", msg, self.index))
+                }
+                Error::Exception(exception) => Error::new(
+                    ExceptionClass::from_value(exception.class().into()).unwrap_or_else(arg_error),
+                    format!("{} (param at index {})", exception, self.index),
+                ),
+                _ => error,
+            })
     }
 }
 
@@ -48,10 +59,55 @@ impl<'a> Params<'a> {
     pub fn to_vec(&self) -> Result<Vec<wasmtime::Val>, Error> {
         let mut vals = Vec::with_capacity(self.0.params().len());
         for (i, (param, value)) in self.0.params().zip(self.1.iter()).enumerate() {
-            let param = Param::new(i, param.clone(), value);
+            let i: u32 = i
+                .try_into()
+                .map_err(|_| Error::new(arg_error(), "too many params"))?;
+            let param = Param::new(i, param.clone(), *value);
             vals.push(param.to_wasmtime_val()?);
         }
 
         Ok(vals)
     }
 }
+
+/// A [`wasmtime::ValType`] that is [`Copy`], so it can be stays on the stack
+///
+/// Note: this can be removed in Wasmtime 8.0 (see https://github.com/bytecodealliance/wasmtime/pull/6138)
+#[derive(Debug, Clone, Copy)]
+pub enum ValTypeCopy {
+    I32,
+    I64,
+    F32,
+    F64,
+    V128,
+    FuncRef,
+    ExternRef,
+}
+
+impl From<ValType> for ValTypeCopy {
+    fn from(ty: ValType) -> Self {
+        match ty {
+            ValType::I32 => Self::I32,
+            ValType::I64 => Self::I64,
+            ValType::F32 => Self::F32,
+            ValType::F64 => Self::F64,
+            ValType::V128 => Self::V128,
+            ValType::FuncRef => Self::FuncRef,
+            ValType::ExternRef => Self::ExternRef,
+        }
+    }
+}
+
+impl From<ValTypeCopy> for ValType {
+    fn from(ty: ValTypeCopy) -> Self {
+        match ty {
+            ValTypeCopy::I32 => Self::I32,
+            ValTypeCopy::I64 => Self::I64,
+            ValTypeCopy::F32 => Self::F32,
+            ValTypeCopy::F64 => Self::F64,
+            ValTypeCopy::V128 => Self::V128,
+            ValTypeCopy::FuncRef => Self::FuncRef,
+            ValTypeCopy::ExternRef => Self::ExternRef,
+        }
+    }
+}