feat: lazy memory allocation, fix instruction rewrite order

Signed-off-by: Henry <mail@henrygressmann.de>

3 files changed, 138 insertions, 49 deletions

diff --git a/crates/parser/src/lib.rs b/crates/parser/src/lib.rs
index ad58cc2..0b5af2e 100644
--- a/crates/parser/src/lib.rs
+++ b/crates/parser/src/lib.rs
@@ -43,8 +43,30 @@ pub use tinywasm_types::TinyWasmModule;
 
 /// Parser optimization and lowering options.
 #[non_exhaustive]
-#[derive(Debug, Clone, Default)]
-pub struct ParserOptions {}
+#[derive(Debug, Clone)]
+pub struct ParserOptions {
+    /// Whether to optimize local memory allocation by skipping allocation of unused local memories.
+    pub optimize_local_memory_allocation: bool,
+}
+
+impl Default for ParserOptions {
+    fn default() -> Self {
+        Self { optimize_local_memory_allocation: true }
+    }
+}
+
+impl ParserOptions {
+    /// Enable or disable the optimization that skips allocating unused local memories.
+    pub const fn with_local_memory_allocation_optimization(mut self, enabled: bool) -> Self {
+        self.optimize_local_memory_allocation = enabled;
+        self
+    }
+
+    /// Returns whether unused local memory allocation optimization is enabled.
+    pub const fn optimize_local_memory_allocation(&self) -> bool {
+        self.optimize_local_memory_allocation
+    }
+}
 
 /// A WebAssembly parser
 #[derive(Debug, Default)]
diff --git a/crates/parser/src/module.rs b/crates/parser/src/module.rs
index 461f79f..9c4d6b4 100644
--- a/crates/parser/src/module.rs
+++ b/crates/parser/src/module.rs
@@ -167,7 +167,7 @@ impl ModuleReader {
         Ok(())
     }
 
-    pub(crate) fn into_module(self, _options: &ParserOptions) -> Result<TinyWasmModule> {
+    pub(crate) fn into_module(self, options: &ParserOptions) -> Result<TinyWasmModule> {
         if !self.end_reached {
             return Err(ParseError::EndNotReached);
         }
@@ -176,28 +176,73 @@ impl ModuleReader {
             return Err(ParseError::Other("Code and code type address count mismatch".to_string()));
         }
 
-        let imported_func_count = self.imports.iter().filter(|i| matches!(&i.kind, ImportKind::Function(_))).count();
-        let funcs = self.code.into_iter().zip(self.code_type_addrs).enumerate().map(
-            |(func_idx, ((instructions, mut data, locals), ty_idx))| {
-                let ty = self.func_types.get(ty_idx as usize).expect("No func type for func, this is a bug").clone();
-                let params = ValueCounts::from_iter(ty.params());
-                let self_func = (imported_func_count + func_idx) as u32;
-                let instructions = optimize::optimize_instructions(instructions, &mut data, self_func);
-                WasmFunction { instructions: ArcSlice::from(instructions), data, locals, params, ty }
-            },
-        );
+        let Self {
+            start_func,
+            func_types,
+            code_type_addrs,
+            exports,
+            code,
+            globals,
+            table_types,
+            memory_types,
+            imports,
+            data,
+            elements,
+            ..
+        } = self;
+
+        let imported_func_count = imports.iter().filter(|i| matches!(&i.kind, ImportKind::Function(_))).count();
+        let imported_memory_count = imports.iter().filter(|i| matches!(&i.kind, ImportKind::Memory(_))).count() as u32;
+        let has_local_memory_export =
+            exports.iter().any(|export| export.kind == ExternalKind::Memory && export.index >= imported_memory_count);
+        let has_active_data_segment_on_local_memory = data.iter().any(|data| match &data.kind {
+            DataKind::Active { mem, .. } => *mem >= imported_memory_count,
+            DataKind::Passive => false,
+        });
+        let optimize_local_memory_allocation = options.optimize_local_memory_allocation();
+        let mut local_memory_allocation = if memory_types.is_empty() {
+            LocalMemoryAllocation::Skip
+        } else if !optimize_local_memory_allocation || has_active_data_segment_on_local_memory {
+            LocalMemoryAllocation::Eager
+        } else if has_local_memory_export {
+            LocalMemoryAllocation::Lazy
+        } else {
+            LocalMemoryAllocation::Skip
+        };
+        let mut funcs = Vec::with_capacity(code.len());
+
+        for (func_idx, ((instructions, mut data, locals), ty_idx)) in code.into_iter().zip(code_type_addrs).enumerate()
+        {
+            let ty = func_types.get(ty_idx as usize).expect("No func type for func, this is a bug").clone();
+            let params = ValueCounts::from_iter(ty.params());
+            let self_func = (imported_func_count + func_idx) as u32;
+            let optimized = optimize::optimize_instructions(
+                instructions,
+                &mut data,
+                self_func,
+                imported_memory_count,
+                optimize_local_memory_allocation && local_memory_allocation != LocalMemoryAllocation::Eager,
+            );
+
+            if optimized.uses_local_memory {
+                local_memory_allocation = LocalMemoryAllocation::Eager;
+            }
+
+            funcs.push(WasmFunction { instructions: ArcSlice::from(optimized.instructions), data, locals, params, ty });
+        }
 
         Ok(TinyWasmModule {
-            funcs: funcs.collect(),
-            func_types: self.func_types.into(),
-            globals: self.globals.into(),
-            table_types: self.table_types.into(),
-            imports: self.imports.into(),
-            start_func: self.start_func,
-            data: self.data.into(),
-            exports: self.exports.into(),
-            elements: self.elements.into(),
-            memory_types: self.memory_types.into(),
+            funcs: funcs.into(),
+            func_types: func_types.into(),
+            globals: globals.into(),
+            table_types: table_types.into(),
+            imports: imports.into(),
+            start_func,
+            data: data.into(),
+            exports: exports.into(),
+            elements: elements.into(),
+            memory_types: memory_types.into(),
+            local_memory_allocation,
         })
     }
 }
diff --git a/crates/parser/src/optimize.rs b/crates/parser/src/optimize.rs
index 0539b61..beb944d 100644
--- a/crates/parser/src/optimize.rs
+++ b/crates/parser/src/optimize.rs
@@ -2,18 +2,32 @@ use crate::macros::optimize::*;
 use alloc::vec::Vec;
 use tinywasm_types::{CmpOp, Instruction, WasmFunctionData};
 
+pub(crate) struct OptimizeResult {
+    pub(crate) instructions: Vec<Instruction>,
+    pub(crate) uses_local_memory: bool,
+}
+
 pub(crate) fn optimize_instructions(
     mut instructions: Vec<Instruction>,
     function_data: &mut WasmFunctionData,
     self_func_addr: u32,
-) -> Vec<Instruction> {
-    rewrite(&mut instructions, self_func_addr);
+    imported_memory_count: u32,
+    track_local_memory_usage: bool,
+) -> OptimizeResult {
+    let uses_local_memory = rewrite(&mut instructions, self_func_addr, imported_memory_count, track_local_memory_usage);
     remove_nop(&mut instructions, function_data);
-    instructions
+    OptimizeResult { instructions, uses_local_memory }
 }
 
-fn rewrite(instrs: &mut [Instruction], self_func_addr: u32) {
+fn rewrite(
+    instrs: &mut [Instruction],
+    self_func_addr: u32,
+    imported_memory_count: u32,
+    track_local_memory_usage: bool,
+) -> bool {
     use Instruction::*;
+    let mut uses_local_memory = false;
+
     for i in 0..instrs.len() {
         match instrs[i] {
             LocalCopy32(a, b) if a == b => instrs[i] = Nop,
@@ -22,14 +36,14 @@ fn rewrite(instrs: &mut [Instruction], self_func_addr: u32) {
             Call(addr) if addr == self_func_addr => instrs[i] = CallSelf,
             ReturnCall(addr) if addr == self_func_addr => instrs[i] = ReturnCallSelf,
             I32Add => {
-                rewrite!(instrs, i, [I32Const(c)] => AddConst32(c));
                 rewrite!(instrs, i, [LocalGet32(a), LocalGet32(b)] => AddLocalLocal32(a, b));
                 rewrite!(instrs, i, [LocalGet32(local), I32Const(c)] => [ Nop, LocalGet32(local), AddConst32(c)]);
+                rewrite!(instrs, i, [I32Const(c)] => AddConst32(c));
             }
             I64Add => {
-                rewrite!(instrs, i, [I64Const(c)] => AddConst64(c));
                 rewrite!(instrs, i, [LocalGet64(a), LocalGet64(b)] => AddLocalLocal64(a, b));
                 rewrite!(instrs, i, [LocalGet64(local), I64Const(c)] => [ Nop, LocalGet64(local), AddConst64(c)]);
+                rewrite!(instrs, i, [I64Const(c)] => AddConst64(c));
             }
             I64Rotl => rewrite!(instrs, i, [I64Xor, I64Const(c)] => XorRotlConst64(c)),
             I32Store(memarg) => {
@@ -69,6 +83,7 @@ fn rewrite(instrs: &mut [Instruction], self_func_addr: u32) {
                 rewrite!(instrs, i, [LocalGet32(src)] => if src == dst { Nop } else { LocalCopy32(src, dst) });
                 rewrite!(instrs, i, [I32Const(c)] => SetLocalConst32(dst, c));
                 rewrite!(instrs, i, [F32Const(c)] => SetLocalConst32(dst, i32::from_ne_bytes(c.to_bits().to_ne_bytes())));
+                rewrite!(instrs, i, [AddLocalLocal32(a, b)] => AddLocalLocalSet32(a, b, dst));
                 rewrite!(instrs, i, [LocalGet32(src), AddConst32(c)] if (src == dst) => AddLocalConst32(dst, c));
                 rewrite!(instrs, i, [LoadLocal32(memarg, addr)] if (let Ok(dst) = u8::try_from(dst)) => LoadLocalSet32(memarg, addr, dst));
                 rewrite!(instrs, i,
@@ -81,6 +96,7 @@ fn rewrite(instrs: &mut [Instruction], self_func_addr: u32) {
                 rewrite!(instrs, i, [LocalGet64(src)] => if src == dst { Nop } else { LocalCopy64(src, dst) });
                 rewrite!(instrs, i, [I64Const(c)] => SetLocalConst64(dst, c));
                 rewrite!(instrs, i, [F64Const(c)] => SetLocalConst64(dst, i64::from_ne_bytes(c.to_bits().to_ne_bytes())));
+                rewrite!(instrs, i, [AddLocalLocal64(a, b)] => AddLocalLocalSet64(a, b, dst));
                 rewrite!(instrs, i,
                     [LocalGet64(src), AddConst64(c)] if (src == dst) =>
                     AddLocalConst64(dst, c)
@@ -130,62 +146,68 @@ fn rewrite(instrs: &mut [Instruction], self_func_addr: u32) {
                     replace!(instrs, i, 1 => [Nop, JumpIfNonZero(ip)]);
                     continue;
                 });
-                rewrite!(instrs, i, [cmp, I32Const(imm)] if (let Some(op) = cmp_op(cmp)) =>
-                    JumpCmpStackConst32 { target_ip: ip, imm, op: inverse_cmp_op(op) }
-                );
-                rewrite!(instrs, i, [cmp, I64Const(imm)] if (let Some(op) = cmp_op_64(cmp)) =>
-                    JumpCmpStackConst64 { target_ip: ip, imm, op: inverse_cmp_op(op) }
-                );
                 rewrite!(instrs, i,
-                    [LocalGet32(local), cmp, I32Const(imm)] if (let Some(op) = cmp_op(cmp)) =>
+                    [LocalGet32(local), I32Const(imm), cmp] if (let Some(op) = cmp_op(cmp)) =>
                     JumpCmpLocalConst32 { target_ip: ip, local, imm, op: inverse_cmp_op(op) }
                 );
                 rewrite!(instrs, i,
-                    [LocalGet64(local), cmp, I64Const(imm)] if
+                    [LocalGet64(local), I64Const(imm), cmp] if
                     (let Some(op) = cmp_op_64(cmp) && let Ok(imm) = i32::try_from(imm)) =>
                     JumpCmpLocalConst64 { target_ip: ip, local, imm, op: inverse_cmp_op(op) }
                 );
                 rewrite!(instrs, i,
-                    [LocalGet32(left), cmp, LocalGet32(right)] if (let Some(op) = cmp_op(cmp)) =>
+                    [LocalGet32(left), LocalGet32(right), cmp] if (let Some(op) = cmp_op(cmp)) =>
                     JumpCmpLocalLocal32 { target_ip: ip, left, right, op: inverse_cmp_op(op) }
                 );
                 rewrite!(instrs, i,
-                    [LocalGet64(left), cmp, LocalGet64(right)] if (let Some(op) = cmp_op_64(cmp)) =>
+                    [LocalGet64(left), LocalGet64(right), cmp] if (let Some(op) = cmp_op_64(cmp)) =>
                     JumpCmpLocalLocal64 { target_ip: ip, left, right, op: inverse_cmp_op(op) }
                 );
+                rewrite!(instrs, i, [I32Const(imm), cmp] if (let Some(op) = cmp_op(cmp)) =>
+                    JumpCmpStackConst32 { target_ip: ip, imm, op: inverse_cmp_op(op) }
+                );
+                rewrite!(instrs, i, [I64Const(imm), cmp] if (let Some(op) = cmp_op_64(cmp)) =>
+                    JumpCmpStackConst64 { target_ip: ip, imm, op: inverse_cmp_op(op) }
+                );
             }
             JumpIfNonZero(ip) => {
                 rewrite!(instrs, i, [I32Eqz] => {
                     replace!(instrs, i, 1 => [Nop, JumpIfZero(ip)]);
                     continue;
                 });
-                rewrite!(instrs, i, [cmp, I32Const(imm)] if (let Some(op) = cmp_op(cmp)) =>
-                    JumpCmpStackConst32 { target_ip: ip, imm, op }
-                );
-                rewrite!(instrs, i, [cmp, I64Const(imm)] if (let Some(op) = cmp_op_64(cmp)) =>
-                    JumpCmpStackConst64 { target_ip: ip, imm, op }
-                );
                 rewrite!(instrs, i,
-                    [LocalGet32(local), cmp, I32Const(imm)] if (let Some(op) = cmp_op(cmp)) =>
+                    [LocalGet32(local), I32Const(imm), cmp] if (let Some(op) = cmp_op(cmp)) =>
                     JumpCmpLocalConst32 { target_ip: ip, local, imm, op }
                 );
                 rewrite!(instrs, i,
-                    [LocalGet64(local), cmp, I64Const(imm)] if
+                    [LocalGet64(local), I64Const(imm), cmp] if
                     (let Some(op) = cmp_op_64(cmp) && let Ok(imm) = i32::try_from(imm)) =>
                     JumpCmpLocalConst64 { target_ip: ip, local, imm, op }
                 );
                 rewrite!(instrs, i,
-                    [LocalGet32(left), cmp, LocalGet32(right)] if (let Some(op) = cmp_op(cmp)) =>
+                    [LocalGet32(left), LocalGet32(right), cmp] if (let Some(op) = cmp_op(cmp)) =>
                     JumpCmpLocalLocal32 { target_ip: ip, left, right, op }
                 );
                 rewrite!(instrs, i,
-                    [LocalGet64(left), cmp, LocalGet64(right)] if (let Some(op) = cmp_op_64(cmp)) =>
+                    [LocalGet64(left), LocalGet64(right), cmp] if (let Some(op) = cmp_op_64(cmp)) =>
                     JumpCmpLocalLocal64 { target_ip: ip, left, right, op }
                 );
+                rewrite!(instrs, i, [I32Const(imm), cmp] if (let Some(op) = cmp_op(cmp)) =>
+                    JumpCmpStackConst32 { target_ip: ip, imm, op }
+                );
+                rewrite!(instrs, i, [I64Const(imm), cmp] if (let Some(op) = cmp_op_64(cmp)) =>
+                    JumpCmpStackConst64 { target_ip: ip, imm, op }
+                );
             }
             _ => {}
         }
+
+        if track_local_memory_usage {
+            uses_local_memory |= instrs[i].memory_addr().is_some_and(|mem| mem >= imported_memory_count);
+        }
     }
+
+    uses_local_memory
 }
 
 fn cmp_op(instr: Instruction) -> Option<CmpOp> {