feat: block witness checks and test fixes

src/ethereum_test_base_types/base_types.py

@@ -181,18 +181,20 @@ class Sized(cls):  # type: ignore
 
         return Sized
 
-    def __new__(cls, input: NumberConvertible | N):
+    def __new__(cls, input: NumberConvertible | N | None):
         """
         Creates a new Number object.
         """
-        i = to_number(input)
-        if i > cls.max_value:
-            raise ValueError(f"Value {i} is too large for {cls.byte_length} bytes")
-        if i < 0:
-            i += cls.max_value + 1
-            if i <= 0:
-                raise ValueError(f"Value {i} is too small for {cls.byte_length} bytes")
-        return super(FixedSizeHexNumber, cls).__new__(cls, i)
+        if input is not None:
+            i = to_number(input)
+            if i > cls.max_value:
+                raise ValueError(f"Value {i} is too large for {cls.byte_length} bytes")
+            if i < 0:
+                i += cls.max_value + 1
+                if i <= 0:
+                    raise ValueError(f"Value {i} is too small for {cls.byte_length} bytes")
+            return super(FixedSizeHexNumber, cls).__new__(cls, i)
+        return None
 
     def __str__(self) -> str:
         """

src/ethereum_test_forks/forks/forks.py

@@ -404,7 +404,9 @@ def call_opcodes(
         """
         At Homestead, DELEGATECALL opcode was introduced.
         """
-        return [(Opcodes.DELEGATECALL, EVMCodeType.LEGACY),] + super(
+        return [
+            (Opcodes.DELEGATECALL, EVMCodeType.LEGACY),
+        ] + super(
             Homestead, cls
         ).call_opcodes(block_number, timestamp)
 
@@ -449,7 +451,9 @@ def call_opcodes(
         """
         At Byzantium, STATICCALL opcode was introduced.
         """
-        return [(Opcodes.STATICCALL, EVMCodeType.LEGACY),] + super(
+        return [
+            (Opcodes.STATICCALL, EVMCodeType.LEGACY),
+        ] + super(
             Byzantium, cls
         ).call_opcodes(block_number, timestamp)
 
@@ -483,7 +487,9 @@ def create_opcodes(
         """
         At Constantinople, `CREATE2` opcode is added.
         """
-        return [(Opcodes.CREATE2, EVMCodeType.LEGACY),] + super(
+        return [
+            (Opcodes.CREATE2, EVMCodeType.LEGACY),
+        ] + super(
             Constantinople, cls
         ).create_opcodes(block_number, timestamp)
 
@@ -974,7 +980,7 @@ def pre_allocation_blockchain(cls) -> Mapping:
         type tests.
         """
         new_allocation = {
-            0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFE: {
+            Address(0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFE): {
                 "nonce": 1,
                 "code": (
                     "0x60203611603157600143035f35116029575f35612000014311602957612000"
@@ -1020,7 +1026,10 @@ def evm_code_types(cls, block_number: int = 0, timestamp: int = 0) -> List[EVMCo
         """
         EOF V1 is supported starting from this fork.
         """
-        return super(CancunEIP7692, cls,).evm_code_types(  # noqa: SC200
+        return super(
+            CancunEIP7692,
+            cls,
+        ).evm_code_types(  # noqa: SC200
             block_number,
             timestamp,
         ) + [EVMCodeType.EOF_V1]
@@ -1049,7 +1058,9 @@ def create_opcodes(
         """
         EOF V1 introduces `EOFCREATE`.
         """
-        return [(Opcodes.EOFCREATE, EVMCodeType.EOF_V1),] + super(
+        return [
+            (Opcodes.EOFCREATE, EVMCodeType.EOF_V1),
+        ] + super(
             CancunEIP7692, cls  # noqa: SC200
         ).create_opcodes(block_number, timestamp)
 
@@ -1084,7 +1095,10 @@ def evm_code_types(cls, block_number: int = 0, timestamp: int = 0) -> List[EVMCo
         """
         EOF V1 is supported starting from this fork.
         """
-        return super(PragueEIP7692, cls,).evm_code_types(  # noqa: SC200
+        return super(
+            PragueEIP7692,
+            cls,
+        ).evm_code_types(  # noqa: SC200
             block_number,
             timestamp,
         ) + [EVMCodeType.EOF_V1]

src/ethereum_test_specs/blockchain.py

@@ -2,7 +2,7 @@
 Ethereum blockchain test spec definition and filler.
 """
 
-from pprint import pprint
+from pprint import pformat
 from typing import Any, Callable, ClassVar, Dict, Generator, List, Optional, Tuple, Type
 
 from pydantic import ConfigDict, Field, field_validator
@@ -48,7 +48,7 @@
     Withdrawal,
     WithdrawalRequest,
 )
-from ethereum_test_types.verkle import VerkleTree, Witness
+from ethereum_test_types.verkle import StateDiff, VerkleTree, Witness, WitnessCheck
 from evm_transition_tool import TransitionTool
 
 from .base import BaseTest, verify_result, verify_transactions
@@ -258,6 +258,10 @@ class Block(Header):
     """
     Custom list of requests to embed in this block.
     """
+    witness_check: WitnessCheck | None = None
+    """
+    Verkle execution witness check for the block.
+    """
 
     def set_environment(self, env: Environment) -> Environment:
         """
@@ -452,20 +456,52 @@ def generate_block_data(
             debug_output_path=self.get_next_transition_tool_output_path(),
         )
 
-        try:
+        try:  # General checks for the transition tool output
             rejected_txs = verify_transactions(txs, transition_tool_output.result)
             verify_result(transition_tool_output.result, env)
-            # TODO: add verify witness (against vkt)
-            # verify_witness(transition_tool_output.witness, transition_tool_output.vkt)
+            if block.witness_check:
+                if transition_tool_output.result.state_diff is None:
+                    raise Exception(
+                        "no state diff in transition tool output, cannot verify witness"
+                    )
+                # TODO: hack for now, temp addition to check hist. storage contract
+                block.witness_check.add_storage_slot(
+                    address=Address(0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFE),
+                    storage_slot=env.number - 1,
+                    # value=env.parent_hash,
+                    value=None,
+                )
+                self.verify_witness(
+                    t8n=t8n,
+                    state_diff=transition_tool_output.result.state_diff,
+                    witness_check=block.witness_check,
+                )
         except Exception as e:
             print_traces(t8n.get_traces())
-            pprint(transition_tool_output.result)
-            pprint(previous_alloc)
-            pprint(transition_tool_output.alloc)
+            print(
+                "\nTransition tool output result:\n"
+                f"{pformat(transition_tool_output.result.model_dump_json())}"
+            )
+            print(
+                "\nPrevious transition tool alloc:\n"
+                f"{pformat(previous_alloc.model_dump_json())}"
+            )
+            if transition_tool_output.alloc is not None:
+                print(
+                    "\nTransition tool output alloc:\n"
+                    f"{pformat(transition_tool_output.alloc.model_dump_json())}"
+                )
             if transition_tool_output.vkt is not None:
-                pprint(transition_tool_output.vkt)
-            if transition_tool_output.witness is not None:
-                pprint(transition_tool_output.witness)
+                print(
+                    "\nTransition tools output verkle tree:\n"
+                    f"{pformat(transition_tool_output.vkt.model_dump_json())}"
+                )
+            # TODO: t8n has the witness state diff from the result for now
+            # if transition_tool_output.witness is not None:
+            # print(
+            # "\nTransition tools output witness:\n"
+            # f"{pformat(transition_tool_output.witness.model_dump_json())}"
+            # )
             raise e
 
         if len(rejected_txs) > 0 and block.exception is None:
@@ -538,7 +574,7 @@ def generate_block_data(
                 )
             )
             transition_tool_output.alloc = previous_alloc
-            # TODO: hack for now
+            # TODO: hack for now, replace with actual witness output once available from t8n
             transition_tool_output.witness = Witness(
                 verkle_proof=transition_tool_output.result.verkle_proof,
                 state_diff=transition_tool_output.result.state_diff,
@@ -587,6 +623,87 @@ def verify_post_state(
             print_traces(t8n.get_traces())
             raise e
 
+    def verify_witness(
+        self,
+        t8n: TransitionTool,
+        state_diff: StateDiff,
+        witness_check: WitnessCheck,
+    ) -> None:
+        """
+        Compares the expected witness check allocation account against the values updated
+        in the block execution witness state diff.
+        """
+        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))
+
+        for stem_state_diff in state_diff.root:
+            actual_stem = stem_state_diff.stem
+            address = witness_check_address_mapping.get(actual_stem, None)
+            print(f"\nChecking witness for stem: {actual_stem} at address: {address}")
+            # 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
+            )
+            if not expected_stem_state_diff:
+                raise ValueError(
+                    "Witness check failed - missing stem not found in expected witness check.\n\n"
+                    + pformat(
+                        {
+                            "test_account_address": str(address),
+                            "stem": str(actual_stem),
+                        },
+                        indent=4,
+                    )
+                )
+            for suffix_diff in stem_state_diff.suffix_diffs:
+                actual_suffix = suffix_diff.suffix
+                actual_current_value = suffix_diff.current_value
+                # check for suffix in the expected witness check
+                expected_suffix_state_diff = next(
+                    (
+                        sd
+                        for sd in expected_stem_state_diff.suffix_diffs
+                        if sd.suffix == actual_suffix
+                    ),
+                    None,
+                )
+                if not expected_suffix_state_diff:
+                    raise ValueError(
+                        "Witness check failed - actual suffix not found in expected witness"
+                        " check.\n\n"
+                        + pformat(
+                            {
+                                "test_account_address": str(address),
+                                "stem": str(actual_stem),
+                                "suffix": actual_suffix,
+                                "value_actual": str(actual_current_value),
+                                "value_expected": "value not found",
+                            },
+                            indent=4,
+                        )
+                    )
+                # check the current value of the actual suffix state diff matches the expected
+                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"
+                        + pformat(
+                            {
+                                "test_account_address": str(address),
+                                "stem": str(actual_stem),
+                                "suffix": actual_suffix,
+                                "value_actual": str(actual_current_value),
+                                "value_expected": str(expected_suffix_state_diff.current_value),
+                            },
+                            indent=4,
+                        )
+                    )
+
     def make_fixture(
         self,
         t8n: TransitionTool,
@@ -841,4 +958,5 @@ def generate(
 
 
 BlockchainTestSpec = Callable[[str], Generator[BlockchainTest, None, None]]
+
 BlockchainTestFiller = Type[BlockchainTest]

src/ethereum_test_tools/__init__.py

@@ -46,8 +46,11 @@
     Storage,
     TestParameterGroup,
     Transaction,
+    VerkleTree,
     Withdrawal,
     WithdrawalRequest,
+    Witness,
+    WitnessCheck,
     add_kzg_version,
     ceiling_division,
     compute_create2_address,
@@ -136,6 +139,9 @@
     "TestPrivateKey2",
     "Transaction",
     "TransactionException",
+    "VerkleTree",
+    "Witness",
+    "WitnessCheck",
     "Withdrawal",
     "WithdrawalRequest",
     "Yul",

src/ethereum_test_types/__init__.py

@@ -30,6 +30,7 @@
     Withdrawal,
     WithdrawalRequest,
 )
+from .verkle import VerkleTree, Witness, WitnessCheck
 
 __all__ = (
     "AccessList",
@@ -53,8 +54,11 @@
     "TestPrivateKey",
     "TestPrivateKey2",
     "Transaction",
+    "VerkleTree",
     "Withdrawal",
     "WithdrawalRequest",
+    "Witness",
+    "WitnessCheck",
     "ZeroPaddedHexNumber",
     "add_kzg_version",
     "ceiling_division",

src/ethereum_test_types/verkle/__init__.py

@@ -2,13 +2,28 @@
 Ethereum Verkle Test Types.
 """
 
-from .types import IpaProof, StateDiff, SuffixStateDiff, VerkleProof, VerkleTree, Witness
+from .helpers import chunkify_code
+from .types import (
+    IpaProof,
+    StateDiff,
+    Stem,
+    StemStateDiff,
+    SuffixStateDiff,
+    VerkleProof,
+    VerkleTree,
+    Witness,
+    WitnessCheck,
+)
 
 __all__ = (
     "IpaProof",
     "StateDiff",
+    "StemStateDiff",
     "SuffixStateDiff",
+    "Stem",
     "VerkleProof",
     "VerkleTree",
     "Witness",
+    "WitnessCheck",
+    "chunkify_code",
 )