fix(stack): stack management

codegen/src/asm.rs

@@ -53,13 +53,16 @@ impl Assembler {
             self.sp,
             self.sp + items
         );
-        self.sp = self
-            .sp
-            .checked_add(items)
-            .ok_or(Error::StackOverflow(self.sp, items))?;
+        self.sp = self.sp.checked_add(items).ok_or(Error::StackOverflow {
+            expected: self.sp,
+            found: self.sp + items,
+        })?;
 
         if self.sp > MAX_STACK_SIZE {
-            return Err(Error::StackOverflow(self.sp, items));
+            return Err(Error::StackOverflow {
+                expected: MAX_STACK_SIZE,
+                found: self.sp,
+            });
         }
 
         Ok(())
@@ -79,9 +82,10 @@ impl Assembler {
         self.sp = if self.sp == items {
             0
         } else {
-            self.sp
-                .checked_sub(items)
-                .ok_or(Error::StackUnderflow(self.sp, items))?
+            self.sp.checked_sub(items).ok_or(Error::StackUnderflow {
+                expected: items,
+                found: self.sp,
+            })?
         };
 
         Ok(())

codegen/src/codegen/dispatcher.rs

@@ -1,7 +1,7 @@
 //! Code generator for EVM dispatcher.
 
 use crate::{
-    wasm::{self, Env, Functions},
+    wasm::{self, Env, Functions, ToLSBytes},
     JumpTable, MacroAssembler, Result,
 };
 use std::collections::BTreeMap;
@@ -61,6 +61,9 @@ impl Dispatcher {
 
     /// Emit selector to buffer.
     fn emit_selector(&mut self, selector: &wasm::Function<'_>, last: bool) -> Result<()> {
+        const RETURN_OFFSET: u8 = 0;
+        const RETURN_SIZE: u8 = 32;
+
         let abi = self.env.load_abi(selector)?;
         self.abi.push(abi.clone());
 
@@ -71,22 +74,37 @@ impl Dispatcher {
             abi.signature(),
         );
 
-        let func = self.env.query_func(&abi.name)?;
-        self.asm.increment_sp(1)?;
+        // Compare selectors.
+        self.asm.push(&selector_bytes)?; // Stack: [selector, selector_bytes]
+        self.asm._eq()?; // Stack: [result]
 
-        // Prepare the `PC` of the callee function.
+        // Conditional jump to function.
+        let func = self.env.query_func(&abi.name)?;
         self.table.call(self.asm.pc(), func);
+        self.asm._jumpi()?; // Jump to func if result != 0
 
+        // Skip to next selector or stop.
         if last {
-            self.asm._swap1()?;
+            self.asm._stop()?;
         } else {
-            self.asm._dup2()?;
+            // Drop result of failed selector match
+            self.asm._pop()?;
         }
 
-        self.asm.push(&selector_bytes)?;
-        self.asm._eq()?;
-        self.asm._swap1()?;
-        self.asm._jumpi()?;
+        // Function return handling.
+        let has_return = self
+            .funcs
+            .get(&func)
+            .map(|ty| !ty.results().is_empty())
+            .unwrap_or(false);
+        if has_return {
+            self.asm._jumpdest()?;
+            self.asm.push(&RETURN_OFFSET.to_ls_bytes())?;
+            self.asm._mstore()?;
+            self.asm.push(&RETURN_SIZE.to_ls_bytes())?;
+            self.asm.push(&RETURN_OFFSET.to_ls_bytes())?;
+            self.asm._return()?;
+        }
 
         Ok(())
     }

codegen/src/codegen/function.rs

@@ -138,7 +138,11 @@ impl Function {
     pub fn finish(self, jump_table: &mut JumpTable, pc: u16) -> Result<Buffer> {
         let sp = self.masm.sp();
         if !self.is_main && self.abi.is_none() && self.masm.sp() != self.ty.results().len() as u16 {
-            return Err(Error::StackNotBalanced(sp));
+            return Err(Error::StackNotBalanced {
+                func_index: self.env.index,
+                expected: self.ty.results().len() as u16,
+                found: sp,
+            });
         }
 
         jump_table.merge(self.table, pc)?;

codegen/src/masm/cmp.rs

@@ -13,7 +13,9 @@ impl MacroAssembler {
         self.push(&[1])?;
         // NOTE: this is the overridden sub but not `self.asm.sub`
         self._sub()?;
-        self.asm._lt()
+        self.asm._lt()?; // a b-1 lt -> a < b-1 -> a <= b
+        self.asm._iszero()?; // Invert: a >= b
+        Ok(())
     }
 
     /// Greater than or equal comparison.
@@ -25,22 +27,26 @@ impl MacroAssembler {
         self.push(&[1])?;
         // NOTE: this is the overridden sub but not `self.asm.sub`
         self._sub()?;
-        self.asm._slt()
+        self.asm._slt()?; // a b-1 slt -> a < b-1 (signed) -> a <= b
+        self.asm._iszero()?; // Invert: a >= b
+        Ok(())
     }
 
-    /// Greater than or equal comparison.
+    /// Less than or equal comparison.
     ///
-    /// a b sge -> a b-1 sgt(slt)
+    /// a b sle -> a b-1 sgt(slt)
     ///
-    /// Using lt due to order of stack.
+    /// Using gt due to order of stack.
     pub fn _sle(&mut self) -> Result<()> {
         self.push(&[1])?;
         // NOTE: this is the overridden sub but not `self.asm.sub`
         self._sub()?;
-        self.asm._slt()
+        self.asm._sgt()?; // a b-1 sgt -> a > b-1 (signed) -> a >= b
+        self.asm._iszero()?; // Invert: a <= b
+        Ok(())
     }
 
-    /// Greater than or equal comparison.
+    /// Less than or equal comparison.
     ///
     /// a b le -> a b-1 lt(gt)
     ///
@@ -49,35 +55,41 @@ impl MacroAssembler {
         self.push(&[1])?;
         // NOTE: this is the overridden sub but not `self.asm.sub`
         self._sub()?;
-        self.asm._lt()
+        self.asm._gt()?; // a b-1 gt -> a > b-1 -> a >= b
+        self.asm._iszero()?; // Invert: a <= b
+        Ok(())
     }
 
-    /// Greater than and equal comparison.
+    /// Signed greater than comparison.
     ///
-    /// Using slt due to order of stack.
+    /// Using sgt due to order of stack.
     pub fn _sgt(&mut self) -> Result<()> {
-        self.asm._slt()
+        self.asm._sgt()?; // Correct: SGT (0x13)
+        Ok(())
     }
 
     /// Greater than comparison.
     ///
-    /// Using lt due to order of stack.
+    /// Using gt due to order of stack.
     pub fn _gt(&mut self) -> Result<()> {
-        self.asm._lt()
+        self.asm._gt()?; // Correct: GT (0x11)
+        Ok(())
     }
 
-    /// less than comparison.
+    /// Less than comparison.
     ///
-    /// Using gt due to order of stack.
+    /// Using lt due to order of stack.
     pub fn _lt(&mut self) -> Result<()> {
-        self.asm._gt()
+        self.asm._lt()?; // Correct: LT (0x10)
+        Ok(())
     }
 
-    /// less than or equal comparison.
+    /// Signed less than comparison.
     ///
-    /// Using gt due to order of stack.
+    /// Using slt due to order of stack.
     pub fn _slt(&mut self) -> Result<()> {
-        self.asm._sgt()
+        self.asm._slt()?; // Correct: SLT (0x12)
+        Ok(())
     }
 
     /// Sign-agnostic compare unequal.

codegen/src/masm/mod.rs

@@ -267,4 +267,17 @@ impl MacroAssembler {
 
         Ok(())
     }
+
+    /// Set the stack pointer to a specific value.
+    pub fn set_sp(&mut self, value: u16) -> Result<()> {
+        if value > 1024 {
+            return Err(Error::StackOverflow {
+                expected: self.sp(),
+                found: value,
+            });
+        }
+        tracing::trace!("set stack pointer {} -> {}", self.sp(), value);
+        self.asm.sp = value;
+        Ok(())
+    }
 }

codegen/src/masm/ret.rs

@@ -39,11 +39,16 @@ impl MacroAssembler {
             self._drop()?;
         }
 
+        // skipping SWAP1 for len=0. for results=[], only JUMP is executed, consuming return PC (sp=1 → sp=0).
+        // this maintains behavior for len>0 (e.g., $func2 in ../stack/dispatcher.wat).
+        if len > 0 {
+            self.shift_stack(len, false)?;
+        }
+
         // Shift stack to prompt the jump instruction,
         // what about just dup it?
         //
         // TODO: handle the length of results > u8::MAX.
-        self.shift_stack(len, false)?;
         self._jump()
     }
 }

codegen/src/result.rs

@@ -69,6 +69,9 @@ pub enum Error {
     /// Failed to parse function selector.
     #[error("Invalid function selector")]
     InvalidSelector,
+    /// Failed to get correct stack value size or format
+    #[error("Invalid stack value")]
+    InvalidStackValue,
     /// Failed to patch jump destination.
     #[error("Invalid frame label")]
     LabelMismatch,
@@ -88,14 +91,43 @@ pub enum Error {
     #[error("Stack index is out of range {0}, max is 255 (0x400)")]
     StackIndexOutOfRange(u16),
     /// Failed to increment stack pointer.
-    #[error("Stack overflow, max is 1024 stack items, but add {1} to {0}")]
-    StackOverflow(u16, u16),
+    #[error("Stack overflow, max is 1024 stack items, attempted {found} (current {expected})")]
+    StackOverflow {
+        /// Expected stack items
+        expected: u16,
+        /// Actual stack items found
+        found: u16,
+    },
     /// Failed to decrement stack pointer.
-    #[error("Stack underflow, current stack items {0}, expect at least {1}")]
-    StackUnderflow(u16, u16),
+    #[error("Stack underflow, current stack items {found}, expect at least {expected}")]
+    StackUnderflow {
+        /// Expected stack items
+        expected: u16,
+        /// Actual stack items found
+        found: u16,
+    },
     /// Failed to pop stack.
-    #[error("Stack not balanced, current stack items {0}")]
-    StackNotBalanced(u16),
+    #[error("Stack not balanced in function {func_index:?}, current stack items {found}, expected {expected}")]
+    StackNotBalanced {
+        /// Function index where imbalance occurred
+        func_index: Option<u32>,
+        /// Expected stack items
+        expected: u16,
+        /// Actual stack items found
+        found: u16,
+    },
+    /// Stack mismatch between expected and actual items.
+    #[error(
+        "Stack mismatch in function {func_index:?}: expected {expected} items, found {found} items"
+    )]
+    StackMismatch {
+        /// Function index where mismatch occurred
+        func_index: Option<u32>,
+        /// Expected stack items
+        expected: u16,
+        /// Actual stack items found
+        found: u16,
+    },
     /// Failed to queue host functions.
     #[error("Unsupported host function {0:?}")]
     UnsupportedHostFunc(crate::wasm::HostFunc),

codegen/src/visitor/call.rs

@@ -67,29 +67,41 @@ impl Function {
         let reserved = self.env.slots.get(&index).unwrap_or(&0);
         let (params, results) = self.env.funcs.get(&index).unwrap_or(&(0, 0));
 
-        // TODO This is a temporary fix to avoid stack underflow.
-        // We need to find a more elegant solution for this.
-        self.masm.increment_sp(1)?;
+        if self.masm.sp() < *params as u16 {
+            return Err(Error::StackUnderflow {
+                expected: *params as u16,
+                found: self.masm.sp(),
+            });
+        }
 
-        // Store parameters in memory and register the call index in the jump table.
+        // Store parameters in memory.
         for i in (0..*params).rev() {
             tracing::trace!("Storing local at {} for function {index}", i + reserved);
             self.masm.push(&((i + reserved) * 0x20).to_ls_bytes())?;
             self.masm._mstore()?;
         }
 
-        // Register the label to jump back.
-        let return_pc = self.masm.pc() + 2;
+        // Emit JUMPDEST to mark the return point.
+        self.masm._jumpdest()?;
+        let return_pc = self.masm.pc(); // return PC is the current PC after JUMPDEST.
+
+        // Register the return PC as a label in the JumpTable.
         self.table.label(self.masm.pc(), return_pc);
-        self.masm._jumpdest()?; // TODO: support same pc different label
 
-        // Register the call index in the jump table.
-        self.table.call(self.masm.pc(), index); // [PUSHN, CALL_PC]
+        // Push the return PC onto the stack.
+        self.masm.push(&return_pc.to_ls_bytes())?;
+
+        // Register the function call in the JumpTable and emit JUMP.
+        self.table.call(self.masm.pc(), index);
         self.masm._jump()?;
 
-        // Adjust the stack pointer for the results.
+        // Drop any excess values to ensure stack contains exactly the expected return values.
+        // Assumes that the callee may leave extra values, but never fewer than expected.
         self.masm._jumpdest()?;
-        self.masm.increment_sp(*results as u16)?;
+        while self.masm.sp() > *results as u16 {
+            self.masm._drop()?;
+        }
+
         Ok(())
     }
 

codegen/src/visitor/control.rs

@@ -157,18 +157,21 @@ impl Function {
     /// - End of function.
     /// - End of program.
     pub fn _end(&mut self) -> Result<()> {
+        tracing::trace!("ENTERING _end, sp: {}", self.masm.sp());
         if let Ok(frame) = self.control.pop() {
             return self.handle_frame_popping(frame);
         }
 
         let results = self.ty.results();
         if self.is_main || self.abi.is_some() {
             tracing::trace!("end of main function");
-            self.masm.main_return(results)
+            self.masm.main_return(results)?;
         } else {
             tracing::trace!("end of call");
-            self.masm.call_return(results)
+            self.masm.call_return(results)?;
         }
+
+        Ok(())
     }
 
     /// Mark as invalid for now.