wip

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

@@ -46,40 +46,42 @@ fn generate_module(
     target_conf: &TargetConf,
     missing_mem_conf: &MissingMemoryConf,
 ) -> TokenStream2 {
-    let mut fields = Vec::new();
+    let fields = module.funcs().map(|f| {
-    let mut get_exports = Vec::new();
+        let name_ident = names.func(&f.name);
-    let mut ctor_externs = Vec::new();
+        quote! { pub #name_ident: wasmtime::Func }
-    let mut ctor_fields = Vec::new();
+    });
-    let mut linker_add = Vec::new();
+    let get_exports = module.funcs().map(|f| {
+        let func_name = f.name.as_str();
+        let name_ident = names.func(&f.name);
+        quote! { #func_name => Some(&self.#name_ident) }
+    });
+    let ctor_fields = module.funcs().map(|f| names.func(&f.name));
+
+    let module_name = module.name.as_str();
+
+    let linker_add = module.funcs().map(|f| {
+        let func_name = f.name.as_str();
+        let name_ident = names.func(&f.name);
+        quote! {
+            linker.define(#module_name, #func_name, self.#name_ident.clone())?;
+        }
+    });
 
     let runtime = names.runtime_mod();
     let target_path = &target_conf.path;
     let missing_mem_err = &missing_mem_conf.err;
-
-    let module_name = module.name.as_str();
     let module_id = names.module(&module.name);
-    for func in module.funcs() {
-        let func_name = func.name.as_str();
-        let name_ident = names.func(&func.name);
-        fields.push(quote! { pub #name_ident: wasmtime::Func });
-        get_exports.push(quote! { #func_name => Some(&self.#name_ident) });
-        ctor_fields.push(name_ident.clone());
-        linker_add.push(quote! {
-            linker.define(#module_name, #func_name, self.#name_ident.clone())?;
-        });
-
-        if let Some(func_override) = module_conf.function_override.find(func_name) {
-            ctor_externs.push(quote! {
-                let #name_ident = wasmtime::Func::wrap(store, #func_override);
-            });
-            continue;
-        }
 
+    let ctor_externs = module.funcs().map(|f| {
+        let name_ident = names.func(&f.name);
+        if let Some(func_override) = module_conf.function_override.find(&f.name.as_str()) {
+            quote! { let #name_ident = wasmtime::Func::wrap(store, #func_override); }
+        } else {
             let mut shim_arg_decls = Vec::new();
             let mut params = Vec::new();
             let mut hostcall_args = Vec::new();
 
-        for param in func.params.iter() {
+            for param in f.params.iter() {
                 let name = names.func_param(&param.name);
 
                 // Registers a new parameter to the shim we're making with the
@@ -161,15 +163,19 @@ fn generate_module(
                         add_param(&len, Abi::I32);
                     }
 
-                witx::Type::ConstPointer(_) | witx::Type::Handle(_) | witx::Type::Pointer(_) => {
+                    witx::Type::ConstPointer(_)
+                    | witx::Type::Handle(_)
+                    | witx::Type::Pointer(_) => {
                         add_param(&name, Abi::I32);
                     }
 
-                witx::Type::Struct(_) | witx::Type::Union(_) => panic!("unsupported argument type"),
+                    witx::Type::Struct(_) | witx::Type::Union(_) => {
+                        panic!("unsupported argument type")
+                    }
                 }
             }
 
-        let mut results = func.results.iter();
+            let mut results = f.results.iter();
             let mut ret_ty = quote! { () };
             let mut cvt_ret = quote! {};
             let mut returns = Vec::new();
@@ -202,7 +208,7 @@ fn generate_module(
                 hostcall_args.push(quote! { #name });
             }
 
-        ctor_externs.push(quote! {
+            quote! {
                 let my_cx = cx.clone();
                 let #name_ident = wasmtime::Func::wrap(
                     store,
@@ -231,8 +237,9 @@ fn generate_module(
                         }
                     }
                 );
-        });
             }
+        }
+    });
 
     let type_name = module_conf.name.clone();
     let type_docs = module_conf