Add State deserialization

cairo/programs/os.cairo

@@ -8,8 +8,9 @@ from src.model import model
 func main{pedersen_ptr: HashBuiltin*, output_ptr: felt*}() {
     %{ dict_manager %}
     tempvar block: model.Block*;
-    tempvar initial_state: model.State*;
+    tempvar state: model.State*;
     %{ block %}
+    %{ state %}
     // TODO: Compute initial state root hash and compare with block.parent_hash
     // TODO: Loop through transactions and apply them to the initial state
     // TODO: Compute the state root hash after applying all transactions

cairo/programs/os.json

@@ -129,9 +129,25 @@
           },
           "reference_ids": {
             "__main__.main.block": 10,
-            "__main__.main.initial_state": 11,
             "__main__.main.output_ptr": 9,
-            "__main__.main.pedersen_ptr": 8
+            "__main__.main.pedersen_ptr": 8,
+            "__main__.main.state": 11
+          }
+        }
+      },
+      {
+        "accessible_scopes": ["__main__", "__main__.main"],
+        "code": "state",
+        "flow_tracking_data": {
+          "ap_tracking": {
+            "group": 3,
+            "offset": 2
+          },
+          "reference_ids": {
+            "__main__.main.block": 10,
+            "__main__.main.output_ptr": 9,
+            "__main__.main.pedersen_ptr": 8,
+            "__main__.main.state": 11
           }
         }
       }
@@ -199,21 +215,6 @@
       ],
       "type": "reference"
     },
-    "__main__.main.initial_state": {
-      "cairo_type": "src.model.model.State*",
-      "full_name": "__main__.main.initial_state",
-      "references": [
-        {
-          "ap_tracking_data": {
-            "group": 3,
-            "offset": 2
-          },
-          "pc": 25,
-          "value": "[cast(ap + (-1), src.model.model.State**)]"
-        }
-      ],
-      "type": "reference"
-    },
     "__main__.main.output_ptr": {
       "cairo_type": "felt*",
       "full_name": "__main__.main.output_ptr",
@@ -244,6 +245,21 @@
       ],
       "type": "reference"
     },
+    "__main__.main.state": {
+      "cairo_type": "src.model.model.State*",
+      "full_name": "__main__.main.state",
+      "references": [
+        {
+          "ap_tracking_data": {
+            "group": 3,
+            "offset": 2
+          },
+          "pc": 25,
+          "value": "[cast(ap + (-1), src.model.model.State**)]"
+        }
+      ],
+      "type": "reference"
+    },
     "__main__.memcpy": {
       "destination": "starkware.cairo.common.memcpy.memcpy",
       "type": "alias"

cairo/src/account.cairo

@@ -91,9 +91,9 @@ namespace Account {
     // @param key The pointer to the storage key
     // @return The updated Account
     // @return The read value
-    func read_storage{pedersen_ptr: HashBuiltin*, range_check_ptr}(
-        self: model.Account*, key: Uint256*
-    ) -> (model.Account*, Uint256*) {
+    func read_storage{pedersen_ptr: HashBuiltin*}(self: model.Account*, key: Uint256*) -> (
+        model.Account*, Uint256*
+    ) {
         alloc_locals;
         let storage = self.storage;
         let (local storage_addr) = Internals._storage_addr(key);
@@ -126,7 +126,7 @@ namespace Account {
     // @notive Read the first value of a given storage slot
     // @dev The storage needs to exists already, to this should be used only
     //      after a storage_read or storage_write has been done
-    func fetch_original_storage{pedersen_ptr: HashBuiltin*, range_check_ptr}(
+    func fetch_original_storage{pedersen_ptr: HashBuiltin*}(
         self: model.Account*, key: Uint256*
     ) -> Uint256 {
         alloc_locals;

cairo/src/state.cairo

@@ -85,9 +85,7 @@ namespace State {
     // @param evm_address The evm address of the Account
     // @return The updated state
     // @return The account
-    func get_account{pedersen_ptr: HashBuiltin*, range_check_ptr, state: model.State*}(
-        evm_address: felt
-    ) -> model.Account* {
+    func get_account{state: model.State*}(evm_address: felt) -> model.Account* {
         alloc_locals;
         local accounts: DictAccess* = state.accounts;
         let (pointer) = dict_read{dict_ptr=accounts}(key=evm_address);

cairo/tests/programs/test_os.cairo

@@ -4,7 +4,8 @@ from starkware.cairo.common.cairo_builtins import HashBuiltin
 
 from programs.os import main
 from src.model import model
-from src.account import Internals
+from src.account import Internals, Account
+from src.state import State
 
 func test_os{pedersen_ptr: HashBuiltin*, output_ptr: felt*}() {
     main();
@@ -36,3 +37,31 @@ func test_account{pedersen_ptr: HashBuiltin*}() -> model.Account* {
 
     return account;
 }
+
+func test_state{pedersen_ptr: HashBuiltin*}() -> model.State* {
+    alloc_locals;
+    tempvar state: model.State*;
+    %{ state %}
+
+    with state {
+        let account = State.get_account(0xc0de);
+        tempvar key = new Uint256(0, 0);
+        let (account, value) = Account.read_storage(account, key);
+        tempvar key = new Uint256(1, 0);
+        let (account, value) = Account.read_storage(account, key);
+        tempvar key = new Uint256(2, 0);
+        let (account, value) = Account.read_storage(account, key);
+        State.update_account(0xc0de, account);
+
+        let account = State.get_account(0x000f3df6d732807ef1319fb7b8bb8522d0beac02);
+        tempvar key = new Uint256(0xf2, 0);
+        let (account, value) = Account.read_storage(account, key);
+        State.update_account(0x000f3df6d732807ef1319fb7b8bb8522d0beac02, account);
+
+        let account = State.get_account(0xa94f5374fce5edbc8e2a8697c15331677e6ebf0b);
+        tempvar key = new Uint256(0, 0);
+        let (account, value) = Account.read_storage(account, key);
+        State.update_account(0xa94f5374fce5edbc8e2a8697c15331677e6ebf0b, account);
+    }
+    return state;
+}

cairo/tests/programs/test_os.py

@@ -1,7 +1,7 @@
 import pytest
 
 from src.utils.uint256 import int_to_uint256
-from tests.utils.models import Account, Block, to_int
+from tests.utils.models import Account, Block, State, to_int
 
 
 @pytest.fixture
@@ -76,10 +76,40 @@ def account():
     )
 
 
+@pytest.fixture
+def state():
+    return State.model_validate(
+        {
+            "0x000000000000000000000000000000000000c0de": {
+                "balance": "0x00",
+                "code": "0x7fa0a1a2a3a4a5a6a7a8a9aaabacadaeafb0b1b2b3b4b5b6b7b8b9babbbcbdbebf5f527fc0c1c2c3c4c5c6c7c8c9cacbcccdcecfd0d1d2d3d4d5d6d7d8d9dadbdcdddedf6020527fe0e1e2e3e4e5e6e7e8e9eaebecedeeeff0f1f2f3f4f5f6f7f8f9fafbfcfdfeff60405260786040356020355f35608a565b5f515f55602051600155604051600255005b5e56",
+                "nonce": "0x01",
+                "storage": {
+                    "0x00": "0xa0a1a2a3a4a5a6a7a8a9aaabacadaeafb0b1b2b3b4b5b6b7b8b9babbbcbdbebf",
+                    "0x01": "0xc0c1c2c3c4c5c6c7c8c9cacbcccdcecfd0d1d2d3d4d5d6d7d8d9dadbdcdddedf",
+                    "0x02": "0xe0e1e2e3e4e5e6e7e8e9eaebecedeeeff0f1f2f3f4f5f6f7f8f9fafbfcfdfeff",
+                },
+            },
+            "0x000f3df6d732807ef1319fb7b8bb8522d0beac02": {
+                "balance": "0x00",
+                "code": "0x3373fffffffffffffffffffffffffffffffffffffffe14604d57602036146024575f5ffd5b5f35801560495762001fff810690815414603c575f5ffd5b62001fff01545f5260205ff35b5f5ffd5b62001fff42064281555f359062001fff015500",
+                "nonce": "0x01",
+                "storage": {"0xf2": "0x64903c57"},
+            },
+            "0xa94f5374fce5edbc8e2a8697c15331677e6ebf0b": {
+                "balance": "0x3b8d6450",
+                "code": "0x",
+                "nonce": "0x01",
+                "storage": {},
+            },
+        }
+    )
+
+
 class TestOs:
 
-    def test_os(self, cairo_run, block):
-        cairo_run("test_os", block=block)
+    def test_os(self, cairo_run, block, state):
+        cairo_run("test_os", block=block, state=state)
 
     def test_block(self, cairo_run, block):
         result = cairo_run("test_block", block=block)
@@ -95,3 +125,6 @@ def test_account(self, cairo_run, account):
         account.storage = {}
 
         assert Account.model_validate(result) == account
+
+    def test_state(self, cairo_run, state):
+        cairo_run("test_state", state=state)

cairo/tests/utils/hints.py

@@ -69,10 +69,18 @@ def gen_arg(
 ids.account = gen_arg(__dict_manager, segments, program_input["account"])
 """
 
+state = f"""
+{dict_manager}
+from tests.utils.hints import gen_arg
+
+ids.state = gen_arg(__dict_manager, segments, program_input["state"])
+"""
+
 hints = {
     "dict_manager": dict_manager,
     "block": block,
     "account": account,
+    "state": state,
 }
 
 

cairo/tests/utils/models.py

@@ -1,5 +1,6 @@
+from collections import defaultdict
 from itertools import chain
-from typing import Union
+from typing import Annotated, DefaultDict, Tuple, Union
 
 from eth_utils import keccak
 from ethereum.cancun.vm.runtime import get_valid_jump_destinations
@@ -8,6 +9,7 @@
     AliasGenerator,
     BaseModel,
     ConfigDict,
+    Field,
     field_validator,
     model_validator,
 )
@@ -195,9 +197,15 @@ class Account(BaseModelIterValuesOnly):
     code_len: int
     code: bytes
     code_hash: tuple[int, int]
-    storage: dict[int, tuple[int, int]]
-    transient_storage: dict[int, tuple[int, int]] = {}
-    valid_jumpdests: dict[int, bool]
+    storage: DefaultDict[
+        int, Annotated[Union[int, Tuple[int, int]], Field(default_factory=int)]
+    ] = defaultdict(int)
+    transient_storage: DefaultDict[
+        int, Annotated[Union[int, Tuple[int, int]], Field(default_factory=int)]
+    ] = defaultdict(int)
+    valid_jumpdests: DefaultDict[int, Annotated[bool, Field(default_factory=bool)]] = (
+        defaultdict(bool)
+    )
     nonce: int
     balance: tuple[int, int]
     selfdestruct: int = 0
@@ -222,15 +230,34 @@ def split_uint256_and_hash_storage(cls, values):
         values["code_hash"] = int.from_bytes(
             keccak(to_bytes(values["code"])), byteorder="big"
         )
-        values["storage"] = {
-            pedersen_hash(*int_to_uint256(to_int(k))): int_to_uint256(to_int(v))
-            for k, v in values["storage"].items()
-        }
-        values["valid_jumpdests"] = {
-            key: True for key in get_valid_jump_destinations(to_bytes(values["code"]))
-        }
+        values["storage"] = defaultdict(
+            int,
+            {
+                pedersen_hash(*int_to_uint256(to_int(k))): int_to_uint256(to_int(v))
+                for k, v in values["storage"].items()
+            },
+        )
+        values["valid_jumpdests"] = defaultdict(
+            int,
+            {
+                key: True
+                for key in get_valid_jump_destinations(to_bytes(values["code"]))
+            },
+        )
         return values
 
 
 class State(BaseModelIterValuesOnly):
-    accounts: dict[int, Account]
+    accounts: DefaultDict[
+        int, Annotated[Union[int, Account], Field(default_factory=int)]
+    ] = defaultdict(int)
+    events_len: int = 0
+    events: list = []
+    transfers_len: int = 0
+    transfers: list = []
+
+    @model_validator(mode="before")
+    def parse_addresses(cls, values):
+        values = values.copy()
+        values["accounts"] = defaultdict(int, {to_int(k): v for k, v in values.items()})
+        return values