feat: control flow opcodes (#383)

Closes #348

---------

Co-authored-by: Mathieu <60658558+enitrat@users.noreply.github.com>

.env.example

@@ -1,6 +1,7 @@
 HYPOTHESIS_PROFILE=dev
 HYPOTHESIS_MAX_ADDRESS_SET_SIZE=50
 HYPOTHESIS_MAX_STORAGE_KEY_SET_SIZE=50
+HYPOTHESIS_MAX_JUMP_DESTINATIONS_SET_SIZE=50
 HYPOTHESIS_MAX_RECURSION_DEPTH=50
 
 ATLANTIC_API_KEY=

cairo/ethereum/cancun/vm/instructions/control_flow.cairo

@@ -0,0 +1,197 @@
+from starkware.cairo.common.bool import TRUE, FALSE
+from starkware.cairo.common.cairo_builtins import BitwiseBuiltin, PoseidonBuiltin
+from starkware.cairo.common.dict import dict_read, DictAccess
+
+from ethereum_types.numeric import (
+    U256,
+    U256Struct,
+    Uint,
+    bool,
+    SetUint,
+    SetUintStruct,
+    SetUintDictAccess,
+)
+
+from ethereum.cancun.vm import Evm, EvmImpl
+from ethereum.cancun.vm.exceptions import ExceptionalHalt, InvalidJumpDestError
+from ethereum.cancun.vm.gas import charge_gas, GasConstants
+from ethereum.cancun.vm.stack import Stack, pop, push
+
+from src.utils.dict import hashdict_read
+
+// @notice Stop further execution of EVM code
+func stop{evm: Evm}() {
+    // STACK
+
+    // GAS
+    // No gas charge for STOP
+
+    // OPERATION
+    EvmImpl.set_running(bool(FALSE));
+
+    // PROGRAM COUNTER
+    EvmImpl.set_pc(Uint(evm.value.pc.value + 1));
+    return ();
+}
+
+// @notice Alter the program counter to the location specified by the top of the stack
+func jump{range_check_ptr, poseidon_ptr: PoseidonBuiltin*, evm: Evm}() -> ExceptionalHalt* {
+    alloc_locals;
+    // STACK
+    let stack = evm.value.stack;
+    with stack {
+        let (jump_dest, err1) = pop();
+        if (cast(err1, felt) != 0) {
+            return err1;
+        }
+    }
+
+    // GAS
+    let err2 = charge_gas(Uint(GasConstants.GAS_MID));
+    if (cast(err2, felt) != 0) {
+        return err2;
+    }
+
+    // OPERATION
+    // Check if jump destination is valid by looking it up in valid_jump_destinations
+    let valid_jump_destinations_ptr = evm.value.valid_jump_destinations.value.dict_ptr;
+    let dict_ptr = cast(valid_jump_destinations_ptr, DictAccess*);
+    let (is_valid_dest) = hashdict_read{dict_ptr=dict_ptr}(1, &jump_dest.value.low);
+    if (is_valid_dest == FALSE) {
+        tempvar err = new ExceptionalHalt(InvalidJumpDestError);
+        return err;
+    }
+
+    let set_dict_ptr = cast(dict_ptr, SetUintDictAccess*);
+    tempvar valid_jumpdests_set = SetUint(
+        new SetUintStruct(evm.value.valid_jump_destinations.value.dict_ptr_start, set_dict_ptr)
+    );
+    EvmImpl.set_valid_jump_destinations(valid_jumpdests_set);
+
+    // PROGRAM COUNTER
+    EvmImpl.set_pc_stack(Uint(jump_dest.value.low), stack);
+    let ok = cast(0, ExceptionalHalt*);
+    return ok;
+}
+
+// @notice Alter the program counter to the specified location if and only if a condition is true
+func jumpi{range_check_ptr, poseidon_ptr: PoseidonBuiltin*, evm: Evm}() -> ExceptionalHalt* {
+    alloc_locals;
+    // STACK
+    let stack = evm.value.stack;
+    with stack {
+        let (jump_dest, err1) = pop();
+        if (cast(err1, felt) != 0) {
+            return err1;
+        }
+        let (condition, err2) = pop();
+        if (cast(err2, felt) != 0) {
+            return err2;
+        }
+    }
+
+    // GAS
+    let err3 = charge_gas(Uint(GasConstants.GAS_HIGH));
+    if (cast(err3, felt) != 0) {
+        return err3;
+    }
+
+    // OPERATION
+    if (condition.value.low == 0 and condition.value.high == 0) {
+        // If condition is false, just increment PC
+        EvmImpl.set_pc_stack(Uint(evm.value.pc.value + 1), stack);
+        let ok = cast(0, ExceptionalHalt*);
+        return ok;
+    }
+
+    let valid_jump_destinations_ptr = evm.value.valid_jump_destinations.value.dict_ptr;
+    let dict_ptr = cast(valid_jump_destinations_ptr, DictAccess*);
+    let (is_valid_dest) = hashdict_read{dict_ptr=dict_ptr}(1, &jump_dest.value.low);
+
+    if (is_valid_dest == FALSE) {
+        tempvar err = new ExceptionalHalt(InvalidJumpDestError);
+        return err;
+    }
+
+    let set_dict_ptr = cast(dict_ptr, SetUintDictAccess*);
+    tempvar valid_jumpdests_set = SetUint(
+        new SetUintStruct(evm.value.valid_jump_destinations.value.dict_ptr_start, set_dict_ptr)
+    );
+    EvmImpl.set_valid_jump_destinations(valid_jumpdests_set);
+
+    // PROGRAM COUNTER
+    EvmImpl.set_pc_stack(Uint(jump_dest.value.low), stack);
+    let ok = cast(0, ExceptionalHalt*);
+    return ok;
+}
+
+// @notice Push the value of the program counter before the increment onto the stack
+func pc{range_check_ptr, evm: Evm}() -> ExceptionalHalt* {
+    alloc_locals;
+    // STACK
+    let stack = evm.value.stack;
+
+    // GAS
+    let err1 = charge_gas(Uint(GasConstants.GAS_BASE));
+    if (cast(err1, felt) != 0) {
+        return err1;
+    }
+
+    // OPERATION
+    with stack {
+        let err2 = push(U256(new U256Struct(evm.value.pc.value, 0)));
+        if (cast(err2, felt) != 0) {
+            return err2;
+        }
+    }
+
+    // PROGRAM COUNTER
+    EvmImpl.set_pc_stack(Uint(evm.value.pc.value + 1), stack);
+    let ok = cast(0, ExceptionalHalt*);
+    return ok;
+}
+
+// @notice Push the amount of available gas onto the stack
+func gas_left{range_check_ptr, evm: Evm}() -> ExceptionalHalt* {
+    alloc_locals;
+    // STACK
+    let stack = evm.value.stack;
+
+    // GAS
+    let err1 = charge_gas(Uint(GasConstants.GAS_BASE));
+    if (cast(err1, felt) != 0) {
+        return err1;
+    }
+
+    // OPERATION
+    with stack {
+        let err2 = push(U256(new U256Struct(evm.value.gas_left.value, 0)));
+        if (cast(err2, felt) != 0) {
+            return err2;
+        }
+    }
+
+    // PROGRAM COUNTER
+    EvmImpl.set_pc_stack(Uint(evm.value.pc.value + 1), stack);
+    let ok = cast(0, ExceptionalHalt*);
+    return ok;
+}
+
+// @notice Mark a valid destination for jumps
+func jumpdest{range_check_ptr, evm: Evm}() -> ExceptionalHalt* {
+    alloc_locals;
+
+    // GAS
+    let err = charge_gas(Uint(GasConstants.GAS_JUMPDEST));
+    if (cast(err, felt) != 0) {
+        return err;
+    }
+
+    // OPERATION
+    // No operation needed, just a marker
+
+    // PROGRAM COUNTER
+    EvmImpl.set_pc(Uint(evm.value.pc.value + 1));
+    let ok = cast(0, ExceptionalHalt*);
+    return ok;
+}

cairo/tests/ethereum/cancun/vm/instructions/test_control_flow.py

@@ -0,0 +1,115 @@
+import pytest
+from ethereum_types.numeric import U256
+from hypothesis import given
+
+from ethereum.cancun.vm.exceptions import ExceptionalHalt
+from ethereum.cancun.vm.instructions.control_flow import (
+    gas_left,
+    jump,
+    jumpdest,
+    jumpi,
+    pc,
+    stop,
+)
+from ethereum.cancun.vm.stack import push
+from tests.utils.args_gen import Evm
+from tests.utils.strategies import evm_lite
+
+pytestmark = pytest.mark.python_vm
+
+
+class TestControlFlow:
+    @given(evm=evm_lite)
+    def test_stop(self, cairo_run, evm: Evm):
+        try:
+            cairo_result = cairo_run("stop", evm)
+        except ExceptionalHalt as cairo_error:
+            with pytest.raises(type(cairo_error)):
+                stop(evm)
+            return
+
+        stop(evm)
+        assert evm == cairo_result
+
+    @given(evm=evm_lite, push_valid_jump_destination=...)
+    def test_jump(self, cairo_run, evm: Evm, push_valid_jump_destination: bool):
+        # Modify the stack to match the valid_jump_destinations generated by hypothesis
+        jump_dest = (
+            next(iter(evm.valid_jump_destinations))
+            if push_valid_jump_destination and evm.valid_jump_destinations
+            else 0
+        )
+        push(evm.stack, U256(jump_dest))
+
+        try:
+            cairo_result = cairo_run("jump", evm)
+        except ExceptionalHalt as cairo_error:
+            with pytest.raises(type(cairo_error)):
+                jump(evm)
+            return
+
+        jump(evm)
+        assert evm == cairo_result
+
+    @given(evm=evm_lite, push_valid_jump_destination=..., jumpi_condition=...)
+    def test_jumpi(
+        self,
+        cairo_run,
+        evm: Evm,
+        push_valid_jump_destination: bool,
+        jumpi_condition: bool,
+    ):
+        # Modify the stack to match the valid_jump_destinations generated by hypothesis
+        push(evm.stack, U256(jumpi_condition))
+        jump_dest = (
+            next(iter(evm.valid_jump_destinations))
+            if push_valid_jump_destination and evm.valid_jump_destinations
+            else 0
+        )
+        push(evm.stack, U256(jump_dest))
+
+        try:
+            cairo_result = cairo_run("jumpi", evm)
+        except ExceptionalHalt as cairo_error:
+            with pytest.raises(type(cairo_error)):
+                jumpi(evm)
+            return
+
+        jumpi(evm)
+        assert evm == cairo_result
+
+    @given(evm=evm_lite)
+    def test_pc(self, cairo_run, evm: Evm):
+        try:
+            cairo_result = cairo_run("pc", evm)
+        except ExceptionalHalt as cairo_error:
+            with pytest.raises(type(cairo_error)):
+                pc(evm)
+            return
+
+        pc(evm)
+        assert evm == cairo_result
+
+    @given(evm=evm_lite)
+    def test_gas_left(self, cairo_run, evm: Evm):
+        try:
+            cairo_result = cairo_run("gas_left", evm)
+        except ExceptionalHalt as cairo_error:
+            with pytest.raises(type(cairo_error)):
+                gas_left(evm)
+            return
+
+        gas_left(evm)
+        assert evm == cairo_result
+
+    @given(evm=evm_lite)
+    def test_jumpdest(self, cairo_run, evm: Evm):
+        try:
+            cairo_result = cairo_run("jumpdest", evm)
+        except ExceptionalHalt as cairo_error:
+            with pytest.raises(type(cairo_error)):
+                jumpdest(evm)
+            return
+
+        jumpdest(evm)
+        assert evm == cairo_result

cairo/tests/utils/args_gen.py

@@ -503,7 +503,7 @@ def _gen_arg(
         dict_ptr = segments.add()
 
         if arg_type_origin is set:
-            arg = {k: True for k in arg}
+            arg = defaultdict(lambda: False, {k: True for k in arg})
             arg_type = Mapping[get_args(arg_type)[0], bool]
 
         data = {

cairo/tests/utils/strategies.py

@@ -77,7 +77,9 @@
 # Maximum size for sets of addresses and tuples of address and bytes32 to avoid heavy memory usage and health check errors
 MAX_ADDRESS_SET_SIZE = int(os.getenv("HYPOTHESIS_MAX_ADDRESS_SET_SIZE", 10))
 MAX_STORAGE_KEY_SET_SIZE = int(os.getenv("HYPOTHESIS_MAX_STORAGE_KEY_SET_SIZE", 10))
-
+MAX_JUMP_DESTINATIONS_SET_SIZE = int(
+    os.getenv("HYPOTHESIS_MAX_JUMP_DESTINATIONS_SET_SIZE", 10)
+)
 # See ethereum.rlp.Simple and ethereum.rlp.Extended for the definition of Simple and Extended
 simple = st.recursive(
     st.one_of(st.binary()),
@@ -205,7 +207,7 @@ def tuple_strategy(thing):
     code=st.just(b""),
     gas_left=st.integers(min_value=0, max_value=BLOCK_GAS_LIMIT).map(Uint),
     env=environment_lite,
-    valid_jump_destinations=st.just(set()),
+    valid_jump_destinations=st.sets(uint, max_size=MAX_ADDRESS_SET_SIZE),
     logs=st.just(()),
     refund_counter=st.just(0),
     running=st.booleans(),