Skip to content

Commit

Permalink
refactor: improve ergonomics of fn replace
Browse files Browse the repository at this point in the history 
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 <[email protected]>
  • Loading branch information
@vados-cosmonic
vados-cosmonic committed Oct 14, 2023
1 parent 632d220 commit 91cc935
Showing 1 changed file with 83 additions and 79 deletions.
162 changes: 83 additions & 79 deletions src/module/functions/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -447,98 +447,110 @@ 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
///
/// 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)
/// });
/// ```
///
/// 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<F>(&mut self, fid: FunctionId, fn_builder: F) -> Result<FunctionId>
where
F: FnOnce((&FuncParams, &FuncResults)) -> Result<LocalFunction>,
F: FnOnce(&mut Self, &FuncParams, &FuncResults) -> Result<LocalFunction>,
{
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((&params, &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, &params, &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<ValType>`) used by the original function
/// - the results (`Vec<ValType>`) 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)
/// });
/// ```
///
/// 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((&FuncParams, &FuncResults)) -> Result<LocalFunction>,
F: FnOnce(&mut Self, &FuncParams, &FuncResults) -> Result<LocalFunction>,
{
// 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((&params, &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, &params, &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");
}
}
}
Expand Down Expand Up @@ -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![]))
})
Expand Down Expand Up @@ -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![]))
})
Expand Down Expand Up @@ -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![]))
})
Expand Down Expand Up @@ -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![]))
})
Expand Down

0 comments on commit 91cc935

Please sign in to comment.