From 73d878924223795ae93fcf75727c821f04f28fe9 Mon Sep 17 00:00:00 2001
From: Chris Fallin <chris@cfallin.org>
Date: Thu, 13 Apr 2023 19:59:38 -0700
Subject: [PATCH] Fix "already-compiled function" functionality: avoid mismatch
 in whether length field is included between CodeSection::raw() and
 function().

---
 src/backend/mod.rs | 31 ++++++++++++++++++++-----------
 src/ir/display.rs  |  4 ++--
 src/ir/func.rs     |  4 ++--
 3 files changed, 24 insertions(+), 15 deletions(-)

diff --git a/src/backend/mod.rs b/src/backend/mod.rs
index 328f92d..994b692 100644
--- a/src/backend/mod.rs
+++ b/src/backend/mod.rs
@@ -48,7 +48,7 @@ impl<'a> WasmFuncBackend<'a> {
         })
     }
 
-    pub fn compile(&self) -> Result<Vec<u8>> {
+    pub fn compile(&self) -> Result<wasm_encoder::Function> {
         let mut func = wasm_encoder::Function::new(
             self.locals
                 .locals
@@ -76,12 +76,7 @@ impl<'a> WasmFuncBackend<'a> {
 
         log::debug!("Compiled to:\n{:?}\n", func);
 
-        let mut bytes = vec![];
-        {
-            use wasm_encoder::Encode;
-            func.encode(&mut bytes);
-        }
-        Ok(bytes)
+        Ok(func)
     }
 
     fn lower_block(&self, block: &WasmBlock<'_>, func: &mut wasm_encoder::Function) {
@@ -703,6 +698,11 @@ pub fn compile(module: &Module<'_>) -> anyhow::Result<Vec<u8>> {
 
     let mut code = wasm_encoder::CodeSection::new();
 
+    enum FuncOrRawBytes<'a> {
+        Raw(&'a [u8]),
+        Func(Cow<'a, wasm_encoder::Function>),
+    }
+
     let bodies = module
         .funcs
         .entries()
@@ -713,14 +713,16 @@ pub fn compile(module: &Module<'_>) -> anyhow::Result<Vec<u8>> {
             match func_decl {
                 FuncDecl::Lazy(_, _name, reader) => {
                     let data = &module.orig_bytes[reader.range()];
-                    Ok(Cow::Borrowed(data))
+                    Ok(FuncOrRawBytes::Raw(data))
+                }
+                FuncDecl::Compiled(_, _name, encoder) => {
+                    Ok(FuncOrRawBytes::Func(Cow::Borrowed(encoder)))
                 }
-                FuncDecl::Compiled(_, _name, bytes) => Ok(Cow::Borrowed(&bytes[..])),
                 FuncDecl::Body(_, name, body) => {
                     log::debug!("Compiling {} \"{}\"", func, name);
                     WasmFuncBackend::new(body)?
                         .compile()
-                        .map(|bytes| Cow::Owned(bytes))
+                        .map(|func| FuncOrRawBytes::Func(Cow::Owned(func)))
                 }
                 FuncDecl::Import(_, _) => unreachable!("Should have skipped imports"),
                 FuncDecl::None => panic!("FuncDecl::None at compilation time"),
@@ -729,7 +731,14 @@ pub fn compile(module: &Module<'_>) -> anyhow::Result<Vec<u8>> {
         .collect::<Result<Vec<_>>>()?;
 
     for body in bodies {
-        code.raw(&body[..]);
+        match body {
+            FuncOrRawBytes::Raw(bytes) => {
+                code.raw(bytes);
+            }
+            FuncOrRawBytes::Func(func) => {
+                code.function(&*func);
+            }
+        }
     }
     into_mod.section(&code);
 
diff --git a/src/ir/display.rs b/src/ir/display.rs
index 61132cf..f377ef7 100644
--- a/src/ir/display.rs
+++ b/src/ir/display.rs
@@ -252,7 +252,7 @@ impl<'a> Display for ModuleDisplay<'a> {
                     )?;
                     writeln!(f, "  # raw bytes (length {})", reader.range().len())?;
                 }
-                FuncDecl::Compiled(sig, name, bytes) => {
+                FuncDecl::Compiled(sig, name, _) => {
                     writeln!(
                         f,
                         "  {} \"{}\": {} = # {}",
@@ -261,7 +261,7 @@ impl<'a> Display for ModuleDisplay<'a> {
                         sig,
                         sig_strs.get(&sig).unwrap()
                     )?;
-                    writeln!(f, "  # already compiled (length {})", bytes.len())?;
+                    writeln!(f, "  # already compiled")?;
                 }
                 FuncDecl::Import(sig, name) => {
                     writeln!(
diff --git a/src/ir/func.rs b/src/ir/func.rs
index f6b7a7a..c107573 100644
--- a/src/ir/func.rs
+++ b/src/ir/func.rs
@@ -18,7 +18,7 @@ pub enum FuncDecl<'a> {
     /// A modified or new function body that requires compilation.
     Body(Signature, String, FunctionBody),
     /// A compiled function body (was IR, has been collapsed back to bytecode).
-    Compiled(Signature, String, Vec<u8>),
+    Compiled(Signature, String, wasm_encoder::Function),
     /// A placeholder.
     #[default]
     None,
@@ -436,7 +436,7 @@ impl FunctionBody {
         Ok(())
     }
 
-    pub fn compile(&self) -> Result<Vec<u8>> {
+    pub fn compile(&self) -> Result<wasm_encoder::Function> {
         let backend = WasmFuncBackend::new(self)?;
         backend.compile()
     }