diff --git a/src/module/exports.rs b/src/module/exports.rs index 66e4bc10..b564fa11 100644 --- a/src/module/exports.rs +++ b/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 { diff --git a/src/module/functions/mod.rs b/src/module/functions/mod.rs index 2de3ab1c..a76a0b67 100644 --- a/src/module/functions/mod.rs +++ b/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) { - // 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 { self.arena.iter().map(|(_, f)| f) @@ -438,12 +427,10 @@ impl Module { pub fn get_exported_func_by_name(&self, name: impl AsRef) -> Result { 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())) } @@ -451,12 +438,10 @@ impl Module { pub fn get_imported_func_by_name(&self, name: impl AsRef) -> Result { 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())) } @@ -464,7 +449,7 @@ impl Module { /// /// 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 { + pub fn get_memory_id(&self) -> Result { 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(&mut self, name: impl AsRef, builder: F) -> Result<()> + pub fn replace_fn_by_id( + &mut self, + fid: FunctionId, + fn_name: impl AsRef, + fn_builder: F, + ) -> Result<()> where - F: FnOnce() -> Result>, + F: FnOnce() -> Result>, { - 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"); } } diff --git a/src/module/imports.rs b/src/module/imports.rs index de2f5bb4..1d3bb32b 100644 --- a/src/module/imports.rs +++ b/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 {