diff --git a/src/ethereum/cancun/state.py b/src/ethereum/cancun/state.py index b067ed2286..3e8a3092af 100644 --- a/src/ethereum/cancun/state.py +++ b/src/ethereum/cancun/state.py @@ -32,6 +32,8 @@ class State: """ Contains all information that is preserved between transactions. + + Now includes optional state tracking """ _main_trie: Trie[Address, Optional[Account]] = field( @@ -47,6 +49,8 @@ class State: ] ] = field(default_factory=list) created_accounts: Set[Address] = field(default_factory=set) + + _state_tracker = field(default=None) @dataclass @@ -71,6 +75,8 @@ def close_state(state: State) -> None: del state._storage_tries del state._snapshots del state.created_accounts + if state._state_tracker is not None: + del state._state_tracker def begin_transaction( @@ -186,6 +192,11 @@ def get_account_optional(state: State, address: Address) -> Optional[Account]: Account at address. """ account = trie_get(state._main_trie, address) + + if state._state_tracker is not None and state._state_tracker.track_reads: + from .state_tracking import log_state_access, ACCOUNT_READ + log_state_access(state, ACCOUNT_READ, address, value_before=account) + return account @@ -205,7 +216,23 @@ def set_account( account : `Account` Account to set at address. """ + # Get old account for tracking + old_account = None + if state._state_tracker is not None and state._state_tracker.track_writes: + old_account = trie_get(state._main_trie, address) + trie_set(state._main_trie, address, account) + + # Log write access if tracking enabled + if state._state_tracker is not None and state._state_tracker.track_writes: + from .state_tracking import log_state_access, ACCOUNT_WRITE + log_state_access( + state, + ACCOUNT_WRITE, + address, + value_before=old_account, + value_after=account + ) def destroy_account(state: State, address: Address) -> None: @@ -284,12 +311,24 @@ def get_storage(state: State, address: Address, key: Bytes32) -> U256: """ trie = state._storage_tries.get(address) if trie is None: - return U256(0) - - value = trie_get(trie, key) - - assert isinstance(value, U256) - return value + result = U256(0) + else: + value = trie_get(trie, key) + assert isinstance(value, U256) + result = value + + # Log read access if tracking enabled + if state._state_tracker is not None and state._state_tracker.track_reads: + from .state_tracking import log_state_access, STORAGE_READ + log_state_access( + state, + STORAGE_READ, + address, + key=key, + value_before=result + ) + + return result def set_storage( @@ -312,6 +351,15 @@ def set_storage( """ assert trie_get(state._main_trie, address) is not None + # Get old value for tracking + old_value = None + if state._state_tracker is not None and state._state_tracker.track_writes: + trie = state._storage_tries.get(address) + if trie is None: + old_value = U256(0) + else: + old_value = trie_get(trie, key) + trie = state._storage_tries.get(address) if trie is None: trie = Trie(secured=True, default=U256(0)) @@ -319,6 +367,18 @@ def set_storage( trie_set(trie, key, value) if trie._data == {}: del state._storage_tries[address] + + # Log write access if tracking enabled + if state._state_tracker is not None and state._state_tracker.track_writes: + from .state_tracking import log_state_access, STORAGE_WRITE + log_state_access( + state, + STORAGE_WRITE, + address, + key=key, + value_before=old_value, + value_after=value + ) def storage_root(state: State, address: Address) -> Root: diff --git a/src/ethereum/cancun/state_tracking.py b/src/ethereum/cancun/state_tracking.py new file mode 100644 index 0000000000..a5a1548527 --- /dev/null +++ b/src/ethereum/cancun/state_tracking.py @@ -0,0 +1,135 @@ +from dataclasses import dataclass, field +from typing import Dict, List, Optional, Set, Tuple, Union + +from ethereum_types.bytes import Bytes32 +from ethereum_types.numeric import U256 + +from .fork_types import Account, Address + +# State access types for tracking +ACCOUNT_READ = "account_read" +ACCOUNT_WRITE = "account_write" +STORAGE_READ = "storage_read" +STORAGE_WRITE = "storage_write" + + +@dataclass +class StateAccess: + """Record of a single state access for proof generation.""" + access_type: str + address: Address + key: Optional[Bytes32] = None + value_before: Optional[Union[Account, U256]] = None + value_after: Optional[Union[Account, U256]] = None + + +@dataclass +class StateTracker: + """Tracks state access for merkle proof generation.""" + accesses: List[StateAccess] = field(default_factory=list) + main_trie_accessed_keys: Set[Address] = field(default_factory=set) + storage_accessed_keys: Dict[Address, Set[Bytes32]] = field(default_factory=dict) + track_reads: bool = True + track_writes: bool = True + + +def enable_state_tracking( + state, + track_reads: bool = True, + track_writes: bool = True +) -> None: + """ + Enable state tracking on a State object. + + Parameters + ---------- + state : State + The state to enable tracking on + track_reads : bool + Whether to track read operations + track_writes : bool + Whether to track write operations + """ + state._state_tracker = StateTracker( + track_reads=track_reads, + track_writes=track_writes + ) + + +def disable_state_tracking(state) -> None: + """ + Disable state tracking on a State object. + + Parameters + ---------- + state : State + The state to disable tracking on + """ + state._state_tracker = None + + +def log_state_access( + state, + access_type: str, + address: Address, + key: Optional[Bytes32] = None, + value_before: Optional[Union[Account, U256]] = None, + value_after: Optional[Union[Account, U256]] = None, +) -> None: + """ + Log a state access if tracking is enabled. + + Parameters + ---------- + state : State + The state (with potential tracker) + access_type : str + Type of access (ACCOUNT_READ, ACCOUNT_WRITE, etc.) + address : Address + Address being accessed + key : Optional[Bytes32] + Storage key (for storage operations) + value_before : Optional[Union[Account, U256]] + Value before the operation + value_after : Optional[Union[Account, U256]] + Value after the operation + """ + if state._state_tracker is None: + return + + tracker = state._state_tracker + access = StateAccess( + access_type=access_type, + address=address, + key=key, + value_before=value_before, + value_after=value_after, + ) + tracker.accesses.append(access) + + if access_type in [ACCOUNT_READ, ACCOUNT_WRITE]: + tracker.main_trie_accessed_keys.add(address) + elif access_type in [STORAGE_READ, STORAGE_WRITE]: + if address not in tracker.storage_accessed_keys: + tracker.storage_accessed_keys[address] = set() + if key is not None: + tracker.storage_accessed_keys[address].add(key) + +# Dummy method +def generate_merkle_proof_requests(state) -> Tuple[List[Address], List[Tuple[Address, Bytes32]]]: + """ + Generate lists of proof requests needed for all tracked accesses. + + Parameters + ---------- + state : State + The state containing tracking logs + + Returns + ------- + account_proofs : List[Address] + List of addresses needing account proofs + storage_proofs : List[Tuple[Address, Bytes32]] + List of (address, storage_key) tuples needing storage proofs + """ + return [], []