From af96691e6b286dd8847474aee5e732764eff7bdd Mon Sep 17 00:00:00 2001
From: Victor Adossi <vadossi@cosmonic.com>
Date: Thu, 19 Oct 2023 03:31:17 +0900
Subject: [PATCH] fix: provide localIds to match args to builder fn

Signed-off-by: Victor Adossi <vadossi@cosmonic.com>
---
 src/module/functions/mod.rs | 26 ++++++++++++++++----------
 1 file changed, 16 insertions(+), 10 deletions(-)

diff --git a/src/module/functions/mod.rs b/src/module/functions/mod.rs
index 1d73ecff..1b26307b 100644
--- 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;
 
@@ -453,7 +453,7 @@ impl Module {
     /// For example, if you wanted to replace an exported function with a no-op,
     ///
     /// ```ignore
-    /// module.replace_exported_func(fid, |body| {
+    /// module.replace_exported_func(fid, |(body, arg_locals)| {
     ///     builder.func_body().unreachable();
     /// });
     /// ```
@@ -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);
 
@@ -507,7 +507,7 @@ impl Module {
     /// For example, if you wanted to replace an imported function with a no-op,
     ///
     /// ```ignore
-    /// module.replace_imported_func(fid, |body| {
+    /// module.replace_imported_func(fid, |(body, arg_locals)| {
     ///     builder.func_body().unreachable();
     /// });
     /// ```
@@ -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`
@@ -740,7 +746,7 @@ mod tests {
 
         // Replace the existing function with a new one with a reversed const value
         let new_fn_id = module
-            .replace_exported_func(original_fn_id, |body| {
+            .replace_exported_func(original_fn_id, |(body, _arg_locals)| {
                 body.unreachable();
             })
             .expect("export function replacement worked");
@@ -819,7 +825,7 @@ mod tests {
 
         // Replace the existing function with a new one with a reversed const value
         let new_fn_id = module
-            .replace_imported_func(original_fn_id, |body| {
+            .replace_imported_func(original_fn_id, |(body, _arg_locals)| {
                 body.unreachable();
             })
             .expect("import fn replacement worked");