From 91cc9351937753cec269cf3ad7eed8d92bef1dd2 Mon Sep 17 00:00:00 2001 From: Victor Adossi Date: Sun, 15 Oct 2023 08:43:37 +0900 Subject: [PATCH] refactor: improve ergonomics of fn replace After trialing `replace_(imported|exported)_func` in WASI-Virt, it's clear that the ergonomics around the builder function need to be improved. `FunctionBuilder` (particularly `FunctionBuilder::new()` is difficult to use without a mutable borrow of the module itself. This commit refactors `replace_(imported|exported)_func` in order to pass through the mutable borrow which makes it easier to use `FunctionBuilder`s. Signed-off-by: Victor Adossi --- src/module/functions/mod.rs | 162 ++++++++++++++++++------------------ 1 file changed, 83 insertions(+), 79 deletions(-) diff --git a/src/module/functions/mod.rs b/src/module/functions/mod.rs index dca3bb49..ce5fa948 100644 --- a/src/module/functions/mod.rs +++ b/src/module/functions/mod.rs @@ -447,16 +447,19 @@ 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 + /// /// For example, if you wanted to replace an exported function with a no-op, /// /// ```ignore - /// // Since `FunctionBuilder` requires a mutable pointer to the module's types, - /// // we must build it *outside* the closure and `move` it in - /// let mut builder = FunctionBuilder::new(&mut module.types, &[], &[]); - /// - /// module.replace_exported_func(fid, move || { + /// module.replace_exported_func(fid, |module, params, results| { + /// let mut builder = FunctionBuilder::new(&mut module.types, params, results); /// builder.func_body().unreachable(); - /// builder.local_func(vec![]) + /// let new_fn = builder.local_func(vec![]); + /// Ok(new_fn) /// }); /// ``` /// @@ -464,46 +467,51 @@ impl Module { /// after it has been inserted into the module as an export. pub fn replace_exported_func(&mut self, fid: FunctionId, fn_builder: F) -> Result where - F: FnOnce((&FuncParams, &FuncResults)) -> Result, + F: FnOnce(&mut Self, &FuncParams, &FuncResults) -> Result, { - match (self.exports.get_exported_func(fid), self.funcs.get(fid)) { - ( - Some(exported_fn), - Function { - kind: FunctionKind::Local(lf), - .. - }, - ) => { - // Retrieve the params & result types for the exported (local) function - let ty = self.types.get(lf.ty()); - let (params, results) = (ty.params().to_vec(), ty.results().to_vec()); - - // Add the function produced by `fn_builder` as a local function, - let new_fid = self.funcs.add_local( - fn_builder((¶ms, &results)).context("export fn builder failed")?, - ); - - // Mutate the existing export to use the new local function - let export = self.exports.get_mut(exported_fn.id()); - export.item = ExportItem::Function(new_fid); - - Ok(new_fid) - } - // The export didn't exist, or the function isn't the kind we expect - _ => bail!("cannot replace function [{fid:?}], it is not an exported function"), + let original_export_id = self + .exports + .get_exported_func(fid) + .map(|e| e.id()) + .context("no exported function with ID [{fid:?}]")?; + + if let Function { + kind: FunctionKind::Local(lf), + .. + } = self.funcs.get(fid) + { + // Retrieve the params & result types for the exported (local) function + let ty = self.types.get(lf.ty()); + let (params, results) = (ty.params().to_vec(), ty.results().to_vec()); + + // Add the function produced by `fn_builder` as a local function + let func = fn_builder(self, ¶ms, &results).context("export fn builder failed")?; + let new_fn_id = self.funcs.add_local(func); + + // Mutate the existing export to use the new local function + let export = self.exports.get_mut(original_export_id); + export.item = ExportItem::Function(new_fn_id); + Ok(new_fn_id) + } else { + bail!("cannot replace function [{fid:?}], it is not an exported function"); } } /// Replace a single imported function with the result of the provided builder function. /// - /// ```ignore - /// // Since `FunctionBuilder` requires a mutable pointer to the module's types, - /// // we must build it *outside* the closure and `move` it in - /// let mut builder = FunctionBuilder::new(&mut module.types, &[], &[]); + /// 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 /// - /// module.replace_imported_func(fid, move || { + /// 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); /// builder.func_body().unreachable(); - /// builder.local_func(vec![]) + /// let new_fn = builder.local_func(vec![]); + /// Ok(new_fn) /// }); /// ``` /// @@ -511,34 +519,38 @@ impl Module { /// 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((&FuncParams, &FuncResults)) -> Result, + F: FnOnce(&mut Self, &FuncParams, &FuncResults) -> Result, { - // If the function is in the imports, replace it - match (self.imports.get_imported_func(fid), self.funcs.get(fid)) { - ( - Some(original_imported_fn), - Function { - kind: FunctionKind::Import(ImportedFunction { ty: tid, .. }), - .. - }, - ) => { - // Retrieve the params & result types for the imported function - let ty = self.types.get(*tid); - let (params, results) = (ty.params().to_vec(), ty.results().to_vec()); - - // Mutate the existing function, changing it from a FunctionKind::ImportedFunction - // to the local function produced by running the provided `fn_builder` - let func = self.funcs.get_mut(fid); - func.kind = FunctionKind::Local( - fn_builder((¶ms, &results)).context("import fn builder failed")?, - ); + let original_import_id = self + .imports + .get_imported_func(fid) + .map(|import| import.id()) + .context("no exported function with ID [{fid:?}]")?; + + if let Function { + kind: FunctionKind::Import(ImportedFunction { ty: tid, .. }), + .. + } = self.funcs.get(fid) + { + // Retrieve the params & result types for the imported function + let ty = self.types.get(*tid); + let (params, results) = (ty.params().to_vec(), ty.results().to_vec()); - self.imports.delete(original_imported_fn.id()); + // Build the new function + let new_func_kind = FunctionKind::Local( + fn_builder(self, ¶ms, &results).context("import fn builder failed")?, + ); - Ok(fid) - } - // The export didn't exist, or the function isn't the kind we expect - _ => bail!("cannot replace function [{fid:?}], it is not an imported function"), + // Mutate the existing function, changing it from a FunctionKind::ImportedFunction + // to the local function produced by running the provided `fn_builder` + let func = self.funcs.get_mut(fid); + func.kind = new_func_kind; + + self.imports.delete(original_import_id); + + Ok(fid) + } else { + bail!("cannot replace function [{fid:?}], it is not an imported function"); } } } @@ -683,12 +695,10 @@ mod tests { let original_fn_id: FunctionId = builder.finish(vec![], &mut module.funcs); let original_export_id = module.exports.add("dummy", original_fn_id); - // Create builder to use inside closure - let mut builder = FunctionBuilder::new(&mut module.types, &[], &[]); - // Replace the existing function with a new one with a reversed const value let new_fn_id = module - .replace_exported_func(original_fn_id, move |_| { + .replace_exported_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![])) }) @@ -728,12 +738,10 @@ mod tests { let original_fn_id: FunctionId = builder.finish(vec![], &mut module.funcs); let original_export_id = module.exports.add("dummy", original_fn_id); - // Create builder to use inside closure - let mut builder = FunctionBuilder::new(&mut module.types, &[], &[]); - // Replace the existing function with a new one with a reversed const value let new_fn_id = module - .replace_exported_func(original_fn_id, move |_| { + .replace_exported_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![])) }) @@ -773,12 +781,10 @@ mod tests { let types = module.types.add(&[], &[]); let (original_fn_id, original_import_id) = module.add_import_func("mod", "dummy", types); - // Create builder to use inside closure - let mut builder = FunctionBuilder::new(&mut module.types, &[], &[]); - // Replace the existing function with a new one with a reversed const value let new_fn_id = module - .replace_imported_func(original_fn_id, |_| { + .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![])) }) @@ -815,12 +821,10 @@ mod tests { let types = module.types.add(&[], &[]); let (original_fn_id, original_import_id) = module.add_import_func("mod", "dummy", types); - // Create builder to use inside closure - let mut builder = FunctionBuilder::new(&mut module.types, &[], &[]); - // Replace the existing function with a new one with a reversed const value let new_fn_id = module - .replace_imported_func(original_fn_id, |_| { + .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![])) })