[WIP] feat(fw): verkle t8n post state verification

src/ethereum_test_tools/spec/base/base_test.py

@@ -44,15 +44,11 @@ def verify_transactions(txs: List[Transaction] | None, result) -> List[int]:
     return list(rejected_txs.keys())
 
 
-def verify_post_alloc(
-    *, expected_post: Mapping, got_alloc: Mapping, got_vkt: Optional[Mapping] = None
-):
+def verify_post_alloc(*, expected_post: Mapping, got_alloc: Mapping):
     """
-    Verify that an allocation matches the expected post in the test.
+    Verify that the final allocation mapping matches the expected post in the test.
     Raises exception on unexpected values.
     """
-    # TODO: If VKT is not None, we have an overlay of the VKT on top of the alloc and we need
-    # to verify that
     got_alloc_normalized: Dict[Address, Any] = {
         Address(address): got_alloc[address] for address in got_alloc
     }
@@ -69,6 +65,37 @@ def verify_post_alloc(
                     raise Exception(f"expected account not found: {address}")
 
 
+def verify_post_vkt(
+    transition_tool: TransitionTool, expected_post: Mapping, got_vkt: Mapping, got_alloc: Mapping
+):
+    """
+    Verify that the final verkle tree mapping matches the expected post in the test.
+    Raises exception on unexpected values.
+    """
+    # TODO: Add this check before filling within the filler cli plugin.
+    if not transition_tool.verkle_subcommand:
+        raise Exception("Only geth's evm tool is supported to verify verkle trees.")
+
+    # Convert the expected post alloc to a verkle tree for comparison.
+    expected_vkt = transition_tool.post_alloc_to_vkt(post_alloc=expected_post)
+
+    # Check for keys that are missing the actual VKT
+    missing_keys = [key for key in expected_vkt if key not in got_vkt]
+    if missing_keys:
+        raise Exception(f"Missing keys in actual VKT: {missing_keys}")
+
+    # TODO: how to determine what is unexpected, i.e TestAddress is expected but not in post state
+    unexpected_keys = [key for key in got_vkt if key not in expected_vkt]
+
+    # Compare the values for each key in the expected VKT
+    for key, expected_value in expected_vkt.items():
+        actual_value = got_vkt.get(key)
+        if expected_value != actual_value:
+            raise Exception(
+                f"VKT mismatch at key {key}: expected {expected_value}, got {actual_value}"
+            )
+
+
 def verify_result(result: Mapping, env: Environment):
     """
     Verify that values in the t8n result match the expected values.

src/ethereum_test_tools/spec/blockchain/blockchain_test.py

@@ -32,6 +32,7 @@
     BaseFixture,
     BaseTest,
     verify_post_alloc,
+    verify_post_vkt,
     verify_result,
     verify_transactions,
 )
@@ -314,10 +315,15 @@ def network_info(self, fork: Fork, eips: Optional[List[int]] = None):
 
     def verify_post_state(self, *, t8n, alloc, vkt=None):
         """
-        Verifies the post alloc after all block/s or payload/s are generated.
+        Verifies the post state after all block/s or payload/s are generated.
         """
         try:
-            verify_post_alloc(expected_post=self.post, got_alloc=alloc, got_vkt=vkt)
+            if vkt is not None:
+                verify_post_vkt(
+                    transition_tool=t8n, expected_post=self.post, got_vkt=vkt, got_alloc=alloc
+                )
+            else:
+                verify_post_alloc(expected_post=self.post, got_alloc=alloc)
         except Exception as e:
             print_traces(t8n.get_traces())
             raise e

src/evm_transition_tool/geth.py

@@ -8,11 +8,13 @@
 import textwrap
 from pathlib import Path
 from re import compile
-from typing import Optional
+from typing import Mapping, Optional
 
 import pytest
+from ethereum.crypto.hash import keccak256
 
 from ethereum_test_forks import Fork
+from ethereum_test_tools import Hash
 
 from .transition_tool import FixtureFormats, TransitionTool, dump_files_to_directory
 
@@ -27,6 +29,7 @@ class GethTransitionTool(TransitionTool):
     t8n_subcommand: Optional[str] = "t8n"
     statetest_subcommand: Optional[str] = "statetest"
     blocktest_subcommand: Optional[str] = "blocktest"
+    verkle_subcommand: Optional[str] = "verkle"
 
     binary: Path
     cached_version: Optional[str] = None
@@ -125,3 +128,63 @@ def verify_fixture(
                 f"Failed to verify fixture via: '{' '.join(command)}'. "
                 f"Error: '{result.stderr.decode()}'"
             )
+
+    def verkle_tree_key(self, account: str, storage_slot: Optional[str] = None) -> str:
+        """
+        Returns the verkle tree key for the input account using the verkle subcommand.
+        Optionally the key for the storage slot if specified.
+        """
+        command = [
+            str(self.binary),
+            str(self.t8n_subcommand),
+            str(self.verkle_subcommand),
+            str(account),
+        ]
+        if storage_slot:
+            command.append(storage_slot)
+
+        result = subprocess.run(
+            command,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+        )
+
+        if result.returncode != 0:
+            raise Exception(
+                f"Failed to run verkle subcommand: '{' '.join(command)}'. "
+                f"Error: '{result.stderr.decode()}'"
+            )
+        return result.stdout.decode().strip()  # strip the newline character
+
+    def post_alloc_to_vkt(self, post_alloc: Mapping) -> Mapping:
+        """
+        Converts the expected post alloc to verkle tree representation using the verkle
+        subcommand.
+        """
+        vkt = {}
+        for address, account in post_alloc.items():
+            # Add the account address: value is simply "0x000...000"
+            address_key = self.verkle_tree_key(address)
+            vkt[address_key] = Hash(0).hex()
+
+            # Add account balance: numbers are little-endian
+            balance_key = address_key[:-2] + "01"
+            balance_value = Hash(account.balance.to_bytes(32, "little")).hex()
+            vkt[balance_key] = balance_value
+
+            # Add account nonce: numbers are little-endian
+            nonce_key = address_key[:-2] + "02"
+            nonce_value = Hash(account.nonce.to_bytes(32, "little")).hex()
+            vkt[nonce_key] = nonce_value
+
+            # Add account code hash: keccak256 hash of the code
+            code_hash_key = address_key[:-2] + "03"
+            code_hash_value = Hash(keccak256(account.code)).hex()
+            vkt[code_hash_key] = code_hash_value
+
+            # Add account storage: each slot has a unique key
+            for slot, value in account.storage.data.items():
+                slot_key = self.verkle_tree_key(address, Hash(slot).hex())
+                vkt[slot_key] = Hash(value).hex()
+
+        return vkt

src/evm_transition_tool/tests/test_alloc_to_vkt.py

@@ -0,0 +1,49 @@
+"""
+Test the verkle tree subcommand from the geth transition tool.
+"""
+
+import pytest
+
+from evm_transition_tool import GethTransitionTool
+
+
+# TODO: Update to use correct types.
+@pytest.mark.parametrize(
+    "post_alloc, expected_vkt",
+    [
+        (
+            {
+                "0x0000000000000000000000000000000000000100": {
+                    "nonce": "0x01",
+                    "balance": "0x01",
+                    "code": "0x60203560003555",
+                    "storage": {"0x0a": "0x0b"},
+                },
+            },
+            {
+                "0x31b64bbd0b09c1d09afea606bebb70bce80a2909189e513c924a0871bbd36300": "0x0000000000000000000000000000000000000000000000000000000000000000",  # noqa: E501
+                "0x31b64bbd0b09c1d09afea606bebb70bce80a2909189e513c924a0871bbd36301": "0x0100000000000000000000000000000000000000000000000000000000000000",  # noqa: E501
+                "0x31b64bbd0b09c1d09afea606bebb70bce80a2909189e513c924a0871bbd36302": "0x0100000000000000000000000000000000000000000000000000000000000000",  # noqa: E501
+                "0x31b64bbd0b09c1d09afea606bebb70bce80a2909189e513c924a0871bbd36303": "0x159c5cfa7fab15702c72a4141ef31b26c42827bd74ffc5671d95033227e662e7",  # noqa: E501
+                "0x31b64bbd0b09c1d09afea606bebb70bce80a2909189e513c924a0871bbd3634a": "0x000000000000000000000000000000000000000000000000000000000000000b",  # noqa: E501
+            },
+        ),
+    ],
+)
+def test_post_alloc_to_vkt(post_alloc, expected_vkt):
+    """
+    Verifies that the `post_alloc_to_vkt` method of the `GethTransitionTool` class.
+    """
+    t8n = GethTransitionTool()
+    result_vkt = t8n.post_alloc_to_vkt(post_alloc)
+
+    assert set(result_vkt.keys()) == set(
+        expected_vkt.keys()
+    ), "Keys in created verkle tree do not match the expected keys."
+
+    for key, expected_value in expected_vkt.items():
+        assert key in result_vkt, f"Key {key} is missing in created verkle tree."
+        assert result_vkt[key] == expected_value, (
+            f"Value for {key} in the created verkle tree does not match expected. "
+            f"Expected {expected_value}, got {result_vkt[key]}."
+        )

src/evm_transition_tool/transition_tool.py

@@ -14,7 +14,7 @@
 from itertools import groupby
 from pathlib import Path
 from re import Pattern
-from typing import Any, Dict, List, Optional, Tuple, Type
+from typing import Any, Dict, List, Mapping, Optional, Tuple, Type
 
 from ethereum_test_forks import Fork
 
@@ -102,6 +102,7 @@ class TransitionTool:
     blocktest_subcommand: Optional[str] = None
     cached_version: Optional[str] = None
     t8n_use_stream: bool = True
+    verkle_subcommand: Optional[str] = None
 
     # Abstract methods that each tool must implement
 
@@ -646,3 +647,13 @@ def verify_fixture(
         raise Exception(
             "The `verify_fixture()` function is not supported by this tool. Use geth's evm tool."
         )
+
+    def post_alloc_to_vkt(self, post_alloc: Mapping) -> Mapping:
+        """
+        Converts the expected post alloc to verkle tree representation using the verkle subcommand.
+
+        Currently only implemented by geth's evm.
+        """
+        raise Exception(
+            "The `create_post_vkt()` function is not supported by this tool. Use geth's evm tool."
+        )