feat: lazy memory allocation, fix instruction rewrite order

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

Author: explodingcamera

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)]

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,
         })
     }
 }

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> {

crates/tinywasm/src/engine.rs

@@ -1,7 +1,7 @@
 use alloc::sync::Arc;
 
 /// Memory backend types and traits.
-pub use crate::store::{LinearMemory, MemoryBackend, PagedMemory, VecMemory};
+pub use crate::store::{LazyLinearMemory, LinearMemory, MemoryBackend, PagedMemory, VecMemory};
 
 /// Global configuration for the WebAssembly interpreter
 ///

crates/tinywasm/src/instance.rs

@@ -133,7 +133,13 @@ impl ModuleInstance {
 
         addrs.funcs.extend(store.init_funcs(&module.0.funcs, idx));
         addrs.tables.extend(store.init_tables(&module.0.table_types, idx));
-        addrs.memories.extend(store.init_memories(&module.0.memory_types, idx)?);
+        match module.0.local_memory_allocation {
+            LocalMemoryAllocation::Skip => {}
+            LocalMemoryAllocation::Lazy => {
+                addrs.memories.extend(store.init_lazy_memories(&module.0.memory_types, idx)?)
+            }
+            LocalMemoryAllocation::Eager => addrs.memories.extend(store.init_memories(&module.0.memory_types, idx)?),
+        }
         let global_addrs = store.init_globals(addrs.globals, &module.0.globals, &addrs.funcs, idx)?;
         let (elem_addrs, elem_trapped) =
             store.init_elements(&addrs.tables, &addrs.funcs, &global_addrs, &module.0.elements, idx)?;

crates/tinywasm/src/interpreter/executor.rs

@@ -172,6 +172,8 @@ impl<'store, const BUDGETED: bool> Executor<'store, BUDGETED> {
             LocalCopy128(from, to) => self.store.value_stack.local_set(&self.cf, *to, self.store.value_stack.local_get::<Value128>(&self.cf, *from)),
             AddLocalLocal32(a, b) => self.store.value_stack.push(self.store.value_stack.local_get::<i32>(&self.cf, *a).wrapping_add(self.store.value_stack.local_get::<i32>(&self.cf, *b)))?,
             AddLocalLocal64(a, b) => self.store.value_stack.push(self.store.value_stack.local_get::<i64>(&self.cf, *a).wrapping_add(self.store.value_stack.local_get::<i64>(&self.cf, *b)))?,
+            AddLocalLocalSet32(a, b, dst) => self.store.value_stack.local_set::<i32>(&self.cf, *dst, self.store.value_stack.local_get::<i32>(&self.cf, *a).wrapping_add(self.store.value_stack.local_get::<i32>(&self.cf, *b))),
+            AddLocalLocalSet64(a, b, dst) => self.store.value_stack.local_set::<i64>(&self.cf, *dst, self.store.value_stack.local_get::<i64>(&self.cf, *a).wrapping_add(self.store.value_stack.local_get::<i64>(&self.cf, *b))),
             AddConst32(c) => stack_op!(unary i32, |v| v.wrapping_add(*c)),
             AddConst64(c) => stack_op!(unary i64, |v| v.wrapping_add(*c)),
             AddLocalConst32(local_index, c) => self.store.value_stack.local_update::<Value32>(&self.cf, *local_index, |local| local.wrapping_add(*c as u32)),

crates/tinywasm/src/lib.rs

@@ -115,7 +115,7 @@ use interpreter::InterpreterRuntime;
 
 /// Global configuration for the WebAssembly interpreter
 pub mod engine;
-pub use engine::{Engine, LinearMemory, MemoryBackend, PagedMemory, VecMemory};
+pub use engine::{Engine, LazyLinearMemory, LinearMemory, MemoryBackend, PagedMemory, VecMemory};
 
 #[cfg(feature = "parser")]
 /// Re-export of [`tinywasm_parser`]. Requires `parser` feature.

crates/tinywasm/src/store/memory/instance.rs

@@ -47,6 +47,22 @@ impl MemoryInstance {
         Ok(Self { kind, inner: storage, page_count: kind.page_count_initial() as usize })
     }
 
+    pub(crate) fn new_lazy(kind: MemoryType, backend: &MemoryBackend) -> Result<Self> {
+        assert!(kind.page_count_initial() <= kind.page_count_max());
+
+        let initial_len = usize::try_from(kind.initial_size())
+            .map_err(|_| Error::UnsupportedFeature("memory size exceeds the host address space"))?;
+
+        crate::log::debug!(
+            "initializing lazy memory with {} pages of {} bytes",
+            kind.page_count_initial(),
+            kind.page_size()
+        );
+
+        let storage = backend.create_lazy(kind, initial_len)?;
+        Ok(Self { kind, inner: storage, page_count: kind.page_count_initial() as usize })
+    }
+
     pub(crate) const fn is_64bit(&self) -> bool {
         matches!(self.kind.arch(), MemoryArch::I64)
     }