WIP: latest system contract changes

Author: nerolation

src/ethereum/osaka/block_access_lists/__init__.py

@@ -15,6 +15,7 @@
 from .tracker import (
     StateChangeTracker,
     set_transaction_index,
+    set_system_transaction_index,
     track_address_access,
     track_balance_change,
     track_code_change,
@@ -40,6 +41,7 @@
     "build",
     "compute_bal_hash",
     "set_transaction_index",
+    "set_system_transaction_index",
     "ssz_encode_block_access_list",
     "track_address_access",
     "track_balance_change",

src/ethereum/osaka/block_access_lists/tracker.py

@@ -63,6 +63,12 @@ class StateChangeTracker:
     """
     The index of the currently executing transaction within the block.
     """
+    
+    system_tx_index: Optional[int] = None
+    """
+    Special transaction index for system contract calls.
+    Set to len(transactions) + 1 to track system contract state changes.
+    """
 
 
 def set_transaction_index(tracker: StateChangeTracker, tx_index: int) -> None:
@@ -80,6 +86,24 @@ def set_transaction_index(tracker: StateChangeTracker, tx_index: int) -> None:
         The index of the transaction about to be processed.
     """
     tracker.current_tx_index = tx_index
+    tracker.system_tx_index = None  # Reset system tx index
+
+
+def set_system_transaction_index(tracker: StateChangeTracker, tx_index: int) -> None:
+    """
+    Set the system transaction index for tracking system contract changes.
+    
+    Used to track state changes from system contracts (pre/post execution)
+    with a special transaction index of len(transactions) + 1.
+    
+    Parameters
+    ----------
+    tracker :
+        The state change tracker instance.
+    tx_index :
+        The system transaction index (typically len(transactions) + 1).
+    """
+    tracker.system_tx_index = tx_index
 
 
 def capture_pre_state(
@@ -200,12 +224,15 @@ def track_storage_write(
 
     value_bytes = new_value.to_be_bytes32()
 
+    # Use system tx index if set, otherwise use current tx index
+    tx_index = tracker.system_tx_index if tracker.system_tx_index is not None else tracker.current_tx_index
+    
     if pre_value != new_value:
         add_storage_write(
             tracker.block_access_list_builder,
             address,
             key,
-            U32(tracker.current_tx_index),
+            U32(tx_index),
             value_bytes
         )
     else:
@@ -239,11 +266,14 @@ def track_balance_change(
     """
     track_address_access(tracker, address)
 
+    # Use system tx index if set, otherwise use current tx index
+    tx_index = tracker.system_tx_index if tracker.system_tx_index is not None else tracker.current_tx_index
+    
     balance_bytes = U128(new_balance).to_be_bytes16()
     add_balance_change(
         tracker.block_access_list_builder,
         address,
-        U32(tracker.current_tx_index),
+        U32(tx_index),
         balance_bytes
     )
 
@@ -276,10 +306,14 @@ def track_nonce_change(
     [`CREATE2`]: ref:ethereum.osaka.vm.instructions.system.create2
     """
     track_address_access(tracker, address)
+    
+    # Use system tx index if set, otherwise use current tx index
+    tx_index = tracker.system_tx_index if tracker.system_tx_index is not None else tracker.current_tx_index
+    
     add_nonce_change(
         tracker.block_access_list_builder,
         address,
-        U32(tracker.current_tx_index),
+        U32(tx_index),
         U64(new_nonce)
     )
 
@@ -313,10 +347,14 @@ def track_code_change(
     [`SETCODE`]: ref:ethereum.osaka.vm.instructions.system.setcode
     """
     track_address_access(tracker, address)
+    
+    # Use system tx index if set, otherwise use current tx index
+    tx_index = tracker.system_tx_index if tracker.system_tx_index is not None else tracker.current_tx_index
+    
     add_code_change(
         tracker.block_access_list_builder,
         address,
-        U32(tracker.current_tx_index),
+        U32(tx_index),
         new_code
     )
 

src/ethereum/osaka/fork.py

@@ -30,7 +30,7 @@
 )
 
 from . import vm
-from .block_access_lists import StateChangeTracker, compute_bal_hash, build
+from .block_access_lists import StateChangeTracker, compute_bal_hash, build, set_system_transaction_index, track_balance_change
 from .blocks import Block, Header, Log, Receipt, Withdrawal, encode_receipt
 from .bloom import logs_bloom
 from .exceptions import (
@@ -582,6 +582,7 @@ def process_system_transaction(
     target_address: Address,
     system_contract_code: Bytes,
     data: Bytes,
+    change_tracker: Optional[StateChangeTracker] = None,
 ) -> MessageCallOutput:
     """
     Process a system transaction with the given code.
@@ -639,6 +640,7 @@ def process_system_transaction(
         disable_precompiles=False,
         warm_code_addresses=set(),
         parent_evm=None,
+        change_tracker=change_tracker,
     )
 
     system_tx_output = process_message_call(system_tx_message)
@@ -650,6 +652,7 @@ def process_checked_system_transaction(
     block_env: vm.BlockEnvironment,
     target_address: Address,
     data: Bytes,
+    change_tracker: Optional[StateChangeTracker] = None,
 ) -> MessageCallOutput:
     """
     Process a system transaction and raise an error if the contract does not
@@ -682,6 +685,7 @@ def process_checked_system_transaction(
         target_address,
         system_contract_code,
         data,
+        change_tracker,
     )
 
     if system_tx_output.error:
@@ -697,6 +701,7 @@ def process_unchecked_system_transaction(
     block_env: vm.BlockEnvironment,
     target_address: Address,
     data: Bytes,
+    change_tracker: Optional[StateChangeTracker] = None,
 ) -> MessageCallOutput:
     """
     Process a system transaction without checking if the contract contains code
@@ -722,6 +727,7 @@ def process_unchecked_system_transaction(
         target_address,
         system_contract_code,
         data,
+        change_tracker,
     )
 
 
@@ -759,16 +765,23 @@ def apply_body(
     # Initialize Block Access List state change tracker
     change_tracker = StateChangeTracker(block_output.block_access_list_builder)
 
+    # Set system transaction index for pre-execution system contracts
+    # Using len(transactions) + 1 as specified
+    system_tx_index = len(transactions) + 1
+    set_system_transaction_index(change_tracker, system_tx_index)
+
     process_unchecked_system_transaction(
         block_env=block_env,
         target_address=BEACON_ROOTS_ADDRESS,
         data=block_env.parent_beacon_block_root,
+        change_tracker=change_tracker,
     )
 
     process_unchecked_system_transaction(
         block_env=block_env,
         target_address=HISTORY_STORAGE_ADDRESS,
         data=block_env.block_hashes[-1],  # The parent hash
+        change_tracker=change_tracker,
     )
 
     for i, tx in enumerate(map(decode_transaction, transactions)):
@@ -777,9 +790,13 @@ def apply_body(
 
     process_withdrawals(block_env, block_output, withdrawals, change_tracker)
 
+    # Set system transaction index for post-execution system contracts
+    set_system_transaction_index(change_tracker, system_tx_index)
+    
     process_general_purpose_requests(
         block_env=block_env,
         block_output=block_output,
+        change_tracker=change_tracker,
     )
 
     return block_output
@@ -788,6 +805,7 @@ def apply_body(
 def process_general_purpose_requests(
     block_env: vm.BlockEnvironment,
     block_output: vm.BlockOutput,
+    change_tracker: StateChangeTracker,
 ) -> None:
     """
     Process all the requests in the block.
@@ -809,6 +827,7 @@ def process_general_purpose_requests(
         block_env=block_env,
         target_address=WITHDRAWAL_REQUEST_PREDEPLOY_ADDRESS,
         data=b"",
+        change_tracker=change_tracker,
     )
 
     if len(system_withdrawal_tx_output.return_data) > 0:
@@ -820,6 +839,7 @@ def process_general_purpose_requests(
         block_env=block_env,
         target_address=CONSOLIDATION_REQUEST_PREDEPLOY_ADDRESS,
         data=b"",
+        change_tracker=change_tracker,
     )
 
     if len(system_consolidation_tx_output.return_data) > 0:
@@ -1022,9 +1042,9 @@ def increase_recipient_balance(recipient: Account) -> None:
 
         modify_state(block_env.state, wd.address, increase_recipient_balance)
 
-        # Track balance change for BAL
+        # Track balance change for BAL (withdrawals are tracked as system contract changes)
         new_balance = get_account(block_env.state, wd.address).balance
-        change_tracker.track_balance_change(wd.address, U256(new_balance), block_env.state)
+        track_balance_change(change_tracker, wd.address, U256(new_balance), block_env.state)
 
         if account_exists_and_is_empty(block_env.state, wd.address):
             destroy_account(block_env.state, wd.address)