refactor: streamline functions, more consistent naming

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

wip: changes to facilitate by-id usage

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

refactor: more consistent naming, concise fns

- use find() for easier to read logic

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

refactor: simplify contains_function()

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

refactor: more consistent naming

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

refactor: use find_map()

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

wip: remove unnecessary comments, add TODO

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

src/module/exports.rs

@@ -123,6 +123,19 @@ impl ModuleExports {
             _ => false,
         })
     }
+
+    /// 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()
+    }
 }
 
 impl Module {

src/module/functions/mod.rs

@@ -21,7 +21,7 @@ use crate::parse::IndicesToIds;
 use crate::tombstone_arena::{Id, Tombstone, TombstoneArena};
 use crate::ty::TypeId;
 use crate::ty::ValType;
-use crate::{ExportItem, ImportKind, Memory, MemoryId};
+use crate::{ExportItem, FunctionBuilder, ImportKind, Memory, MemoryId};
 
 pub use self::local_function::LocalFunction;
 
@@ -210,17 +210,6 @@ impl ModuleFunctions {
         self.arena.delete(id);
     }
 
-    /// Remove a function from this module by name that corresponds to it's import or export.
-    ///
-    /// Unlike [`by_name()`](ModuleFunctions::by_name), "name" corresponds to either a functions import name or it's export name.
-    pub fn delete_by_name(&mut self, name: impl AsRef<str>) {
-        // TODO: How do we differentiate between imports/exports here?
-        // add some sort of `FunctionLocation` enum (ex. FunctionLocation::ModuleImports/FunctionLocation::ModuleExports?)
-        //
-        // FunctionKind is almost there as a marker but there's no "Export"
-        todo!()
-    }
-
     /// Get a shared reference to this module's functions.
     pub fn iter(&self) -> impl Iterator<Item = &Function> {
         self.arena.iter().map(|(_, f)| f)
@@ -438,33 +427,29 @@ impl Module {
     pub fn get_exported_func_by_name(&self, name: impl AsRef<str>) -> Result<FunctionId> {
         self.exports
             .iter()
-            // Find the export with the correct name and internal type
-            .filter_map(|expt| match expt.item {
+            .find_map(|expt| match expt.item {
                 ExportItem::Function(fid) if expt.name == name.as_ref() => Some(fid),
                 _ => None,
             })
-            .nth(0)
             .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 the export with the correct name and internal type
-            .filter_map(|impt| match impt.kind {
+            .find_map(|impt| match impt.kind {
                 ImportKind::Function(fid) if impt.name == name.as_ref() => Some(fid),
                 _ => None,
             })
-            .nth(0)
             .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),
     /// and will error if more than one memory is present.
-    pub fn get_primary_memory_id(&self) -> Result<MemoryId> {
+    pub fn get_memory_id(&self) -> Result<MemoryId> {
         if self.memories.len() > 1 {
             bail!("multiple memories unsupported")
         }
@@ -480,18 +465,56 @@ 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_name<F>(&mut self, name: impl AsRef<str>, builder: F) -> Result<()>
+    pub fn replace_fn_by_id<F>(
+        &mut self,
+        fid: FunctionId,
+        fn_name: impl AsRef<str>,
+        fn_builder: F,
+    ) -> Result<()>
     where
-        F: FnOnce() -> Result<Option<Function>>,
+        F: FnOnce() -> Result<Option<LocalFunction>>,
     {
-        let built_fn = builder().context("fn builder failed")?;
+        // 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);
+
+            // TODO: Replace an imported function
+
+            return Ok(());
+        }
 
-        // TODO: Hmnn.... how should we distinguish between which one to go for?
-        // this is a similar problem to the more general delete_by_name() above.
-        let imported_fn_id = self.get_imported_func_by_name(name.as_ref());
-        let exported_fn_id = self.get_exported_func_by_name(name.as_ref());
+        // Replace a function export
+        if self.exports.contains_function_by_id(fid) {
+            self.funcs.add_local(new_local_fn);
 
-        Ok(())
+            // 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(());
+        }
+
+        bail!("function with ID [{fid:?}] is neither an import or an export");
     }
 }
 

src/module/imports.rs

@@ -103,6 +103,19 @@ impl ModuleImports {
 
         Some(import?.0)
     }
+
+    /// Retrieve an imported function by ID
+    pub fn get_function_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,
+        })
+    }
+
+    /// 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()
+    }
 }
 
 impl Module {