Fix Aggregation Pool for Electra (sigp#5754)

* Fix Aggregation Pool for Electra

* Remove Outdated Interface

beacon_node/beacon_chain/src/beacon_chain.rs

@@ -1677,17 +1677,14 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
         }
     }
 
+    // TODO(electra): call this function from the new beacon API method
     pub fn get_aggregated_attestation_electra(
         &self,
-        slot: Slot,
-        attestation_data_root: &Hash256,
+        data: &AttestationData,
         committee_index: CommitteeIndex,
     ) -> Result<Option<Attestation<T::EthSpec>>, Error> {
-        let attestation_key = crate::naive_aggregation_pool::AttestationKey::new_electra(
-            slot,
-            *attestation_data_root,
-            committee_index,
-        );
+        let attestation_key =
+            crate::naive_aggregation_pool::AttestationKey::new_electra(data, committee_index);
         if let Some(attestation) = self.naive_aggregation_pool.read().get(&attestation_key) {
             self.filter_optimistic_attestation(attestation)
                 .map(Option::Some)

beacon_node/beacon_chain/src/naive_aggregation_pool.rs

@@ -1,5 +1,7 @@
 use crate::metrics;
 use crate::observed_aggregates::AsReference;
+use itertools::Itertools;
+use smallvec::SmallVec;
 use std::collections::HashMap;
 use tree_hash::{MerkleHasher, TreeHash, TreeHashType};
 use types::consts::altair::SYNC_COMMITTEE_SUBNET_COUNT;
@@ -16,7 +18,7 @@ type SyncDataRoot = Hash256;
 #[derive(Debug, Clone, PartialEq)]
 pub struct AttestationKey {
     data_root: Hash256,
-    committee_index: Option<CommitteeIndex>,
+    committee_index: Option<u64>,
     slot: Slot,
 }
 
@@ -93,9 +95,10 @@ impl AttestationKey {
         }
     }
 
-    pub fn new_electra(slot: Slot, data_root: Hash256, committee_index: CommitteeIndex) -> Self {
+    pub fn new_electra(data: &AttestationData, committee_index: u64) -> Self {
+        let slot = data.slot;
         Self {
-            data_root,
+            data_root: data.tree_hash_root(),
             committee_index: Some(committee_index),
             slot,
         }
@@ -238,36 +241,31 @@ impl<E: EthSpec> AggregateMap for AggregatedAttestationMap<E> {
     fn insert(&mut self, a: AttestationRef<E>) -> Result<InsertOutcome, Error> {
         let _timer = metrics::start_timer(&metrics::ATTESTATION_PROCESSING_AGG_POOL_CORE_INSERT);
 
-        let set_bits = match a {
+        let aggregation_bit = match a {
             AttestationRef::Base(att) => att
                 .aggregation_bits
                 .iter()
                 .enumerate()
-                .filter(|(_i, bit)| *bit)
-                .map(|(i, _bit)| i)
-                .collect::<Vec<_>>(),
+                .filter_map(|(i, bit)| if bit { Some(i) } else { None })
+                .at_most_one()
+                .map_err(|iter| Error::MoreThanOneAggregationBitSet(iter.count()))?
+                .ok_or(Error::NoAggregationBitsSet)?,
             AttestationRef::Electra(att) => att
                 .aggregation_bits
                 .iter()
                 .enumerate()
-                .filter(|(_i, bit)| *bit)
-                .map(|(i, _bit)| i)
-                .collect::<Vec<_>>(),
+                .filter_map(|(i, bit)| if bit { Some(i) } else { None })
+                .at_most_one()
+                .map_err(|iter| Error::MoreThanOneAggregationBitSet(iter.count()))?
+                .ok_or(Error::NoAggregationBitsSet)?,
         };
 
-        let committee_index = set_bits
-            .first()
-            .copied()
-            .ok_or(Error::NoAggregationBitsSet)?;
-
         let attestation_key = AttestationKey::from_attestation_ref(a)?;
-        let attestation_key_root = attestation_key.tree_hash_root();
-
-        let attestation_data_root = a.data().tree_hash_root();
+        let attestation_data_root = attestation_key.tree_hash_root();
 
         if let Some(existing_attestation) = self.map.get_mut(&attestation_data_root) {
             if existing_attestation
-                .get_aggregation_bit(committee_index)
+                .get_aggregation_bit(aggregation_bit)
                 .map_err(|_| Error::InconsistentBitfieldLengths)?
             {
                 Ok(InsertOutcome::SignatureAlreadyKnown {
@@ -288,7 +286,9 @@ impl<E: EthSpec> AggregateMap for AggregatedAttestationMap<E> {
 
             self.map
                 .insert(attestation_data_root, a.clone_as_attestation());
-            Ok(InsertOutcome::NewItemInserted { committee_index })
+            Ok(InsertOutcome::NewItemInserted {
+                committee_index: aggregation_bit,
+            })
         }
     }
 
@@ -599,12 +599,22 @@ mod tests {
 
     type E = types::MainnetEthSpec;
 
-    fn get_attestation(slot: Slot) -> Attestation<E> {
-        let mut a: Attestation<E> = test_random_instance();
-        a.data_mut().slot = slot;
-        *a.aggregation_bits_base_mut().unwrap() =
-            BitList::with_capacity(4).expect("should create bitlist");
-        a
+    fn get_attestation_base(slot: Slot) -> Attestation<E> {
+        let mut a: AttestationBase<E> = test_random_instance();
+        a.data.slot = slot;
+        a.aggregation_bits = BitList::with_capacity(4).expect("should create bitlist");
+        Attestation::Base(a)
+    }
+
+    fn get_attestation_electra(slot: Slot) -> Attestation<E> {
+        let mut a: AttestationElectra<E> = test_random_instance();
+        a.data.slot = slot;
+        a.aggregation_bits = BitList::with_capacity(4).expect("should create bitlist");
+        a.committee_bits = BitVector::new();
+        a.committee_bits
+            .set(0, true)
+            .expect("should set committee bit");
+        Attestation::Electra(a)
     }
 
     fn get_sync_contribution(slot: Slot) -> SyncCommitteeContribution<E> {
@@ -647,10 +657,16 @@ mod tests {
     }
 
     fn unset_attestation_bit(a: &mut Attestation<E>, i: usize) {
-        a.aggregation_bits_base_mut()
-            .unwrap()
-            .set(i, false)
-            .expect("should unset aggregation bit")
+        match a {
+            Attestation::Base(ref mut att) => att
+                .aggregation_bits
+                .set(i, false)
+                .expect("should unset aggregation bit"),
+            Attestation::Electra(ref mut att) => att
+                .aggregation_bits
+                .set(i, false)
+                .expect("should unset aggregation bit"),
+        }
     }
 
     fn unset_sync_contribution_bit(a: &mut SyncCommitteeContribution<E>, i: usize) {
@@ -671,8 +687,8 @@ mod tests {
         a.data().beacon_block_root == block_root
     }
 
-    fn key_from_attestation(a: &Attestation<E>) -> AttestationData {
-        a.data().clone()
+    fn key_from_attestation(a: &Attestation<E>) -> AttestationKey {
+        AttestationKey::from_attestation_ref(a.to_ref()).expect("should create attestation key")
     }
 
     fn mutate_sync_contribution_block_root(

beacon_node/beacon_chain/src/test_utils.rs

@@ -1398,7 +1398,7 @@ where
                         // aggregate locally.
                         let aggregate = self
                             .chain
-                            .get_aggregated_attestation(attestation.data())
+                            .get_aggregated_attestation_base(attestation.data())
                             .unwrap()
                             .unwrap_or_else(|| {
                                 committee_attestations.iter().skip(1).fold(

beacon_node/beacon_chain/tests/payload_invalidation.rs

@@ -1229,10 +1229,7 @@ async fn attesting_to_optimistic_head() {
     let get_aggregated_by_slot_and_root = || {
         rig.harness
             .chain
-            .get_aggregated_attestation_by_slot_and_root(
-                attestation.data().slot,
-                &attestation.data().tree_hash_root(),
-            )
+            .get_aggregated_attestation_base(attestation.data())
     };
 
     /*

consensus/types/src/attestation.rs

@@ -97,6 +97,7 @@ impl<E: EthSpec> Decode for Attestation<E> {
     }
 }
 
+// TODO(electra): think about how to handle fork variants here
 impl<E: EthSpec> TestRandom for Attestation<E> {
     fn random_for_test(rng: &mut impl RngCore) -> Self {
         let aggregation_bits: BitList<E::MaxValidatorsPerCommittee> = BitList::random_for_test(rng);