Merge pull request #1591 from hwwhww/get_fork

Add `get_fork_version` and `get_domain` helpers

Author: hwwhww

eth/beacon/enums.py

@@ -0,0 +1,21 @@
+from enum import IntEnum
+
+
+class ValidatorStatusCode(IntEnum):
+    PENDING_ACTIVATION = 0
+    ACTIVE = 1
+    ACTIVE_PENDING_EXIT = 2
+    EXITED_WITHOUT_PENALTY = 3
+    EXITED_WITH_PENALTY = 4
+
+
+class ValidatorRegistryDeltaFlag(IntEnum):
+    ACTIVATION = 0
+    EXIT = 1
+
+
+class SignatureDomain(IntEnum):
+    DOMAIN_DEPOSIT = 0
+    DOMAIN_ATTESTATION = 1
+    DOMAIN_PROPOSAL = 2
+    DOMAIN_EXIT = 3

eth/beacon/enums/signature_domain.py

@@ -28,7 +28,7 @@
 )
 
 from eth.beacon.block_committees_info import BlockCommitteesInfo
-from eth.beacon.enums.validator_status_codes import (
+from eth.beacon.enums import (
     ValidatorStatusCode,
 )
 from eth.beacon.types.shard_committees import (
@@ -41,9 +41,11 @@
 
 
 if TYPE_CHECKING:
+    from eth.beacon.enums import SignatureDomain  # noqa: F401
     from eth.beacon.types.attestation_records import AttestationRecord  # noqa: F401
     from eth.beacon.types.blocks import BaseBeaconBlock  # noqa: F401
     from eth.beacon.types.states import BeaconState  # noqa: F401
+    from eth.beacon.types.fork_data import ForkData  # noqa: F401
     from eth.beacon.types.validator_records import ValidatorRecord  # noqa: F401
 
 
@@ -154,7 +156,7 @@ def get_active_validator_indices(validators: Sequence['ValidatorRecord']) -> Tup
     """
     return tuple(
         i for i, v in enumerate(validators)
-        if v.status in [ValidatorStatusCode.ACTIVE, ValidatorStatusCode.PENDING_EXIT]
+        if v.status in [ValidatorStatusCode.ACTIVE, ValidatorStatusCode.ACTIVE_PENDING_EXIT]
     )
 
 
@@ -393,3 +395,27 @@ def get_new_validator_registry_delta_chain_tip(current_validator_registry_delta_
         # TODO: currently, we use 256-bit pubkey which is different form the spec
         pubkey.to_bytes(32, 'big')
     )
+
+
+def get_fork_version(fork_data: 'ForkData',
+                     slot: int) -> int:
+    """
+    Return the current ``fork_version`` from the given ``fork_data`` and ``slot``.
+    """
+    if slot < fork_data.fork_slot:
+        return fork_data.pre_fork_version
+    else:
+        return fork_data.post_fork_version
+
+
+def get_domain(fork_data: 'ForkData',
+               slot: int,
+               domain_type: 'SignatureDomain') -> int:
+    """
+    Return the domain number of the current fork and ``domain_type``.
+    """
+    # 2 ** 32 = 4294967296
+    return get_fork_version(
+        fork_data,
+        slot,
+    ) * 4294967296 + domain_type

eth/beacon/enums/special_record_types.py

@@ -9,7 +9,7 @@
 import eth.utils.bls as bls
 from eth.utils.blake import blake
 
-from eth.beacon.enums.validator_status_codes import (
+from eth.beacon.enums import (
     ValidatorStatusCode,
 )
 from eth.beacon.state_machines.forks.serenity.configs import SERENITY_CONFIG

eth/beacon/enums/validator_registry_delta_flags.py

@@ -1,4 +1,4 @@
-from eth.beacon.enums.validator_status_codes import (
+from eth.beacon.enums import (
     ValidatorStatusCode,
 )
 from eth.beacon.types.validator_records import (

eth/beacon/enums/validator_status_codes.py

@@ -10,10 +10,11 @@
 )
 
 
-from eth.beacon.enums.validator_status_codes import (
+from eth.beacon.enums import (
     ValidatorStatusCode,
 )
 from eth.beacon.types.blocks import BaseBeaconBlock
+from eth.beacon.types.fork_data import ForkData
 from eth.beacon.types.shard_committees import ShardCommittee
 from eth.beacon.types.states import BeaconState
 from eth.beacon.types.validator_records import ValidatorRecord
@@ -24,6 +25,8 @@
     get_beacon_proposer_index,
     get_block_hash,
     get_effective_balance,
+    get_domain,
+    get_fork_version,
     get_hashes_from_latest_block_hashes,
     get_new_shuffling,
     get_new_validator_registry_delta_chain_tip,
@@ -500,14 +503,14 @@ def test_get_active_validator_indices(sample_validator_record_params):
     active_validator_indices = get_active_validator_indices(validators)
     assert len(active_validator_indices) == 3
 
-    # Make one validator becomes PENDING_EXIT.
+    # Make one validator becomes ACTIVE_PENDING_EXIT.
     validators[0] = validators[0].copy(
-        status=ValidatorStatusCode.PENDING_EXIT,
+        status=ValidatorStatusCode.ACTIVE_PENDING_EXIT,
     )
     active_validator_indices = get_active_validator_indices(validators)
     assert len(active_validator_indices) == 3
 
-    # Make one validator becomes PENDING_EXIT.
+    # Make one validator becomes EXITED_WITHOUT_PENALTY.
     validators[0] = validators[0].copy(
         status=ValidatorStatusCode.EXITED_WITHOUT_PENALTY,
     )
@@ -661,3 +664,68 @@ def test_get_new_validator_registry_delta_chain_tip(index,
         flag=flag,
     )
     assert result == expected
+
+
+@pytest.mark.parametrize(
+    (
+        'pre_fork_version,'
+        'post_fork_version,'
+        'fork_slot,'
+        'current_slot,'
+        'expected'
+    ),
+    [
+        (0, 0, 0, 0, 0),
+        (0, 0, 0, 1, 0),
+        (0, 1, 20, 10, 0),
+        (0, 1, 20, 20, 1),
+        (0, 1, 10, 20, 1),
+    ]
+)
+def test_get_fork_version(pre_fork_version,
+                          post_fork_version,
+                          fork_slot,
+                          current_slot,
+                          expected):
+    fork_data = ForkData(
+        pre_fork_version=pre_fork_version,
+        post_fork_version=post_fork_version,
+        fork_slot=fork_slot,
+    )
+    assert expected == get_fork_version(
+        fork_data,
+        current_slot,
+    )
+
+
+@pytest.mark.parametrize(
+    (
+        'pre_fork_version,'
+        'post_fork_version,'
+        'fork_slot,'
+        'current_slot,'
+        'domain_type,'
+        'expected'
+    ),
+    [
+        (1, 2, 20, 10, 10, 1 * 2 ** 32 + 10),
+        (1, 2, 20, 20, 11, 2 * 2 ** 32 + 11),
+        (1, 2, 10, 20, 12, 2 * 2 ** 32 + 12),
+    ]
+)
+def test_get_domain(pre_fork_version,
+                    post_fork_version,
+                    fork_slot,
+                    current_slot,
+                    domain_type,
+                    expected):
+    fork_data = ForkData(
+        pre_fork_version=pre_fork_version,
+        post_fork_version=post_fork_version,
+        fork_slot=fork_slot,
+    )
+    assert expected == get_domain(
+        fork_data=fork_data,
+        slot=current_slot,
+        domain_type=domain_type,
+    )