feat: upstream wasi virt walrus ops (#248)

* feat: upstream ops from WASI-virt

WASI-virt contains functions that are helpful for manipulating modules
and dealing with exports/imports, which would be helpful to an even
wider group if upstreamed here to walrus.

This commit copies and upstreams some operations that were introduced
in WASI-virt for wider use via walrus.

See also: bytecodealliance/WASI-Virt#20

Signed-off-by: Victor Adossi <[email protected]>

* fix(tests): update WAT used in tests

The excellent pre-existing test suite contains a lot of WAT files as
test-cases that seem to have older syntax in them.

The breakage is introduced because the current test suite runs stable,
latest and nightly versions of Rust while conducting tests. While the
rest of the Wasm toolchain (binaryen, wabt, etc) are hard-coded to versions that should
stay compatible with the WAT as they were written, the upstreaming of
various changes to Rust itself seems to be causing failures when using
the existing WAT for tests.

This PR updates the WAT (mostly operation renaming) that have changed
in order to get tests running again.

Signed-off-by: Victor Adossi <[email protected]>

---------

Signed-off-by: Victor Adossi <[email protected]>

src/module/exports.rs

@@ -1,5 +1,7 @@
 //! Exported items in a wasm module.
 
+use anyhow::Context;
+
 use crate::emit::{Emit, EmitContext};
 use crate::parse::IndicesToIds;
 use crate::tombstone_arena::{Id, Tombstone, TombstoneArena};
@@ -100,6 +102,16 @@ impl ModuleExports {
         })
     }
 
+    /// Retrieve an exported function by name
+    pub fn get_func_by_name(&self, name: impl AsRef<str>) -> Result<FunctionId> {
+        self.iter()
+            .find_map(|expt| match expt.item {
+                ExportItem::Function(fid) if expt.name == name.as_ref() => Some(fid),
+                _ => None,
+            })
+            .with_context(|| format!("unable to find function export '{}'", name.as_ref()))
+    }
+
     /// Get a reference to a table export given its table id.
     pub fn get_exported_table(&self, t: TableId) -> Option<&Export> {
         self.iter().find(|e| match e.item {
@@ -354,6 +366,16 @@ mod tests {
         assert!(module.exports.get_exported_func(fn_id).is_none());
     }
 
+    #[test]
+    fn get_func_by_name() {
+        let mut module = Module::default();
+        let fn_id: FunctionId = always_the_same_id();
+        let export_id: ExportId = module.exports.add("dummy", fn_id);
+        assert!(module.exports.get_func_by_name("dummy").is_ok());
+        module.exports.delete(export_id);
+        assert!(module.exports.get_func_by_name("dummy").is_err());
+    }
+
     #[test]
     fn iter_mut_can_update_export_item() {
         let mut module = Module::default();

src/module/functions/mod.rs

@@ -1,5 +1,15 @@
 //! Functions within a wasm module.
 
+use std::cmp;
+use std::collections::BTreeMap;
+
+use anyhow::{bail, Context};
+use wasm_encoder::Encode;
+use wasmparser::{FuncValidator, FunctionBody, Range, ValidatorResources};
+
+#[cfg(feature = "parallel")]
+use rayon::prelude::*;
+
 mod local_function;
 
 use crate::emit::{Emit, EmitContext};
@@ -11,19 +21,19 @@ use crate::parse::IndicesToIds;
 use crate::tombstone_arena::{Id, Tombstone, TombstoneArena};
 use crate::ty::TypeId;
 use crate::ty::ValType;
-use std::cmp;
-use std::collections::BTreeMap;
-use wasm_encoder::Encode;
-use wasmparser::{FuncValidator, FunctionBody, Range, ValidatorResources};
-
-#[cfg(feature = "parallel")]
-use rayon::prelude::*;
+use crate::{ExportItem, Memory, MemoryId};
 
 pub use self::local_function::LocalFunction;
 
 /// A function identifier.
 pub type FunctionId = Id<Function>;
 
+/// Parameter(s) to a function
+pub type FuncParams = Vec<ValType>;
+
+/// Result(s) of a given function
+pub type FuncResults = Vec<ValType>;
+
 /// A wasm function.
 ///
 /// Either defined locally or externally and then imported; see `FunctionKind`.
@@ -418,6 +428,119 @@ impl Module {
 
         Ok(())
     }
+
+    /// Retrieve the ID for the first exported memory.
+    ///
+    /// This method does not work in contexts with [multi-memory enabled](https://github.com/WebAssembly/multi-memory),
+    /// and will error if more than one memory is present.
+    pub fn get_memory_id(&self) -> Result<MemoryId> {
+        if self.memories.len() > 1 {
+            bail!("multiple memories unsupported")
+        }
+
+        self.memories
+            .iter()
+            .next()
+            .map(Memory::id)
+            .context("module does not export a memory")
+    }
+
+    /// Replace a single exported function with the result of the provided builder 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 || {
+    ///     builder.func_body().unreachable();
+    ///     builder.local_func(vec![])
+    /// });
+    /// ```
+    ///
+    /// 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>,
+    {
+        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"),
+        }
+    }
+
+    /// 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, &[], &[]);
+    ///
+    /// module.replace_imported_func(fid, move || {
+    ///     builder.func_body().unreachable();
+    ///     builder.local_func(vec![])
+    /// });
+    /// ```
+    ///
+    /// 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>,
+    {
+        // 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")?,
+                );
+
+                self.imports.delete(original_imported_fn.id());
+
+                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"),
+        }
+    }
 }
 
 fn used_local_functions<'a>(cx: &mut EmitContext<'a>) -> Vec<(FunctionId, &'a LocalFunction, u64)> {
@@ -535,3 +658,192 @@ impl Emit for ModuleFunctions {
         cx.code_transform.instruction_map = instruction_map.into_iter().collect();
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::{Export, FunctionBuilder, Module};
+
+    #[test]
+    fn get_memory_id() {
+        let mut module = Module::default();
+        let expected_id = module.memories.add_local(false, 0, None);
+        assert!(module.get_memory_id().is_ok_and(|id| id == expected_id));
+    }
+
+    /// Running `replace_exported_func` with a closure that builds
+    /// a function should replace the existing function with the new one
+    #[test]
+    fn replace_exported_func() {
+        let mut module = Module::default();
+
+        // Create original function
+        let mut builder = FunctionBuilder::new(&mut module.types, &[], &[]);
+        builder.func_body().i32_const(1234).drop();
+        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 |_| {
+                builder.func_body().i32_const(4321).drop();
+                Ok(builder.local_func(vec![]))
+            })
+            .expect("function replacement worked");
+
+        assert!(
+            module.exports.get_exported_func(original_fn_id).is_none(),
+            "replaced function cannot be gotten by ID"
+        );
+
+        // Ensure the function was replaced
+        match module
+            .exports
+            .get_exported_func(new_fn_id)
+            .expect("failed to unwrap exported func")
+        {
+            exp @ Export {
+                item: ExportItem::Function(fid),
+                ..
+            } => {
+                assert_eq!(*fid, new_fn_id, "retrieved function ID matches");
+                assert_eq!(exp.id(), original_export_id, "export ID is unchanged");
+            }
+            _ => panic!("expected an Export with a Function inside"),
+        }
+    }
+
+    /// Running `replace_exported_func` with a closure that returns None
+    /// should replace the function with a generated no-op function
+    #[test]
+    fn replace_exported_func_generated_no_op() {
+        let mut module = Module::default();
+
+        // Create original function
+        let mut builder = FunctionBuilder::new(&mut module.types, &[], &[]);
+        builder.func_body().i32_const(1234).drop();
+        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 |_| {
+                builder.func_body().unreachable();
+                Ok(builder.local_func(vec![]))
+            })
+            .expect("export function replacement worked");
+
+        assert!(
+            module.exports.get_exported_func(original_fn_id).is_none(),
+            "replaced export function cannot be gotten by ID"
+        );
+
+        // Ensure the function was replaced
+        match module
+            .exports
+            .get_exported_func(new_fn_id)
+            .expect("failed to unwrap exported func")
+        {
+            exp @ Export {
+                item: ExportItem::Function(fid),
+                name,
+                ..
+            } => {
+                assert_eq!(name, "dummy", "function name on export is unchanged");
+                assert_eq!(*fid, new_fn_id, "retrieved function ID matches");
+                assert_eq!(exp.id(), original_export_id, "export ID is unchanged");
+            }
+            _ => panic!("expected an Export with a Function inside"),
+        }
+    }
+
+    /// Running `replace_imported_func` with a closure that builds
+    /// a function should replace the existing function with the new one
+    #[test]
+    fn replace_imported_func() {
+        let mut module = Module::default();
+
+        // Create original import function
+        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, |_| {
+                builder.func_body().i32_const(4321).drop();
+                Ok(builder.local_func(vec![]))
+            })
+            .expect("import fn replacement worked");
+
+        assert!(
+            !module.imports.iter().any(|i| i.id() == original_import_id),
+            "original import is missing",
+        );
+
+        assert!(
+            module.imports.get_imported_func(original_fn_id).is_none(),
+            "replaced import function cannot be gotten by ID"
+        );
+
+        assert!(
+            module.imports.get_imported_func(new_fn_id).is_none(),
+            "new import function cannot be gotten by ID (it is now local)"
+        );
+
+        assert!(
+            matches!(module.funcs.get(new_fn_id).kind, FunctionKind::Local(_)),
+            "new local function has the right kind"
+        );
+    }
+
+    /// Running `replace_imported_func` with a closure that returns None
+    /// should replace the function with a generated no-op function
+    #[test]
+    fn replace_imported_func_generated_no_op() {
+        let mut module = Module::default();
+
+        // Create original import function
+        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, |_| {
+                builder.func_body().unreachable();
+                Ok(builder.local_func(vec![]))
+            })
+            .expect("import fn replacement worked");
+
+        assert!(
+            !module.imports.iter().any(|i| i.id() == original_import_id),
+            "original import is missing",
+        );
+
+        assert!(
+            module.imports.get_imported_func(original_fn_id).is_none(),
+            "replaced import function cannot be gotten by ID"
+        );
+
+        assert!(
+            module.imports.get_imported_func(new_fn_id).is_none(),
+            "new import function cannot be gotten by ID (it is now local)"
+        );
+
+        assert!(
+            matches!(module.funcs.get(new_fn_id).kind, FunctionKind::Local(_)),
+            "new local function has the right kind"
+        );
+    }
+}