feat: make rust API more consistent with Cairo and accept compound di…

…ct keys

Cargo.toml

@@ -59,4 +59,4 @@ cairo-vm = { git = "https://github.com/lambdaclass/cairo-vm.git", tag = "v2.0.0-
 ] }
 
 [patch."https://github.com/lambdaclass/cairo-vm.git"]
-cairo-vm = { git = "https://github.com/ClementWalter/cairo-rs" }
+cairo-vm = { git = "https://github.com/kkrt-labs/cairo-vm", rev = "dd3d3f6e76248fc02395b31cbefe5fe8183222f1" }

cairo/tests/ethereum/cancun/test_fork_types.py

@@ -1,9 +1,5 @@
-import pytest
-
 from ethereum.cancun.fork_types import EMPTY_ACCOUNT
 
-pytestmark = pytest.mark.python_vm
-
 
 class TestForkTypes:
     def test_account_default(self, cairo_run):

cairo/tests/ethereum/cancun/utils/test_address.py

@@ -1,6 +1,5 @@
 from typing import Union
 
-import pytest
 from ethereum_types.bytes import Bytes32
 from ethereum_types.numeric import U256, Uint
 from hypothesis import given
@@ -12,8 +11,6 @@
     to_address,
 )
 
-pytestmark = pytest.mark.python_vm
-
 
 class TestAddress:
 

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

@@ -17,8 +17,6 @@
 from tests.utils.args_gen import Evm
 from tests.utils.strategies import evm_lite
 
-pytestmark = pytest.mark.python_vm
-
 
 class TestArithmetic:
     @given(evm=evm_lite)

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

@@ -15,8 +15,6 @@
 from tests.utils.args_gen import Evm
 from tests.utils.strategies import evm_lite
 
-pytestmark = pytest.mark.python_vm
-
 
 class TestBitwise:
     @given(evm=evm_lite)

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

@@ -14,8 +14,6 @@
 from tests.utils.args_gen import Evm
 from tests.utils.strategies import evm_lite
 
-pytestmark = pytest.mark.python_vm
-
 
 class TestBlock:
     @given(evm=evm_lite)

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

@@ -13,8 +13,6 @@
 from tests.utils.args_gen import Evm
 from tests.utils.strategies import evm_lite
 
-pytestmark = pytest.mark.python_vm
-
 
 class TestComparison:
     @given(evm=evm_lite)

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

@@ -15,8 +15,6 @@
 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)

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

@@ -8,8 +8,6 @@
 from tests.utils.args_gen import Evm
 from tests.utils.strategies import evm_lite, memory_access_size, memory_start_position
 
-pytestmark = pytest.mark.python_vm
-
 
 class TestKeccak:
     @given(

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

@@ -8,8 +8,6 @@
 from tests.utils.args_gen import Evm
 from tests.utils.strategies import evm_lite, memory_access_size, memory_start_position
 
-pytestmark = pytest.mark.python_vm
-
 
 class TestLog:
     @given(

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

@@ -6,8 +6,6 @@
 from tests.utils.args_gen import Evm
 from tests.utils.strategies import evm_lite
 
-pytestmark = pytest.mark.python_vm
-
 
 class TestMemory:
     @given(evm=evm_lite)

cairo/tests/ethereum/cancun/vm/test_memory.py

@@ -1,4 +1,3 @@
-import pytest
 from ethereum_types.bytes import Bytes
 from ethereum_types.numeric import U256
 from hypothesis import given
@@ -38,9 +37,6 @@ def memory_read_strategy(draw):
     return memory, start_position, size
 
 
-pytestmark = pytest.mark.python_vm
-
-
 class TestMemory:
     @given(memory_write_strategy())
     def test_memory_write(self, cairo_run, params):

cairo/tests/ethereum/cancun/vm/test_runtime.py

@@ -1,11 +1,8 @@
-import pytest
 from ethereum_types.bytes import Bytes
 from hypothesis import given
 
 from ethereum.cancun.vm.runtime import get_valid_jump_destinations
 
-pytestmark = pytest.mark.python_vm
-
 
 class TestRuntime:
     @given(code=...)

cairo/tests/ethereum/cancun/vm/test_stack.py

@@ -7,8 +7,6 @@
 from ethereum.cancun.vm.exceptions import StackOverflowError, StackUnderflowError
 from ethereum.cancun.vm.stack import pop, push
 
-pytestmark = pytest.mark.python_vm
-
 
 class TestStack:
     def test_pop_underflow(self, cairo_run):

cairo/tests/ethereum/crypto/test_hash.py

@@ -1,11 +1,8 @@
-import pytest
 from ethereum_types.bytes import Bytes
 from hypothesis import assume, given
 
 from ethereum.crypto.hash import keccak256
 
-pytestmark = pytest.mark.python_vm
-
 
 class TestHash:
 

cairo/tests/ethereum/test_rlp.py

@@ -23,8 +23,6 @@
 )
 from tests.utils.errors import cairo_error
 
-pytestmark = pytest.mark.python_vm
-
 
 class TestRlp:
     class TestEncode:

cairo/tests/src/test_stack.py

@@ -1,8 +1,3 @@
-import pytest
-
-pytestmark = pytest.mark.python_vm
-
-
 class TestStack:
     class TestPeek:
         def test_should_return_stack_at_given_index__when_value_is_0(self, cairo_run):

cairo/tests/test_serde.py

@@ -166,9 +166,6 @@ def single_evm_parent(b: Union[Message, Evm]) -> bool:
     return True
 
 
-pytestmark = pytest.mark.python_vm
-
-
 class TestSerde:
     @given(b=...)
     # 20 examples per type

cairo/tests/utils/args_gen.py

@@ -74,6 +74,7 @@
 )
 
 from cairo_addons.vm import DictTracker as RustDictTracker
+from cairo_addons.vm import MemorySegmentManager as RustMemorySegmentManager
 from cairo_addons.vm import Relocatable as RustRelocatable
 from ethereum_types.bytes import (
     Bytes,
@@ -100,6 +101,7 @@
 from starkware.cairo.lang.compiler.program import Program
 from starkware.cairo.lang.compiler.scoped_name import ScopedName
 from starkware.cairo.lang.vm.crypto import poseidon_hash_many
+from starkware.cairo.lang.vm.memory_segments import MemorySegmentManager
 from starkware.cairo.lang.vm.relocatable import RelocatableValue
 
 from ethereum.cancun.blocks import Header, Log, Receipt, Withdrawal
@@ -374,7 +376,7 @@ def gen_arg(dict_manager, segments):
 
 def _gen_arg(
     dict_manager,
-    segments,
+    segments: Union[MemorySegmentManager, RustMemorySegmentManager],
     arg_type: Type,
     arg: Any,
     annotations: Optional[Any] = None,
@@ -545,19 +547,15 @@ def _gen_arg(
                 data=data, current_ptr=current_ptr
             )
         else:
-            dict_manager.insert(
-                dict_ptr.segment_index,
-                RustDictTracker(
-                    keys=list(data.keys()),
-                    values=list(data.values()),
-                    current_ptr=current_ptr,
-                    default_value=(
-                        data.default_factory()
-                        if isinstance(data, defaultdict)
-                        else None
-                    ),
-                ),
+            default_value = (
+                data.default_factory() if isinstance(data, defaultdict) else None
+            )
+            dict_manager.trackers[dict_ptr.segment_index] = RustDictTracker(
+                data=data,
+                current_ptr=current_ptr,
+                default_value=default_value,
             )
+
         base = segments.add()
         segments.load_data(base, [dict_ptr, current_ptr])
         return base
@@ -589,11 +587,17 @@ def _gen_arg(
 
         if arg_type_origin is Trie:
             # In case of a Trie, we need the dict to be a defaultdict with the trie.default as the default value.
-            dict_ptr = segments.memory[data[2]]
-            dict_manager.trackers[dict_ptr.segment_index].data = defaultdict(
-                lambda: data[1], dict_manager.trackers[dict_ptr.segment_index].data
-            )
-
+            dict_ptr = segments.memory.get(data[2])
+            if isinstance(dict_manager, DictManager):
+                dict_manager.trackers[dict_ptr.segment_index].data = defaultdict(
+                    lambda: data[1], dict_manager.trackers[dict_ptr.segment_index].data
+                )
+            else:
+                dict_manager.trackers[dict_ptr.segment_index] = RustDictTracker(
+                    data=dict_manager.trackers[dict_ptr.segment_index].data,
+                    current_ptr=dict_ptr,
+                    default_value=data[1],
+                )
         return struct_ptr
 
     if arg_type in (U256, Hash32, Bytes32, Bytes256):

cairo/tests/utils/serde.py

@@ -105,9 +105,7 @@ def __init__(
         cairo_file=None,
     ):
         self.segments = segments
-        self.memory = (
-            segments.memory if isinstance(segments, MemorySegmentManager) else segments
-        )
+        self.memory = segments.memory
         self.program = program
         self.dict_manager = dict_manager
         self.cairo_file = cairo_file or Path()