feat: make init stack size configurable (#44)

7 files changed, 176 insertions, 21 deletions

diff --git a/crates/tinywasm/src/config.rs b/crates/tinywasm/src/config.rs
new file mode 100644
index 0000000..70ca4db
--- /dev/null
+++ b/crates/tinywasm/src/config.rs
@@ -0,0 +1,115 @@
+use core::fmt;
+
+/// Default initial size for the 32-bit value stack (i32, f32 values).
+pub const DEFAULT_VALUE_STACK_32_INIT_SIZE: usize = 32 * 1024; // 32KB
+
+/// Default initial size for the 64-bit value stack (i64, f64 values).
+pub const DEFAULT_VALUE_STACK_64_INIT_SIZE: usize = 16 * 1024; // 16KB
+
+/// Default initial size for the 128-bit value stack (v128 values).
+pub const DEFAULT_VALUE_STACK_128_INIT_SIZE: usize = 8 * 1024; // 8KB
+
+/// Default initial size for the reference value stack (funcref, externref values).
+pub const DEFAULT_VALUE_STACK_REF_INIT_SIZE: usize = 1024; // 1KB
+
+/// Default initial size for the block stack.
+pub const DEFAULT_BLOCK_STACK_INIT_SIZE: usize = 128;
+
+/// Configuration for the WebAssembly interpreter's stack preallocation.
+///
+/// This struct allows you to configure how much space is preallocated for the
+/// different parts of the stack that the interpreter uses to store values.
+#[derive(Debug, Clone)]
+pub struct StackConfig {
+    value_stack_32_init_size: Option<usize>,
+    value_stack_64_init_size: Option<usize>,
+    value_stack_128_init_size: Option<usize>,
+    value_stack_ref_init_size: Option<usize>,
+    block_stack_init_size: Option<usize>,
+}
+
+impl StackConfig {
+    /// Create a new stack configuration with default settings.
+    pub fn new() -> Self {
+        Self {
+            value_stack_32_init_size: None,
+            value_stack_64_init_size: None,
+            value_stack_128_init_size: None,
+            value_stack_ref_init_size: None,
+            block_stack_init_size: None,
+        }
+    }
+
+    /// Get the initial size for the 32-bit value stack.
+    pub fn value_stack_32_init_size(&self) -> usize {
+        self.value_stack_32_init_size.unwrap_or(DEFAULT_VALUE_STACK_32_INIT_SIZE)
+    }
+
+    /// Get the initial size for the 64-bit value stack.
+    pub fn value_stack_64_init_size(&self) -> usize {
+        self.value_stack_64_init_size.unwrap_or(DEFAULT_VALUE_STACK_64_INIT_SIZE)
+    }
+
+    /// Get the initial size for the 128-bit value stack.
+    pub fn value_stack_128_init_size(&self) -> usize {
+        self.value_stack_128_init_size.unwrap_or(DEFAULT_VALUE_STACK_128_INIT_SIZE)
+    }
+
+    /// Get the initial size for the reference value stack.
+    pub fn value_stack_ref_init_size(&self) -> usize {
+        self.value_stack_ref_init_size.unwrap_or(DEFAULT_VALUE_STACK_REF_INIT_SIZE)
+    }
+
+    /// Get the initial size for the block stack.
+    pub fn block_stack_init_size(&self) -> usize {
+        self.block_stack_init_size.unwrap_or(DEFAULT_BLOCK_STACK_INIT_SIZE)
+    }
+
+    /// Set the initial capacity for the 32-bit value stack.
+    pub fn with_value_stack_32_init_size(mut self, capacity: usize) -> Self {
+        self.value_stack_32_init_size = Some(capacity);
+        self
+    }
+
+    /// Set the initial capacity for the 64-bit value stack.
+    pub fn with_value_stack_64_init_size(mut self, capacity: usize) -> Self {
+        self.value_stack_64_init_size = Some(capacity);
+        self
+    }
+
+    /// Set the initial capacity for the 128-bit value stack.
+    pub fn with_value_stack_128_init_size(mut self, capacity: usize) -> Self {
+        self.value_stack_128_init_size = Some(capacity);
+        self
+    }
+
+    /// Set the initial capacity for the reference value stack.
+    pub fn with_value_stack_ref_init_size(mut self, capacity: usize) -> Self {
+        self.value_stack_ref_init_size = Some(capacity);
+        self
+    }
+
+    /// Set the initial capacity for the block stack.
+    pub fn with_block_stack_init_size(mut self, capacity: usize) -> Self {
+        self.block_stack_init_size = Some(capacity);
+        self
+    }
+}
+
+impl Default for StackConfig {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl fmt::Display for StackConfig {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "StackConfig {{ ")?;
+        write!(f, "value_stack_32: {}, ", self.value_stack_32_init_size())?;
+        write!(f, "value_stack_64: {}, ", self.value_stack_64_init_size())?;
+        write!(f, "value_stack_128: {}, ", self.value_stack_128_init_size())?;
+        write!(f, "value_stack_ref: {}, ", self.value_stack_ref_init_size())?;
+        write!(f, "block_stack: {} }}", self.block_stack_init_size())?;
+        Ok(())
+    }
+}
diff --git a/crates/tinywasm/src/func.rs b/crates/tinywasm/src/func.rs
index 97bf1e8..e8055d1 100644
--- a/crates/tinywasm/src/func.rs
+++ b/crates/tinywasm/src/func.rs
@@ -63,7 +63,7 @@ impl FuncHandle {
 
         // 7. Push the frame f to the call stack
         // & 8. Push the values to the stack (Not needed since the call frame owns the values)
-        let mut stack = Stack::new(call_frame);
+        let mut stack = Stack::new(call_frame, &store.config);
 
         // 9. Invoke the function instance
         let runtime = store.runtime();
diff --git a/crates/tinywasm/src/interpreter/stack/block_stack.rs b/crates/tinywasm/src/interpreter/stack/block_stack.rs
index 2267194..e5d87a8 100644
--- a/crates/tinywasm/src/interpreter/stack/block_stack.rs
+++ b/crates/tinywasm/src/interpreter/stack/block_stack.rs
@@ -1,4 +1,4 @@
-use crate::unlikely;
+use crate::{StackConfig, unlikely};
 use alloc::vec::Vec;
 
 use crate::interpreter::values::{StackHeight, StackLocation};
@@ -6,13 +6,11 @@ use crate::interpreter::values::{StackHeight, StackLocation};
 #[derive(Debug)]
 pub(crate) struct BlockStack(Vec<BlockFrame>);
 
-impl Default for BlockStack {
-    fn default() -> Self {
-        Self(Vec::with_capacity(128))
+impl BlockStack {
+    pub(crate) fn new(config: &StackConfig) -> Self {
+        Self(Vec::with_capacity(config.block_stack_init_size()))
     }
-}
 
-impl BlockStack {
     #[inline(always)]
     pub(crate) fn len(&self) -> usize {
         self.0.len()
diff --git a/crates/tinywasm/src/interpreter/stack/mod.rs b/crates/tinywasm/src/interpreter/stack/mod.rs
index c526c4a..a4706f1 100644
--- a/crates/tinywasm/src/interpreter/stack/mod.rs
+++ b/crates/tinywasm/src/interpreter/stack/mod.rs
@@ -6,6 +6,8 @@ pub(crate) use block_stack::{BlockFrame, BlockStack, BlockType};
 pub(crate) use call_stack::{CallFrame, CallStack, Locals};
 pub(crate) use value_stack::ValueStack;
 
+use crate::StackConfig;
+
 /// A WebAssembly Stack
 #[derive(Debug)]
 pub(crate) struct Stack {
@@ -15,7 +17,11 @@ pub(crate) struct Stack {
 }
 
 impl Stack {
-    pub(crate) fn new(call_frame: CallFrame) -> Self {
-        Self { values: ValueStack::new(), blocks: BlockStack::default(), call_stack: CallStack::new(call_frame) }
+    pub(crate) fn new(call_frame: CallFrame, config: &StackConfig) -> Self {
+        Self {
+            values: ValueStack::new(config),
+            blocks: BlockStack::new(config),
+            call_stack: CallStack::new(call_frame),
+        }
     }
 }
diff --git a/crates/tinywasm/src/interpreter/stack/value_stack.rs b/crates/tinywasm/src/interpreter/stack/value_stack.rs
index e8f5e82..054ffe3 100644
--- a/crates/tinywasm/src/interpreter/stack/value_stack.rs
+++ b/crates/tinywasm/src/interpreter/stack/value_stack.rs
@@ -1,13 +1,9 @@
 use alloc::vec::Vec;
 use tinywasm_types::{ExternRef, FuncRef, ValType, ValueCounts, ValueCountsSmall, WasmValue};
 
-use crate::{Result, interpreter::*};
+use crate::{Result, StackConfig, interpreter::*};
 
 use super::Locals;
-pub(crate) const STACK_32_SIZE: usize = 1024 * 32;
-pub(crate) const STACK_64_SIZE: usize = 1024 * 16;
-pub(crate) const STACK_128_SIZE: usize = 1024 * 8;
-pub(crate) const STACK_REF_SIZE: usize = 1024;
 
 #[derive(Debug)]
 pub(crate) struct ValueStack {
@@ -18,12 +14,12 @@ pub(crate) struct ValueStack {
 }
 
 impl ValueStack {
-    pub(crate) fn new() -> Self {
+    pub(crate) fn new(config: &StackConfig) -> Self {
         Self {
-            stack_32: Vec::with_capacity(STACK_32_SIZE),
-            stack_64: Vec::with_capacity(STACK_64_SIZE),
-            stack_128: Vec::with_capacity(STACK_128_SIZE),
-            stack_ref: Vec::with_capacity(STACK_REF_SIZE),
+            stack_32: Vec::with_capacity(config.value_stack_32_init_size()),
+            stack_64: Vec::with_capacity(config.value_stack_64_init_size()),
+            stack_128: Vec::with_capacity(config.value_stack_128_init_size()),
+            stack_ref: Vec::with_capacity(config.value_stack_ref_init_size()),
         }
     }
 
diff --git a/crates/tinywasm/src/lib.rs b/crates/tinywasm/src/lib.rs
index c23b595..0037829 100644
--- a/crates/tinywasm/src/lib.rs
+++ b/crates/tinywasm/src/lib.rs
@@ -69,6 +69,29 @@
 //! and other modules to be linked into the module when it is instantiated.
 //!
 //! See the [`Imports`] documentation for more information.
+//!
+//! ## Runtime Configuration
+//!
+//! For resource-constrained targets, you can configure the initial memory allocation:
+//!
+//! ```rust
+//! use tinywasm::{Store, StackConfig};
+//!
+//! // Create a store with minimal initial allocation (90% reduction in pre-allocated memory)
+//! let config = StackConfig::new()
+//!     .with_value_stack_32_init_size(1024)  // 1KB instead of 32KB
+//!     .with_value_stack_64_init_size(512)   // 512B instead of 16KB
+//!     .with_value_stack_128_init_size(256)  // 256B instead of 8KB
+//!     .with_value_stack_ref_init_size(128)  // 128B instead of 1KB
+
+//!     .with_block_stack_init_size(32);      // 32 instead of 128
+//! let store = Store::with_config(config);
+//!
+//! // Or create a partial configuration (only override what you need)
+//! let config = StackConfig::new()
+//!     .with_value_stack_32_init_size(2048); // Only override 32-bit stack size
+//! let store = Store::with_config(config);
+//! ```
 
 mod std;
 extern crate alloc;
@@ -110,6 +133,10 @@ mod store;
 pub mod interpreter;
 pub use interpreter::InterpreterRuntime;
 
+/// Configuration for the WebAssembly interpreter's stack preallocation.
+pub mod config;
+pub use config::StackConfig;
+
 #[cfg(feature = "parser")]
 /// Re-export of [`tinywasm_parser`]. Requires `parser` feature.
 pub mod parser {
diff --git a/crates/tinywasm/src/store/mod.rs b/crates/tinywasm/src/store/mod.rs
index 634fdfe..8446297 100644
--- a/crates/tinywasm/src/store/mod.rs
+++ b/crates/tinywasm/src/store/mod.rs
@@ -4,7 +4,7 @@ use core::sync::atomic::{AtomicUsize, Ordering};
 use tinywasm_types::*;
 
 use crate::interpreter::{self, InterpreterRuntime, TinyWasmValue};
-use crate::{Error, Function, ModuleInstance, Result, Trap, cold};
+use crate::{Error, Function, ModuleInstance, Result, StackConfig, Trap, cold};
 
 mod data;
 mod element;
@@ -33,6 +33,7 @@ pub struct Store {
 
     pub(crate) data: StoreData,
     pub(crate) runtime: Runtime,
+    pub(crate) config: StackConfig,
 }
 
 impl Debug for Store {
@@ -57,6 +58,12 @@ impl Store {
         Self::default()
     }
 
+    /// Create a new store with the given stack configuration
+    pub fn with_config(config: StackConfig) -> Self {
+        let id = STORE_ID.fetch_add(1, Ordering::Relaxed);
+        Self { id, module_instances: Vec::new(), data: StoreData::default(), runtime: Runtime::Default, config }
+    }
+
     /// Get a module instance by the internal id
     pub fn get_module_instance(&self, addr: ModuleInstanceAddr) -> Option<&ModuleInstance> {
         self.module_instances.get(addr as usize)
@@ -83,7 +90,13 @@ impl PartialEq for Store {
 impl Default for Store {
     fn default() -> Self {
         let id = STORE_ID.fetch_add(1, Ordering::Relaxed);
-        Self { id, module_instances: Vec::new(), data: StoreData::default(), runtime: Runtime::Default }
+        Self {
+            id,
+            module_instances: Vec::new(),
+            data: StoreData::default(),
+            runtime: Runtime::Default,
+            config: StackConfig::default(),
+        }
     }
 }