From 265788be58cfbcdd7280d17eedc309eb0bef0347 Mon Sep 17 00:00:00 2001
From: Hsiao-Wei Wang <hsiaowei.eth@gmail.com>
Date: Thu, 4 Apr 2024 18:01:08 +0900
Subject: [PATCH] refactor

---
 .../eth2spec/test/helpers/attestations.py     | 75 +++++++++++++------
 1 file changed, 53 insertions(+), 22 deletions(-)

diff --git a/tests/core/pyspec/eth2spec/test/helpers/attestations.py b/tests/core/pyspec/eth2spec/test/helpers/attestations.py
index e57eb8318..567bbd4c4 100644
--- a/tests/core/pyspec/eth2spec/test/helpers/attestations.py
+++ b/tests/core/pyspec/eth2spec/test/helpers/attestations.py
@@ -225,6 +225,51 @@ def get_valid_attestations_at_slot(state, spec, slot_to_attest, participation_fn
         )
 
 
+def _get_aggregate_committee_indices(spec, attestations):
+    """
+    Aggregate all unique committee indices from the given attestations.
+    """
+    all_committee_indices = set()
+    for attestation in attestations:
+        committee_indices = spec.get_committee_indices(attestation.committee_bits)
+        assert len(committee_indices) == 1
+        all_committee_indices.add(committee_indices[0])
+
+    return all_committee_indices
+
+
+def _aggregate_aggregation_bits_and_signatures(spec, state, slot, aggregate, attestations):
+    """
+    Aggregate the aggregation bits and signatures from the attestations,
+    incorporating the calculation of aggregation bits offset directly.
+    """
+    # initialize aggregation bits for the aggregate attestation
+    aggregate.aggregation_bits = get_empty_eip7549_aggregation_bits(
+        spec, state, aggregate.committee_bits, slot)
+
+    signatures = []
+
+    offset = 0
+    attestations = sorted(attestations, key=lambda att: spec.get_committee_indices(att.committee_bits)[0])
+    for attestation in attestations:
+        # retrieve the single committee index for the attestation.
+        committee_index = spec.get_committee_indices(attestation.committee_bits)[0]
+
+        # update the aggregate's aggregation bits based on each attestation.
+        for i, bit in enumerate(attestation.aggregation_bits):
+            aggregate.aggregation_bits[offset + i] = bit
+
+        # collect signatures for aggregation.
+        signatures.append(attestation.signature)
+
+        # update offset
+        committee = spec.get_beacon_committee(state, slot, committee_index)
+        offset += len(committee)
+
+    # aggregate signatures from all attestations.
+    aggregate.signature = bls.Aggregate(signatures)
+
+
 def get_valid_attestation_at_slot(state, spec, slot_to_attest, participation_fn=None, beacon_block_root=None):
     """
     Return the aggregate attestation
@@ -235,34 +280,18 @@ def get_valid_attestation_at_slot(state, spec, slot_to_attest, participation_fn=
         participation_fn=participation_fn,
         beacon_block_root=beacon_block_root,
     ))
-    if len(attestations) == 0:
+    if not attestations:
         return None
+
+    # initialize the aggregate attestation.
     aggregate = spec.Attestation(data=attestations[0].data)
 
     # fill in committee_bits
-    all_committee_indices = []
-    for attestation in attestations:
-        committee_indices_of_attestation = spec.get_committee_indices(attestation.committee_bits)
-        assert len(committee_indices_of_attestation) == 1
-        all_committee_indices += committee_indices_of_attestation
-
-    all_committee_indices = set(all_committee_indices)
+    all_committee_indices = _get_aggregate_committee_indices(spec, attestations)
     for committee_index in all_committee_indices:
         aggregate.committee_bits[committee_index] = True
 
-    # aggregate the attestation data and sigs
-    aggregate.aggregation_bits = get_empty_eip7549_aggregation_bits(
-        spec, state, aggregate.committee_bits, slot_to_attest)
-    for attestation in attestations:
-        committee_indices_of_attestation = spec.get_committee_indices(attestation.committee_bits)
-        assert len(committee_indices_of_attestation) == 1
-        committee_index = committee_indices_of_attestation[0]
-        offset = get_eip7549_aggregation_bits_offset(
-            spec, state, attestation.data.slot, aggregate.committee_bits, committee_index)
-        for i in range(len(attestation.aggregation_bits)):
-            aggregation_bits_index = offset + i
-            aggregate.aggregation_bits[aggregation_bits_index] = attestation.aggregation_bits[i]
-    aggregate.signature = bls.Aggregate([attestation.signature for attestation in attestations])
+    _aggregate_aggregation_bits_and_signatures(spec, state, slot_to_attest, aggregate, attestations)
 
     return aggregate
 
@@ -471,7 +500,9 @@ def get_empty_eip7549_aggregation_bits(spec, state, committee_bits, slot):
 
 
 def get_eip7549_aggregation_bits_offset(spec, state, slot, committee_bits, committee_index):
-    # FIXME: it's not efficient to use it as an aggregator
+    """
+    Calculate the offset for the aggregation bits based on the committee index.
+    """
     committee_indices = spec.get_committee_indices(committee_bits)
     assert committee_index in committee_indices
     offset = 0