wip: fill in implementation of swap function

Signed-off-by: Victor Adossi <vadossi@cosmonic.com>

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<&Export> {
+        self.iter()
+            .find(|expt| match expt.item {
+                ExportItem::Function(_) => expt.name == name.as_ref(),
+                _ => false,
+            })
+            .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 {
@@ -124,17 +136,9 @@ impl ModuleExports {
         })
     }
 
-    /// Retrieve an exported function by ID
-    pub fn get_function_by_id(&self, id: FunctionId) -> Option<&Export> {
-        self.arena.iter().find_map(|(_, export)| match export.item {
-            ExportItem::Function(fid) if fid == id => Some(export),
-            _ => None,
-        })
-    }
-
     /// Check whether exports include the given function by ID
-    pub fn contains_function_by_id(&self, fid: FunctionId) -> bool {
-        self.get_function_by_id(fid).is_some()
+    pub fn contains_func_by_id(&self, fid: FunctionId) -> bool {
+        self.get_exported_func(fid).is_some()
     }
 }
 

src/module/functions/mod.rs

@@ -21,13 +21,19 @@ use crate::parse::IndicesToIds;
 use crate::tombstone_arena::{Id, Tombstone, TombstoneArena};
 use crate::ty::TypeId;
 use crate::ty::ValType;
-use crate::{ExportItem, FunctionBuilder, ImportKind, Memory, MemoryId};
+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`.
@@ -423,28 +429,6 @@ impl Module {
         Ok(())
     }
 
-    /// Retrieve an exported function by name
-    pub fn get_exported_func_by_name(&self, name: impl AsRef<str>) -> Result<FunctionId> {
-        self.exports
-            .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()))
-    }
-
-    /// Retrieve an imported function by name
-    pub fn get_imported_func_by_name(&self, name: impl AsRef<str>) -> Result<FunctionId> {
-        self.imports
-            .iter()
-            .find_map(|impt| match impt.kind {
-                ImportKind::Function(fid) if impt.name == name.as_ref() => Some(fid),
-                _ => None,
-            })
-            .with_context(|| format!("unable to find function export '{}'", name.as_ref()))
-    }
-
     /// 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),
@@ -465,52 +449,60 @@ impl Module {
     ///
     /// 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_fn_by_id<F>(
-        &mut self,
-        fid: FunctionId,
-        fn_name: impl AsRef<str>,
-        fn_builder: F,
-    ) -> Result<()>
+    pub fn replace_fn_by_id<F>(&mut self, fid: FunctionId, fn_builder: F) -> Result<()>
     where
-        F: FnOnce() -> Result<Option<LocalFunction>>,
+        F: FnOnce(Option<(&FuncParams, &FuncResults)>) -> Result<Option<LocalFunction>>,
     {
-        // Run the builder function to produce a new local function, or create a stub that is unreachable
-        let new_local_fn = if let Some(new_local_fn) = fn_builder().context("fn builder failed")? {
-            new_local_fn
-        } else {
-            let mut builder = FunctionBuilder::new(&mut self.types, &Vec::new(), &Vec::new());
-            builder.func_body().unreachable();
-            builder.local_func(vec![])
-        };
-
         // Replace a function import
-        if self.imports.contains_function_by_id(fid) {
-            self.funcs.add_local(new_local_fn);
+        if let Some(original_imported_fn) = self.imports.get_func_by_id(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 = if let Some(new_local_fn) =
+                    fn_builder(Some((&params, &results))).context("fn builder failed")?
+                {
+                    new_local_fn
+                } else {
+                    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(());
+            }
+        }
 
-            // TODO: Replace an imported function
+        // Replace the existing exported function if present
+        if let Some(exported_fn) = self.exports.get_exported_func(fid) {
+            // Run the builder function to produce a new local function, or create a stub that is unreachable
+            let new_local_fn = if let Some(new_local_fn) =
+                fn_builder(None).context("fn builder failed")?
+            {
+                new_local_fn
+            } else {
+                let mut builder = FunctionBuilder::new(&mut self.types, &Vec::new(), &Vec::new());
+                builder.func_body().unreachable();
+                builder.local_func(vec![])
+            };
 
-            return Ok(());
-        }
+            let new_fid = self.funcs.add_local(new_local_fn);
 
-        // Replace a function export
-        if self.exports.contains_function_by_id(fid) {
-            self.funcs.add_local(new_local_fn);
+            let export = self.exports.get_mut(exported_fn.id());
+            export.item = ExportItem::Function(new_fid);
 
-            // Create or replace export
-            if let Some(exported_fn) = self.exports.get_function_by_id(fid) {
-                // If the function exists, replace it
-                let export = self.exports.get_mut(exported_fn.id());
-                export.item = ExportItem::Function(fid);
-            } else {
-                // If the function doesn't exist, create a new export
-
-                // TODO: we should never get here, because we check for existing function up front.
-                // how can we add the function if it doesn't exist *AND* know if we wanted to 
-                // add an export or an import specifically?
-
-                self.exports
-                    .add(fn_name.as_ref(), ExportItem::Function(fid));
-            }
             return Ok(());
         }
 

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>;
@@ -104,8 +105,24 @@ 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<&Import> {
+        self.iter()
+            .find(|impt| match impt.kind {
+                ImportKind::Function(_) => {
+                    impt.module == module.as_ref() && impt.name == name.as_ref()
+                }
+                _ => false,
+            })
+            .with_context(|| format!("unable to find function export '{}'", name.as_ref()))
+    }
+
     /// Retrieve an imported function by ID
-    pub fn get_function_by_id(&self, id: FunctionId) -> Option<&Import> {
+    pub fn get_func_by_id(&self, id: FunctionId) -> Option<&Import> {
         self.arena.iter().find_map(|(_, import)| match import.kind {
             ImportKind::Function(fid) if fid == id => Some(import),
             _ => None,
@@ -114,7 +131,7 @@ impl ModuleImports {
 
     /// Check whether imports include a function with the given ID
     pub fn contains_function_by_id(&self, fid: FunctionId) -> bool {
-        self.get_function_by_id(fid).is_some()
+        self.get_func_by_id(fid).is_some()
     }
 }