Add a new Func::wrap_cabi API

This commit adds a new API to the `wasmtime::Func` type for wrapping a
C-ABI function with a well-defined signature derived from a wasm type
signature. The purpose of this API is to add the-most-optimized-we-can
path for using the C API and having wasm call host functions. Previously
when wasm called a host function it would perform these steps:

1. Using a trampoline, place all arguments into a `u128*` array on the
   stack.
2. Call `Func::invoke` which uses the type of the function (dynamically)
   to read values from this `u128*`.
3. Values are placed within a `Vec<Val>` after being read.
4. The C API receives `&[Val]` and translates this to
   `&[wasmtime_val_t]`, iterating over each value and copying its
   contents into a new vector.
5. Then the host function is actually called.
6. The above argument-shuffling steps are all performed in reverse for
   the results, shipping everything through `wasmtime_val_t` and `Val`.

PRs such as bytecodealliance#3319 and related attempts have made this sequence faster,
but the numbers on bytecodealliance#3319 show that even after we get all the allocation
and such bits out of the way we're still spending quite a lot of time
shuffling arguments back-and-forth relative to the `Func::wrap` API that
Rust can use.

This commit fixes the issue by eliminating all steps except 1/5 above.
Although we still place all arguments on the stack and read them out
again to call the C-defined function with it's much faster than pushing
this all through the `Val` and `wasmtime_val_t` machinery. This overall
gets the cost of a wasm->host call basically on-par with `Func::wrap`,
although it's still not as optimized. Benchmarking the overhead of
wasm->host calls, where `i64` returns one i64 value and `many` takes 5
`i32` parameters and returns one `i64` value, the numbers I get are:

| Import | Rust | C before | C after |
|--------|------|----------|---------|
| `i64`  | 1ns  | 40ns     | 7ns     |
| `many` | 1ns  | 91ns     | 10ns    |

This commit is a clear win over the previous implementation, but it's
even still somewhat far away from Rust. That being said I think I'm out
of ideas of how to make this better. Without open-coding much larger
portions of `wasmtime` I'm not sure how much better we can get here. The
time in C after this commit is almost entirely spent in trampolines
storing the arguments to the stack and loading them later, and at this
point I'm not sure how much more optimized we can get than that since
Rust needs to enter the picture here somehow to handle the Wasmtime
fiddly-bits of calling back into C. I'm hopeful, though, that this is
such a large improvement from before that the question of optimizing
this further in C is best left for another day.

The new `Func::wrap_cabi` method is unlikely to ever get used from Rust,
but a `wasmtime_func_wrap` API was added to the C API to mirror
`Func::wrap` where if the host function pointer has a specific ABI this
function can be called instead of `wasmtime_func_new`.

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

@@ -87,6 +87,61 @@ WASM_API_EXTERN void wasmtime_func_new(
   wasmtime_func_t *ret
 );
 
+/**
+ * \brief Creates a new host-defined function with a custom signature.
+ *
+ * This function is the same as #wasmtime_func_new except that the signature of
+ * `callback` is different. The signature of `callback` must correspond to the
+ * function `type` specified, if it's not then calling the function is UB.
+ *
+ * \param store the store in which to create the function
+ * \param type the wasm type of the function that's being created
+ * \param callback the host-defined callback to invoke
+ * \param env host-specific data passed to the callback invocation, can be
+ * `NULL`
+ * \param finalizer optional finalizer for `env`, can be `NULL`
+ * \param ret the #wasmtime_func_t return value to be filled in.
+ *
+ * The returned function can only be used with the specified `store`.
+ *
+ * The signature of `callback` is as follows:
+ *
+ * * First it takes a `void*` parameter which is the `env` provided here.
+ * * Next it receives a `wasmtime_caller_t*` which is only valid for the
+ *   duration of that one call.
+ * * Next it receives each WebAssembly argument individually. Mapping from wasm
+ *   types to host types is provided below.
+ * * Next it receives an out-pointer for each WebAssembly result individually.
+ * * Finally the function must return a `wasm_trap_t*` indicating whether the
+ *   host function should trap or not.
+ *
+ * The mapping of WebAssembly types to host types is:
+ *
+ * * `i32 <=> int32_t`
+ * * `i64 <=> int64_t`
+ * * `f32 <=> float`
+ * * `f64 <=> double`
+ * * `v128`
+ *   * Received indirectly as `const uint8_t*` when passed as an argument
+ *   * Also passed as `const uint8_t*` for an out-pointer.
+ * * `extenref` - not supported
+ * * `funcref` - not supported
+ *
+ * All arguments are passed by-value except for `v128`, which is passed
+ * indirectly to the host function.
+ *
+ * This function will return an error if the `type` specified uses either
+ * `externref` or `funcref`.
+ */
+WASM_API_EXTERN wasmtime_error_t *wasmtime_func_wrap(
+  wasmtime_context_t *store,
+  const wasm_functype_t* type,
+  size_t callback,
+  void *env,
+  void (*finalizer)(void*),
+  wasmtime_func_t *ret
+);
+
 /**
  * \brief Returns the type of the function specified
  *

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

@@ -81,7 +81,8 @@ WASM_API_EXTERN wasmtime_error_t* wasmtime_linker_define(
 );
 
 /**
- * \brief Defines a new function in this linker.
+ * \brief Defines a new function in this linker along the lines of
+ * #wasmtime_func_new.
  *
  * \param linker the linker the name is being defined in.
  * \param module the module name the item is defined under.
@@ -115,6 +116,46 @@ WASM_API_EXTERN wasmtime_error_t* wasmtime_linker_define_func(
     void (*finalizer)(void*)
 );
 
+/**
+ * \brief Defines a new function in this linker along the lines of
+ * #wasmtime_func_wrap.
+ *
+ * \param linker the linker the name is being defined in.
+ * \param module the module name the item is defined under.
+ * \param module_len the byte length of `module`
+ * \param name the field name the item is defined under
+ * \param name_len the byte length of `name`
+ * \param ty the type of the function that's being defined
+ * \param cb the host callback to invoke when the function is called
+ * \param data the host-provided data to provide as the first argument to the callback
+ * \param finalizer an optional finalizer for the `data` argument.
+ *
+ * \return On success `NULL` is returned, otherwise an error is returned which
+ * describes why the definition failed.
+ *
+ * For more information about name resolution consult the [Rust
+ * documentation](https://bytecodealliance.github.io/wasmtime/api/wasmtime/struct.Linker.html#name-resolution).
+ *
+ * Note that this function does not create a #wasmtime_func_t. This creates a
+ * store-independent function within the linker, allowing this function
+ * definition to be used with multiple stores.
+ *
+ * Also note that the documentation of #wasmtime_func_wrap should be consulted
+ * to determine the correct function signature for the function provided in
+ * `cb`.
+ */
+WASM_API_EXTERN wasmtime_error_t* wasmtime_linker_func_wrap(
+    wasmtime_linker_t *linker,
+    const char *module,
+    size_t module_len,
+    const char *name,
+    size_t name_len,
+    const wasm_functype_t *ty,
+    size_t cb,
+    void *data,
+    void (*finalizer)(void*)
+);
+
 /**
  * \brief Defines WASI functions in this linker.
  *

crates/c-api/src/func.rs

@@ -1,7 +1,8 @@
 use crate::wasm_trap_t;
 use crate::{
-    wasm_extern_t, wasm_functype_t, wasm_store_t, wasm_val_t, wasm_val_vec_t, wasmtime_error_t,
-    wasmtime_extern_t, wasmtime_val_t, wasmtime_val_union, CStoreContext, CStoreContextMut,
+    handle_result, wasm_extern_t, wasm_functype_t, wasm_store_t, wasm_val_t, wasm_val_vec_t,
+    wasmtime_error_t, wasmtime_extern_t, wasmtime_val_t, wasmtime_val_union, CStoreContext,
+    CStoreContextMut,
 };
 use anyhow::anyhow;
 use std::ffi::c_void;
@@ -182,7 +183,10 @@ pub extern "C" fn wasm_func_as_extern_const(f: &wasm_func_t) -> &wasm_extern_t {
     &(*f).ext
 }
 
-#[repr(C)]
+// Note the `transparent` annotation which is required due to the
+// `wasmtime_func_wrap` function where the C code uses `*mut wasmtime_caller_t`
+// to represent the `&mut Caller<'_, T>` type in Rust.
+#[repr(transparent)]
 pub struct wasmtime_caller_t<'a> {
     caller: Caller<'a, crate::StoreData>,
 }
@@ -211,6 +215,22 @@ pub unsafe extern "C" fn wasmtime_func_new(
     *func = f;
 }
 
+#[no_mangle]
+pub unsafe extern "C" fn wasmtime_func_wrap(
+    store: CStoreContextMut<'_>,
+    ty: &wasm_functype_t,
+    callback: usize,
+    data: usize,
+    finalizer: Option<extern "C" fn(usize)>,
+    func: &mut Func,
+) -> Option<Box<wasmtime_error_t>> {
+    let ty = ty.ty().ty.clone();
+
+    handle_result(Func::wrap_cabi(store, ty, callback, data, finalizer), |f| {
+        *func = f;
+    })
+}
+
 pub(crate) unsafe fn c_callback_to_rust_fn(
     callback: wasmtime_func_callback_t,
     data: *mut c_void,

crates/c-api/src/linker.rs

@@ -75,6 +75,29 @@ pub unsafe extern "C" fn wasmtime_linker_define_func(
     handle_result(linker.linker.func_new(module, name, ty, cb), |_linker| ())
 }
 
+#[no_mangle]
+pub unsafe extern "C" fn wasmtime_linker_func_wrap(
+    linker: &mut wasmtime_linker_t,
+    module: *const u8,
+    module_len: usize,
+    name: *const u8,
+    name_len: usize,
+    ty: &wasm_functype_t,
+    callback: usize,
+    data: usize,
+    finalizer: Option<extern "C" fn(usize)>,
+) -> Option<Box<wasmtime_error_t>> {
+    let ty = ty.ty().ty.clone();
+    let module = to_str!(module, module_len);
+    let name = to_str!(name, name_len);
+    handle_result(
+        linker
+            .linker
+            .func_wrap_cabi(module, name, ty, callback, data, finalizer),
+        |_linker| (),
+    )
+}
+
 #[cfg(feature = "wasi")]
 #[no_mangle]
 pub extern "C" fn wasmtime_linker_define_wasi(

crates/c-api/src/trap.rs

@@ -2,7 +2,12 @@ use crate::{wasm_frame_vec_t, wasm_instance_t, wasm_name_t, wasm_store_t};
 use once_cell::unsync::OnceCell;
 use wasmtime::{Trap, TrapCode};
 
-#[repr(C)]
+// Note the `transparent` representation here which is required by the
+// `wasmtime_func_wrap` API since the C-ABI function pointers return `*mut
+// wasm_trap_t` which is interpreted as `Option<Box<Trap>>` in Rust. Note that
+// `wasm_trap_t` is opaque in C, though, so it's always asking this library to
+// allocate.
+#[repr(transparent)]
 #[derive(Clone)]
 pub struct wasm_trap_t {
     pub(crate) trap: Trap,

crates/cranelift/src/compiler.rs

@@ -8,7 +8,7 @@ use crate::{
 };
 use anyhow::{Context as _, Result};
 use cranelift_codegen::ir::{self, ExternalName, InstBuilder, MemFlags};
-use cranelift_codegen::isa::TargetIsa;
+use cranelift_codegen::isa::{CallConv, TargetIsa};
 use cranelift_codegen::print_errors::pretty_error;
 use cranelift_codegen::settings;
 use cranelift_codegen::MachSrcLoc;
@@ -29,7 +29,7 @@ use std::sync::Mutex;
 use wasmtime_environ::{
     AddressMapSection, CompileError, FilePos, FlagValue, FunctionBodyData, FunctionInfo,
     InstructionAddressMap, Module, ModuleTranslation, StackMapInformation, Trampoline, TrapCode,
-    TrapEncodingBuilder, TrapInformation, Tunables, TypeTables, VMOffsets,
+    TrapEncodingBuilder, TrapInformation, Tunables, TypeTables, VMOffsets, WasmType,
 };
 
 /// A compiler that compiles a WebAssembly module with Compiler, translating
@@ -325,16 +325,18 @@ impl wasmtime_environ::Compiler for Compiler {
         ty: &WasmFuncType,
         host_fn: usize,
         obj: &mut Object,
-    ) -> Result<(Trampoline, Trampoline)> {
+    ) -> Result<(Trampoline, Trampoline, Trampoline)> {
         let host_to_wasm = self.host_to_wasm_trampoline(ty)?;
         let wasm_to_host = self.wasm_to_host_trampoline(ty, host_fn)?;
+        let host_to_c = self.host_to_c_trampoline(ty)?;
         let module = Module::new();
         let mut builder = ObjectBuilder::new(obj, &module, &*self.isa);
         let a = builder.trampoline(SignatureIndex::new(0), &host_to_wasm);
         let b = builder.trampoline(SignatureIndex::new(1), &wasm_to_host);
+        let c = builder.trampoline(SignatureIndex::new(2), &host_to_c);
         builder.unwind_info();
         builder.finish()?;
-        Ok((a, b))
+        Ok((a, b, c))
     }
 
     fn triple(&self) -> &target_lexicon::Triple {
@@ -525,6 +527,116 @@ impl Compiler {
         Ok(func)
     }
 
+    fn host_to_c_trampoline(&self, ty: &WasmFuncType) -> Result<CompiledFunction, CompileError> {
+        let isa = &*self.isa;
+        let pointer_type = isa.pointer_type();
+
+        // This is the signature of the C stub that this trampoline calls.
+        // Effectively this is the definition of how wasmtime translates a wasm
+        // signature to a C signature. Currently the signature parameters are:
+        //
+        // * a void* pointer argument (user-provided)
+        // * a pointer argument for a "caller" structure (`wasmtime::Caller`)
+        // * each wasm type as an individual parameter
+        // * each wasm return as an out-pointer
+        //
+        // And then the function returns a pointer which is a trap object. The
+        // trap is a nullable pointer.
+        let mut c_sig = ir::Signature::new(CallConv::triple_default(isa.triple()));
+        c_sig.params.push(ir::AbiParam::new(pointer_type));
+        c_sig.params.push(ir::AbiParam::new(pointer_type));
+        for ty in ty.params.iter() {
+            let ty = match ty {
+                WasmType::I32 | WasmType::I64 | WasmType::F32 | WasmType::F64 => {
+                    value_type(isa, *ty)
+                }
+                WasmType::V128 => pointer_type,
+                WasmType::FuncRef | WasmType::ExternRef | WasmType::ExnRef => unimplemented!(),
+            };
+            c_sig.params.push(ir::AbiParam::new(ty));
+        }
+        for _ in ty.returns.iter() {
+            c_sig.params.push(ir::AbiParam::new(pointer_type));
+        }
+        c_sig.returns.push(ir::AbiParam::new(pointer_type));
+
+        // This trampoline itself has a fixed signature so it can be called from
+        // Rust. The parameters here are:
+        //
+        // * a void* pionter argument (user-provided)
+        // * a pointer argument for a "caller" structure (`wasmtime::Caller`)
+        // * the `*const u128` values_vec allocation for arguments/returns
+        // * the C function that we'll be calling
+        //
+        // This function returns a pointer which is the trap the C function
+        // returns.
+        let mut stub_signature = ir::Signature::new(CallConv::triple_default(isa.triple()));
+        stub_signature.params.push(ir::AbiParam::new(pointer_type));
+        stub_signature.params.push(ir::AbiParam::new(pointer_type));
+        stub_signature.params.push(ir::AbiParam::new(pointer_type));
+        stub_signature.params.push(ir::AbiParam::new(pointer_type));
+        stub_signature.returns.push(ir::AbiParam::new(pointer_type));
+
+        let value_size = mem::size_of::<u128>();
+        let mut context = Context::new();
+        context.func =
+            ir::Function::with_name_signature(ir::ExternalName::user(0, 0), stub_signature);
+        let mut func_translator = self.take_translator();
+        let mut builder = FunctionBuilder::new(&mut context.func, func_translator.context());
+        let block0 = builder.create_block();
+        builder.append_block_params_for_function_params(block0);
+        builder.switch_to_block(block0);
+        builder.seal_block(block0);
+
+        // Build the arguments to the C function that we're going to call. The
+        // first two arguments are simply forwarding our own first two
+        // arguments. Next comes all the wasm parameters which are passed
+        // by-value. Finally all returns get a return pointer which points back
+        // into our `values_vec` parameter.
+        let mut args = Vec::new();
+        args.push(builder.func.dfg.block_params(block0)[0]);
+        args.push(builder.func.dfg.block_params(block0)[1]);
+        let values_vec_ptr_val = builder.func.dfg.block_params(block0)[2];
+        for (i, ty) in ty.params.iter().enumerate() {
+            let offset = (i * value_size) as i32;
+            let arg = match ty {
+                WasmType::I32 | WasmType::I64 | WasmType::F32 | WasmType::F64 => {
+                    builder.ins().load(
+                        value_type(isa, *ty),
+                        MemFlags::trusted(),
+                        values_vec_ptr_val,
+                        offset,
+                    )
+                }
+                WasmType::V128 => builder
+                    .ins()
+                    .iadd_imm(values_vec_ptr_val, i64::from(offset)),
+                WasmType::FuncRef | WasmType::ExternRef | WasmType::ExnRef => unimplemented!(),
+            };
+            args.push(arg);
+        }
+        for (i, _ty) in ty.returns.iter().enumerate() {
+            let addr = builder
+                .ins()
+                .iadd_imm(values_vec_ptr_val, (i * value_size) as i64);
+            args.push(addr);
+        }
+
+        // With all the arguments we call the indirect function passed as a
+        // parameter to our function, then we're good to go and we simply return
+        // that function's value.
+        let new_sig = builder.import_signature(c_sig);
+        let callee_value = builder.func.dfg.block_params(block0)[3];
+        let call = builder.ins().call_indirect(new_sig, callee_value, &args);
+        let trap = builder.inst_results(call)[0];
+        builder.ins().return_(&[trap]);
+        builder.finalize();
+
+        let func = self.finish_trampoline(context, isa)?;
+        self.save_translator(func_translator);
+        Ok(func)
+    }
+
     fn finish_trampoline(
         &self,
         mut context: Context,

crates/environ/src/compilation.rs

@@ -175,19 +175,23 @@ pub trait Compiler: Send + Sync {
         obj: &mut Object,
     ) -> Result<(PrimaryMap<DefinedFuncIndex, FunctionInfo>, Vec<Trampoline>)>;
 
-    /// Inserts two functions for host-to-wasm and wasm-to-host trampolines into
-    /// the `obj` provided.
+    /// Inserts various trampolines for interacting with the wasm function
+    /// siganture `ty` into the `obj` provided.
     ///
     /// This will configure the same sections as `emit_obj`, but will likely be
-    /// much smaller. The two returned `Trampoline` structures describe where to
-    /// find the host-to-wasm and wasm-to-host trampolines in the text section,
-    /// respectively.
+    /// much smaller. The returned `Trampoline` structures at this time are:
+    ///
+    /// * Where to find the host-to-wasm trampoline.
+    /// * Where to find the wasm-to-host trampoline.
+    /// * Where to fidn the host-to-c trampoline.
+    ///
+    /// All trampoline offsets are relative to the text section.
     fn emit_trampoline_obj(
         &self,
         ty: &WasmFuncType,
         host_fn: usize,
         obj: &mut Object,
-    ) -> Result<(Trampoline, Trampoline)>;
+    ) -> Result<(Trampoline, Trampoline, Trampoline)>;
 
     /// Creates a new `Object` file which is used to build the results of a
     /// compilation into.