feat: test_balance exhaustive checks example.

src/ethereum_test_specs/blockchain.py

@@ -466,6 +466,7 @@ def generate_block_data(
                     )
                 self.verify_witness(
                     t8n=t8n,
+                    fork=fork,
                     state_diff=transition_tool_output.result.state_diff,
                     witness_check=block.witness_check,
                 )
@@ -619,6 +620,7 @@ def verify_post_state(
     def verify_witness(
         self,
         t8n: TransitionTool,
+        fork: Fork,
         state_diff: StateDiff,
         witness_check: WitnessCheck,
     ) -> None:
@@ -629,13 +631,19 @@ def verify_witness(
         witness_check_state_diff, witness_check_address_mapping = t8n.get_witness_check_mapping(
             witness_check
         )
+        print("\nExpected witness check state diff:")
+        print(witness_check_state_diff.model_dump_json(indent=4))
+        system_contract_accounts = fork.pre_allocation_blockchain().items()
+        addresses_to_skip = [Address(address) for address, _ in system_contract_accounts]
+
         for stem_state_diff in state_diff.root:
             actual_stem = stem_state_diff.stem
             address = witness_check_address_mapping.get(actual_stem, None)
-            print("Stem we are checking from actual witness:")
-            print(str(actual_stem))
-            print("Address we are checking from actual witness:")
-            print(str(address))
+            print(f"\nChecking witness for stem: {actual_stem} at address: {address}")
+            # TODO: skip system contract addresses for now
+            if address in addresses_to_skip:
+                print(f"Skipping system contract address in witness check: {address}")
+                continue
             # check for stem in the expected witness check
             expected_stem_state_diff = next(
                 (sd for sd in witness_check_state_diff.root if sd.stem == actual_stem), None
@@ -679,7 +687,9 @@ def verify_witness(
                         )
                     )
                 # check the current value of the actual suffix state diff matches the expected
-                if actual_current_value != expected_suffix_state_diff.current_value:
+                if str(actual_current_value) != str(
+                    expected_suffix_state_diff.current_value
+                ):  # TODO: temp fix str casting
                     raise ValueError(
                         "Witness check failed - current value mismatch. The stem and suffix"
                         " exist.\n\n"

src/ethereum_test_types/verkle/types.py

@@ -10,6 +10,7 @@
 from pydantic.functional_serializers import model_serializer
 
 from ethereum_test_base_types import Address, CamelModel, HexNumber, PaddedFixedSizeBytes
+from ethereum_test_forks import Fork, Verkle
 from ethereum_test_types import Account
 
 IPA_PROOF_DEPTH = 8
@@ -171,16 +172,22 @@ class AccountHeaderEntry(Enum):
         BASIC_DATA = 0
         CODEHASH = 1
 
-    def __init__(self) -> None:
+    def __init__(self, fork: Fork) -> None:
         """
         Initializes a WitnessCheck instance.
         """
+        assert fork >= Verkle, "WitnessCheck is only supported for Verkle fork and later"
         self.account_entries: List[
             Tuple[Address, WitnessCheck.AccountHeaderEntry, Optional[Hash]]
         ] = []
         self.storage_slots: List[Tuple[Address, int, Optional[Hash]]] = []
         self.code_chunks: List[Tuple[Address, int, Optional[Hash]]] = []
 
+        # Add the pre-allocation accounts by default
+        pre_allocations = fork.pre_allocation_blockchain()
+        for address, account_data in pre_allocations.items():
+            self.add_account_full(Address(address), Account(**account_data))
+
     def __repr__(self) -> str:
         """
         Provides a detailed string representation of the WitnessCheck object for debugging.
@@ -197,13 +204,17 @@ def add_account_full(self, address: Address, account: Account | None) -> None:
         """
         Adds the address, nonce, balance, and code. Delays actual key computation until later.
         """
-        if account and account.code:
-            code_hash = Hash(keccak256(account.code))
-        else:
-            # keccak256 of empty byte array
-            code_hash = Hash(0xC5D2460186F7233C927E7DB2DCC703C0E500B653CA82273B7BFAD8045D85A470)
         self.add_account_basic_data(address, account)
-        self.add_account_codehash(address, code_hash)
+        if account:
+            if account.code:
+                code_hash = Hash(keccak256(account.code))
+            else:  # keccak256 of empty byte array
+                code_hash = Hash(
+                    0xC5D2460186F7233C927E7DB2DCC703C0E500B653CA82273B7BFAD8045D85A470
+                )
+            self.add_account_codehash(address, code_hash)
+        else:
+            self.add_account_codehash(address, None)
 
     def add_account_basic_data(self, address: Address, account: Account | None) -> None:
         """

tests/verkle/eip4762_verkle_gas_witness/test_balance.py

@@ -7,18 +7,21 @@
 
 import pytest
 
+from ethereum_test_forks import Verkle
 from ethereum_test_tools import (
     Account,
     Address,
     Block,
     BlockchainTestFiller,
     Environment,
+    Hash,
     TestAddress,
     TestAddress2,
     Transaction,
     WitnessCheck,
 )
 from ethereum_test_tools.vm.opcode import Opcodes as Op
+from ethereum_test_types.verkle.helpers import chunkify_code
 
 # TODO(verkle): Update reference spec version
 REFERENCE_SPEC_GIT_PATH = "EIPS/eip-4762.md"
@@ -88,12 +91,15 @@ def _balance(
         gas_price=10,
     )
 
-    witness_check = WitnessCheck()
-    for address in [env.fee_recipient, TestAddress, TestAddress2]:
-        witness_check.add_account_full(
-            address=address,
-            account=(None if address == env.fee_recipient else pre[address]),
-        )
+    witness_check = WitnessCheck(fork=Verkle)
+    witness_check.add_account_full(address=TestAddress, account=pre[TestAddress])
+    witness_check.add_account_full(address=TestAddress2, account=pre[TestAddress2])
+    witness_check.add_storage_slot(address=TestAddress2, storage_slot=0, value=None)
+    code_chunks = chunkify_code(pre[TestAddress2].code)
+    for i, chunk in enumerate(code_chunks, start=0):
+        witness_check.add_code_chunk(address=TestAddress2, chunk_number=i, value=chunk)
+
+    witness_check.add_account_full(address=env.fee_recipient, account=None)
     for address in exp_addr_basic_data:
         witness_check.add_account_basic_data(
             address=address,