WIP: Don't return code transforms, apply them to custom sections && s…

…moke test for code transforms

crates/fuzz/src/lib.rs

@@ -75,7 +75,6 @@ impl Config {
         walrus::Module::from_buffer(&wasm)
             .unwrap()
             .emit_wasm()
-            .unwrap()
             .into()
     }
 

crates/tests/tests/custom_sections.rs

@@ -1,7 +1,7 @@
 //! Tests for working with custom sections that `walrus` doesn't know about.
 
 use std::borrow::Cow;
-use walrus::{CustomSection, Module, ModuleConfig};
+use walrus::{CodeTransform, CustomSection, Module, ModuleConfig, ValType};
 
 #[derive(Clone, Debug, Default, PartialEq, Eq)]
 struct HelloCustomSection(String);
@@ -48,7 +48,7 @@ fn round_trip_unkown_custom_sections() {
         [(world_id.into(), world.data())]
     );
 
-    let wasm = module.emit_wasm().unwrap().wasm;
+    let wasm = module.emit_wasm();
     let mut module = config.parse(&wasm).unwrap();
 
     let world_round_tripped = module.customs.remove_raw("hello").unwrap();
@@ -58,6 +58,95 @@ fn round_trip_unkown_custom_sections() {
     assert_eq!(new_world.data(), world.data());
     module.customs.add(new_world);
 
-    let new_wasm = module.emit_wasm().unwrap().wasm;
+    let new_wasm = module.emit_wasm();
     assert_eq!(wasm, new_wasm);
 }
+
+// Insert a `(drop (i32.const 0))` at the start of the function and assert that
+// all instructions are pushed down by the size of a `(drop (i32.const 0))`,
+// which is 3.
+#[test]
+fn smoke_test_code_transform() {
+    use std::sync::atomic::{AtomicUsize, Ordering};
+
+    static APPLIED_CODE_TRANSFORM: AtomicUsize = AtomicUsize::new(0);
+
+    #[derive(Debug)]
+    struct CheckCodeTransform;
+    impl CustomSection for CheckCodeTransform {
+        fn name(&self) -> &str {
+            "check-code-transform"
+        }
+
+        fn data(&self) -> Cow<[u8]> {
+            vec![].into()
+        }
+
+        fn apply_code_transform(&mut self, transform: &CodeTransform) {
+            APPLIED_CODE_TRANSFORM.store(1, Ordering::SeqCst);
+            assert!(!transform.is_empty());
+            for (input_offset, output_offset) in transform.iter().cloned() {
+                println!("0x{:x} -> 0x{:x}", input_offset, output_offset);
+
+                // TODO: this assert currently fails
+                //
+                // assert_eq!(input_offset + 3, output_offset);
+            }
+        }
+    }
+
+    let mut config = ModuleConfig::new();
+    config.generate_producers_section(false);
+
+    let wasm = {
+        let mut module = Module::with_config(config.clone());
+
+        let ty = module.types.add(&[], &[ValType::I32]);
+
+        let mut builder = walrus::FunctionBuilder::new();
+        let exprs = vec![builder.i32_const(1337)];
+        let locals = vec![];
+        let f_id = builder.finish(ty, locals, exprs, &mut module);
+
+        module.exports.add("f", f_id);
+
+        module.emit_wasm()
+    };
+
+    use std::fs;
+    use std::path::Path;
+
+    fs::write("input.wasm", &wasm).unwrap();
+    println!(
+        "=== input.wasm ===\n {}",
+        walrus_tests_utils::wasm2wat(Path::new("input.wasm"))
+    );
+
+    config.preserve_code_transform(true);
+
+    let mut module = config.parse(&wasm).unwrap();
+    module.customs.add(CheckCodeTransform);
+
+    for (_id, f) in module.funcs.iter_local_mut() {
+        let builder = f.builder_mut();
+        let zero = builder.i32_const(0);
+        let drop = builder.drop(zero);
+
+        let entry = f.entry_block();
+        f.block_mut(entry).exprs.insert(0, drop);
+    }
+
+    // Emit the new, transformed wasm. This should trigger the
+    // `apply_code_transform` method to be called.
+    let wasm = module.emit_wasm();
+
+    fs::write("output.wasm", &wasm).unwrap();
+    println!(
+        "=== output.wasm ===\n {}",
+        walrus_tests_utils::wasm2wat(Path::new("output.wasm"))
+    );
+
+    assert_eq!(APPLIED_CODE_TRANSFORM.load(Ordering::SeqCst), 1);
+
+    panic!("TODO: make the commented out assertion in `apply_code_transform` pass");
+}

crates/tests/tests/spec-tests.rs

@@ -84,19 +84,15 @@ fn run(wast: &Path) -> Result<(), failure::Error> {
             }
             cmd => {
                 let wasm = fs::read(&path)?;
-                let wasm = config
+                let mut wasm = config
                     .parse(&wasm)
                     .context(format!("error parsing wasm (line {})", line))?;
-                let wasm1 = wasm
-                    .emit_wasm()
-                    .context(format!("error emitting wasm (line {})", line))?
-                    .wasm;
+                let wasm1 = wasm.emit_wasm();
                 fs::write(&path, &wasm1)?;
                 let wasm2 = config
                     .parse(&wasm1)
-                    .and_then(|m| m.emit_wasm())
-                    .context(format!("error re-parsing wasm (line {})", line))?
-                    .wasm;
+                    .map(|mut m| m.emit_wasm())
+                    .context(format!("error re-parsing wasm (line {})", line))?;
                 if wasm1 != wasm2 {
                     panic!("wasm module at line {} isn't deterministic", line);
                 }
@@ -124,7 +120,7 @@ fn run(wast: &Path) -> Result<(), failure::Error> {
             walrus::passes::gc::run(&mut module);
         }
 
-        let wasm = module.emit_wasm()?.wasm;
+        let wasm = module.emit_wasm();
         fs::write(&file, wasm)?;
     }
 

examples/round-trip.rs

@@ -3,8 +3,8 @@
 fn main() {
     env_logger::init();
     let a = std::env::args().nth(1).unwrap();
-    let m = walrus::Module::from_file(&a).unwrap();
-    let walrus::EmitResult { wasm, .. } = m.emit_wasm().unwrap();
+    let mut m = walrus::Module::from_file(&a).unwrap();
+    let wasm = m.emit_wasm();
     if let Some(destination) = std::env::args().nth(2) {
         std::fs::write(destination, wasm).unwrap();
     }

src/emit.rs

@@ -5,8 +5,8 @@
 use crate::encode::{Encoder, MAX_U32_LENGTH};
 use crate::ir::Local;
 use crate::map::{IdHashMap, IdHashSet};
+use crate::{CodeTransform, Global, GlobalId, Memory, MemoryId, Module, Table, TableId};
 use crate::{Data, DataId, Element, ElementId, Function, FunctionId};
-use crate::{Global, GlobalId, Memory, MemoryId, Module, Table, TableId};
 use crate::{Type, TypeId};
 use std::ops::{Deref, DerefMut};
 
@@ -15,7 +15,7 @@ pub struct EmitContext<'a> {
     pub indices: &'a mut IdsToIndices,
     pub encoder: Encoder<'a>,
     pub locals: IdHashMap<Function, IdHashSet<Local>>,
-    pub code_transform: Vec<(usize, usize)>,
+    pub code_transform: CodeTransform,
 }
 
 pub struct SubContext<'a, 'cx> {

src/module/custom.rs

@@ -1,6 +1,7 @@
 //! Working with custom sections.
 
 use crate::tombstone_arena::{Id, Tombstone, TombstoneArena};
+use crate::CodeTransform;
 use std::any::Any;
 use std::borrow::Cow;
 use std::fmt::{self, Debug};
@@ -28,6 +29,26 @@ pub trait CustomSection: Any + Debug + Send + Sync {
     /// section's name, or the count of how many bytes are in the
     /// payload. `walrus` will handle these for you.
     fn data(&self) -> Cow<[u8]>;
+
+    /// Apply the given code transformations to this custom section.
+    ///
+    /// If the module was not configured with `preserve_code_transform = true`,
+    /// then this method is never called.
+    ///
+    /// This method is called after we have emitted the non-custom Wasm
+    /// sections, just before a custom section's data is emitted into the Wasm
+    /// binary. If this custom section references offsets in the Wasm code, this
+    /// is a chance to update them so they are valid for the new, transformed
+    /// Wasm code that is being emitted.
+    ///
+    /// For example, DWARF debug info references Wasm instructions via offsets
+    /// into the code section, and we can use these transforms to fix those
+    /// offsets after having transformed various functions and instructions.
+    ///
+    /// The default provided method does nothing.
+    fn apply_code_transform(&mut self, transform: &CodeTransform) {
+        let _ = transform;
+    }
 }
 
 // We have to have this so that we can convert `CustomSection` trait objects

src/module/mod.rs

@@ -36,6 +36,7 @@ pub use crate::module::types::ModuleTypes;
 use crate::parse::IndicesToIds;
 use failure::{bail, ResultExt};
 use std::fs;
+use std::mem;
 use std::path::Path;
 
 pub use self::config::ModuleConfig;
@@ -66,25 +67,17 @@ pub struct Module {
     /// The name of this module, used for debugging purposes in the `name`
     /// custom section.
     pub name: Option<String>,
-    pub(crate) config: ModuleConfig,
-}
-
-/// Emitted module information.
-#[derive(Debug)]
-pub struct EmitResult {
-    /// The wasm file bytecode.
-    pub wasm: Vec<u8>,
 
-    /// Optional code transform (if module generated from
-    /// the existing wasm file).
-    pub code_transform: Vec<(usize, usize)>,
+    pub(crate) config: ModuleConfig,
 }
 
-impl Into<Vec<u8>> for EmitResult {
-    fn into(self) -> Vec<u8> {
-        self.wasm
-    }
-}
+/// Maps from an offset of an instruction in the input Wasm to its offset in the
+/// output Wasm.
+///
+/// Note that an input offset may be mapped to multiple output offsets, and vice
+/// versa, due to transformations like function inlinining or constant
+/// propagation.
+pub type CodeTransform = Vec<(usize, usize)>;
 
 impl Module {
     /// Create a default, empty module that uses the given configuration.
@@ -243,24 +236,26 @@ impl Module {
     }
 
     /// Emit this module into a `.wasm` file at the given path.
-    pub fn emit_wasm_file<P>(&self, path: P) -> Result<()>
+    pub fn emit_wasm_file<P>(&mut self, path: P) -> Result<()>
     where
         P: AsRef<Path>,
     {
-        let EmitResult { wasm: buffer, .. } = self.emit_wasm()?;
+        let buffer = self.emit_wasm();
         fs::write(path, buffer).context("failed to write wasm module")?;
         Ok(())
     }
 
     /// Emit this module into an in-memory wasm buffer.
-    pub fn emit_wasm(&self) -> Result<Vec<u8>> {
+    pub fn emit_wasm(&mut self) -> Vec<u8> {
         log::debug!("start emit");
 
         let indices = &mut IdsToIndices::default();
         let mut wasm = Vec::new();
         wasm.extend(&[0x00, 0x61, 0x73, 0x6d]); // magic
         wasm.extend(&[0x01, 0x00, 0x00, 0x00]); // version
 
+        let mut customs = mem::replace(&mut self.customs, ModuleCustomSections::default());
+
         let mut cx = EmitContext {
             module: self,
             indices,
@@ -291,24 +286,27 @@ impl Module {
             self.producers.emit(&mut cx);
         }
 
-        for (_id, section) in self.customs.iter() {
+        for (_id, section) in customs.iter_mut() {
             if !self.config.generate_dwarf && section.name().starts_with(".debug") {
                 log::debug!("skipping DWARF custom section {}", section.name());
                 continue;
             }
 
-            section.
-
             log::debug!("emitting custom section {}", section.name());
+
+            if self.config.preserve_code_transform {
+                section.apply_code_transform(&cx.code_transform);
+            }
+
             cx.custom_section(&section.name())
                 .encoder
                 .raw(&section.data());
         }
 
-        let code_transform = cx.code_transform;
+        self.customs = customs;
 
         log::debug!("emission finished");
-        Ok(wasm)
+        wasm
     }
 
     /// Returns an iterator over all functions in this module