Add get_fork_version and get_domain

Author: hwwhww

eth/beacon/enums/__init__.py

@@ -1,7 +1,7 @@
 from enum import IntEnum
 
 
-class BLSDomain(IntEnum):
+class SignatureDomain(IntEnum):
     DOMAIN_DEPOSIT = 0
     DOMAIN_ATTESTATION = 1
     DOMAIN_PROPOSAL = 2

eth/beacon/enums/signature_domain.py

@@ -41,9 +41,11 @@
 
 
 if TYPE_CHECKING:
+    from eth.beacon.enums.signature_domain 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
 
 
@@ -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/helpers.py

@@ -14,6 +14,7 @@
     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,
@@ -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, 0, 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,
+    )