refactor: build local func after builer_fn finishes

Signed-off-by: Victor Adossi <vadossi@cosmonic.com>

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, ModuleTypes};
+use crate::{ExportItem, FunctionBuilder, InstrSeqBuilder, Memory, MemoryId};
 
 pub use self::local_function::LocalFunction;
 
@@ -447,28 +447,26 @@ impl Module {
 
     /// Replace a single exported function with the result of the provided builder function.
     ///
-    /// The builder function is provided:
-    /// - a mutable borrow (`&mut Module`) to the module,
-    /// - the parameters (`Vec<ValType>`) used by the original function
-    /// - the results (`Vec<ValType>`) returned by the original function
+    /// The builder function is provided a mutable reference to an [`InstrSeqBuilder`] which can be
+    /// used to build the function as necessary.
     ///
     /// For example, if you wanted to replace an exported function with a no-op,
     ///
     /// ```ignore
-    /// module.replace_exported_func(fid, |module, params, results| {
-    ///     let mut builder = FunctionBuilder::new(&mut module.types, params, results);
+    /// module.replace_exported_func(fid, |body| {
     ///     builder.func_body().unreachable();
-    ///     let new_fn = builder.local_func(vec![]);
-    ///     Ok(new_fn)
     /// });
     /// ```
     ///
+    /// The arguments passed to the original function will be passed to the
+    /// new exported function that was built in your closure.
+    ///
     /// This function returns the function ID of the *new* function,
     /// after it has been inserted into the module as an export.
     pub fn replace_exported_func(
         &mut self,
         fid: FunctionId,
-        builder_fn: impl FnOnce(InstrSeqBuilder) -> LocalFunction,
+        builder_fn: impl FnOnce(&mut InstrSeqBuilder),
     ) -> Result<FunctionId> {
         let original_export_id = self
             .exports
@@ -487,9 +485,9 @@ impl Module {
 
             // Add the function produced by `fn_builder` as a local function
             let mut builder = FunctionBuilder::new(&mut self.types, &params, &results);
-            let seq_builder = builder.func_body();
-            // let func = builder_fn(seq_builder).context("export fn builder failed")?;
-            let func = builder_fn(seq_builder);
+            let mut new_fn_body = builder.func_body();
+            builder_fn(&mut new_fn_body);
+            let func = builder.local_func(lf.args.clone());
             let new_fn_id = self.funcs.add_local(func);
 
             // Mutate the existing export to use the new local function
@@ -503,28 +501,27 @@ impl Module {
 
     /// Replace a single imported function with the result of the provided builder function.
     ///
-    /// The builder function is provided:
-    /// - a mutable borrow (`&mut Module`) to the module,
-    /// - the parameters (`Vec<ValType>`) used by the original function
-    /// - the results (`Vec<ValType>`) returned by the original function
+    /// The builder function is provided a mutable reference to an [`InstrSeqBuilder`] which can be
+    /// used to build the function as necessary.
     ///
     /// For example, if you wanted to replace an imported function with a no-op,
     ///
     /// ```ignore
-    /// module.replace_imported_func(fid, |module, params, results| {
-    ///     let mut builder = FunctionBuilder::new(&mut module.types, params, results);
+    /// module.replace_imported_func(fid, |body| {
     ///     builder.func_body().unreachable();
-    ///     let new_fn = builder.local_func(vec![]);
-    ///     Ok(new_fn)
     /// });
     /// ```
     ///
+    /// The arguments passed to the original function will be passed to the
+    /// new exported function that was built in your closure.
+    ///
     /// This function returns the function ID of the *new* function, and
     /// removes the existing import that has been replaced (the function will become local).
-    pub fn replace_imported_func<F>(&mut self, fid: FunctionId, fn_builder: F) -> Result<FunctionId>
-    where
-        F: FnOnce(&mut Self, &FuncParams, &FuncResults) -> Result<LocalFunction>,
-    {
+    pub fn replace_imported_func(
+        &mut self,
+        fid: FunctionId,
+        builder_fn: impl FnOnce(&mut InstrSeqBuilder),
+    ) -> Result<FunctionId> {
         let original_import_id = self
             .imports
             .get_imported_func(fid)
@@ -541,9 +538,10 @@ impl Module {
             let (params, results) = (ty.params().to_vec(), ty.results().to_vec());
 
             // Build the new function
-            let new_func_kind = FunctionKind::Local(
-                fn_builder(self, &params, &results).context("import fn builder failed")?,
-            );
+            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![]));
 
             // Mutate the existing function, changing it from a FunctionKind::ImportedFunction
             // to the local function produced by running the provided `fn_builder`
@@ -701,9 +699,8 @@ 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, |mut body| {
+            .replace_exported_func(original_fn_id, |body| {
                 body.i32_const(4321).drop();
-                body.local_func(vec![]) // TODO(FIX): Can't move LocalFunction produced here out of deref to body!
             })
             .expect("function replacement worked");
 
@@ -743,9 +740,8 @@ 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, |mut body| {
+            .replace_exported_func(original_fn_id, |body| {
                 body.unreachable();
-                body.local_func(vec![])
             })
             .expect("export function replacement worked");
 
@@ -785,10 +781,8 @@ 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, |module, params, results| {
-                let mut builder = FunctionBuilder::new(&mut module.types, params, results);
-                builder.func_body().i32_const(4321).drop();
-                Ok(builder.local_func(vec![]))
+            .replace_imported_func(original_fn_id, |body| {
+                body.i32_const(4321).drop();
             })
             .expect("import fn replacement worked");
 
@@ -825,10 +819,8 @@ 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, |module, params, results| {
-                let mut builder = FunctionBuilder::new(&mut module.types, params, results);
-                builder.func_body().unreachable();
-                Ok(builder.local_func(vec![]))
+            .replace_imported_func(original_fn_id, |body| {
+                body.unreachable();
             })
             .expect("import fn replacement worked");