fix(specs): EIP-7928 - Refactor state.state_tracker argument pattern

Author: nerolation

src/ethereum/forks/amsterdam/block_access_lists/__init__.py

@@ -21,6 +21,9 @@
     StateChangeTracker,
     begin_call_frame,
     commit_call_frame,
+    handle_in_transaction_selfdestruct,
+    normalize_balance_changes,
+    prepare_balance_tracking,
     rollback_call_frame,
     set_block_access_index,
     track_address_access,
@@ -44,6 +47,9 @@
     "build_block_access_list",
     "commit_call_frame",
     "compute_block_access_list_hash",
+    "handle_in_transaction_selfdestruct",
+    "normalize_balance_changes",
+    "prepare_balance_tracking",
     "rollback_call_frame",
     "set_block_access_index",
     "rlp_encode_block_access_list",

src/ethereum/forks/amsterdam/block_access_lists/tracker.py

@@ -113,9 +113,7 @@ class StateChangeTracker:
     """
 
 
-def set_block_access_index(
-    tracker: StateChangeTracker, block_access_index: Uint
-) -> None:
+def set_block_access_index(state: "State", block_access_index: Uint) -> None:
     """
     Set the current block access index for tracking changes.
 
@@ -129,13 +127,14 @@ def set_block_access_index(
 
     Parameters
     ----------
-    tracker :
-        The state change tracker instance.
+    state :
+        The current execution state.
     block_access_index :
         The block access index (0 for pre-execution,
         1..n for transactions, n+1 for post-execution).
 
     """
+    tracker = state.change_tracker
     tracker.current_block_access_index = block_access_index
     # Clear the pre-storage cache for each new transaction to ensure
     # no-op writes are detected relative to the transaction start
@@ -181,9 +180,7 @@ def capture_pre_state(
     return tracker.pre_storage_cache[cache_key]
 
 
-def track_address_access(
-    tracker: StateChangeTracker, address: Address
-) -> None:
+def track_address_access(state: "State", address: Address) -> None:
     """
     Track that an address was accessed.
 
@@ -192,18 +189,18 @@ def track_address_access(
 
     Parameters
     ----------
-    tracker :
-        The state change tracker instance.
+    state :
+        The current execution state.
     address :
         The account address that was accessed.
 
     """
-    add_touched_account(tracker.block_access_list_builder, address)
+    add_touched_account(
+        state.change_tracker.block_access_list_builder, address
+    )
 
 
-def track_storage_read(
-    tracker: StateChangeTracker, address: Address, key: Bytes32, state: "State"
-) -> None:
+def track_storage_read(state: "State", address: Address, key: Bytes32) -> None:
     """
     Track a storage read operation.
 
@@ -213,29 +210,28 @@ def track_storage_read(
 
     Parameters
     ----------
-    tracker :
-        The state change tracker instance.
+    state :
+        The current execution state.
     address :
         The account address whose storage is being read.
     key :
         The storage slot being read.
-    state :
-        The current execution state.
 
     """
-    track_address_access(tracker, address)
+    track_address_access(state, address)
 
-    capture_pre_state(tracker, address, key, state)
+    capture_pre_state(state.change_tracker, address, key, state)
 
-    add_storage_read(tracker.block_access_list_builder, address, key)
+    add_storage_read(
+        state.change_tracker.block_access_list_builder, address, key
+    )
 
 
 def track_storage_write(
-    tracker: StateChangeTracker,
+    state: "State",
     address: Address,
     key: Bytes32,
     new_value: U256,
-    state: "State",
 ) -> None:
     """
     Track a storage write operation.
@@ -246,22 +242,21 @@ def track_storage_write(
 
     Parameters
     ----------
-    tracker :
-        The state change tracker instance.
+    state :
+        The current execution state.
     address :
         The account address whose storage is being modified.
     key :
         The storage slot being written to.
     new_value :
         The new value to write.
-    state :
-        The current execution state.
 
     [EIP-7928]: https://eips.ethereum.org/EIPS/eip-7928
 
     """
-    track_address_access(tracker, address)
+    track_address_access(state, address)
 
+    tracker = state.change_tracker
     pre_value = capture_pre_state(tracker, address, key, state)
 
     value_bytes = new_value.to_be_bytes32()
@@ -322,8 +317,33 @@ def capture_pre_balance(
     return tracker.pre_balance_cache[address]
 
 
+def prepare_balance_tracking(state: "State", address: Address) -> None:
+    """
+    Prepare for tracking balance changes by caching the pre-transaction
+    balance.
+
+    This should be called before any balance modifications when you need to
+    ensure the pre-balance is captured for later normalization. This is
+    particularly important for operations like withdrawals where the balance
+    might not actually change.
+
+    Parameters
+    ----------
+    state :
+        The current execution state.
+    address :
+        The account address whose balance will be tracked.
+
+    """
+    # Ensure the address is tracked
+    track_address_access(state, address)
+
+    # Cache the pre-balance for later normalization
+    capture_pre_balance(state.change_tracker, address, state)
+
+
 def track_balance_change(
-    tracker: StateChangeTracker,
+    state: "State",
     address: Address,
     new_balance: U256,
 ) -> None:
@@ -335,16 +355,17 @@ def track_balance_change(
 
     Parameters
     ----------
-    tracker :
-        The state change tracker instance.
+    state :
+        The current execution state.
     address :
         The account address whose balance changed.
     new_balance :
         The new balance value.
 
     """
-    track_address_access(tracker, address)
+    track_address_access(state, address)
 
+    tracker = state.change_tracker
     block_access_index = BlockAccessIndex(tracker.current_block_access_index)
     add_balance_change(
         tracker.block_access_list_builder,
@@ -362,7 +383,7 @@ def track_balance_change(
 
 
 def track_nonce_change(
-    tracker: StateChangeTracker, address: Address, new_nonce: Uint
+    state: "State", address: Address, new_nonce: Uint
 ) -> None:
     """
     Track a nonce change for an account.
@@ -373,20 +394,19 @@ def track_nonce_change(
 
     Parameters
     ----------
-    tracker :
-        The state change tracker instance.
+    state :
+        The current execution state.
     address :
         The account address whose nonce changed.
     new_nonce :
         The new nonce value.
-    state :
-        The current execution state.
 
     [`CREATE`]: ref:ethereum.forks.amsterdam.vm.instructions.system.create
     [`CREATE2`]: ref:ethereum.forks.amsterdam.vm.instructions.system.create2
 
     """
-    track_address_access(tracker, address)
+    track_address_access(state, address)
+    tracker = state.change_tracker
     block_access_index = BlockAccessIndex(tracker.current_block_access_index)
     nonce_u64 = U64(new_nonce)
     add_nonce_change(
@@ -403,7 +423,7 @@ def track_nonce_change(
 
 
 def track_code_change(
-    tracker: StateChangeTracker, address: Address, new_code: Bytes
+    state: "State", address: Address, new_code: Bytes
 ) -> None:
     """
     Track a code change for contract deployment.
@@ -414,8 +434,8 @@ def track_code_change(
 
     Parameters
     ----------
-    tracker :
-        The state change tracker instance.
+    state :
+        The current execution state.
     address :
         The address receiving the contract code.
     new_code :
@@ -425,7 +445,8 @@ def track_code_change(
     [`CREATE2`]: ref:ethereum.forks.amsterdam.vm.instructions.system.create2
 
     """
-    track_address_access(tracker, address)
+    track_address_access(state, address)
+    tracker = state.change_tracker
     block_access_index = BlockAccessIndex(tracker.current_block_access_index)
     add_code_change(
         tracker.block_access_list_builder,
@@ -441,7 +462,7 @@ def track_code_change(
 
 
 def handle_in_transaction_selfdestruct(
-    tracker: StateChangeTracker, address: Address
+    state: "State", address: Address
 ) -> None:
     """
     Handle an account that self-destructed in the same transaction it was
@@ -456,12 +477,13 @@ def handle_in_transaction_selfdestruct(
 
     Parameters
     ----------
-    tracker :
-        The state change tracker instance.
+    state :
+        The current execution state.
     address :
         The address that self-destructed.
 
     """
+    tracker = state.change_tracker
     builder = tracker.block_access_list_builder
     if address not in builder.accounts:
         return
@@ -493,9 +515,7 @@ def handle_in_transaction_selfdestruct(
     ]
 
 
-def normalize_balance_changes(
-    tracker: StateChangeTracker, state: "State"
-) -> None:
+def normalize_balance_changes(state: "State") -> None:
     """
     Normalize balance changes for the current block access index.
 
@@ -515,15 +535,14 @@ def normalize_balance_changes(
 
     Parameters
     ----------
-    tracker :
-        The state change tracker instance.
     state :
         The current execution state.
 
     """
     # Import locally to avoid circular import
     from ..state import get_account
 
+    tracker = state.change_tracker
     builder = tracker.block_access_list_builder
     current_index = tracker.current_block_access_index
 
@@ -548,7 +567,7 @@ def normalize_balance_changes(
             ]
 
 
-def begin_call_frame(tracker: StateChangeTracker) -> None:
+def begin_call_frame(state: "State") -> None:
     """
     Begin a new call frame for tracking reverts.
 
@@ -557,14 +576,14 @@ def begin_call_frame(tracker: StateChangeTracker) -> None:
 
     Parameters
     ----------
-    tracker :
-        The state change tracker instance.
+    state :
+        The current execution state.
 
     """
-    tracker.call_frame_snapshots.append(CallFrameSnapshot())
+    state.change_tracker.call_frame_snapshots.append(CallFrameSnapshot())
 
 
-def rollback_call_frame(tracker: StateChangeTracker) -> None:
+def rollback_call_frame(state: "State") -> None:
     """
     Rollback changes from the current call frame.
 
@@ -578,10 +597,11 @@ def rollback_call_frame(tracker: StateChangeTracker) -> None:
 
     Parameters
     ----------
-    tracker :
-        The state change tracker instance.
+    state :
+        The current execution state.
 
     """
+    tracker = state.change_tracker
     if not tracker.call_frame_snapshots:
         return
 
@@ -651,7 +671,7 @@ def rollback_call_frame(tracker: StateChangeTracker) -> None:
     # All touched addresses remain in the access list (already tracked)
 
 
-def commit_call_frame(tracker: StateChangeTracker) -> None:
+def commit_call_frame(state: "State") -> None:
     """
     Commit changes from the current call frame.
 
@@ -660,9 +680,9 @@ def commit_call_frame(tracker: StateChangeTracker) -> None:
 
     Parameters
     ----------
-    tracker :
-        The state change tracker instance.
+    state :
+        The current execution state.
 
     """
-    if tracker.call_frame_snapshots:
-        tracker.call_frame_snapshots.pop()
+    if state.change_tracker.call_frame_snapshots:
+        state.change_tracker.call_frame_snapshots.pop()