From 684d5a04a928ea4132d0c5e61bb4086fac9feb22 Mon Sep 17 00:00:00 2001 From: Henry Date: Sun, 19 Apr 2026 15:12:25 +0200 Subject: feat: lazy memory allocation, fix instruction rewrite order Signed-off-by: Henry --- crates/parser/src/lib.rs | 26 ++++- crates/parser/src/module.rs | 87 ++++++++++++---- crates/parser/src/optimize.rs | 74 +++++++++----- crates/tinywasm/src/engine.rs | 2 +- crates/tinywasm/src/instance.rs | 8 +- crates/tinywasm/src/interpreter/executor.rs | 2 + crates/tinywasm/src/lib.rs | 2 +- crates/tinywasm/src/store/memory/instance.rs | 16 +++ crates/tinywasm/src/store/memory/lazy.rs | 143 +++++++++++++++++++++++++++ crates/tinywasm/src/store/memory/mod.rs | 7 +- crates/tinywasm/src/store/mod.rs | 16 ++- crates/tinywasm/tests/memory_backends.rs | 134 +++++++++++++++++++++++++ crates/tinywasm/tests/store_ownership.rs | 6 +- crates/types/src/instructions.rs | 69 +++++++++++++ crates/types/src/lib.rs | 17 ++++ 15 files changed, 552 insertions(+), 57 deletions(-) create mode 100644 crates/tinywasm/src/store/memory/lazy.rs (limited to 'crates') 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 { + pub(crate) fn into_module(self, options: &ParserOptions) -> Result { 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, + pub(crate) uses_local_memory: bool, +} + pub(crate) fn optimize_instructions( mut instructions: Vec, function_data: &mut WasmFunctionData, self_func_addr: u32, -) -> Vec { - 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 { diff --git a/crates/tinywasm/src/engine.rs b/crates/tinywasm/src/engine.rs index 916e493..36220fe 100644 --- a/crates/tinywasm/src/engine.rs +++ b/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 /// diff --git a/crates/tinywasm/src/instance.rs b/crates/tinywasm/src/instance.rs index 11c4191..6b4f407 100644 --- a/crates/tinywasm/src/instance.rs +++ b/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)?; diff --git a/crates/tinywasm/src/interpreter/executor.rs b/crates/tinywasm/src/interpreter/executor.rs index 0c85d02..acd14f6 100644 --- a/crates/tinywasm/src/interpreter/executor.rs +++ b/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::(&self.cf, *from)), AddLocalLocal32(a, b) => self.store.value_stack.push(self.store.value_stack.local_get::(&self.cf, *a).wrapping_add(self.store.value_stack.local_get::(&self.cf, *b)))?, AddLocalLocal64(a, b) => self.store.value_stack.push(self.store.value_stack.local_get::(&self.cf, *a).wrapping_add(self.store.value_stack.local_get::(&self.cf, *b)))?, + AddLocalLocalSet32(a, b, dst) => self.store.value_stack.local_set::(&self.cf, *dst, self.store.value_stack.local_get::(&self.cf, *a).wrapping_add(self.store.value_stack.local_get::(&self.cf, *b))), + AddLocalLocalSet64(a, b, dst) => self.store.value_stack.local_set::(&self.cf, *dst, self.store.value_stack.local_get::(&self.cf, *a).wrapping_add(self.store.value_stack.local_get::(&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::(&self.cf, *local_index, |local| local.wrapping_add(*c as u32)), diff --git a/crates/tinywasm/src/lib.rs b/crates/tinywasm/src/lib.rs index 7806016..d352042 100644 --- a/crates/tinywasm/src/lib.rs +++ b/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. diff --git a/crates/tinywasm/src/store/memory/instance.rs b/crates/tinywasm/src/store/memory/instance.rs index 0ea03eb..56d1915 100644 --- a/crates/tinywasm/src/store/memory/instance.rs +++ b/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 { + 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) } diff --git a/crates/tinywasm/src/store/memory/lazy.rs b/crates/tinywasm/src/store/memory/lazy.rs new file mode 100644 index 0000000..54c7ddc --- /dev/null +++ b/crates/tinywasm/src/store/memory/lazy.rs @@ -0,0 +1,143 @@ +use alloc::boxed::Box; +use alloc::vec::Vec; +use core::cell::RefCell; + +use tinywasm_types::MemoryType; + +use crate::{Error, MemoryBackend, Result}; + +use super::LinearMemory; + +/// A linear memory wrapper that materializes its inner backend on first access. +/// +/// If the wrapped backend fails to create the inner memory during first access, +/// this wrapper will panic. +pub struct LazyLinearMemory { + ty: MemoryType, + initial_len: usize, + backend: MemoryBackend, + inner: RefCell>>, +} + +impl LazyLinearMemory { + /// Creates a lazy memory for `ty` using `backend` for the eventual materialized storage. + pub fn new(ty: MemoryType, backend: MemoryBackend) -> Result { + let initial_len = usize::try_from(ty.initial_size()) + .map_err(|_| Error::UnsupportedFeature("memory size exceeds the host address space"))?; + Ok(Self::new_with_initial_len(ty, initial_len, backend)) + } + + pub(crate) fn new_with_initial_len(ty: MemoryType, initial_len: usize, backend: MemoryBackend) -> Self { + Self { ty, initial_len, backend, inner: RefCell::new(None) } + } + + fn with_inner(&self, f: impl FnOnce(&dyn LinearMemory) -> R) -> R { + self.ensure_materialized(); + let inner = self.inner.borrow(); + f(inner.as_deref().expect("lazy memory should be materialized")) + } + + fn with_inner_mut(&self, f: impl FnOnce(&mut dyn LinearMemory) -> R) -> R { + self.ensure_materialized(); + let mut inner = self.inner.borrow_mut(); + f(inner.as_deref_mut().expect("lazy memory should be materialized")) + } + + fn ensure_materialized(&self) { + if self.inner.borrow().is_some() { + return; + } + + // Lazy materialization happens from trait methods that cannot surface backend creation errors. + let storage = self.backend.create(self.ty, self.initial_len).expect("lazy memory materialization failed"); + *self.inner.borrow_mut() = Some(storage.0); + } +} + +impl LinearMemory for LazyLinearMemory { + fn len(&self) -> usize { + self.with_inner(|inner| inner.len()) + } + + fn grow_to(&mut self, new_len: usize) -> Option<()> { + self.with_inner_mut(|inner| inner.grow_to(new_len)) + } + + fn read(&self, addr: usize, dst: &mut [u8]) -> usize { + self.with_inner(|inner| inner.read(addr, dst)) + } + + fn write(&mut self, addr: usize, src: &[u8]) -> usize { + self.with_inner_mut(|inner| inner.write(addr, src)) + } + + fn write_all(&mut self, addr: usize, src: &[u8]) -> Option<()> { + self.with_inner_mut(|inner| inner.write_all(addr, src)) + } + + fn fill(&mut self, addr: usize, len: usize, val: u8) -> Option<()> { + self.with_inner_mut(|inner| inner.fill(addr, len, val)) + } + + fn copy_within(&mut self, dst: usize, src: usize, len: usize) -> Option<()> { + self.with_inner_mut(|inner| inner.copy_within(dst, src, len)) + } + + fn read_exact(&self, addr: usize, dst: &mut [u8]) -> Option<()> { + self.with_inner(|inner| inner.read_exact(addr, dst)) + } + + fn read_vec(&self, addr: usize, len: usize) -> Option> { + self.with_inner(|inner| inner.read_vec(addr, len)) + } + + fn read_8(&self, base: u64, offset: u64) -> core::result::Result { + self.with_inner(|inner| inner.read_8(base, offset)) + } + + fn read_16(&self, base: u64, offset: u64) -> core::result::Result<[u8; 2], crate::Trap> { + self.with_inner(|inner| inner.read_16(base, offset)) + } + + fn read_32(&self, base: u64, offset: u64) -> core::result::Result<[u8; 4], crate::Trap> { + self.with_inner(|inner| inner.read_32(base, offset)) + } + + fn read_64(&self, base: u64, offset: u64) -> core::result::Result<[u8; 8], crate::Trap> { + self.with_inner(|inner| inner.read_64(base, offset)) + } + + fn read_128(&self, base: u64, offset: u64) -> core::result::Result<[u8; 16], crate::Trap> { + self.with_inner(|inner| inner.read_128(base, offset)) + } + + fn write_8(&mut self, base: u64, offset: u64, byte: u8) -> core::result::Result<(), crate::Trap> { + self.with_inner_mut(|inner| inner.write_8(base, offset, byte)) + } + + fn write_16(&mut self, base: u64, offset: u64, bytes: [u8; 2]) -> core::result::Result<(), crate::Trap> { + self.with_inner_mut(|inner| inner.write_16(base, offset, bytes)) + } + + fn write_32(&mut self, base: u64, offset: u64, bytes: [u8; 4]) -> core::result::Result<(), crate::Trap> { + self.with_inner_mut(|inner| inner.write_32(base, offset, bytes)) + } + + fn write_64(&mut self, base: u64, offset: u64, bytes: [u8; 8]) -> core::result::Result<(), crate::Trap> { + self.with_inner_mut(|inner| inner.write_64(base, offset, bytes)) + } + + fn write_128(&mut self, base: u64, offset: u64, bytes: [u8; 16]) -> core::result::Result<(), crate::Trap> { + self.with_inner_mut(|inner| inner.write_128(base, offset, bytes)) + } +} + +#[cfg(feature = "debug")] +impl core::fmt::Debug for LazyLinearMemory { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + f.debug_struct("LazyLinearMemory") + .field("ty", &self.ty) + .field("materialized", &self.inner.borrow().is_some()) + .finish() + } +} diff --git a/crates/tinywasm/src/store/memory/mod.rs b/crates/tinywasm/src/store/memory/mod.rs index d62b80e..2c6558f 100644 --- a/crates/tinywasm/src/store/memory/mod.rs +++ b/crates/tinywasm/src/store/memory/mod.rs @@ -12,13 +12,14 @@ use tinywasm_types::MemoryType; use crate::{Error, Result}; mod instance; +mod lazy; mod paged; #[path = "vec.rs"] mod vec_memory; pub(crate) use instance::MemoryInstance; -pub use {paged::PagedMemory, vec_memory::VecMemory}; +pub use {lazy::LazyLinearMemory, paged::PagedMemory, vec_memory::VecMemory}; /// Backend storage for a linear memory. pub trait LinearMemory { @@ -322,6 +323,10 @@ impl MemoryBackend { Ok(MemoryStorage(storage)) } + + pub(crate) fn create_lazy(&self, ty: MemoryType, initial_len: usize) -> Result { + Ok(MemoryStorage(Box::new(LazyLinearMemory::new_with_initial_len(ty, initial_len, self.clone())))) + } } #[cfg(feature = "debug")] diff --git a/crates/tinywasm/src/store/mod.rs b/crates/tinywasm/src/store/mod.rs index 9c3d0ad..d26cb8a 100644 --- a/crates/tinywasm/src/store/mod.rs +++ b/crates/tinywasm/src/store/mod.rs @@ -15,7 +15,7 @@ mod global; mod memory; mod table; -pub use memory::{LinearMemory, MemoryBackend, PagedMemory, VecMemory}; +pub use memory::{LazyLinearMemory, LinearMemory, MemoryBackend, PagedMemory, VecMemory}; pub(crate) use memory::{MemValue, MemoryInstance}; pub(crate) use {data::*, element::*, function::*, global::*, table::*}; @@ -294,6 +294,20 @@ impl Store { Ok(mem_addrs) } + pub(crate) fn init_lazy_memories( + &mut self, + memories: &[MemoryType], + _idx: ModuleInstanceAddr, + ) -> Result> { + let mem_count = self.state.memories.len(); + let mut mem_addrs = Vec::with_capacity(mem_count); + for (i, mem) in memories.iter().enumerate() { + self.state.memories.push(MemoryInstance::new_lazy(*mem, &self.engine.config().memory_backend)?); + mem_addrs.push((i + mem_count) as MemAddr); + } + Ok(mem_addrs) + } + /// Add globals to the store, returning their addresses in the store pub(crate) fn init_globals( &mut self, diff --git a/crates/tinywasm/tests/memory_backends.rs b/crates/tinywasm/tests/memory_backends.rs index 0efbee0..cd9f770 100644 --- a/crates/tinywasm/tests/memory_backends.rs +++ b/crates/tinywasm/tests/memory_backends.rs @@ -8,6 +8,45 @@ use eyre::Result; use tinywasm::engine::Config; use tinywasm::types::{MemoryArch, MemoryType}; use tinywasm::{Engine, Memory, MemoryBackend, Module, PagedMemory, Store}; +use tinywasm_parser::{Parser, ParserOptions}; + +fn instantiate_module_with_counting_backend(module: Module) -> Result { + let created = Arc::new(AtomicUsize::new(0)); + let factory_calls = created.clone(); + let backend = MemoryBackend::custom(move |ty| { + factory_calls.fetch_add(1, Ordering::Relaxed); + Ok(PagedMemory::new(ty.initial_size() as usize, 16)) + }); + let engine = Engine::new(Config::new().with_memory_backend(backend)); + let mut store = Store::new(engine); + + let _ = module.instantiate(&mut store, None)?; + + Ok(created.load(Ordering::Relaxed)) +} + +fn instantiate_with_counting_backend(wat: &str) -> Result { + let wasm = wat::parse_str(wat)?; + let module = Module::parse_bytes(&wasm)?; + instantiate_module_with_counting_backend(module) +} + +fn instantiate_exported_memory_with_counting_backend( + wat: &str, +) -> Result<(Store, tinywasm::ModuleInstance, Arc)> { + let wasm = wat::parse_str(wat)?; + let module = Module::parse_bytes(&wasm)?; + let created = Arc::new(AtomicUsize::new(0)); + let factory_calls = created.clone(); + let backend = MemoryBackend::custom(move |ty| { + factory_calls.fetch_add(1, Ordering::Relaxed); + Ok(PagedMemory::new(ty.initial_size() as usize, 16)) + }); + let engine = Engine::new(Config::new().with_memory_backend(backend)); + let mut store = Store::new(engine); + let instance = module.instantiate(&mut store, None)?; + Ok((store, instance, created)) +} #[test] fn paged_backend_works_for_module_memories() -> Result<()> { @@ -58,6 +97,101 @@ fn custom_backend_factory_is_used_for_host_memories() -> Result<()> { Ok(()) } +#[test] +fn local_memory_without_observable_use_is_not_allocated() -> Result<()> { + let created = instantiate_with_counting_backend( + r#" + (module + (memory 1) + (func (export "run")) + ) + "#, + )?; + + assert_eq!(created, 0); + Ok(()) +} + +#[test] +fn exported_local_memory_is_not_eagerly_allocated() -> Result<()> { + let created = instantiate_with_counting_backend( + r#" + (module + (memory (export "memory") 1) + ) + "#, + )?; + + assert_eq!(created, 0); + Ok(()) +} + +#[test] +fn exported_local_memory_materializes_on_first_method_call() -> Result<()> { + let (store, instance, created) = instantiate_exported_memory_with_counting_backend( + r#" + (module + (memory (export "memory") 1) + ) + "#, + )?; + + let memory = instance.memory("memory")?; + assert_eq!(created.load(Ordering::Relaxed), 0); + assert_eq!(memory.len(&store)?, 65536); + assert_eq!(created.load(Ordering::Relaxed), 1); + Ok(()) +} + +#[test] +fn active_data_segment_on_local_memory_is_allocated() -> Result<()> { + let created = instantiate_with_counting_backend( + r#" + (module + (memory 1) + (data (i32.const 0) "hi") + ) + "#, + )?; + + assert_eq!(created, 1); + Ok(()) +} + +#[test] +fn local_memory_instruction_is_allocated() -> Result<()> { + let created = instantiate_with_counting_backend( + r#" + (module + (memory 1) + (func (export "run") (drop (memory.size))) + ) + "#, + )?; + + assert_eq!(created, 1); + Ok(()) +} + +#[test] +fn disabled_local_memory_allocation_optimization_keeps_old_behavior() -> Result<()> { + let wasm = wat::parse_str( + r#" + (module + (memory 1) + (func (export "run")) + ) + "#, + )?; + let parser = Parser::with_options(ParserOptions::default().with_local_memory_allocation_optimization(false)); + let module = Module::from(parser.parse_module_bytes(&wasm)?); + + let created = instantiate_module_with_counting_backend(module)?; + + assert_eq!(created, 1); + Ok(()) +} + #[test] fn read_returns_short_count_at_end_of_memory() -> Result<()> { let mut store = Store::default(); diff --git a/crates/tinywasm/tests/store_ownership.rs b/crates/tinywasm/tests/store_ownership.rs index f1bd117..0f81ea3 100644 --- a/crates/tinywasm/tests/store_ownership.rs +++ b/crates/tinywasm/tests/store_ownership.rs @@ -1,5 +1,5 @@ use eyre::Result; -use tinywasm::{Error, Module, Store}; +use tinywasm::{Module, Store}; const MODULE_WAT: &str = r#" (module @@ -22,7 +22,7 @@ fn func_handle_rejects_wrong_store() -> Result<()> { let mut other_store = Store::default(); let err = func.call(&mut other_store, &[1.into(), 2.into()]).unwrap_err(); - assert!(matches!(err, Error::InvalidStore)); + assert!(err.to_string().contains("invalid store")); Ok(()) } @@ -38,7 +38,7 @@ fn memory_access_rejects_wrong_store() -> Result<()> { let memory = instance.memory("memory")?; let other_store = Store::default(); let err = memory.len(&other_store).unwrap_err(); - assert!(matches!(err, Error::InvalidStore)); + assert!(err.to_string().contains("invalid store")); Ok(()) } diff --git a/crates/types/src/instructions.rs b/crates/types/src/instructions.rs index 696477a..b26a51a 100644 --- a/crates/types/src/instructions.rs +++ b/crates/types/src/instructions.rs @@ -78,6 +78,7 @@ pub enum CmpOp { pub enum Instruction { LocalCopy32(LocalAddr, LocalAddr), LocalCopy64(LocalAddr, LocalAddr), LocalCopy128(LocalAddr, LocalAddr), AddLocalLocal32(LocalAddr, LocalAddr), AddLocalLocal64(LocalAddr, LocalAddr), + AddLocalLocalSet32(LocalAddr, LocalAddr, LocalAddr), AddLocalLocalSet64(LocalAddr, LocalAddr, LocalAddr), AddConst32(i32), AddConst64(i64), AddLocalConst32(LocalAddr, i32), AddLocalConst64(LocalAddr, i64), SetLocalConst32(LocalAddr, i32), SetLocalConst64(LocalAddr, i64), @@ -298,6 +299,74 @@ pub enum Instruction { I32x4RelaxedDotI8x16I7x16AddS } +impl Instruction { + #[inline] + pub const fn memory_addr(&self) -> Option { + match self { + Self::StoreLocalLocal32(arg, ..) + | Self::StoreLocalLocal64(arg, ..) + | Self::StoreLocalLocal128(arg, ..) + | Self::LoadLocal32(arg, ..) + | Self::LoadLocalTee32(arg, ..) + | Self::LoadLocalSet32(arg, ..) + | Self::LoadLocalTee128(arg, ..) + | Self::LoadLocalSet128(arg, ..) + | Self::I32Load(arg) + | Self::I64Load(arg) + | Self::F32Load(arg) + | Self::F64Load(arg) + | Self::I32Load8S(arg) + | Self::I32Load8U(arg) + | Self::I32Load16S(arg) + | Self::I32Load16U(arg) + | Self::I64Load8S(arg) + | Self::I64Load8U(arg) + | Self::I64Load16S(arg) + | Self::I64Load16U(arg) + | Self::I64Load32S(arg) + | Self::I64Load32U(arg) + | Self::I32Store(arg) + | Self::I64Store(arg) + | Self::F32Store(arg) + | Self::F64Store(arg) + | Self::I32Store8(arg) + | Self::I32Store16(arg) + | Self::I64Store8(arg) + | Self::I64Store16(arg) + | Self::I64Store32(arg) + | Self::V128Load(arg) + | Self::V128Load8x8S(arg) + | Self::V128Load8x8U(arg) + | Self::V128Load16x4S(arg) + | Self::V128Load16x4U(arg) + | Self::V128Load32x2S(arg) + | Self::V128Load32x2U(arg) + | Self::V128Load8Splat(arg) + | Self::V128Load16Splat(arg) + | Self::V128Load32Splat(arg) + | Self::V128Load64Splat(arg) + | Self::V128Load8Lane(arg, ..) + | Self::V128Load16Lane(arg, ..) + | Self::V128Load32Lane(arg, ..) + | Self::V128Load64Lane(arg, ..) + | Self::V128Load32Zero(arg) + | Self::V128Load64Zero(arg) + | Self::V128Store(arg) + | Self::V128Store8Lane(arg, ..) + | Self::V128Store16Lane(arg, ..) + | Self::V128Store32Lane(arg, ..) + | Self::V128Store64Lane(arg, ..) => Some(arg.mem_addr()), + Self::MemorySize(mem) + | Self::MemoryGrow(mem) + | Self::MemoryInit(mem, ..) + | Self::MemoryFill(mem) + | Self::MemoryFillImm(mem, ..) => Some(*mem), + Self::MemoryCopy { dst_mem, src_mem } => Some(if *dst_mem >= *src_mem { *dst_mem } else { *src_mem }), + _ => None, + } + } +} + #[cfg(test)] mod tests { use super::Instruction; diff --git a/crates/types/src/lib.rs b/crates/types/src/lib.rs index 30cae5d..914e8c4 100644 --- a/crates/types/src/lib.rs +++ b/crates/types/src/lib.rs @@ -118,6 +118,23 @@ pub struct TinyWasmModule { /// /// Corresponds to the `elem` section of the original WebAssembly module. pub elements: ArcSlice, + + /// How instantiation should prepare the module's local memories. + pub local_memory_allocation: LocalMemoryAllocation, +} + +/// How instantiation should prepare local memories declared by the module. +#[derive(Clone, Copy, PartialEq, Eq, Default)] +#[cfg_attr(feature = "debug", derive(Debug))] +#[cfg_attr(feature = "archive", derive(serde::Serialize, serde::Deserialize))] +pub enum LocalMemoryAllocation { + /// The module's local memories are unobservable and can be skipped entirely. + #[default] + Skip, + /// The module's local memories may be observed through exports, but can be delayed until first use. + Lazy, + /// The module's local memories must be allocated during instantiation. + Eager, } /// A WebAssembly External Kind. -- cgit v1.3.1