Lazy function parsing and recompilation when roundtripping

fuzz/fuzz_targets/roundtrip.rs

@@ -24,6 +24,7 @@ fuzz_target!(|module: wasm_smith::Module| {
             }
         }
     };
+    parsed_module.expand_all_funcs().unwrap();
     parsed_module.optimize();
     let _ = parsed_module.to_wasm_bytes();
 });

src/backend/mod.rs

@@ -585,7 +585,7 @@ pub fn compile(module: &Module<'_>) -> anyhow::Result<Vec<u8>> {
     for (func, func_decl) in module.funcs().skip(num_func_imports) {
         match func_decl {
             FuncDecl::Import(_) => anyhow::bail!("Import comes after func with body: {}", func),
-            FuncDecl::Body(sig, _) => {
+            FuncDecl::Lazy(sig, _) | FuncDecl::Body(sig, _) => {
                 funcs.function(sig.index() as u32);
             }
         }
@@ -689,20 +689,42 @@ pub fn compile(module: &Module<'_>) -> anyhow::Result<Vec<u8>> {
     into_mod.section(&elem);
 
     let mut code = wasm_encoder::CodeSection::new();
+    enum FuncOrRawBytes<'a> {
+        Func(wasm_encoder::Function),
+        Raw(&'a [u8]),
+    }
+
     let bodies = module
         .funcs()
         .skip(num_func_imports)
         .collect::<Vec<_>>()
         .par_iter()
-        .map(|(func, func_decl)| -> Result<wasm_encoder::Function> {
-            let body = func_decl.body().unwrap();
-            log::debug!("Compiling {}", func);
-            WasmFuncBackend::new(body)?.compile()
+        .map(|(func, func_decl)| -> Result<FuncOrRawBytes> {
+            match func_decl {
+                FuncDecl::Lazy(_, reader) => {
+                    let data = &module.orig_bytes[reader.range()];
+                    Ok(FuncOrRawBytes::Raw(data))
+                }
+                FuncDecl::Body(_, body) => {
+                    log::debug!("Compiling {}", func);
+                    WasmFuncBackend::new(body)?
+                        .compile()
+                        .map(|f| FuncOrRawBytes::Func(f))
+                }
+                FuncDecl::Import(_) => unreachable!("Should have skipped imports"),
+            }
         })
-        .collect::<Result<Vec<_>>>()?;
+        .collect::<Result<Vec<FuncOrRawBytes<'_>>>>()?;
 
     for body in bodies {
-        code.function(&body);
+        match body {
+            FuncOrRawBytes::Func(f) => {
+                code.function(&f);
+            }
+            FuncOrRawBytes::Raw(bytes) => {
+                code.raw(bytes);
+            }
+        }
     }
     into_mod.section(&code);
 

src/frontend.rs

@@ -143,10 +143,7 @@ fn handle_payload<'a>(
             *next_func += 1;
 
             let my_sig = module.func(func_idx).sig();
-            let body = parse_body(module, my_sig, body)?;
-
-            let existing_body = module.func_mut(func_idx).body_mut().unwrap();
-            *existing_body = body;
+            *module.func_mut(func_idx) = FuncDecl::Lazy(my_sig, body);
         }
         Payload::ExportSection(reader) => {
             for export in reader {
@@ -267,10 +264,10 @@ fn handle_payload<'a>(
     Ok(())
 }
 
-fn parse_body<'a>(
+pub(crate) fn parse_body<'a>(
     module: &'a Module,
     my_sig: Signature,
-    body: wasmparser::FunctionBody,
+    body: &mut wasmparser::FunctionBody,
 ) -> Result<FunctionBody> {
     let mut ret: FunctionBody = FunctionBody::default();
 

src/ir/display.rs

@@ -217,6 +217,10 @@ impl<'a> Display for ModuleDisplay<'a> {
                     writeln!(f, "  {}: {} = # {}", func, sig, sig_strs.get(&sig).unwrap())?;
                     writeln!(f, "{}", body.display("    "))?;
                 }
+                FuncDecl::Lazy(sig, reader) => {
+                    writeln!(f, "  {}: {} = # {}", func, sig, sig_strs.get(&sig).unwrap())?;
+                    writeln!(f, "  # raw bytes (length {})", reader.range().len())?;
+                }
                 FuncDecl::Import(sig) => {
                     writeln!(f, "  {}: {} # {}", func, sig, sig_strs.get(&sig).unwrap())?;
                 }

src/ir/func.rs

@@ -1,21 +1,36 @@
 use super::{Block, FunctionBodyDisplay, Local, Module, Signature, Type, Value, ValueDef};
 use crate::cfg::CFGInfo;
 use crate::entity::{EntityRef, EntityVec, PerEntity};
+use crate::frontend::parse_body;
+use anyhow::Result;
 
 #[derive(Clone, Debug)]
-pub enum FuncDecl {
+pub enum FuncDecl<'a> {
     Import(Signature),
+    Lazy(Signature, wasmparser::FunctionBody<'a>),
     Body(Signature, FunctionBody),
 }
 
-impl FuncDecl {
+impl<'a> FuncDecl<'a> {
     pub fn sig(&self) -> Signature {
         match self {
             FuncDecl::Import(sig) => *sig,
+            FuncDecl::Lazy(sig, ..) => *sig,
             FuncDecl::Body(sig, ..) => *sig,
         }
     }
 
+    pub fn parse(&mut self, module: &Module) -> Result<()> {
+        match self {
+            FuncDecl::Lazy(sig, body) => {
+                let body = parse_body(module, *sig, body)?;
+                *self = FuncDecl::Body(*sig, body);
+                Ok(())
+            }
+            _ => Ok(()),
+        }
+    }
+
     pub fn body(&self) -> Option<&FunctionBody> {
         match self {
             FuncDecl::Body(_, body) => Some(body),

src/ir/module.rs

@@ -6,8 +6,8 @@ use anyhow::Result;
 
 #[derive(Clone, Debug)]
 pub struct Module<'a> {
-    orig_bytes: &'a [u8],
-    funcs: EntityVec<Func, FuncDecl>,
+    pub orig_bytes: &'a [u8],
+    funcs: EntityVec<Func, FuncDecl<'a>>,
     signatures: EntityVec<Signature, SignatureData>,
     globals: EntityVec<Global, GlobalData>,
     tables: EntityVec<Table, TableData>,
@@ -142,13 +142,13 @@ impl<'a> Module<'a> {
 }
 
 impl<'a> Module<'a> {
-    pub fn func<'b>(&'b self, id: Func) -> &'b FuncDecl {
+    pub fn func<'b>(&'b self, id: Func) -> &'b FuncDecl<'a> {
         &self.funcs[id]
     }
-    pub fn func_mut<'b>(&'b mut self, id: Func) -> &'b mut FuncDecl {
+    pub fn func_mut<'b>(&'b mut self, id: Func) -> &'b mut FuncDecl<'a> {
         &mut self.funcs[id]
     }
-    pub fn funcs<'b>(&'b self) -> impl Iterator<Item = (Func, &'b FuncDecl)> {
+    pub fn funcs<'b>(&'b self) -> impl Iterator<Item = (Func, &'b FuncDecl<'a>)> {
         self.funcs.entries()
     }
     pub fn signature<'b>(&'b self, id: Signature) -> &'b SignatureData {
@@ -192,7 +192,7 @@ impl<'a> Module<'a> {
     pub(crate) fn frontend_add_signature(&mut self, ty: SignatureData) {
         self.signatures.push(ty);
     }
-    pub(crate) fn frontend_add_func(&mut self, body: FuncDecl) -> Func {
+    pub(crate) fn frontend_add_func(&mut self, body: FuncDecl<'a>) -> Func {
         self.funcs.push(body)
     }
     pub(crate) fn frontend_add_table(&mut self, ty: Type, max: Option<u32>) -> Table {
@@ -236,6 +236,23 @@ impl<'a> Module<'a> {
         }
     }
 
+    pub fn expand_func<'b>(&'b mut self, id: Func) -> Result<&'b FuncDecl<'a>> {
+        let mut funcs = std::mem::take(&mut self.funcs);
+        let ret = funcs[id].parse(self);
+        self.funcs = funcs;
+        ret.and(Ok(&self.funcs[id]))
+    }
+
+    pub fn expand_all_funcs(&mut self) -> Result<()> {
+        let mut funcs = std::mem::take(&mut self.funcs);
+        let mut ret = Ok(());
+        for func_decl in funcs.values_mut() {
+            ret = ret.and_then(|_| func_decl.parse(self));
+        }
+        self.funcs = funcs;
+        ret
+    }
+
     pub fn optimize(&mut self) {
         self.per_func_body(|body| {
             let cfg = crate::cfg::CFGInfo::new(body);