Allow merging modules and skipping DWARF reading on load

src/module/config.rs

@@ -18,6 +18,7 @@ pub struct ModuleConfig {
     pub(crate) on_parse:
         Option<Box<dyn Fn(&mut Module, &IndicesToIds) -> Result<()> + Sync + Send + 'static>>,
     pub(crate) on_instr_loc: Option<Box<dyn Fn(&usize) -> InstrLocId + Sync + Send + 'static>>,
+    pub(crate) no_read_dwarf: bool,
 }
 
 impl Clone for ModuleConfig {
@@ -32,6 +33,7 @@ impl Clone for ModuleConfig {
             skip_producers_section: self.skip_producers_section,
             skip_name_section: self.skip_name_section,
             preserve_code_transform: self.preserve_code_transform,
+            no_read_dwarf: self.no_read_dwarf,
 
             // ... and this is left empty.
             on_parse: None,
@@ -54,6 +56,7 @@ impl fmt::Debug for ModuleConfig {
             ref preserve_code_transform,
             ref on_parse,
             ref on_instr_loc,
+            ref no_read_dwarf,
         } = self;
 
         f.debug_struct("ModuleConfig")
@@ -69,6 +72,7 @@ impl fmt::Debug for ModuleConfig {
             .field("preserve_code_transform", preserve_code_transform)
             .field("on_parse", &on_parse.as_ref().map(|_| ".."))
             .field("on_instr_loc", &on_instr_loc.as_ref().map(|_| ".."))
+            .field("no_read_dwarf", no_read_dwarf)
             .finish()
     }
 }
@@ -79,6 +83,11 @@ impl ModuleConfig {
         ModuleConfig::default()
     }
 
+    /// Sets a flag to whether DWARF debug sections are read for this module.
+    pub fn read_dwarf(&mut self, dw: bool) -> &mut ModuleConfig {
+        self.no_read_dwarf = !dw;
+        return self;
+    }
     /// Sets a flag to whether DWARF debug sections are generated for this
     /// module.
     ///
@@ -208,6 +217,12 @@ impl ModuleConfig {
         Module::parse(wasm, self)
     }
 
+    /// Parses an in-memory WebAssembly file into an existing `Module` using this
+    /// configuration.
+    pub fn parse_into(&self, wasm: &[u8], pre: &mut Module) -> Result<()> {
+        Module::parse_in(wasm, self, pre)
+    }
+
     /// Parses a WebAssembly file into a `Module` using this configuration.
     pub fn parse_file<P>(&self, path: P) -> Result<Module>
     where

src/module/functions/mod.rs

@@ -350,7 +350,7 @@ impl Module {
         // necessary and it's a bottleneck!
         let mut bodies = Vec::with_capacity(functions.len());
         for (i, (body, mut validator)) in functions.into_iter().enumerate() {
-            let index = (num_imports + i) as u32;
+            let index = (indices.num_fun_imports + i) as u32;
             let id = indices.get_func(index)?;
             let ty = match self.funcs.arena[id].kind {
                 FunctionKind::Uninitialized(ty) => ty,

src/module/imports.rs

@@ -160,6 +160,7 @@ impl Module {
                         entry.field.expect("module linking not supported"),
                         ty,
                     );
+                    ids.num_fun_imports += 1;
                     ids.push_func(id.0);
                 }
                 wasmparser::ImportSectionEntryType::Table(t) => {
@@ -171,6 +172,7 @@ impl Module {
                         t.maximum,
                         ty,
                     );
+                    ids.num_table_imports += 1;
                     ids.push_table(id.0);
                 }
                 wasmparser::ImportSectionEntryType::Memory(m) => {
@@ -184,6 +186,7 @@ impl Module {
                         m.initial as u32,
                         m.maximum.map(|m| m as u32),
                     );
+                    ids.num_mem_imports += 1;
                     ids.push_memory(id.0);
                 }
                 wasmparser::ImportSectionEntryType::Global(g) => {
@@ -193,6 +196,7 @@ impl Module {
                         ValType::parse(&g.content_type)?,
                         g.mutable,
                     );
+                    ids.num_global_imports += 1;
                     ids.push_global(id.0);
                 }
                 wasmparser::ImportSectionEntryType::Module(_)

src/module/mod.rs

@@ -38,6 +38,7 @@ pub use crate::module::producers::ModuleProducers;
 pub use crate::module::tables::{ModuleTables, Table, TableId};
 pub use crate::module::types::ModuleTypes;
 use crate::parse::IndicesToIds;
+use crate::FunctionBuilder;
 use anyhow::{bail, Context};
 use id_arena::Id;
 use log::warn;
@@ -128,9 +129,21 @@ impl Module {
         ModuleConfig::new().parse(wasm)
     }
 
+    /// Merge a module from the in-memory wasm buffer with the default configuration
+    pub fn merge_buffer(&mut self, wasm: &[u8]) -> Result<()> {
+        return Self::parse_in(wasm, &ModuleConfig::default(), self);
+    }
+
     fn parse(wasm: &[u8], config: &ModuleConfig) -> Result<Module> {
         let mut ret = Module::default();
         ret.config = config.clone();
+        Self::parse_in(wasm, config, &mut ret)?;
+        return Ok(ret);
+    }
+    fn parse_in(wasm: &[u8], config: &ModuleConfig, ret: &mut Module) -> Result<()> {
+        // let mut ret = Module::default();
+        // ret.config = config.clone();
+        let old_start = ret.start.take();
         let mut indices = IndicesToIds::default();
         let mut validator = Validator::new();
         validator.wasm_features(WasmFeatures {
@@ -299,19 +312,31 @@ impl Module {
             config.on_instr_loc.as_ref().map(|f| f.as_ref()),
         )
         .context("failed to parse code section")?;
-
-        ret.parse_debug_sections(debug_sections)
-            .context("failed to parse debug data section")?;
-
+        if !config.no_read_dwarf {
+            ret.parse_debug_sections(debug_sections)
+                .context("failed to parse debug data section")?;
+        }
         ret.producers
             .add_processed_by("walrus", env!("CARGO_PKG_VERSION"));
 
+        match ret.start {
+            Some(a) => {
+                if let Some(r) = old_start {
+                    let mut new_builder = FunctionBuilder::new(&mut ret.types, &[], &[]);
+                    new_builder.func_body().call(r).call(a).return_();
+                    let new = new_builder.local_func(vec![]);
+                    let new = ret.funcs.add_local(new);
+                    ret.start = Some(new);
+                }
+            }
+            None => ret.start = old_start,
+        }
         if let Some(on_parse) = &config.on_parse {
-            on_parse(&mut ret, &indices)?;
+            on_parse(ret, &indices)?;
         }
 
         log::debug!("parse complete");
-        Ok(ret)
+        Ok(())
     }
 
     /// Emit this module into a `.wasm` file at the given path.

src/parse.rs

@@ -26,6 +26,10 @@ pub struct IndicesToIds {
     elements: Vec<ElementId>,
     data: Vec<DataId>,
     locals: IdHashMap<Function, Vec<LocalId>>,
+    pub(crate) num_fun_imports: usize,
+    pub(crate) num_mem_imports: usize,
+    pub(crate) num_global_imports: usize,
+    pub(crate) num_table_imports: usize,
 }
 
 macro_rules! define_push_get {