feat(codegen): handle backward jumps in jump table (#286)

* feat(codegen): handle backward jumps in jump table

* feat(zink): introduce equal for u256

* feat(examples): introduce example br_balance

* chore(examples): mark the control flow problems

codegen/src/jump/target.rs

@@ -33,7 +33,14 @@ impl JumpTable {
         // First pass: calculate all target sizes and accumulate offsets
         for (original_pc, jump) in jumps.iter() {
             let pc = original_pc + total_offset;
-            let target = self.target(jump)? + total_offset;
+            let raw_target = self.target(jump)?;
+
+            // Calculate the absolute target including future offsets
+            let target = if raw_target > *original_pc {
+                raw_target + total_offset
+            } else {
+                raw_target
+            };
 
             // Calculate instruction size based on absolute target value
             let instr_size = if target > 0xff {

codegen/src/visitor/control.rs

@@ -113,10 +113,17 @@ impl Function {
     /// Conditional branch to a given label in an enclosing construct.
     pub fn _br_if(&mut self, depth: u32) -> Result<()> {
         let label = self.control.label_from_depth(depth)?;
+
+        // Register the jump target for breaking out
         self.table.label(self.masm.pc(), label);
-        self.masm.asm.increment_sp(1)?;
+        self.masm.increment_sp(1)?;
+
+        // JUMPI will check condition and jump if true
         self.masm._jumpi()?;
 
+        // If we don't jump, we continue in the current frame
+        // No need for explicit ISZERO since JUMPI handles the condition
+
         Ok(())
     }
 

codegen/src/wasm/host.rs

@@ -75,6 +75,7 @@ impl TryFrom<(&str, &str)> for HostFunc {
             ("zinkc", "u256_add") => Ok(Self::Evm(OpCode::ADD)),
             ("zinkc", "u256_sub") => Ok(Self::Evm(OpCode::SUB)),
             ("zinkc", "u256_lt") => Ok(Self::Evm(OpCode::LT)),
+            ("zinkc", "u256_eq") => Ok(Self::Evm(OpCode::EQ)),
             ("zinkc", "u256_max") => Ok(Self::U256MAX),
             ("zinkc", "u256_addmod") => Ok(Self::Evm(OpCode::ADDMOD)),
             ("zinkc", "u256_mulmod") => Ok(Self::Evm(OpCode::MULMOD)),

examples/br_balance.rs

@@ -0,0 +1,60 @@
+#![cfg_attr(target_arch = "wasm32", no_std)]
+#![cfg_attr(target_arch = "wasm32", no_main)]
+
+extern crate zink;
+use zink::{storage, Storage};
+
+#[storage(i32)]
+struct Balance;
+
+#[zink::external]
+fn check_and_update(value: i32) -> bool {
+    let current = Balance::get();
+
+    // This mimics the ERC20 balance check
+    if current < value {
+        zink::revert!("Not enough balance");
+        // TODO: #287
+        // return false;
+    }
+
+    Balance::set(current - value);
+    return true;
+}
+
+// TODO: identify if the problem is caused by control flow of incorrect opcode mapping. (issue #287)
+#[test]
+fn test_balance_check() -> anyhow::Result<()> {
+    use zint::{Bytes32, Contract, EVM};
+
+    let mut evm = EVM::default().commit(true);
+    let mut contract = Contract::search("br_balance")?.compile()?;
+
+    // Initialize with balance of 42
+    let info = evm.deploy(
+        &contract
+            .construct(
+                [(
+                    Balance::STORAGE_KEY.to_bytes32().into(),
+                    vec![42].try_into()?,
+                )]
+                .into_iter()
+                .collect(),
+            )?
+            .bytecode()?,
+    )?;
+
+    // Try to transfer 21 (should succeed)
+    let info = evm
+        .calldata(&contract.encode(&[
+            b"check_and_update(int32)".to_vec(),
+            21i32.to_bytes32().to_vec(),
+        ])?)
+        .call(info.address)?;
+    assert_eq!(info.ret, true.to_bytes32(), "{info:?}");
+
+    Ok(())
+}
+
+#[cfg(not(target_arch = "wasm32"))]
+fn main() {}

examples/erc20.rs

@@ -78,6 +78,7 @@ fn _update(from: Address, to: Address, value: U256) {
         TotalSupply::set(TotalSupply::get().add(value));
     } else {
         let from_balance = Balances::get(from);
+
         if from_balance.lt(value) {
             zink::revert!("Insufficient balance");
         }
@@ -143,13 +144,16 @@ fn deploy() -> anyhow::Result<()> {
     use zint::{Bytes32, Contract, EVM};
 
     let caller_bytes = hex::decode("be862ad9abfe6f22bcb087716c7d89a26051f74c")?;
+    let spender = [42; 20];
     let mut caller = [0; 20];
     caller.copy_from_slice(&caller_bytes);
 
     let mut evm = EVM::default().commit(true).caller(caller);
     let mut contract = Contract::search("erc20")?.compile()?;
     let name = "The Zink Language";
     let symbol = "zink";
+    let value = 42;
+    //let half_value = 21;
 
     // 1. deploy
     let info = evm.deploy(
@@ -165,6 +169,10 @@ fn deploy() -> anyhow::Result<()> {
                         TotalSupply::STORAGE_KEY.to_bytes32().into(),
                         vec![42].try_into()?,
                     ),
+                    (
+                        Balances::storage_key(Address(caller)).into(),
+                        vec![42].try_into()?,
+                    ),
                 ]
                 .into_iter()
                 .collect(),
@@ -173,58 +181,90 @@ fn deploy() -> anyhow::Result<()> {
     )?;
     let address = info.address;
 
-    // 2. get name
-    let info = evm
-        .calldata(&contract.encode(&[b"name()".to_vec()])?)
-        .call(address)?;
-    assert_eq!(info.ret, name.to_bytes32(), "{info:?}");
-
-    // 3. get symbol
-    let info = evm
-        .calldata(&contract.encode(&[b"symbol()".to_vec()])?)
-        .call(address)?;
-    assert_eq!(info.ret, symbol.to_bytes32(), "{info:?}");
-
-    // 4. get total supply
-    let info = evm
-        .calldata(&contract.encode(&[b"total_supply()".to_vec()])?)
-        .call(address)?;
-    assert_eq!(info.ret, 42u64.to_bytes32(), "{info:?}");
-
-    // 5. check decimals
-    let info = evm
-        .calldata(&contract.encode(&[b"decimals()".to_vec()])?)
-        .call(address)?;
-    assert_eq!(info.ret, 8u64.to_bytes32(), "{info:?}");
-
-    // TODO: refactor offset handling (#280)
-    // 6. check approval
-    /*   let value = 42;
-    let spender = [42; 20];
-    let info = evm
-        .calldata(&contract.encode(&[
-            b"approve(address,uint256)".to_vec(),
-            spender.to_bytes32().to_vec(),
-            value.to_bytes32().to_vec(),
-        ])?)
-        .call(address)?;
-    assert_eq!(info.ret, true.to_bytes32(), "{info:?}"); */
-
-    // let allowance = evm.storage(
-    //     address,
-    //     Allowance::storage_key(Address(evm.caller), Address(spender)),
-    // )?;
-    // assert_eq!(value.to_bytes32(), allowance);
-    //
-    // // 7. check approval results
-    // let info = evm
-    //     .calldata(&contract.encode(&[
-    //         b"allowance(address,address)".to_vec(),
-    //         evm.caller.to_bytes32().to_vec(),
-    //         spender.to_bytes32().to_vec(),
-    //     ])?)
-    //     .call(address)?;
-    // assert_eq!(info.ret, allowance);
+    // 2. get static data
+    {
+        // 2.1. get name
+        let info = evm
+            .calldata(&contract.encode(&[b"name()".to_vec()])?)
+            .call(address)?;
+        assert_eq!(info.ret, name.to_bytes32(), "{info:?}");
+
+        // 2.2. get symbol
+        let info = evm
+            .calldata(&contract.encode(&[b"symbol()".to_vec()])?)
+            .call(address)?;
+        assert_eq!(info.ret, symbol.to_bytes32(), "{info:?}");
+
+        // 2.3. get total supply
+        let info = evm
+            .calldata(&contract.encode(&[b"total_supply()".to_vec()])?)
+            .call(address)?;
+        assert_eq!(info.ret, 42u64.to_bytes32(), "{info:?}");
+
+        // 2.4. check decimals
+        let info = evm
+            .calldata(&contract.encode(&[b"decimals()".to_vec()])?)
+            .call(address)?;
+        assert_eq!(info.ret, 8u64.to_bytes32(), "{info:?}");
+
+        // 2.5. check balance of the caller
+        let balance = evm.storage(address, Balances::storage_key(Address(caller)))?;
+        assert_eq!(value.to_bytes32(), balance);
+    }
+
+    // 3. check approval
+    {
+        // 3.1. approve
+        let info = evm
+            .calldata(&contract.encode(&[
+                b"approve(address,uint256)".to_vec(),
+                spender.to_bytes32().to_vec(),
+                value.to_bytes32().to_vec(),
+            ])?)
+            .call(address)?;
+        assert_eq!(info.ret, true.to_bytes32(), "{info:?}");
+
+        let allowance = evm.storage(
+            address,
+            Allowance::storage_key(Address(evm.caller), Address(spender)),
+        )?;
+        assert_eq!(value.to_bytes32(), allowance);
+
+        // 3.2. check approval results
+        let info = evm
+            .calldata(&contract.encode(&[
+                b"allowance(address,address)".to_vec(),
+                evm.caller.to_bytes32().to_vec(),
+                spender.to_bytes32().to_vec(),
+            ])?)
+            .call(address)?;
+        assert_eq!(info.ret, allowance);
+    }
+
+    // 4. check transfer
+    {
+        // 4.1. verify balance of the caller
+        let info = evm
+            .calldata(&contract.encode(&[
+                b"balances(address)".to_vec(),
+                evm.caller.to_bytes32().to_vec(),
+            ])?)
+            .call(address)?;
+        assert_eq!(info.ret, value.to_bytes32(), "{info:?}");
+
+        //  TODO: see br_balance.rs (#287)
+        // 4.2. check transfer
+        /* evm = evm.commit(false);
+        let info = evm
+            .calldata(&contract.encode(&[
+                b"transfer(address,uint256)".to_vec(),
+                spender.to_bytes32().to_vec(),
+                half_value.to_bytes32().to_vec(),
+            ])?)
+            .call(address)?;
+        println!("{info:?}");
+        assert_eq!(info.ret, true.to_bytes32(), "{info:?}"); */
+    }
 
     Ok(())
 }

zink/src/ffi/mod.rs

@@ -23,6 +23,12 @@ extern "C" {
     /// Less than operation for addresses
     pub fn u256_lt(this: U256, other: U256) -> bool;
 
+    /// Equal operation for addresses
+    pub fn u256_eq(this: U256, other: U256) -> bool;
+
+    /// Returns zero value
+    pub fn u256_zero() -> U256;
+
     /// Equal operation for addresses
     pub fn u256_max() -> U256;
 

zink/src/primitives/u256.rs

@@ -16,6 +16,18 @@ impl U256 {
         U256([0; 32])
     }
 
+    /// Returns zero value
+    #[cfg(not(target_family = "wasm"))]
+    pub const fn zero() -> Self {
+        Self([0; 32])
+    }
+
+    /// Returns zero value
+    #[cfg(target_family = "wasm")]
+    pub const fn zero() -> Self {
+        Self(0)
+    }
+
     /// u256 add
     #[inline(always)]
     pub fn add(self, other: Self) -> Self {
@@ -28,6 +40,12 @@ impl U256 {
         unsafe { ffi::u256_lt(other, self) }
     }
 
+    /// u256 eq
+    #[inline(always)]
+    pub fn eq(self, other: Self) -> bool {
+        unsafe { ffi::u256_eq(self, other) }
+    }
+
     /// u256 sub
     #[inline(always)]
     pub fn sub(self, other: Self) -> Self {