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]>

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 {

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, FunctionBuilder, 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,103 @@ 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 function that is either an import or an export
+    /// with the result of the provided builder function.
+    ///
+    /// When called, if `builder` produces a None value, the function in question will be
+    /// replaced with a stub that does nothing (more precisely, a function with an unreachable body).
+    pub fn replace_func_by_id<F>(&mut self, fid: FunctionId, fn_builder: F) -> Result<()>
+    where
+        F: FnOnce((&FuncParams, &FuncResults)) -> Result<Option<LocalFunction>>,
+    {
+        // If the function is in the imports, replace it
+        if let Some(original_imported_fn) = self.imports.get_imported_func(fid) {
+            // Change types of existing function
+            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());
+
+                // Run the builder function to produce a new local function, or create a stub that is unreachable
+                let new_local_fn = match fn_builder((&params, &results))
+                    .context("import fn builder failed")?
+                {
+                    Some(f) => f,
+                    None => {
+                        let mut builder = FunctionBuilder::new(&mut self.types, &params, &results);
+                        builder.func_body().unreachable();
+                        builder.local_func(vec![])
+                    }
+                };
+
+                let func = self.funcs.get_mut(fid);
+                func.kind = FunctionKind::Local(new_local_fn);
+
+                self.imports.delete(original_imported_fn.id());
+
+                return Ok(());
+            } else {
+                bail!("ID [{fid:?}] did not resolve to an imported function");
+            }
+        }
+
+        // If the function is in the exports, replace it
+        if let Some(exported_fn) = self.exports.get_exported_func(fid) {
+            // Change types of existing function
+            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());
+
+                // Run the builder function to produce a new local function, or create a stub that is unreachable
+                let new_local_fn = match fn_builder((&params, &results))
+                    .context("export fn builder failed")?
+                {
+                    Some(f) => f,
+                    None => {
+                        let mut builder = FunctionBuilder::new(&mut self.types, &params, &results);
+                        builder.func_body().unreachable();
+                        builder.local_func(vec![])
+                    }
+                };
+
+                let new_fid = self.funcs.add_local(new_local_fn);
+                let export = self.exports.get_mut(exported_fn.id());
+                export.item = ExportItem::Function(new_fid);
+
+                return Ok(());
+            } else {
+                bail!("ID [{fid:?}] did not resolve to an exported function");
+            }
+        }
+
+        bail!("function with ID [{fid:?}] is neither an import or an export");
+    }
 }
 
 fn used_local_functions<'a>(cx: &mut EmitContext<'a>) -> Vec<(FunctionId, &'a LocalFunction, u64)> {

src/module/imports.rs

@@ -1,11 +1,12 @@
 //! A wasm module's imports.
 
+use anyhow::{bail, Context};
+
 use crate::emit::{Emit, EmitContext};
 use crate::parse::IndicesToIds;
 use crate::tombstone_arena::{Id, Tombstone, TombstoneArena};
 use crate::{FunctionId, GlobalId, MemoryId, Result, TableId};
 use crate::{Module, TypeId, ValType};
-use anyhow::bail;
 
 /// The id of an import.
 pub type ImportId = Id<Import>;
@@ -103,6 +104,32 @@ impl ModuleImports {
 
         Some(import?.0)
     }
+
+    /// Retrieve an imported function by name, including the module in which it resides
+    pub fn get_func_by_name(
+        &self,
+        module: impl AsRef<str>,
+        name: impl AsRef<str>,
+    ) -> Result<FunctionId> {
+        self.iter()
+            .find_map(|impt| match impt.kind {
+                ImportKind::Function(fid)
+                    if impt.module == module.as_ref() && impt.name == name.as_ref() =>
+                {
+                    Some(fid)
+                }
+                _ => None,
+            })
+            .with_context(|| format!("unable to find function export '{}'", name.as_ref()))
+    }
+
+    /// Retrieve an imported function by ID
+    pub fn get_imported_func(&self, id: FunctionId) -> Option<&Import> {
+        self.arena.iter().find_map(|(_, import)| match import.kind {
+            ImportKind::Function(fid) if fid == id => Some(import),
+            _ => None,
+        })
+    }
 }
 
 impl Module {

src/module/memories.rs

@@ -115,6 +115,11 @@ impl ModuleMemories {
     pub fn iter_mut(&mut self) -> impl Iterator<Item = &mut Memory> {
         self.arena.iter_mut().map(|(_, f)| f)
     }
+
+    /// Get the number of memories in this module
+    pub fn len(&self) -> usize {
+        self.arena.len()
+    }
 }
 
 impl Module {