Add add_to_config to Wasi.

This commit adds `add_to_config` to the code generation for Wasmtime's `Wasi`
type.

The new method adds the WASI functions to the given config as host functions.

crates/wiggle/wasmtime/macro/src/lib.rs

@@ -1,7 +1,7 @@
 use proc_macro::TokenStream;
 use proc_macro2::TokenStream as TokenStream2;
-use quote::quote;
-use syn::parse_macro_input;
+use quote::{format_ident, quote};
+use syn::{parse_macro_input, Ident};
 use wiggle_generate::Names;
 
 mod config;
@@ -88,14 +88,28 @@ fn generate_module(
     let module_id = names.module(&module.name);
     let target_module = quote! { #target_path::#module_id };
 
-    let ctor_externs = module.funcs().map(|f| {
+    let mut fns = Vec::new();
+    let mut ctor_externs = Vec::new();
+    let mut host_funcs = Vec::new();
+
+    for f in module.funcs() {
         if let Some(func_override) = module_conf.function_override.find(&f.name.as_str()) {
             let name_ident = names.func(&f.name);
-            quote! { let #name_ident = wasmtime::Func::wrap(store, #func_override); }
+            ctor_externs
+                .push(quote! { let #name_ident = wasmtime::Func::wrap(store, #func_override); });
+            host_funcs.push(quote! { config.wrap_host_func(#func_override); });
         } else {
-            generate_func(&f, names, &target_module)
+            generate_func(
+                &module_id,
+                &f,
+                names,
+                &target_module,
+                &mut fns,
+                &mut ctor_externs,
+                &mut host_funcs,
+            );
         }
-    });
+    }
 
     let type_name = module_conf.name.clone();
     let type_docs = module_conf
@@ -148,25 +162,44 @@ contained in the `cx` parameter.",
                 #(#linker_add)*
                 Ok(())
             }
+
+            /// Adds the WASI host functions to the given `Config`.
+            pub fn add_to_config(config: &mut wasmtime::Config) {
+                #(#host_funcs)*
+            }
+
+            #(#fns)*
         }
     }
 }
 
 fn generate_func(
+    module_ident: &Ident,
     func: &witx::InterfaceFunc,
     names: &Names,
     target_module: &TokenStream2,
-) -> TokenStream2 {
+    fns: &mut Vec<TokenStream2>,
+    ctors: &mut Vec<TokenStream2>,
+    host_funcs: &mut Vec<TokenStream2>,
+) {
     let name_ident = names.func(&func.name);
 
     let coretype = func.core_type();
 
-    let arg_decls = coretype.args.iter().map(|arg| {
-        let name = names.func_core_arg(arg);
-        let atom = names.atom_type(arg.repr());
-        quote! { #name: #atom }
-    });
-    let arg_names = coretype.args.iter().map(|arg| names.func_core_arg(arg));
+    let arg_decls = coretype
+        .args
+        .iter()
+        .map(|arg| {
+            let name = names.func_core_arg(arg);
+            let atom = names.atom_type(arg.repr());
+            quote! { #name: #atom }
+        })
+        .collect::<Vec<_>>();
+    let arg_names = coretype
+        .args
+        .iter()
+        .map(|arg| names.func_core_arg(arg))
+        .collect::<Vec<_>>();
 
     let ret_ty = if let Some(ret) = &coretype.ret {
         let ret_ty = match ret.signifies {
@@ -179,32 +212,51 @@ fn generate_func(
     };
 
     let runtime = names.runtime_mod();
+    let fn_ident = format_ident!("{}_{}", module_ident, name_ident);
+    let ctx_type = names.ctx_type();
 
-    quote! {
+    fns.push(quote! {
+        fn #fn_ident(caller: &wasmtime::Caller<'_>, ctx: &mut #ctx_type #(, #arg_decls)*) -> Result<#ret_ty, wasmtime::Trap> {
+            unsafe {
+                let mem = match caller.get_export("memory") {
+                    Some(wasmtime::Extern::Memory(m)) => m,
+                    _ => {
+                        return Err(wasmtime::Trap::new("missing required memory export"));
+                    }
+                };
+                let mem = #runtime::WasmtimeGuestMemory::new(mem);
+                match #target_module::#name_ident(ctx, &mem #(, #arg_names)*) {
+                    Ok(r) => Ok(r.into()),
+                    Err(wasmtime_wiggle::Trap::String(err)) => Err(wasmtime::Trap::new(err)),
+                    Err(wasmtime_wiggle::Trap::I32Exit(err)) => Err(wasmtime::Trap::i32_exit(err)),
+                }
+            }
+        }
+    });
+
+    ctors.push(quote! {
         let my_cx = cx.clone();
         let #name_ident = wasmtime::Func::wrap(
             store,
-            move |caller: wasmtime::Caller<'_> #(,#arg_decls)*| -> Result<#ret_ty, wasmtime::Trap> {
-                unsafe {
-                    let mem = match caller.get_export("memory") {
-                        Some(wasmtime::Extern::Memory(m)) => m,
-                        _ => {
-                            return Err(wasmtime::Trap::new("missing required memory export"));
-                        }
-                    };
-                    let mem = #runtime::WasmtimeGuestMemory::new(mem);
-                    let result = #target_module::#name_ident(
-                        &mut my_cx.borrow_mut(),
-                        &mem,
-                        #(#arg_names),*
-                    );
-                    match result {
-                        Ok(r) => Ok(r.into()),
-                        Err(wasmtime_wiggle::Trap::String(err)) => Err(wasmtime::Trap::new(err)),
-                        Err(wasmtime_wiggle::Trap::I32Exit(err)) => Err(wasmtime::Trap::i32_exit(err)),
-                    }
-                }
+            move |caller: wasmtime::Caller<'_> #(, #arg_decls)*| -> Result<#ret_ty, wasmtime::Trap> {
+                Self::#fn_ident(&caller, &mut my_cx.borrow_mut() #(, #arg_names)*)
             }
         );
-    }
+    });
+
+    host_funcs.push(quote! {
+        config.wrap_host_func(
+            stringify!(#module_ident),
+            stringify!(#name_ident),
+            move |caller: wasmtime::Caller<'_> #(, #arg_decls)*| -> Result<#ret_ty, wasmtime::Trap> {
+                let store = caller.store().clone();
+                store.with_context(move |ctx: Option<&mut #ctx_type>| {
+                    match ctx {
+                        Some(ctx) => Self::#fn_ident(&caller, ctx #(, #arg_names)*),
+                        None => Err(wasmtime::Trap::new("WasiContext is missing in store context"))
+                    }
+                })
+            },
+        );
+    });
 }

tests/all/host_funcs.rs

@@ -1,6 +1,7 @@
 use anyhow::Result;
 use std::sync::atomic::{AtomicUsize, Ordering::SeqCst};
 use wasmtime::*;
+use wasmtime_wasi::Wasi;
 
 #[test]
 fn wrap_func() {
@@ -545,3 +546,64 @@ fn store_with_context() -> Result<()> {
 
     Ok(())
 }
+
+#[test]
+fn wasi_imports_missing_context() -> Result<()> {
+    let mut config = Config::default();
+    Wasi::add_to_config(&mut config);
+
+    let wasm = wat::parse_str(
+        r#"
+        (import "wasi_snapshot_preview1" "proc_exit" (func $__wasi_proc_exit (param i32)))
+        (memory (export "memory") 0)
+        (func (export "_start")
+            (call $__wasi_proc_exit (i32.const 123))
+        )
+        "#,
+    )?;
+
+    let engine = Engine::new(&config);
+    let module = Module::new(&engine, wasm)?;
+    let store = Store::new(&engine);
+    let linker = Linker::new(&store);
+    let instance = linker.instantiate(&module)?;
+
+    let start = instance.get_func("_start").unwrap().get0::<()>()?;
+
+    let trap = start().unwrap_err();
+
+    assert!(trap.to_string().contains("WasiContext is missing"));
+    assert!(trap.i32_exit_status().is_none());
+
+    Ok(())
+}
+
+#[test]
+fn wasi_imports() -> Result<()> {
+    let mut config = Config::default();
+    Wasi::add_to_config(&mut config);
+
+    let wasm = wat::parse_str(
+        r#"
+        (import "wasi_snapshot_preview1" "proc_exit" (func $__wasi_proc_exit (param i32)))
+        (memory (export "memory") 0)
+        (func (export "_start")
+            (call $__wasi_proc_exit (i32.const 123))
+        )
+        "#,
+    )?;
+
+    let engine = Engine::new(&config);
+    let module = Module::new(&engine, wasm)?;
+    let store = Store::new(&engine);
+    store.insert_context(wasmtime_wasi::WasiCtxBuilder::new().build()?);
+    let linker = Linker::new(&store);
+    let instance = linker.instantiate(&module)?;
+
+    let start = instance.get_func("_start").unwrap().get0::<()>()?;
+
+    let trap = start().unwrap_err();
+    assert_eq!(trap.i32_exit_status(), Some(123));
+
+    Ok(())
+}