From 01e2356832d606beb81ba560c70a6100f6b60dab Mon Sep 17 00:00:00 2001 From: Victor Adossi Date: Wed, 18 Oct 2023 09:29:28 +0900 Subject: [PATCH] refactor: build local func after builer_fn finishes Signed-off-by: Victor Adossi --- src/module/functions/mod.rs | 72 +++++++++++++++++-------------------- 1 file changed, 32 insertions(+), 40 deletions(-) diff --git a/src/module/functions/mod.rs b/src/module/functions/mod.rs index 76abb866..1f547b34 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, 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`) used by the original function - /// - the results (`Vec`) 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 { 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, ¶ms, &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`) used by the original function - /// - the results (`Vec`) 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(&mut self, fid: FunctionId, fn_builder: F) -> Result - where - F: FnOnce(&mut Self, &FuncParams, &FuncResults) -> Result, - { + pub fn replace_imported_func( + &mut self, + fid: FunctionId, + builder_fn: impl FnOnce(&mut InstrSeqBuilder), + ) -> Result { 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, ¶ms, &results).context("import fn builder failed")?, - ); + let mut builder = FunctionBuilder::new(&mut self.types, ¶ms, &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");