Track original wasm instructions locations.

crates/fuzz/src/lib.rs

@@ -76,6 +76,7 @@ impl Config {
             .unwrap()
             .emit_wasm()
             .unwrap()
+            .into()
     }
 
     fn run_one(&self, wat: &str) -> Option<FailingTestCase> {

examples/round-trip.rs

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

src/emit.rs

@@ -15,6 +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 struct SubContext<'a, 'cx> {

src/function_builder.rs

@@ -1,4 +1,5 @@
 use crate::ir::*;
+use crate::map::IdHashMap;
 use crate::tombstone_arena::TombstoneArena;
 use crate::{FunctionId, LocalFunction, Module, TypeId, ValType};
 use crate::{ModuleFunctions, ModuleTypes};
@@ -141,7 +142,7 @@ impl FunctionBuilder {
             results: ty.results().to_vec().into_boxed_slice(),
             exprs,
         });
-        let func = LocalFunction::new(ty_id, args, self, entry);
+        let func = LocalFunction::new(ty_id, args, self, entry, IdHashMap::default());
         funcs.add_local(func)
     }
 }

src/module/config.rs

@@ -11,6 +11,7 @@ pub struct ModuleConfig {
     pub(crate) skip_strict_validate: bool,
     pub(crate) skip_producers_section: bool,
     pub(crate) skip_name_section: bool,
+    pub(crate) preserve_code_transform: bool,
 }
 
 impl ModuleConfig {
@@ -97,6 +98,14 @@ impl ModuleConfig {
         self
     }
 
+    /// Sets a flag to whether code transform is preverved during parsing.
+    ///
+    /// By default this flag is `false`.
+    pub fn preserve_code_transform(&mut self, preserve: bool) -> &mut ModuleConfig {
+        self.preserve_code_transform = preserve;
+        self
+    }
+
     /// Parses an in-memory WebAssembly file into a `Module` using this
     /// configuration.
     pub fn parse(&self, wasm: &[u8]) -> Result<Module> {

src/module/functions/local_function/emit.rs

@@ -11,6 +11,7 @@ pub(crate) fn run(
     indices: &IdsToIndices,
     local_indices: &IdHashMap<Local, u32>,
     encoder: &mut Encoder,
+    map: Option<&mut IdHashMap<Expr, usize>>,
 ) {
     let mut v = Emit {
         func,
@@ -19,6 +20,7 @@ pub(crate) fn run(
         blocks: vec![],
         encoder,
         local_indices,
+        map,
     };
     v.visit(func.entry_block());
 }
@@ -40,6 +42,9 @@ struct Emit<'a, 'b> {
 
     // The instruction sequence we are building up to emit.
     encoder: &'a mut Encoder<'b>,
+
+    // Encoded ExprId -> offset map.
+    map: Option<&'a mut IdHashMap<Expr, usize>>,
 }
 
 impl Emit<'_, '_> {
@@ -56,6 +61,10 @@ impl Emit<'_, '_> {
         let old = self.id;
         self.id = id;
 
+        if let Some(map) = self.map.as_mut() {
+            map.insert(old, self.encoder.pos());
+        }
+
         match self.func.get(id) {
             Const(e) => e.value.emit(self.encoder),
             Block(e) => self.visit_block(e),