chore: optimize block stack, remove EndFunc (#12)

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

18 files changed, 72 insertions, 68 deletions

diff --git a/crates/cli/src/bin.rs b/crates/cli/src/bin.rs
index 34d4bcd..6508c5f 100644
--- a/crates/cli/src/bin.rs
+++ b/crates/cli/src/bin.rs
@@ -4,7 +4,7 @@ use argh::FromArgs;
 use args::WasmArg;
 use color_eyre::eyre::Result;
 use log::{debug, info};
-use tinywasm::{self, types::WasmValue, Module};
+use tinywasm::{types::WasmValue, Module};
 
 use crate::args::to_wasm_args;
 mod args;
diff --git a/crates/parser/src/visit.rs b/crates/parser/src/visit.rs
index b438fc1..c5743b8 100644
--- a/crates/parser/src/visit.rs
+++ b/crates/parser/src/visit.rs
@@ -392,7 +392,7 @@ impl<'a> wasmparser::VisitOperator<'a> for FunctionBuilder {
 
     fn visit_end(&mut self) -> Self::Output {
         let Some(label_pointer) = self.label_ptrs.pop() else {
-            return self.visit(Instruction::EndFunc);
+            return self.visit(Instruction::Return);
         };
 
         let current_instr_ptr = self.instructions.len();
diff --git a/crates/tinywasm/Cargo.toml b/crates/tinywasm/Cargo.toml
index 3f42ca0..43f459a 100644
--- a/crates/tinywasm/Cargo.toml
+++ b/crates/tinywasm/Cargo.toml
@@ -13,7 +13,7 @@ name="tinywasm"
 path="src/lib.rs"
 
 [dependencies]
-log={version="0.4", optional=true}
+_log={version="0.4", optional=true, package="log"}
 tinywasm-parser={version="0.5.0", path="../parser", default-features=false, optional=true}
 tinywasm-types={version="0.5.0", path="../types", default-features=false}
 libm={version="0.2", default-features=false}
@@ -29,7 +29,7 @@ pretty_env_logger="0.5"
 
 [features]
 default=["std", "parser", "logging", "archive"]
-logging=["log", "tinywasm-parser?/logging", "tinywasm-types/logging"]
+logging=["_log", "tinywasm-parser?/logging", "tinywasm-types/logging"]
 std=["tinywasm-parser?/std", "tinywasm-types/std"]
 parser=["tinywasm-parser"]
 unsafe=["tinywasm-types/unsafe"]
diff --git a/crates/tinywasm/src/imports.rs b/crates/tinywasm/src/imports.rs
index 949c5fe..67ac360 100644
--- a/crates/tinywasm/src/imports.rs
+++ b/crates/tinywasm/src/imports.rs
@@ -191,6 +191,7 @@ impl From<&Import> for ExternName {
 ///
 /// ## Example
 /// ```rust
+/// # use _log as log;
 /// # fn main() -> tinywasm::Result<()> {
 /// use tinywasm::{Imports, Extern};
 /// use tinywasm::types::{ValType, TableType, MemoryType, WasmValue};
diff --git a/crates/tinywasm/src/lib.rs b/crates/tinywasm/src/lib.rs
index 583a68d..e2d57fc 100644
--- a/crates/tinywasm/src/lib.rs
+++ b/crates/tinywasm/src/lib.rs
@@ -79,7 +79,7 @@ extern crate alloc;
 // log for logging (optional).
 #[cfg(feature = "logging")]
 #[allow(clippy::single_component_path_imports)]
-use log;
+use _log as log;
 
 // noop fallback if logging is disabled.
 #[cfg(not(feature = "logging"))]
diff --git a/crates/tinywasm/src/runtime/interpreter/mod.rs b/crates/tinywasm/src/runtime/interpreter/mod.rs
index b014a85..2f0328b 100644
--- a/crates/tinywasm/src/runtime/interpreter/mod.rs
+++ b/crates/tinywasm/src/runtime/interpreter/mod.rs
@@ -80,7 +80,7 @@ enum ExecResult {
 fn exec_one(cf: &mut CallFrame, stack: &mut Stack, store: &mut Store, module: &ModuleInstance) -> Result<ExecResult> {
     let instrs = &cf.func_instance.0.instructions;
 
-    if unlikely(cf.instr_ptr >= instrs.len() || instrs.is_empty()) {
+    if unlikely(cf.instr_ptr as usize >= instrs.len() || instrs.is_empty()) {
         log::error!("instr_ptr out of bounds: {} >= {}", cf.instr_ptr, instrs.len());
         return Err(Error::Other(format!("instr_ptr out of bounds: {} >= {}", cf.instr_ptr, instrs.len())));
     }
@@ -128,7 +128,7 @@ fn exec_one(cf: &mut CallFrame, stack: &mut Stack, store: &mut Store, module: &M
             };
 
             let params = stack.values.pop_n_rev(wasm_func.ty.params.len())?;
-            let call_frame = CallFrame::new(wasm_func, func_inst.owner, params, stack.blocks.len());
+            let call_frame = CallFrame::new(wasm_func, func_inst.owner, params, stack.blocks.len() as u32);
 
             // push the call frame
             cf.instr_ptr += 1; // skip the call instruction
@@ -181,7 +181,7 @@ fn exec_one(cf: &mut CallFrame, stack: &mut Stack, store: &mut Store, module: &M
             }
 
             let params = stack.values.pop_n_rev(wasm_func.ty.params.len())?;
-            let call_frame = CallFrame::new(wasm_func, func_inst.owner, params, stack.blocks.len());
+            let call_frame = CallFrame::new(wasm_func, func_inst.owner, params, stack.blocks.len() as u32);
 
             // push the call frame
             cf.instr_ptr += 1; // skip the call instruction
@@ -198,8 +198,8 @@ fn exec_one(cf: &mut CallFrame, stack: &mut Stack, store: &mut Store, module: &M
                 cf.enter_block(
                     BlockFrame::new(
                         cf.instr_ptr,
-                        cf.instr_ptr + *end_offset as usize,
-                        stack.values.len(),
+                        cf.instr_ptr + *end_offset,
+                        stack.values.len() as u32,
                         BlockType::If,
                         &args.unpack(),
                         module,
@@ -213,17 +213,17 @@ fn exec_one(cf: &mut CallFrame, stack: &mut Stack, store: &mut Store, module: &M
             // falsy value is on the top of the stack
             if *else_offset != 0 {
                 let label = BlockFrame::new(
-                    cf.instr_ptr + *else_offset as usize,
-                    cf.instr_ptr + *end_offset as usize,
-                    stack.values.len(),
+                    cf.instr_ptr + *else_offset,
+                    cf.instr_ptr + *end_offset,
+                    stack.values.len() as u32,
                     BlockType::Else,
                     &args.unpack(),
                     module,
                 );
-                cf.instr_ptr += *else_offset as usize;
+                cf.instr_ptr += *else_offset;
                 cf.enter_block(label, &mut stack.values, &mut stack.blocks);
             } else {
-                cf.instr_ptr += *end_offset as usize;
+                cf.instr_ptr += *end_offset;
             }
         }
 
@@ -231,8 +231,8 @@ fn exec_one(cf: &mut CallFrame, stack: &mut Stack, store: &mut Store, module: &M
             cf.enter_block(
                 BlockFrame::new(
                     cf.instr_ptr,
-                    cf.instr_ptr + *end_offset as usize,
-                    stack.values.len(),
+                    cf.instr_ptr + *end_offset,
+                    stack.values.len() as u32,
                     BlockType::Loop,
                     args,
                     module,
@@ -246,8 +246,8 @@ fn exec_one(cf: &mut CallFrame, stack: &mut Stack, store: &mut Store, module: &M
             cf.enter_block(
                 BlockFrame::new(
                     cf.instr_ptr,
-                    cf.instr_ptr + *end_offset as usize,
-                    stack.values.len(),
+                    cf.instr_ptr + *end_offset,
+                    stack.values.len() as u32,
                     BlockType::Block,
                     args,
                     module,
@@ -258,7 +258,9 @@ fn exec_one(cf: &mut CallFrame, stack: &mut Stack, store: &mut Store, module: &M
         }
 
         BrTable(default, len) => {
-            let instr = cf.instructions()[cf.instr_ptr + 1..cf.instr_ptr + 1 + *len as usize]
+            let start = cf.instr_ptr + 1;
+            let end = cf.instr_ptr + 1 + *len;
+            let instr = cf.instructions()[start as usize..end as usize]
                 .iter()
                 .map(|i| match i {
                     BrLabel(l) => Ok(*l),
@@ -294,24 +296,13 @@ fn exec_one(cf: &mut CallFrame, stack: &mut Stack, store: &mut Store, module: &M
             false => return Ok(ExecResult::Call),
         },
 
-        EndFunc => {
-            if unlikely(stack.blocks.len() != cf.block_ptr) {
-                panic!("endfunc: block frames not empty, this should have been validated by the parser");
-            }
-
-            match stack.call_stack.is_empty() {
-                true => return Ok(ExecResult::Return),
-                false => return Ok(ExecResult::Call),
-            }
-        }
-
         // We're essentially using else as a EndBlockFrame instruction for if blocks
         Else(end_offset) => {
             let block =
                 stack.blocks.pop().expect("else: no label to end, this should have been validated by the parser");
 
-            stack.values.truncate_keep(block.stack_ptr, block.results);
-            cf.instr_ptr += *end_offset as usize;
+            stack.values.truncate_keep(block.stack_ptr, block.results as u32);
+            cf.instr_ptr += *end_offset;
         }
 
         // remove the label from the label stack
@@ -321,7 +312,7 @@ fn exec_one(cf: &mut CallFrame, stack: &mut Stack, store: &mut Store, module: &M
                 .pop()
                 .expect("end blockframe: no label to end, this should have been validated by the parser");
 
-            stack.values.truncate_keep(block.stack_ptr, block.results);
+            stack.values.truncate_keep(block.stack_ptr, block.results as u32);
         }
 
         LocalGet(local_index) => stack.values.push(cf.get_local(*local_index as usize)),
diff --git a/crates/tinywasm/src/runtime/stack/block_stack.rs b/crates/tinywasm/src/runtime/stack/block_stack.rs
index edaf2d1..719dc00 100644
--- a/crates/tinywasm/src/runtime/stack/block_stack.rs
+++ b/crates/tinywasm/src/runtime/stack/block_stack.rs
@@ -18,15 +18,15 @@ impl BlockStack {
 
     #[inline]
     /// get the label at the given index, where 0 is the top of the stack
-    pub(crate) fn get_relative_to(&self, index: usize, offset: usize) -> Option<&BlockFrame> {
-        let len = self.0.len() - offset;
+    pub(crate) fn get_relative_to(&self, index: u32, offset: u32) -> Option<&BlockFrame> {
+        let len = (self.0.len() as u32) - offset;
 
         // the vast majority of wasm functions don't use break to return
         if unlikely(index >= len) {
             return None;
         }
 
-        Some(&self.0[self.0.len() - index - 1])
+        Some(&self.0[self.0.len() - index as usize - 1])
     }
 
     #[inline]
@@ -36,31 +36,32 @@ impl BlockStack {
 
     /// keep the top `len` blocks and discard the rest
     #[inline]
-    pub(crate) fn truncate(&mut self, len: usize) {
-        self.0.truncate(len);
+    pub(crate) fn truncate(&mut self, len: u32) {
+        self.0.truncate(len as usize);
     }
 }
 
 #[derive(Debug, Clone)]
 pub(crate) struct BlockFrame {
     // position of the instruction pointer when the block was entered
-    pub(crate) instr_ptr: usize,
+    pub(crate) instr_ptr: u32,
     // position of the end instruction of the block
-    pub(crate) end_instr_ptr: usize,
+    pub(crate) end_instr_ptr: u32,
 
     // position of the stack pointer when the block was entered
-    pub(crate) stack_ptr: usize,
-    pub(crate) results: usize,
-    pub(crate) params: usize,
+    pub(crate) stack_ptr: u32,
+
+    pub(crate) results: u8,
+    pub(crate) params: u8,
     pub(crate) ty: BlockType,
 }
 
 impl BlockFrame {
     #[inline]
     pub(crate) fn new(
-        instr_ptr: usize,
-        end_instr_ptr: usize,
-        stack_ptr: usize,
+        instr_ptr: u32,
+        end_instr_ptr: u32,
+        stack_ptr: u32,
         ty: BlockType,
         args: &BlockArgs,
         module: &ModuleInstance,
@@ -70,7 +71,7 @@ impl BlockFrame {
             BlockArgs::Type(_) => (0, 1),
             BlockArgs::FuncType(t) => {
                 let ty = module.func_ty(*t);
-                (ty.params.len(), ty.results.len())
+                (ty.params.len() as u8, ty.results.len() as u8)
             }
         };
 
diff --git a/crates/tinywasm/src/runtime/stack/call_stack.rs b/crates/tinywasm/src/runtime/stack/call_stack.rs
index 12270d1..c242b74 100644
--- a/crates/tinywasm/src/runtime/stack/call_stack.rs
+++ b/crates/tinywasm/src/runtime/stack/call_stack.rs
@@ -48,8 +48,8 @@ impl CallStack {
 
 #[derive(Debug, Clone)]
 pub(crate) struct CallFrame {
-    pub(crate) instr_ptr: usize,
-    pub(crate) block_ptr: usize,
+    pub(crate) instr_ptr: u32,
+    pub(crate) block_ptr: u32,
     pub(crate) func_instance: (Rc<WasmFunction>, ModuleInstanceAddr),
     pub(crate) locals: Box<[RawWasmValue]>,
 }
@@ -63,7 +63,9 @@ impl CallFrame {
         blocks: &mut super::BlockStack,
     ) {
         if block_frame.params > 0 {
-            values.extend_from_within((block_frame.stack_ptr - block_frame.params)..block_frame.stack_ptr);
+            let start = (block_frame.stack_ptr - block_frame.params as u32) as usize;
+            let end = block_frame.stack_ptr as usize;
+            values.extend_from_within(start..end);
         }
 
         blocks.push(block_frame);
@@ -77,7 +79,7 @@ impl CallFrame {
         values: &mut super::ValueStack,
         blocks: &mut super::BlockStack,
     ) -> Option<()> {
-        let break_to = blocks.get_relative_to(break_to_relative as usize, self.block_ptr)?;
+        let break_to = blocks.get_relative_to(break_to_relative, self.block_ptr)?;
 
         // instr_ptr points to the label instruction, but the next step
         // will increment it by 1 since we're changing the "current" instr_ptr
@@ -92,7 +94,7 @@ impl CallFrame {
                 // check if we're breaking to the loop
                 if break_to_relative != 0 {
                     // we also want to trim the label stack to the loop (but not including the loop)
-                    blocks.truncate(blocks.len() - break_to_relative as usize);
+                    blocks.truncate(blocks.len() as u32 - break_to_relative);
                     return Some(());
                 }
             }
@@ -106,7 +108,7 @@ impl CallFrame {
                 self.instr_ptr = break_to.end_instr_ptr;
 
                 // we also want to trim the label stack, including the block
-                blocks.truncate(blocks.len() - (break_to_relative as usize + 1));
+                blocks.truncate(blocks.len() as u32 - (break_to_relative + 1));
             }
         }
 
@@ -118,8 +120,8 @@ impl CallFrame {
     pub(crate) fn new(
         wasm_func_inst: Rc<WasmFunction>,
         owner: ModuleInstanceAddr,
-        params: impl Iterator<Item = RawWasmValue> + ExactSizeIterator,
-        block_ptr: usize,
+        params: impl ExactSizeIterator<Item = RawWasmValue>,
+        block_ptr: u32,
     ) -> Self {
         let locals = {
             let local_types = &wasm_func_inst.locals;
@@ -150,6 +152,6 @@ impl CallFrame {
 
     #[inline(always)]
     pub(crate) fn current_instruction(&self) -> &Instruction {
-        &self.func_instance.0.instructions[self.instr_ptr]
+        &self.func_instance.0.instructions[self.instr_ptr as usize]
     }
 }
diff --git a/crates/tinywasm/src/runtime/stack/value_stack.rs b/crates/tinywasm/src/runtime/stack/value_stack.rs
index fed043f..2fa03e9 100644
--- a/crates/tinywasm/src/runtime/stack/value_stack.rs
+++ b/crates/tinywasm/src/runtime/stack/value_stack.rs
@@ -5,6 +5,7 @@ use alloc::vec::Vec;
 use tinywasm_types::{ValType, WasmValue};
 
 pub(crate) const MIN_VALUE_STACK_SIZE: usize = 1024;
+// pub(crate) const MAX_VALUE_STACK_SIZE: usize = 1024 * 1024;
 
 #[derive(Debug)]
 pub(crate) struct ValueStack {
@@ -48,9 +49,9 @@ impl ValueStack {
         self.stack.len()
     }
 
-    pub(crate) fn truncate_keep(&mut self, n: usize, end_keep: usize) {
+    pub(crate) fn truncate_keep(&mut self, n: u32, end_keep: u32) {
         let total_to_keep = n + end_keep;
-        let len = self.stack.len();
+        let len = self.stack.len() as u32;
         assert!(len >= total_to_keep, "Total to keep should be less than or equal to self.top");
 
         if len <= total_to_keep {
@@ -58,8 +59,8 @@ impl ValueStack {
         }
 
         let items_to_remove = len - total_to_keep;
-        let remove_start_index = len - items_to_remove - end_keep;
-        let remove_end_index = len - end_keep;
+        let remove_start_index = (len - items_to_remove - end_keep) as usize;
+        let remove_end_index = (len - end_keep) as usize;
         self.stack.drain(remove_start_index..remove_end_index);
     }
 
@@ -108,8 +109,10 @@ impl ValueStack {
     }
 
     #[inline]
-    pub(crate) fn break_to(&mut self, new_stack_size: usize, result_count: usize) {
-        self.stack.drain(new_stack_size..(self.stack.len() - result_count));
+    pub(crate) fn break_to(&mut self, new_stack_size: u32, result_count: u8) {
+        let start = new_stack_size as usize;
+        let end = self.stack.len() - result_count as usize;
+        self.stack.drain(start..end);
     }
 
     #[inline]
diff --git a/crates/tinywasm/src/store/memory.rs b/crates/tinywasm/src/store/memory.rs
index 9866b6b..00bcc97 100644
--- a/crates/tinywasm/src/store/memory.rs
+++ b/crates/tinywasm/src/store/memory.rs
@@ -211,7 +211,7 @@ impl_mem_loadable_for_primitive!(
 #[cfg(test)]
 mod memory_instance_tests {
     use super::*;
-    use tinywasm_types::{MemoryArch, MemoryType, ModuleInstanceAddr};
+    use tinywasm_types::MemoryArch;
 
     fn create_test_memory() -> MemoryInstance {
         let kind = MemoryType { arch: MemoryArch::I32, page_count_initial: 1, page_count_max: Some(2) };
diff --git a/crates/tinywasm/tests/test-mvp.rs b/crates/tinywasm/tests/test-mvp.rs
index 445b7fa..a1c2067 100644
--- a/crates/tinywasm/tests/test-mvp.rs
+++ b/crates/tinywasm/tests/test-mvp.rs
@@ -1,4 +1,5 @@
 mod testsuite;
+use _log as log;
 use eyre::{eyre, Result};
 use owo_colors::OwoColorize;
 use testsuite::TestSuite;
diff --git a/crates/tinywasm/tests/test-two.rs b/crates/tinywasm/tests/test-two.rs
index e710d1a..b5ed9c8 100644
--- a/crates/tinywasm/tests/test-two.rs
+++ b/crates/tinywasm/tests/test-two.rs
@@ -1,4 +1,5 @@
 mod testsuite;
+use _log as log;
 use eyre::{eyre, Result};
 use owo_colors::OwoColorize;
 use testsuite::TestSuite;
diff --git a/crates/tinywasm/tests/test-wast.rs b/crates/tinywasm/tests/test-wast.rs
index 1d3fbe3..fa3af91 100644
--- a/crates/tinywasm/tests/test-wast.rs
+++ b/crates/tinywasm/tests/test-wast.rs
@@ -1,5 +1,6 @@
 use std::path::PathBuf;
 
+use _log as log;
 use eyre::{bail, eyre, Result};
 use owo_colors::OwoColorize;
 use testsuite::TestSuite;
diff --git a/crates/tinywasm/tests/testsuite/mod.rs b/crates/tinywasm/tests/testsuite/mod.rs
index 35f0607..6223982 100644
--- a/crates/tinywasm/tests/testsuite/mod.rs
+++ b/crates/tinywasm/tests/testsuite/mod.rs
@@ -1,5 +1,6 @@
 #![allow(dead_code)] // rust analyzer doesn't recognize that code is used by tests without harness
 
+use _log as log;
 use eyre::Result;
 use owo_colors::OwoColorize;
 use std::io::{BufRead, Seek, SeekFrom};
diff --git a/crates/tinywasm/tests/testsuite/run.rs b/crates/tinywasm/tests/testsuite/run.rs
index 84efc3a..2b2bbc4 100644
--- a/crates/tinywasm/tests/testsuite/run.rs
+++ b/crates/tinywasm/tests/testsuite/run.rs
@@ -3,6 +3,7 @@ use crate::testsuite::util::*;
 use std::{borrow::Cow, collections::HashMap};
 
 use super::TestSuite;
+use _log as log;
 use eyre::{eyre, Result};
 use log::{debug, error, info};
 use tinywasm::{Extern, Imports, ModuleInstance};
diff --git a/crates/types/src/archive.rs b/crates/types/src/archive.rs
index 273d6ed..9bf095b 100644
--- a/crates/types/src/archive.rs
+++ b/crates/types/src/archive.rs
@@ -48,6 +48,9 @@ impl Display for TwasmError {
 }
 
 #[cfg(feature = "std")]
+extern crate std;
+
+#[cfg(feature = "std")]
 impl std::error::Error for TwasmError {}
 
 impl TinyWasmModule {
diff --git a/crates/types/src/instructions.rs b/crates/types/src/instructions.rs
index c932704..402cc9a 100644
--- a/crates/types/src/instructions.rs
+++ b/crates/types/src/instructions.rs
@@ -73,8 +73,7 @@ pub enum ConstInstruction {
 ///
 /// # Differences to the spec
 /// * `br_table` stores the jump labels in the following `br_label` instructions to keep this enum small.
-/// * Lables/Blocks: we store the label end offset in the instruction itself and
-///   have seperate EndBlockFrame and EndFunc instructions to mark the end of a block or function.
+/// * Lables/Blocks: we store the label end offset in the instruction itself and use `EndBlockFrame` to mark the end of a block.
 ///   This makes it easier to implement the label stack iteratively.
 ///
 /// See <https://webassembly.github.io/spec/core/binary/instructions.html>
@@ -121,7 +120,6 @@ pub enum Instruction {
     If(BlockArgsPacked, ElseOffset, EndOffset), // If else offset is 0 if there is no else block
     Else(EndOffset),
     EndBlockFrame,
-    EndFunc,
     Br(LabelAddr),
     BrIf(LabelAddr),
     BrTable(BrTableDefault, BrTableLen), // has to be followed by multiple BrLabel instructions
diff --git a/crates/types/src/lib.rs b/crates/types/src/lib.rs
index 679f3bf..c5c8071 100644
--- a/crates/types/src/lib.rs
+++ b/crates/types/src/lib.rs
@@ -3,7 +3,7 @@
     attr(deny(warnings, rust_2018_idioms), allow(dead_code, unused_assignments, unused_variables))
 ))]
 #![warn(missing_debug_implementations, rust_2018_idioms, unreachable_pub)]
-#![cfg_attr(not(feature = "std"), no_std)]
+#![no_std]
 #![cfg_attr(not(feature = "unsafe"), forbid(unsafe_code))]
 #![cfg_attr(feature = "unsafe", deny(unused_unsafe))]