fix: proivde localIds to match args to builder fn

Signed-off-by: Victor Adossi <[email protected]>
rustwasm · Oct 18, 2023 · 4549c65 · 4549c65
1 parent fcac3ea
commit 4549c65
Showing 1 changed file with 12 additions and 6 deletions.
diff --git a/src/module/functions/mod.rs b/src/module/functions/mod.rs
@@ -21,7 +21,7 @@ use crate::parse::IndicesToIds;
 use crate::tombstone_arena::{Id, Tombstone, TombstoneArena};
 use crate::ty::TypeId;
 use crate::ty::ValType;
-use crate::{ExportItem, FunctionBuilder, InstrSeqBuilder, Memory, MemoryId};
+use crate::{ExportItem, FunctionBuilder, InstrSeqBuilder, LocalId, Memory, MemoryId};
 
 pub use self::local_function::LocalFunction;
 
@@ -466,7 +466,7 @@ impl Module {
     pub fn replace_exported_func(
         &mut self,
         fid: FunctionId,
-        builder_fn: impl FnOnce(&mut InstrSeqBuilder),
+        builder_fn: impl FnOnce((&mut InstrSeqBuilder, &Vec<LocalId>)),
     ) -> Result<FunctionId> {
         let original_export_id = self
             .exports
@@ -486,7 +486,7 @@ impl Module {
             // Add the function produced by `fn_builder` as a local function
             let mut builder = FunctionBuilder::new(&mut self.types, &params, &results);
             let mut new_fn_body = builder.func_body();
-            builder_fn(&mut new_fn_body);
+            builder_fn((&mut new_fn_body, &lf.args));
             let func = builder.local_func(lf.args.clone());
             let new_fn_id = self.funcs.add_local(func);
 
@@ -520,7 +520,7 @@ impl Module {
     pub fn replace_imported_func(
         &mut self,
         fid: FunctionId,
-        builder_fn: impl FnOnce(&mut InstrSeqBuilder),
+        builder_fn: impl FnOnce((&mut InstrSeqBuilder, &Vec<LocalId>)),
     ) -> Result<FunctionId> {
         let original_import_id = self
             .imports
@@ -537,11 +537,17 @@ impl Module {
             let ty = self.types.get(*tid);
             let (params, results) = (ty.params().to_vec(), ty.results().to_vec());
 
+            // Build the list LocalIds used by args to match the original function
+            let args = params
+                .iter()
+                .map(|ty| self.locals.add(*ty))
+                .collect::<Vec<_>>();
+
             // Build the new function
             let mut builder = FunctionBuilder::new(&mut self.types, &params, &results);
             let mut new_fn_body = builder.func_body();
-            builder_fn(&mut new_fn_body);
-            let new_func_kind = FunctionKind::Local(builder.local_func(vec![]));
+            builder_fn((&mut new_fn_body, &args));
+            let new_func_kind = FunctionKind::Local(builder.local_func(args));
 
             // Mutate the existing function, changing it from a FunctionKind::ImportedFunction
             // to the local function produced by running the provided `fn_builder`