Merge pull request #3563 from etan-status/lc-branchtypes

Use types for representing LC Merkle branches

Author: hwwhww

pysetup/helpers.py

@@ -22,8 +22,11 @@ def collect_prev_forks(fork: str) -> list[str]:
         forks.append(fork)
 
 
-def is_byte_vector(value: str) -> bool:
-    return value.startswith(('ByteVector'))
+def requires_mypy_type_ignore(value: str) -> bool:
+    return (
+        value.startswith(('ByteVector'))
+        or (value.startswith(('Vector')) and 'floorlog2' in value)
+    )
 
 
 def make_function_abstract(protocol_def: ProtocolDefinition, key: str):
@@ -41,7 +44,8 @@ def objects_to_spec(preset_name: str,
     new_type_definitions = (
         '\n\n'.join(
             [
-                f"class {key}({value}):\n    pass\n" if not is_byte_vector(value) else f"class {key}({value}):  # type: ignore\n    pass\n"
+                f"class {key}({value}):\n    pass\n" if not requires_mypy_type_ignore(value)
+                else f"class {key}({value}):  # type: ignore\n    pass\n"
                 for key, value in spec_object.custom_types.items()
             ]
         )
@@ -108,7 +112,7 @@ def format_constant(name: str, vardef: VariableDefinition) -> str:
         if vardef.comment is not None:
             out += f'  # {vardef.comment}'
         return out
-    
+
     # Merge all constant objects
     hardcoded_ssz_dep_constants =         reduce(lambda obj, builder: {**obj, **builder.hardcoded_ssz_dep_constants()}, builders, {})
     hardcoded_custom_type_dep_constants = reduce(lambda obj, builder: {**obj, **builder.hardcoded_custom_type_dep_constants(spec_object)}, builders, {})
@@ -131,12 +135,13 @@ def format_constant(name: str, vardef: VariableDefinition) -> str:
         imports,
         preparations,
         f"fork = \'{fork}\'\n",
+        # The helper functions that some SSZ containers require. Need to be defined before `custom_type_dep_constants`
+        CONSTANT_DEP_SUNDRY_CONSTANTS_FUNCTIONS,
         # The constants that some SSZ containers require. Need to be defined before `new_type_definitions`
         custom_type_dep_constants,
-        new_type_definitions,
-        CONSTANT_DEP_SUNDRY_CONSTANTS_FUNCTIONS,
         # The constants that some SSZ containers require. Need to be defined before `constants_spec`
         ssz_dep_constants,
+        new_type_definitions,
         constant_vars_spec,
         preset_vars_spec,
         config_spec,

pysetup/spec_builders/capella.py

@@ -16,5 +16,5 @@ def imports(cls, preset_name: str):
     @classmethod
     def hardcoded_ssz_dep_constants(cls) -> Dict[str, str]:
         return {
-            'EXECUTION_PAYLOAD_INDEX': 'GeneralizedIndex(25)',
+            'EXECUTION_PAYLOAD_GINDEX': 'GeneralizedIndex(25)',
         }

pysetup/spec_builders/whisk.py

@@ -27,6 +27,6 @@ def hardcoded_custom_type_dep_constants(cls, spec_object) -> str:
     @classmethod
     def hardcoded_ssz_dep_constants(cls) -> Dict[str, str]:
         constants = {
-            'EXECUTION_PAYLOAD_INDEX': 'GeneralizedIndex(41)',
+            'EXECUTION_PAYLOAD_GINDEX': 'GeneralizedIndex(41)',
         }
         return {**super().hardcoded_ssz_dep_constants(), **constants}

setup.py

@@ -93,6 +93,8 @@ def _get_class_info_from_source(source: str) -> Tuple[str, Optional[str]]:
     base = class_def.bases[0]
     if isinstance(base, ast.Name):
         parent_class = base.id
+    elif isinstance(base, ast.Subscript):
+        parent_class = base.value.id
     else:
         # NOTE: SSZ definition derives from earlier phase...
         # e.g. `phase0.SignedBeaconBlock`

specs/altair/light-client/full-node.md

@@ -75,7 +75,8 @@ def create_light_client_bootstrap(state: BeaconState,
     return LightClientBootstrap(
         header=block_to_light_client_header(block),
         current_sync_committee=state.current_sync_committee,
-        current_sync_committee_branch=compute_merkle_proof(state, CURRENT_SYNC_COMMITTEE_GINDEX),
+        current_sync_committee_branch=CurrentSyncCommitteeBranch(
+            compute_merkle_proof(state, CURRENT_SYNC_COMMITTEE_GINDEX)),
     )
 ```
 
@@ -122,7 +123,8 @@ def create_light_client_update(state: BeaconState,
     # `next_sync_committee` is only useful if the message is signed by the current sync committee
     if update_attested_period == update_signature_period:
         update.next_sync_committee = attested_state.next_sync_committee
-        update.next_sync_committee_branch = compute_merkle_proof(attested_state, NEXT_SYNC_COMMITTEE_GINDEX)
+        update.next_sync_committee_branch = NextSyncCommitteeBranch(
+            compute_merkle_proof(attested_state, NEXT_SYNC_COMMITTEE_GINDEX))
 
     # Indicate finality whenever possible
     if finalized_block is not None:
@@ -131,7 +133,8 @@ def create_light_client_update(state: BeaconState,
             assert hash_tree_root(update.finalized_header.beacon) == attested_state.finalized_checkpoint.root
         else:
             assert attested_state.finalized_checkpoint.root == Bytes32()
-        update.finality_branch = compute_merkle_proof(attested_state, FINALIZED_ROOT_GINDEX)
+        update.finality_branch = FinalityBranch(
+            compute_merkle_proof(attested_state, FINALIZED_ROOT_GINDEX))
 
     update.sync_aggregate = block.message.body.sync_aggregate
     update.signature_slot = block.message.slot

specs/altair/light-client/sync-protocol.md

@@ -9,6 +9,7 @@
 <!-- DON'T EDIT THIS SECTION, INSTEAD RE-RUN doctoc TO UPDATE -->
 
 - [Introduction](#introduction)
+- [Custom types](#custom-types)
 - [Constants](#constants)
 - [Preset](#preset)
   - [Misc](#misc)
@@ -56,6 +57,14 @@ Additional documents describe how the light client sync protocol can be used:
 - [Light client](./light-client.md)
 - [Networking](./p2p-interface.md)
 
+## Custom types
+
+| Name | SSZ equivalent | Description |
+| - | - | - |
+| `FinalityBranch` | `Vector[Bytes32, floorlog2(FINALIZED_ROOT_GINDEX)]` | Merkle branch of `finalized_checkpoint.root` within `BeaconState` |
+| `CurrentSyncCommitteeBranch` | `Vector[Bytes32, floorlog2(CURRENT_SYNC_COMMITTEE_GINDEX)]` | Merkle branch of `current_sync_committee` within `BeaconState` |
+| `NextSyncCommitteeBranch` | `Vector[Bytes32, floorlog2(NEXT_SYNC_COMMITTEE_GINDEX)]` | Merkle branch of `next_sync_committee` within `BeaconState` |
+
 ## Constants
 
 | Name | Value |
@@ -93,7 +102,7 @@ class LightClientBootstrap(Container):
     header: LightClientHeader
     # Current sync committee corresponding to `header.beacon.state_root`
     current_sync_committee: SyncCommittee
-    current_sync_committee_branch: Vector[Bytes32, floorlog2(CURRENT_SYNC_COMMITTEE_GINDEX)]
+    current_sync_committee_branch: CurrentSyncCommitteeBranch
 ```
 
 ### `LightClientUpdate`
@@ -104,10 +113,10 @@ class LightClientUpdate(Container):
     attested_header: LightClientHeader
     # Next sync committee corresponding to `attested_header.beacon.state_root`
     next_sync_committee: SyncCommittee
-    next_sync_committee_branch: Vector[Bytes32, floorlog2(NEXT_SYNC_COMMITTEE_GINDEX)]
+    next_sync_committee_branch: NextSyncCommitteeBranch
     # Finalized header corresponding to `attested_header.beacon.state_root`
     finalized_header: LightClientHeader
-    finality_branch: Vector[Bytes32, floorlog2(FINALIZED_ROOT_GINDEX)]
+    finality_branch: FinalityBranch
     # Sync committee aggregate signature
     sync_aggregate: SyncAggregate
     # Slot at which the aggregate signature was created (untrusted)
@@ -122,7 +131,7 @@ class LightClientFinalityUpdate(Container):
     attested_header: LightClientHeader
     # Finalized header corresponding to `attested_header.beacon.state_root`
     finalized_header: LightClientHeader
-    finality_branch: Vector[Bytes32, floorlog2(FINALIZED_ROOT_GINDEX)]
+    finality_branch: FinalityBranch
     # Sync committee aggregate signature
     sync_aggregate: SyncAggregate
     # Slot at which the aggregate signature was created (untrusted)
@@ -174,14 +183,14 @@ def is_valid_light_client_header(header: LightClientHeader) -> bool:
 
 ```python
 def is_sync_committee_update(update: LightClientUpdate) -> bool:
-    return update.next_sync_committee_branch != [Bytes32() for _ in range(floorlog2(NEXT_SYNC_COMMITTEE_GINDEX))]
+    return update.next_sync_committee_branch != NextSyncCommitteeBranch()
 ```
 
 ### `is_finality_update`
 
 ```python
 def is_finality_update(update: LightClientUpdate) -> bool:
-    return update.finality_branch != [Bytes32() for _ in range(floorlog2(FINALIZED_ROOT_GINDEX))]
+    return update.finality_branch != FinalityBranch()
 ```
 
 ### `is_better_update`
@@ -493,7 +502,7 @@ def process_light_client_finality_update(store: LightClientStore,
     update = LightClientUpdate(
         attested_header=finality_update.attested_header,
         next_sync_committee=SyncCommittee(),
-        next_sync_committee_branch=[Bytes32() for _ in range(floorlog2(NEXT_SYNC_COMMITTEE_GINDEX))],
+        next_sync_committee_branch=NextSyncCommitteeBranch(),
         finalized_header=finality_update.finalized_header,
         finality_branch=finality_update.finality_branch,
         sync_aggregate=finality_update.sync_aggregate,
@@ -512,9 +521,9 @@ def process_light_client_optimistic_update(store: LightClientStore,
     update = LightClientUpdate(
         attested_header=optimistic_update.attested_header,
         next_sync_committee=SyncCommittee(),
-        next_sync_committee_branch=[Bytes32() for _ in range(floorlog2(NEXT_SYNC_COMMITTEE_GINDEX))],
+        next_sync_committee_branch=NextSyncCommitteeBranch(),
         finalized_header=LightClientHeader(),
-        finality_branch=[Bytes32() for _ in range(floorlog2(FINALIZED_ROOT_GINDEX))],
+        finality_branch=FinalityBranch(),
         sync_aggregate=optimistic_update.sync_aggregate,
         signature_slot=optimistic_update.signature_slot,
     )

specs/capella/light-client/full-node.md

@@ -46,14 +46,15 @@ def block_to_light_client_header(block: SignedBeaconBlock) -> LightClientHeader:
             transactions_root=hash_tree_root(payload.transactions),
             withdrawals_root=hash_tree_root(payload.withdrawals),
         )
-        execution_branch = compute_merkle_proof(block.message.body, EXECUTION_PAYLOAD_INDEX)
+        execution_branch = ExecutionBranch(
+            compute_merkle_proof(block.message.body, EXECUTION_PAYLOAD_GINDEX))
     else:
         # Note that during fork transitions, `finalized_header` may still point to earlier forks.
         # While Bellatrix blocks also contain an `ExecutionPayload` (minus `withdrawals_root`),
         # it was not included in the corresponding light client data. To ensure compatibility
         # with legacy data going through `upgrade_lc_header_to_capella`, leave out execution data.
         execution_header = ExecutionPayloadHeader()
-        execution_branch = [Bytes32() for _ in range(floorlog2(EXECUTION_PAYLOAD_INDEX))]
+        execution_branch = ExecutionBranch()
 
     return LightClientHeader(
         beacon=BeaconBlockHeader(

specs/capella/light-client/sync-protocol.md

@@ -9,6 +9,7 @@
 <!-- DON'T EDIT THIS SECTION, INSTEAD RE-RUN doctoc TO UPDATE -->
 
 - [Introduction](#introduction)
+- [Custom types](#custom-types)
 - [Constants](#constants)
 - [Containers](#containers)
   - [Modified `LightClientHeader`](#modified-lightclientheader)
@@ -27,11 +28,17 @@ Additional documents describes the impact of the upgrade on certain roles:
 - [Full node](./full-node.md)
 - [Networking](./p2p-interface.md)
 
+## Custom types
+
+| Name | SSZ equivalent | Description |
+| - | - | - |
+| `ExecutionBranch` | `Vector[Bytes32, floorlog2(EXECUTION_PAYLOAD_GINDEX)]` | Merkle branch of `execution_payload` within `BeaconBlockBody` |
+
 ## Constants
 
 | Name | Value |
 | - | - |
-| `EXECUTION_PAYLOAD_INDEX` | `get_generalized_index(BeaconBlockBody, 'execution_payload')` (= 25) |
+| `EXECUTION_PAYLOAD_GINDEX` | `get_generalized_index(BeaconBlockBody, 'execution_payload')` (= 25) |
 
 ## Containers
 
@@ -43,7 +50,7 @@ class LightClientHeader(Container):
     beacon: BeaconBlockHeader
     # Execution payload header corresponding to `beacon.body_root` (from Capella onward)
     execution: ExecutionPayloadHeader
-    execution_branch: Vector[Bytes32, floorlog2(EXECUTION_PAYLOAD_INDEX)]
+    execution_branch: ExecutionBranch
 ```
 
 ## Helper functions
@@ -69,14 +76,14 @@ def is_valid_light_client_header(header: LightClientHeader) -> bool:
     if epoch < CAPELLA_FORK_EPOCH:
         return (
             header.execution == ExecutionPayloadHeader()
-            and header.execution_branch == [Bytes32() for _ in range(floorlog2(EXECUTION_PAYLOAD_INDEX))]
+            and header.execution_branch == ExecutionBranch()
         )
 
     return is_valid_merkle_branch(
         leaf=get_lc_execution_root(header),
         branch=header.execution_branch,
-        depth=floorlog2(EXECUTION_PAYLOAD_INDEX),
-        index=get_subtree_index(EXECUTION_PAYLOAD_INDEX),
+        depth=floorlog2(EXECUTION_PAYLOAD_GINDEX),
+        index=get_subtree_index(EXECUTION_PAYLOAD_GINDEX),
         root=header.beacon.body_root,
     )
 ```

specs/deneb/light-client/full-node.md

@@ -52,14 +52,15 @@ def block_to_light_client_header(block: SignedBeaconBlock) -> LightClientHeader:
             execution_header.blob_gas_used = payload.blob_gas_used
             execution_header.excess_blob_gas = payload.excess_blob_gas
 
-        execution_branch = compute_merkle_proof(block.message.body, EXECUTION_PAYLOAD_INDEX)
+        execution_branch = ExecutionBranch(
+            compute_merkle_proof(block.message.body, EXECUTION_PAYLOAD_GINDEX))
     else:
         # Note that during fork transitions, `finalized_header` may still point to earlier forks.
         # While Bellatrix blocks also contain an `ExecutionPayload` (minus `withdrawals_root`),
         # it was not included in the corresponding light client data. To ensure compatibility
         # with legacy data going through `upgrade_lc_header_to_capella`, leave out execution data.
         execution_header = ExecutionPayloadHeader()
-        execution_branch = [Bytes32() for _ in range(floorlog2(EXECUTION_PAYLOAD_INDEX))]
+        execution_branch = ExecutionBranch()
 
     return LightClientHeader(
         beacon=BeaconBlockHeader(

specs/deneb/light-client/sync-protocol.md

@@ -74,14 +74,14 @@ def is_valid_light_client_header(header: LightClientHeader) -> bool:
     if epoch < CAPELLA_FORK_EPOCH:
         return (
             header.execution == ExecutionPayloadHeader()
-            and header.execution_branch == [Bytes32() for _ in range(floorlog2(EXECUTION_PAYLOAD_INDEX))]
+            and header.execution_branch == ExecutionBranch()
         )
 
     return is_valid_merkle_branch(
         leaf=get_lc_execution_root(header),
         branch=header.execution_branch,
-        depth=floorlog2(EXECUTION_PAYLOAD_INDEX),
-        index=get_subtree_index(EXECUTION_PAYLOAD_INDEX),
+        depth=floorlog2(EXECUTION_PAYLOAD_GINDEX),
+        index=get_subtree_index(EXECUTION_PAYLOAD_GINDEX),
         root=header.beacon.body_root,
     )
 ```