Change wasm-to-host trampolines to take the values_vec size (#4192)

* Change wasm-to-host trampolines to take the values_vec size

This commit changes the ABI of wasm-to-host trampolines, which are
only used right now for functions created with `Func::new`, to pass
along the size of the `values_vec` argument. Previously the trampoline
simply received `*mut ValRaw` and assumed that it was the appropriate
size. By receiving a size as well we can thread through `&mut [ValRaw]`
internally instead of `*mut ValRaw`.

The original motivation for this is that I'm planning to leverage these
trampolines for the component model for host-defined functions. Out of
an abundance of caution of making sure that everything lines up I wanted
to be able to write down asserts about the size received at runtime
compared to the size expected. This overall led me to the desire to
thread this size parameter through on the assumption that it would not
impact performance all that much.

I ran two benchmarks locally from the `call.rs` benchmark and got:

* `sync/no-hook/wasm-to-host - nop - unchecked` - no change
* `sync/no-hook/wasm-to-host - nop-params-and-results - unchecked` - 5%
  slower

This is what I roughly expected in that if nothing actually reads the
new parameter (e.g. no arguments) then threading through the parameter
is effectively otherwise free. Otherwise though accesses to the `ValRaw`
storage is now bounds-checked internally in Wasmtime instead of assuming
it's valid, leading to the 5% slowdown (~9.6ns to ~10.3ns). If this
becomes a peformance bottleneck for a particular use case then we should
be fine to remove the bounds checking here or otherwise only bounds
check in debug mode, otherwise I plan on leaving this as-is.

Of particular note this also changes the C API for `*_unchecked`
functions where the C callback now receives the size of the array as
well.

* Add docs

benches/call.rs

@@ -314,15 +314,15 @@ fn wasm_to_host(c: &mut Criterion) {
             let ty = FuncType::new([ValType::I32, ValType::I64], [ValType::F32]);
             unchecked
                 .func_new_unchecked("", "nop-params-and-results", ty, |mut caller, space| {
-                    match Val::from_raw(&mut caller, *space, ValType::I32) {
+                    match Val::from_raw(&mut caller, space[0], ValType::I32) {
                         Val::I32(0) => {}
                         _ => unreachable!(),
                     }
-                    match Val::from_raw(&mut caller, *space.add(1), ValType::I64) {
+                    match Val::from_raw(&mut caller, space[1], ValType::I64) {
                         Val::I64(0) => {}
                         _ => unreachable!(),
                     }
-                    *space = Val::F32(0).to_raw(&mut caller);
+                    space[0] = Val::F32(0).to_raw(&mut caller);
                     Ok(())
                 })
                 .unwrap();

crates/c-api/include/wasmtime/func.h

@@ -101,6 +101,8 @@ WASM_API_EXTERN void wasmtime_func_new(
  *        array depends on the function type that the host function is created
  *        with, but it will be the maximum of the number of parameters and
  *        number of results.
+ * \param num_args_and_results the size of the `args_and_results` parameter in
+ *        units of #wasmtime_val_raw_t.
  *
  * This callback can optionally return a #wasm_trap_t indicating that a trap
  * should be raised in WebAssembly. It's expected that in this case the caller
@@ -121,7 +123,8 @@ WASM_API_EXTERN void wasmtime_func_new(
 typedef wasm_trap_t* (*wasmtime_func_unchecked_callback_t)(
     void *env,
     wasmtime_caller_t* caller,
-    wasmtime_val_raw_t *args_and_results);
+    wasmtime_val_raw_t *args_and_results,
+    size_t num_args_and_results);
 
 /**
  * \brief Creates a new host function in the same manner of #wasmtime_func_new,

crates/c-api/src/func.rs

@@ -209,8 +209,12 @@ pub type wasmtime_func_callback_t = extern "C" fn(
     usize,
 ) -> Option<Box<wasm_trap_t>>;
 
-pub type wasmtime_func_unchecked_callback_t =
-    extern "C" fn(*mut c_void, *mut wasmtime_caller_t, *mut ValRaw) -> Option<Box<wasm_trap_t>>;
+pub type wasmtime_func_unchecked_callback_t = extern "C" fn(
+    *mut c_void,
+    *mut wasmtime_caller_t,
+    *mut ValRaw,
+    usize,
+) -> Option<Box<wasm_trap_t>>;
 
 #[no_mangle]
 pub unsafe extern "C" fn wasmtime_func_new(
@@ -295,12 +299,12 @@ pub(crate) unsafe fn c_unchecked_callback_to_rust_fn(
     callback: wasmtime_func_unchecked_callback_t,
     data: *mut c_void,
     finalizer: Option<extern "C" fn(*mut std::ffi::c_void)>,
-) -> impl Fn(Caller<'_, crate::StoreData>, *mut ValRaw) -> Result<(), Trap> {
+) -> impl Fn(Caller<'_, crate::StoreData>, &mut [ValRaw]) -> Result<(), Trap> {
     let foreign = crate::ForeignData { data, finalizer };
     move |caller, values| {
         drop(&foreign); // move entire foreign into this closure
         let mut caller = wasmtime_caller_t { caller };
-        match callback(foreign.data, &mut caller, values) {
+        match callback(foreign.data, &mut caller, values.as_mut_ptr(), values.len()) {
             None => Ok(()),
             Some(trap) => Err(trap.trap),
         }

crates/cranelift/src/compiler.rs

@@ -516,14 +516,18 @@ impl Compiler {
         let isa = &*self.isa;
         let pointer_type = isa.pointer_type();
         let wasm_signature = indirect_signature(isa, ty);
-        // The host signature has an added parameter for the `values_vec` input
-        // and output.
         let mut host_signature = blank_sig(isa, wasmtime_call_conv(isa));
+        // The host signature has an added parameter for the `values_vec`
+        // input/output buffer in addition to the size of the buffer, in units
+        // of `ValRaw`.
+        host_signature.params.push(ir::AbiParam::new(pointer_type));
         host_signature.params.push(ir::AbiParam::new(pointer_type));
 
         // Compute the size of the values vector. The vmctx and caller vmctx are passed separately.
         let value_size = mem::size_of::<u128>();
-        let values_vec_len = (value_size * cmp::max(ty.params().len(), ty.returns().len())) as u32;
+        let values_vec_len = cmp::max(ty.params().len(), ty.returns().len());
+        let values_vec_byte_size = u32::try_from(value_size * values_vec_len).unwrap();
+        let values_vec_len = u32::try_from(values_vec_len).unwrap();
 
         let CompilerContext {
             mut func_translator,
@@ -535,7 +539,7 @@ impl Compiler {
 
         let ss = context.func.create_stack_slot(ir::StackSlotData::new(
             ir::StackSlotKind::ExplicitSlot,
-            values_vec_len,
+            values_vec_byte_size,
         ));
 
         let mut builder = FunctionBuilder::new(&mut context.func, func_translator.context());
@@ -559,7 +563,14 @@ impl Compiler {
         let vmctx_ptr_val = block_params[0];
         let caller_vmctx_ptr_val = block_params[1];
 
-        let callee_args = vec![vmctx_ptr_val, caller_vmctx_ptr_val, values_vec_ptr_val];
+        let callee_args = vec![
+            vmctx_ptr_val,
+            caller_vmctx_ptr_val,
+            values_vec_ptr_val,
+            builder
+                .ins()
+                .iconst(pointer_type, i64::from(values_vec_len)),
+        ];
 
         let new_sig = builder.import_signature(host_signature);
 

crates/wasmtime/src/func.rs

@@ -367,7 +367,7 @@ impl Func {
     pub unsafe fn new_unchecked<T>(
         mut store: impl AsContextMut<Data = T>,
         ty: FuncType,
-        func: impl Fn(Caller<'_, T>, *mut ValRaw) -> Result<(), Trap> + Send + Sync + 'static,
+        func: impl Fn(Caller<'_, T>, &mut [ValRaw]) -> Result<(), Trap> + Send + Sync + 'static,
     ) -> Self {
         let store = store.as_context_mut().0;
         let host = HostFunc::new_unchecked(store.engine(), ty, func);
@@ -1024,7 +1024,7 @@ impl Func {
     fn invoke<T>(
         mut caller: Caller<'_, T>,
         ty: &FuncType,
-        values_vec: *mut ValRaw,
+        values_vec: &mut [ValRaw],
         func: &dyn Fn(Caller<'_, T>, &[Val], &mut [Val]) -> Result<(), Trap>,
     ) -> Result<(), Trap> {
         // Translate the raw JIT arguments in `values_vec` into a `Val` which
@@ -1041,7 +1041,7 @@ impl Func {
         let nparams = ty.params().len();
         val_vec.reserve(nparams + ty.results().len());
         for (i, ty) in ty.params().enumerate() {
-            val_vec.push(unsafe { Val::from_raw(&mut caller.store, *values_vec.add(i), ty) })
+            val_vec.push(unsafe { Val::from_raw(&mut caller.store, values_vec[i], ty) })
         }
 
         val_vec.extend((0..ty.results().len()).map(|_| Val::null()));
@@ -1075,7 +1075,7 @@ impl Func {
                 ));
             }
             unsafe {
-                *values_vec.add(i) = ret.to_raw(&mut caller.store);
+                values_vec[i] = ret.to_raw(&mut caller.store);
             }
         }
 
@@ -2045,9 +2045,9 @@ impl HostFunc {
     pub unsafe fn new_unchecked<T>(
         engine: &Engine,
         ty: FuncType,
-        func: impl Fn(Caller<'_, T>, *mut ValRaw) -> Result<(), Trap> + Send + Sync + 'static,
+        func: impl Fn(Caller<'_, T>, &mut [ValRaw]) -> Result<(), Trap> + Send + Sync + 'static,
     ) -> Self {
-        let func = move |caller_vmctx, values: *mut ValRaw| {
+        let func = move |caller_vmctx, values: &mut [ValRaw]| {
             Caller::<T>::with(caller_vmctx, |mut caller| {
                 caller.store.0.call_hook(CallHook::CallingHost)?;
                 let result = func(caller.sub_caller(), values)?;

crates/wasmtime/src/linker.rs

@@ -324,7 +324,7 @@ impl<T> Linker<T> {
         module: &str,
         name: &str,
         ty: FuncType,
-        func: impl Fn(Caller<'_, T>, *mut ValRaw) -> Result<(), Trap> + Send + Sync + 'static,
+        func: impl Fn(Caller<'_, T>, &mut [ValRaw]) -> Result<(), Trap> + Send + Sync + 'static,
     ) -> Result<&mut Self> {
         let func = HostFunc::new_unchecked(&self.engine, ty, func);
         let key = self.import_key(module, Some(name));

crates/wasmtime/src/trampoline/func.rs

@@ -26,8 +26,9 @@ unsafe extern "C" fn stub_fn<F>(
     vmctx: *mut VMContext,
     caller_vmctx: *mut VMContext,
     values_vec: *mut ValRaw,
+    values_vec_len: usize,
 ) where
-    F: Fn(*mut VMContext, *mut ValRaw) -> Result<(), Trap> + 'static,
+    F: Fn(*mut VMContext, &mut [ValRaw]) -> Result<(), Trap> + 'static,
 {
     // Here we are careful to use `catch_unwind` to ensure Rust panics don't
     // unwind past us. The primary reason for this is that Rust considers it UB
@@ -49,6 +50,7 @@ unsafe extern "C" fn stub_fn<F>(
         let state = (*vmctx).host_state();
         debug_assert!(state.is::<TrampolineState<F>>());
         let state = &*(state as *const _ as *const TrampolineState<F>);
+        let values_vec = std::slice::from_raw_parts_mut(values_vec, values_vec_len);
         (state.func)(caller_vmctx, values_vec)
     }));
 
@@ -109,7 +111,7 @@ pub fn create_function<F>(
     engine: &Engine,
 ) -> Result<(InstanceHandle, VMTrampoline)>
 where
-    F: Fn(*mut VMContext, *mut ValRaw) -> Result<(), Trap> + Send + Sync + 'static,
+    F: Fn(*mut VMContext, &mut [ValRaw]) -> Result<(), Trap> + Send + Sync + 'static,
 {
     let mut obj = engine.compiler().object()?;
     let (t1, t2) = engine.compiler().emit_trampoline_obj(

tests/all/call_hook.rs

@@ -52,10 +52,10 @@ fn call_wrapped_func() -> Result<(), Error> {
             |caller: Caller<State>, space| {
                 verify(caller.data());
 
-                assert_eq!((*space.add(0)).get_i32(), 1i32);
-                assert_eq!((*space.add(1)).get_i64(), 2i64);
-                assert_eq!((*space.add(2)).get_f32(), 3.0f32.to_bits());
-                assert_eq!((*space.add(3)).get_f64(), 4.0f64.to_bits());
+                assert_eq!(space[0].get_i32(), 1i32);
+                assert_eq!(space[1].get_i64(), 2i64);
+                assert_eq!(space[2].get_f32(), 3.0f32.to_bits());
+                assert_eq!(space[3].get_f64(), 4.0f64.to_bits());
                 Ok(())
             },
         )