Support compilation of individual functions before serializing whole module, to keep memory usage down

src/backend/mod.rs

@@ -48,7 +48,7 @@ impl<'a> WasmFuncBackend<'a> {
         })
     }
 
-    pub fn compile(&self) -> Result<wasm_encoder::Function> {
+    pub fn compile(&self) -> Result<Vec<u8>> {
         let mut func = wasm_encoder::Function::new(
             self.locals
                 .locals
@@ -76,7 +76,12 @@ impl<'a> WasmFuncBackend<'a> {
 
         log::debug!("Compiled to:\n{:?}\n", func);
 
-        Ok(func)
+        let mut bytes = vec![];
+        {
+            use wasm_encoder::Encode;
+            func.encode(&mut bytes);
+        }
+        Ok(bytes)
     }
 
     fn lower_block(&self, block: &WasmBlock<'_>, func: &mut wasm_encoder::Function) {
@@ -590,7 +595,9 @@ pub fn compile(module: &Module<'_>) -> anyhow::Result<Vec<u8>> {
     for (func, func_decl) in module.funcs.entries().skip(num_func_imports) {
         match func_decl {
             FuncDecl::Import(_, _) => anyhow::bail!("Import comes after func with body: {}", func),
-            FuncDecl::Lazy(sig, _, _) | FuncDecl::Body(sig, _, _) => {
+            FuncDecl::Lazy(sig, _, _)
+            | FuncDecl::Body(sig, _, _)
+            | FuncDecl::Compiled(sig, _, _) => {
                 funcs.function(sig.index() as u32);
             }
             FuncDecl::None => panic!("FuncDecl::None at compilation time"),
@@ -695,10 +702,6 @@ 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
@@ -706,33 +709,27 @@ pub fn compile(module: &Module<'_>) -> anyhow::Result<Vec<u8>> {
         .skip(num_func_imports)
         .collect::<Vec<_>>()
         .par_iter()
-        .map(|(func, func_decl)| -> Result<FuncOrRawBytes> {
+        .map(|(func, func_decl)| -> Result<_> {
             match func_decl {
                 FuncDecl::Lazy(_, _name, reader) => {
                     let data = &module.orig_bytes[reader.range()];
-                    Ok(FuncOrRawBytes::Raw(data))
+                    Ok(Cow::Borrowed(data))
                 }
+                FuncDecl::Compiled(_, _name, bytes) => Ok(Cow::Borrowed(&bytes[..])),
                 FuncDecl::Body(_, name, body) => {
                     log::debug!("Compiling {} \"{}\"", func, name);
                     WasmFuncBackend::new(body)?
                         .compile()
-                        .map(|f| FuncOrRawBytes::Func(f))
+                        .map(|bytes| Cow::Owned(bytes))
                 }
                 FuncDecl::Import(_, _) => unreachable!("Should have skipped imports"),
                 FuncDecl::None => panic!("FuncDecl::None at compilation time"),
             }
         })
-        .collect::<Result<Vec<FuncOrRawBytes<'_>>>>()?;
+        .collect::<Result<Vec<_>>>()?;
 
     for body in bodies {
-        match body {
-            FuncOrRawBytes::Func(f) => {
-                code.function(&f);
-            }
-            FuncOrRawBytes::Raw(bytes) => {
-                code.raw(bytes);
-            }
-        }
+        code.raw(&body[..]);
     }
     into_mod.section(&code);
 

src/interp.rs

@@ -90,6 +90,7 @@ impl InterpContext {
     pub fn call(&mut self, module: &Module<'_>, func: Func, args: &[ConstVal]) -> InterpResult {
         let body = match &module.funcs[func] {
             FuncDecl::Lazy(..) => panic!("Un-expanded function"),
+            FuncDecl::Compiled(..) => panic!("Already-compiled function"),
             FuncDecl::Import(..) => {
                 let import = &module.imports[func.index()];
                 assert_eq!(import.kind, ImportKind::Func(func));

src/ir/display.rs

@@ -252,6 +252,17 @@ impl<'a> Display for ModuleDisplay<'a> {
                     )?;
                     writeln!(f, "  # raw bytes (length {})", reader.range().len())?;
                 }
+                FuncDecl::Compiled(sig, name, bytes) => {
+                    writeln!(
+                        f,
+                        "  {} \"{}\": {} = # {}",
+                        func,
+                        name,
+                        sig,
+                        sig_strs.get(&sig).unwrap()
+                    )?;
+                    writeln!(f, "  # already compiled (length {})", bytes.len())?;
+                }
                 FuncDecl::Import(sig, name) => {
                     writeln!(
                         f,

src/ir/func.rs

@@ -1,4 +1,5 @@
 use super::{Block, FunctionBodyDisplay, Local, Module, Signature, Type, Value, ValueDef};
+use crate::backend::WasmFuncBackend;
 use crate::cfg::CFGInfo;
 use crate::entity::{EntityRef, EntityVec, PerEntity};
 use crate::frontend::parse_body;
@@ -16,6 +17,8 @@ pub enum FuncDecl<'a> {
     Lazy(Signature, String, wasmparser::FunctionBody<'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>),
     /// A placeholder.
     #[default]
     None,
@@ -27,6 +30,7 @@ impl<'a> FuncDecl<'a> {
             FuncDecl::Import(sig, ..) => *sig,
             FuncDecl::Lazy(sig, ..) => *sig,
             FuncDecl::Body(sig, ..) => *sig,
+            FuncDecl::Compiled(sig, ..) => *sig,
             FuncDecl::None => panic!("No signature for FuncDecl::None"),
         }
     }
@@ -76,18 +80,20 @@ impl<'a> FuncDecl<'a> {
 
     pub fn name(&self) -> &str {
         match self {
-            FuncDecl::Body(_, name, _) | FuncDecl::Lazy(_, name, _) | FuncDecl::Import(_, name) => {
-                &name[..]
-            }
+            FuncDecl::Body(_, name, _)
+            | FuncDecl::Lazy(_, name, _)
+            | FuncDecl::Import(_, name)
+            | FuncDecl::Compiled(_, name, _) => &name[..],
             FuncDecl::None => panic!("No name for FuncDecl::None"),
         }
     }
 
     pub fn set_name(&mut self, new_name: &str) {
         match self {
-            FuncDecl::Body(_, name, _) | FuncDecl::Lazy(_, name, _) | FuncDecl::Import(_, name) => {
-                *name = new_name.to_owned()
-            }
+            FuncDecl::Body(_, name, _)
+            | FuncDecl::Lazy(_, name, _)
+            | FuncDecl::Import(_, name)
+            | FuncDecl::Compiled(_, name, _) => *name = new_name.to_owned(),
             FuncDecl::None => panic!("No name for FuncDecl::None"),
         }
     }
@@ -429,6 +435,11 @@ impl FunctionBody {
 
         Ok(())
     }
+
+    pub fn compile(&self) -> Result<Vec<u8>> {
+        let backend = WasmFuncBackend::new(self)?;
+        backend.compile()
+    }
 }
 
 #[derive(Clone, Debug, Default)]