From 4795e4e5df25ab3745c314d892f3f588856cd009 Mon Sep 17 00:00:00 2001 From: g4titanx Date: Tue, 15 Apr 2025 16:58:40 +0100 Subject: [PATCH 1/3] fix(stack): stack management wrt dispatcher test --- codegen/src/asm.rs | 20 ++++++---- codegen/src/codegen/function.rs | 6 ++- codegen/src/masm/mod.rs | 13 +++++++ codegen/src/masm/ret.rs | 7 +++- codegen/src/result.rs | 41 ++++++++++++++++++--- codegen/src/visitor/call.rs | 31 +++++++++++----- codegen/src/visitor/control.rs | 7 +++- compiler/filetests/wat/stack/dispatcher.wat | 27 ++++++++++++++ tests/stack.rs | 21 +++++++++++ 9 files changed, 145 insertions(+), 28 deletions(-) create mode 100644 compiler/filetests/wat/stack/dispatcher.wat create mode 100644 tests/stack.rs diff --git a/codegen/src/asm.rs b/codegen/src/asm.rs index 8ab558e43..984107557 100644 --- a/codegen/src/asm.rs +++ b/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(()) diff --git a/codegen/src/codegen/function.rs b/codegen/src/codegen/function.rs index 3b5596adb..29c45f91b 100644 --- a/codegen/src/codegen/function.rs +++ b/codegen/src/codegen/function.rs @@ -138,7 +138,11 @@ impl Function { pub fn finish(self, jump_table: &mut JumpTable, pc: u16) -> Result { 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)?; diff --git a/codegen/src/masm/mod.rs b/codegen/src/masm/mod.rs index 10d686c4f..301a886a6 100644 --- a/codegen/src/masm/mod.rs +++ b/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(()) + } } diff --git a/codegen/src/masm/ret.rs b/codegen/src/masm/ret.rs index ec5b65a04..07a992cce 100644 --- a/codegen/src/masm/ret.rs +++ b/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() } } diff --git a/codegen/src/result.rs b/codegen/src/result.rs index 12816f3b4..31b687e45 100644 --- a/codegen/src/result.rs +++ b/codegen/src/result.rs @@ -88,14 +88,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, + /// 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, + /// 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), diff --git a/codegen/src/visitor/call.rs b/codegen/src/visitor/call.rs index cdf9066af..406b3f06f 100644 --- a/codegen/src/visitor/call.rs +++ b/codegen/src/visitor/call.rs @@ -67,9 +67,12 @@ 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. for i in (0..*params).rev() { @@ -78,18 +81,26 @@ impl Function { self.masm._mstore()?; } - // Register the label to jump back. - let return_pc = self.masm.pc() + 2; + let return_pc = self.masm.pc() + 3; + self.masm.push(&return_pc.to_ls_bytes())?; 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] + self.masm._jumpdest()?; + self.table.call(self.masm.pc(), index); self.masm._jump()?; // Adjust the stack pointer for the results. self.masm._jumpdest()?; - self.masm.increment_sp(*results as u16)?; + if *results > 0 { + self.masm._push0()?; + self.masm._mload()?; + while self.masm.sp() > *results as u16 { + self.masm._drop()?; + } + } else { + // Preserve return PC, let caller handle result + self.masm._jumpdest()?; + } + Ok(()) } diff --git a/codegen/src/visitor/control.rs b/codegen/src/visitor/control.rs index d767f5325..3c3e05b6f 100644 --- a/codegen/src/visitor/control.rs +++ b/codegen/src/visitor/control.rs @@ -157,6 +157,7 @@ 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); } @@ -164,11 +165,13 @@ impl Function { 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. diff --git a/compiler/filetests/wat/stack/dispatcher.wat b/compiler/filetests/wat/stack/dispatcher.wat new file mode 100644 index 000000000..92d0e24b2 --- /dev/null +++ b/compiler/filetests/wat/stack/dispatcher.wat @@ -0,0 +1,27 @@ +(module + (func $func1 (param i32) (result i32) + (local.get 0) + (i32.const 1) + (i32.add) + ) + (func $func2 (param i32) (result i32) + (local.get 0) + (i32.const 2) + (i32.add) + ) + (func (export "dispatcher") (param i32) (result i32) + (local.get 0) + (i32.const 1) + (i32.eq) + (if (result i32) + (then + (local.get 0) + (call $func1) + ) + (else + (local.get 0) + (call $func2) + ) + ) + ) +) \ No newline at end of file diff --git a/tests/stack.rs b/tests/stack.rs new file mode 100644 index 000000000..2234e356f --- /dev/null +++ b/tests/stack.rs @@ -0,0 +1,21 @@ +//! Stack stability tests for the Zink compiler. +#![cfg(test)] + +use anyhow::Result; +use filetests::Test; +use zint::{Bytes32, Contract}; + +#[test] +fn dispatcher_stack() -> Result<()> { + let mut contract = Contract::from(Test::STACK_DISPATCHER).pure().compile()?; + + // Test input 1: Should call func1 (input + 1), expect return value 2 + let info = contract.execute(&[1.to_bytes32()])?; + assert_eq!(info.ret, 2.to_bytes32()); + + // Test input 0: Should call func2 (input + 2), expect return value 2 + let info = contract.execute(&[0.to_bytes32()])?; + assert_eq!(info.ret, 2.to_bytes32()); + + Ok(()) +} \ No newline at end of file From 6e7179920e0841c8e1db4562cb27b2c01787f7e3 Mon Sep 17 00:00:00 2001 From: g4titanx Date: Mon, 19 May 2025 00:37:09 +0100 Subject: [PATCH 2/3] fix(stack): register pc to jumptable --- codegen/src/masm/cmp.rs | 50 ++++++++++++++++++------------ codegen/src/result.rs | 3 ++ codegen/src/visitor/call.rs | 20 +++++++----- codegen/src/visitor/log.rs | 7 +++-- zink/abi/src/lib.rs | 2 +- zink/{codegen => abi}/src/utils.rs | 27 +++++++--------- zink/codegen/src/lib.rs | 1 - zink/codegen/src/storage.rs | 2 +- 8 files changed, 63 insertions(+), 49 deletions(-) rename zink/{codegen => abi}/src/utils.rs (79%) diff --git a/codegen/src/masm/cmp.rs b/codegen/src/masm/cmp.rs index a2357c0fb..eeba2096b 100644 --- a/codegen/src/masm/cmp.rs +++ b/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. diff --git a/codegen/src/result.rs b/codegen/src/result.rs index 31b687e45..c4de12b89 100644 --- a/codegen/src/result.rs +++ b/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, diff --git a/codegen/src/visitor/call.rs b/codegen/src/visitor/call.rs index 406b3f06f..525128af6 100644 --- a/codegen/src/visitor/call.rs +++ b/codegen/src/visitor/call.rs @@ -74,21 +74,28 @@ impl Function { }); } - // 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()?; } - let return_pc = self.masm.pc() + 3; - self.masm.push(&return_pc.to_ls_bytes())?; - self.table.label(self.masm.pc(), return_pc); + // 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); + + // 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. + // Adjust the stack for results. self.masm._jumpdest()?; if *results > 0 { self.masm._push0()?; @@ -96,9 +103,6 @@ impl Function { while self.masm.sp() > *results as u16 { self.masm._drop()?; } - } else { - // Preserve return PC, let caller handle result - self.masm._jumpdest()?; } Ok(()) diff --git a/codegen/src/visitor/log.rs b/codegen/src/visitor/log.rs index 0f0e4fe95..43f49a79e 100644 --- a/codegen/src/visitor/log.rs +++ b/codegen/src/visitor/log.rs @@ -1,6 +1,7 @@ //! System instructions use crate::{masm::MemoryInfo, wasm::ToLSBytes, Error, Function, Result}; +use zabi::utils::Bytes32; impl Function { /// Parse log data from the bytecode. @@ -63,14 +64,14 @@ impl Function { let (offset, size) = self.data()?; let data = self.env.data.load(offset, size as usize)?; - // 1. write data to memory + // 1. write data to memory with 32-byte padding let MemoryInfo { offset, size } = self.masm.memory_write_bytes(&data)?; // 3. prepare the offset and size of the data. - self.masm.push(&size.to_ls_bytes())?; + self.masm.push(&size.to_bytes32())?; self.masm.push(&offset)?; - // 4. run log for the data + // 4. emit log opcode match count { 0 => self.masm._log0(), 1 => self.masm._log1(), diff --git a/zink/abi/src/lib.rs b/zink/abi/src/lib.rs index f40b1afb5..b3345ad13 100644 --- a/zink/abi/src/lib.rs +++ b/zink/abi/src/lib.rs @@ -3,10 +3,10 @@ //! Currently just a wrapper of solidity ABI. mod abi; -#[cfg(feature = "encoding")] mod encoding; pub mod result; pub mod selector; +pub mod utils; #[cfg(feature = "encoding")] pub use encoding::{decode, encode, is_dynamic_type, AbiDecode, AbiEncode, DecodeError}; diff --git a/zink/codegen/src/utils.rs b/zink/abi/src/utils.rs similarity index 79% rename from zink/codegen/src/utils.rs rename to zink/abi/src/utils.rs index 7ab578804..722ee5676 100644 --- a/zink/codegen/src/utils.rs +++ b/zink/abi/src/utils.rs @@ -1,6 +1,4 @@ //! Utils for bytes conversion. -//! -//! TODO: move this util to other library /// Trait for converting type to bytes32. pub trait Bytes32: Sized { @@ -13,28 +11,25 @@ pub trait Bytes32: Sized { } } -/// Implement Bytes32 for types. macro_rules! impl_bytes32 { ($($ty:ident),+) => { $( impl Bytes32 for $ty { fn to_bytes32(&self) -> [u8; 32] { let mut bytes = [0u8; 32]; - let ls_bytes = { - self.to_le_bytes() - .into_iter() - .rev() - .skip_while(|b| *b == 0) - .collect::>() - .into_iter() - .rev() - .collect::>() - }; - - bytes[(32 - ls_bytes.len())..].copy_from_slice(&ls_bytes); + let src = self.to_le_bytes(); + + // To prevent empty slices for 0u32. + // zero has no significant bytes, and the EVM expects [0; 32] for a zero value + if *self == 0 { + return bytes; + } + + let significant_bytes = src.len() - (self.leading_zeros() as usize / 8); + let end = significant_bytes.max(1); // Ensures non-empty slice + bytes[(32 - end)..].copy_from_slice(&src[..end]); bytes } - fn to_vec(&self) -> Vec { self.to_le_bytes().to_vec() } diff --git a/zink/codegen/src/lib.rs b/zink/codegen/src/lib.rs index 0537b514b..7c2fa6122 100644 --- a/zink/codegen/src/lib.rs +++ b/zink/codegen/src/lib.rs @@ -13,7 +13,6 @@ mod event; mod revert; mod selector; mod storage; -mod utils; /// Revert with the input message /// diff --git a/zink/codegen/src/storage.rs b/zink/codegen/src/storage.rs index adae244dc..f0d5499eb 100644 --- a/zink/codegen/src/storage.rs +++ b/zink/codegen/src/storage.rs @@ -1,4 +1,3 @@ -use crate::utils::Bytes32; use heck::AsSnakeCase; use proc_macro::TokenStream; use proc_macro2::{Literal, Span, TokenTree}; @@ -9,6 +8,7 @@ use syn::{ parse::{Parse, ParseStream, Result}, parse_quote, Attribute, Ident, ItemFn, ItemStruct, Visibility, }; +use zabi::utils::Bytes32; thread_local! { static STORAGE_REGISTRY: RefCell> = RefCell::new(HashSet::new()); From 2b6d7dc62e0692202b970a530aefcc0e47712c00 Mon Sep 17 00:00:00 2001 From: g4titanx Date: Wed, 21 May 2025 11:20:21 +0100 Subject: [PATCH 3/3] fix(dispatcher): handle return value after jump --- codegen/src/codegen/dispatcher.rs | 38 +++++++++++++++++++++++-------- codegen/src/visitor/call.rs | 11 ++++----- 2 files changed, 32 insertions(+), 17 deletions(-) diff --git a/codegen/src/codegen/dispatcher.rs b/codegen/src/codegen/dispatcher.rs index 628bfcd51..3b224c6d0 100644 --- a/codegen/src/codegen/dispatcher.rs +++ b/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(()) } diff --git a/codegen/src/visitor/call.rs b/codegen/src/visitor/call.rs index 525128af6..4bd5c4c44 100644 --- a/codegen/src/visitor/call.rs +++ b/codegen/src/visitor/call.rs @@ -95,14 +95,11 @@ impl Function { self.table.call(self.masm.pc(), index); self.masm._jump()?; - // Adjust the stack for 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()?; - if *results > 0 { - self.masm._push0()?; - self.masm._mload()?; - while self.masm.sp() > *results as u16 { - self.masm._drop()?; - } + while self.masm.sp() > *results as u16 { + self.masm._drop()?; } Ok(())