PR feedback and move aggregation sig related functions into aggregation.py

Author: hwwhww

eth/beacon/aggregation.py

@@ -0,0 +1,79 @@
+from typing import (
+    Iterable,
+    Tuple,
+)
+from cytoolz import (
+    pipe
+)
+
+
+from eth_typing import (
+    Hash32,
+)
+
+from eth.utils import bls
+from eth.utils.bitfield import (
+    set_voted,
+)
+from eth.utils.blake import blake
+
+
+def create_signing_message(slot: int,
+                           parent_hashes: Iterable[Hash32],
+                           shard_id: int,
+                           shard_block_hash: Hash32,
+                           justified_slot: int) -> bytes:
+    """
+    Return the signining message for attesting.
+    """
+    # TODO: Will be updated with SSZ encoded attestation.
+    return blake(
+        slot.to_bytes(8, byteorder='big') +
+        b''.join(parent_hashes) +
+        shard_id.to_bytes(2, byteorder='big') +
+        shard_block_hash +
+        justified_slot.to_bytes(8, 'big')
+    )
+
+
+def verify_votes(
+        message: bytes,
+        votes: Iterable[Tuple[int, bytes, int]]) -> Tuple[Tuple[bytes, ...], Tuple[int, ...]]:
+    """
+    Verify the given votes.
+
+    vote: (committee_index, sig, public_key)
+    """
+    sigs_with_committe_info = tuple(
+        (sig, committee_index)
+        for (committee_index, sig, public_key)
+        in votes
+        if bls.verify(message, public_key, sig)
+    )
+    try:
+        sigs, committee_indices = zip(*sigs_with_committe_info)
+    except ValueError:
+        sigs = tuple()
+        committee_indices = tuple()
+
+    return sigs, committee_indices
+
+
+def aggregate_votes(bitfield: bytes,
+                    sigs: Iterable[bytes],
+                    voting_sigs: Iterable[bytes],
+                    voting_committee_indices: Iterable[int]) -> Tuple[bytes, Tuple[int]]:
+    """
+    Aggregate the votes.
+    """
+    # Update the bitfield and append the signatures
+    sigs = tuple(sigs) + tuple(voting_sigs)
+    bitfield = pipe(
+        bitfield,
+        *(
+            set_voted(index=committee_index)
+            for committee_index in voting_committee_indices
+        )
+    )
+
+    return bitfield, bls.aggregate_sigs(sigs)

eth/beacon/helpers.py

@@ -2,9 +2,6 @@
     repeat,
 )
 
-from cytoolz import (
-    pipe
-)
 from typing import (
     Any,
     Iterable,
@@ -21,11 +18,6 @@
     Hash32,
 )
 
-from eth.utils import bls
-from eth.utils.bitfield import (
-    set_voted,
-)
-from eth.utils.blake import blake
 from eth.utils.numeric import (
     clamp,
 )
@@ -394,105 +386,3 @@ def get_block_committees_info(parent_block: 'BaseBeaconBlock',
         proposer_committee_size=proposer_committee_size,
         shards_and_committees=shards_and_committees,
     )
-
-
-#
-# Signatures and Aggregation
-#
-def create_signing_message(slot: int,
-                           parent_hashes: Iterable[Hash32],
-                           shard_id: int,
-                           shard_block_hash: Hash32,
-                           justified_slot: int) -> bytes:
-    # TODO: will be updated to hashed encoded attestation
-    return blake(
-        slot.to_bytes(8, byteorder='big') +
-        b''.join(parent_hashes) +
-        shard_id.to_bytes(2, byteorder='big') +
-        shard_block_hash +
-        justified_slot.to_bytes(8, 'big')
-    )
-
-
-def aggregate_attestation_record(last_justified_slot: int,
-                                 recent_block_hashes: Iterable[Hash32],
-                                 block: 'BaseBeaconBlock',
-                                 votes: Iterable[Tuple[int, bytes, int]],
-                                 proposer_attestation: 'AttestationRecord',
-                                 cycle_length: int) -> 'AttestationRecord':
-    """
-    Aggregate the votes.
-
-    TODO: Write tests
-    """
-    # Get signing message
-    parent_hashes = get_hashes_to_sign(
-        recent_block_hashes,
-        block,
-        cycle_length,
-    )
-    message = create_signing_message(
-        block.slot_number,
-        parent_hashes,
-        proposer_attestation.shard_id,
-        proposer_attestation.shard_block_hash,
-        last_justified_slot,
-    )
-
-    # Verify
-    # TODO: Verification is very slow, needs to compute in parallel
-    voting_sigs, committee_indices = verify_votes(message, votes)
-
-    # Aggregate the votes
-    bitfield, sigs = aggregate_votes(
-        bitfield=proposer_attestation.bitfield,
-        sigs=proposer_attestation.aggregate_sig,
-        voting_sigs=voting_sigs,
-        voting_committee_indices=committee_indices
-    )
-
-    return proposer_attestation.copy(
-        bitfield=bitfield,
-        sigs=bls.aggregate_sigs(sigs),
-    )
-
-
-def verify_votes(
-        message: bytes,
-        votes: Iterable[Tuple[int, bytes, int]]) -> Tuple[Tuple[bytes, ...], Tuple[int, ...]]:
-    """
-    Verify the given votes
-    """
-    sigs_with_committe_info = tuple(
-        (sig, committee_index)
-        for (committee_index, sig, public_key)
-        in votes
-        if bls.verify(message, public_key, sig)
-    )
-    try:
-        sigs, committee_indices = zip(*sigs_with_committe_info)
-    except ValueError:
-        sigs = tuple()
-        committee_indices = tuple()
-
-    return sigs, committee_indices
-
-
-def aggregate_votes(bitfield: bytes,
-                    sigs: Iterable[bytes],
-                    voting_sigs: Iterable[bytes],
-                    voting_committee_indices: Iterable[int]) -> Tuple[bytes, Tuple[int]]:
-    """
-    Aggregate the votes
-    """
-    # Update the bitfield and append the signatures
-    sigs = tuple(sigs) + tuple(voting_sigs)
-    bitfield = pipe(
-        bitfield,
-        *(
-            set_voted(index=committee_index)
-            for committee_index in voting_committee_indices
-        )
-    )
-
-    return bitfield, bls.aggregate_sigs(sigs)

eth/beacon/state_machines/base.py

@@ -31,10 +31,12 @@
     Configurable,
 )
 
+from eth.beacon.aggregation import (
+    create_signing_message,
+)
 from eth.beacon.block_proposal import BlockProposal
 from eth.beacon.db.chain import BaseBeaconChainDB
 from eth.beacon.helpers import (
-    create_signing_message,
     get_block_committees_info,
     get_hashes_to_sign,
     get_new_recent_block_hashes,

eth/beacon/state_machines/validation.py

@@ -19,8 +19,10 @@
     has_voted,
 )
 
-from eth.beacon.helpers import (
+from eth.beacon.aggregation import (
     create_signing_message,
+)
+from eth.beacon.helpers import (
     get_attestation_indices,
     get_block_committees_info,
     get_signed_parent_hashes,

tests/beacon/test_aggregation.py

@@ -0,0 +1,68 @@
+import pytest
+from hypothesis import (
+    given,
+    settings,
+    strategies as st,
+)
+
+from eth.utils import bls
+from eth.utils.bitfield import (
+    get_empty_bitfield,
+    has_voted,
+)
+from eth.beacon.aggregation import (
+    aggregate_votes,
+    verify_votes,
+)
+
+
+@settings(max_examples=1)
+@given(random=st.randoms())
+@pytest.mark.parametrize(
+    (
+        'votes_count'
+    ),
+    [
+        (0),
+        (9),
+    ],
+)
+def test_aggregate_votes(votes_count, random, privkeys, pubkeys):
+    bit_count = 10
+    pre_bitfield = get_empty_bitfield(bit_count)
+    pre_sigs = ()
+
+    random_votes = random.sample([i for i in range(bit_count)], votes_count)
+    message = b'hello'
+
+    # Get votes: (committee_index, sig, public_key)
+    votes = [
+        (committee_index, bls.sign(message, privkeys[committee_index]), pubkeys[committee_index])
+        for committee_index in random_votes
+    ]
+
+    # Verify
+    sigs, committee_indices = verify_votes(message, votes)
+
+    # Aggregate the votes
+    bitfield, sigs = aggregate_votes(
+        bitfield=pre_bitfield,
+        sigs=pre_sigs,
+        voting_sigs=sigs,
+        voting_committee_indices=committee_indices
+    )
+
+    try:
+        _, _, pubs = zip(*votes)
+    except ValueError:
+        pubs = ()
+
+    voted_index = [
+        committee_index
+        for committee_index in random_votes
+        if has_voted(bitfield, committee_index)
+    ]
+    assert len(voted_index) == len(votes)
+
+    aggregated_pubs = bls.aggregate_pubs(pubs)
+    assert bls.verify(message, aggregated_pubs, sigs)

tests/beacon/test_helpers.py

@@ -14,7 +14,6 @@
 from eth.beacon.types.shard_and_committees import ShardAndCommittee
 from eth.beacon.helpers import (
     _get_element_from_recent_list,
-    aggregate_votes,
     get_attestation_indices,
     get_block_hash,
     get_hashes_from_recent_block_hashes,
@@ -23,7 +22,6 @@
     get_shards_and_committees_for_slot,
     get_signed_parent_hashes,
     get_block_committees_info,
-    verify_votes,
 )
 
 from tests.beacon.helpers import (
@@ -400,55 +398,3 @@ def mock_get_shards_and_committees_for_slot(parent_block,
             block_committees_info.proposer_index_in_committee ==
             result_proposer_index_in_committee
         )
-
-
-@settings(max_examples=1)
-@given(random=st.randoms())
-@pytest.mark.parametrize(
-    (
-        'votes_count'
-    ),
-    [
-        (0),
-        (9),
-    ],
-)
-def test_aggregate_votes(votes_count, random, privkeys, pubkeys):
-    bit_count = 10
-    pre_bitfield = get_empty_bitfield(bit_count)
-    pre_sigs = ()
-
-    random_votes = random.sample([i for i in range(bit_count)], votes_count)
-    message = b'hello'
-
-    # Get votes: (committee_index, sig, public_key)
-    votes = [
-        (committee_index, bls.sign(message, privkeys[committee_index]), pubkeys[committee_index])
-        for committee_index in random_votes
-    ]
-
-    # Verify
-    sigs, committee_indices = verify_votes(message, votes)
-
-    # Aggregate the votes
-    bitfield, sigs = aggregate_votes(
-        bitfield=pre_bitfield,
-        sigs=pre_sigs,
-        voting_sigs=sigs,
-        voting_committee_indices=committee_indices
-    )
-
-    try:
-        _, _, pubs = zip(*votes)
-    except ValueError:
-        pubs = ()
-
-    voted_index = [
-        committee_index
-        for committee_index in random_votes
-        if has_voted(bitfield, committee_index)
-    ]
-    assert len(voted_index) == len(votes)
-
-    aggregated_pubs = bls.aggregate_pubs(pubs)
-    assert bls.verify(message, aggregated_pubs, sigs)