Fix bugs in cross-committee aggregation

beacon_node/operation_pool/src/attestation_storage.rs

@@ -4,7 +4,7 @@ use std::collections::{BTreeMap, HashMap};
 use types::{
     attestation::{AttestationBase, AttestationElectra},
     superstruct, AggregateSignature, Attestation, AttestationData, BeaconState, BitList, BitVector,
-    Checkpoint, Epoch, EthSpec, Hash256, Slot,
+    Checkpoint, Epoch, EthSpec, Hash256, Slot, Unsigned,
 };
 
 #[derive(Debug, PartialEq, Eq, Hash, Clone, Copy)]
@@ -30,7 +30,6 @@ pub struct CompactIndexedAttestation<E: EthSpec> {
     #[superstruct(only(Electra), partial_getter(rename = "aggregation_bits_electra"))]
     pub aggregation_bits: BitList<E::MaxValidatorsPerSlot>,
     pub signature: AggregateSignature,
-    pub index: u64,
     #[superstruct(only(Electra))]
     pub committee_bits: BitVector<E::MaxCommitteesPerSlot>,
 }
@@ -79,7 +78,6 @@ impl<E: EthSpec> SplitAttestation<E> {
                     attesting_indices,
                     aggregation_bits: attn.aggregation_bits,
                     signature: attestation.signature().clone(),
-                    index: data.index,
                 })
             }
             Attestation::Electra(attn) => {
@@ -87,7 +85,6 @@ impl<E: EthSpec> SplitAttestation<E> {
                     attesting_indices,
                     aggregation_bits: attn.aggregation_bits,
                     signature: attestation.signature().clone(),
-                    index: data.index,
                     committee_bits: attn.committee_bits,
                 })
             }
@@ -182,18 +179,11 @@ impl<E: EthSpec> CompactIndexedAttestation<E> {
             (
                 CompactIndexedAttestation::Electra(this),
                 CompactIndexedAttestation::Electra(other),
-            ) => this.aggregate(other),
+            ) => this.aggregate_same_committee(other),
             // TODO(electra) is a mix of electra and base compact indexed attestations an edge case we need to deal with?
             _ => (),
         }
     }
-
-    pub fn committee_index(&self) -> u64 {
-        match self {
-            CompactIndexedAttestation::Base(att) => att.index,
-            CompactIndexedAttestation::Electra(att) => att.committee_index(),
-        }
-    }
 }
 
 impl<E: EthSpec> CompactIndexedAttestationBase<E> {
@@ -225,14 +215,43 @@ impl<E: EthSpec> CompactIndexedAttestationElectra<E> {
                 .is_zero()
     }
 
-    pub fn aggregate(&mut self, other: &Self) {
+    pub fn aggregate_same_committee(&mut self, other: &Self) {
+        // TODO(electra): remove assert in favour of Result
+        assert_eq!(self.committee_bits, other.committee_bits);
+        self.aggregation_bits = self.aggregation_bits.union(&other.aggregation_bits);
+        self.attesting_indices = self
+            .attesting_indices
+            .drain(..)
+            .merge(other.attesting_indices.iter().copied())
+            .dedup()
+            .collect();
+        self.signature.add_assign_aggregate(&other.signature);
+    }
+
+    pub fn aggregate_with_disjoint_committees(&mut self, other: &Self) {
+        // TODO(electra): remove asserts or use Result
+        assert!(
+            self.committee_bits
+                .intersection(&other.committee_bits)
+                .is_zero(),
+            self.committee_bits,
+            other.committee_bits
+        );
+        // The attestation being aggregated in must only have 1 committee bit set.
+        assert_eq!(other.committee_bits.num_set_bits(), 1);
+        // Check we are aggregating in increasing committee index order (so we can append
+        // aggregation bits).
+        assert!(self.committee_bits.highest_set_bit() < other.committee_bits.highest_set_bit());
+
+        self.committee_bits = self.committee_bits.union(&other.committee_bits);
+        self.aggregation_bits =
+            bitlist_extend(&self.aggregation_bits, &other.aggregation_bits).unwrap();
         self.attesting_indices = self
             .attesting_indices
             .drain(..)
             .merge(other.attesting_indices.iter().copied())
             .dedup()
             .collect();
-        self.aggregation_bits = self.aggregation_bits.union(&other.aggregation_bits);
         self.signature.add_assign_aggregate(&other.signature);
     }
 
@@ -249,6 +268,25 @@ impl<E: EthSpec> CompactIndexedAttestationElectra<E> {
     }
 }
 
+// TODO(electra): upstream this or a more efficient implementation
+fn bitlist_extend<N: Unsigned>(list1: &BitList<N>, list2: &BitList<N>) -> Option<BitList<N>> {
+    let new_length = list1.len() + list2.len();
+    let mut list = BitList::<N>::with_capacity(new_length).ok()?;
+
+    // Copy bits from list1.
+    for (i, bit) in list1.iter().enumerate() {
+        list.set(i, bit).ok()?;
+    }
+
+    // Copy bits from list2, starting from the end of list1.
+    let offset = list1.len();
+    for (i, bit) in list2.iter().enumerate() {
+        list.set(offset + i, bit).ok()?;
+    }
+
+    Some(list)
+}
+
 impl<E: EthSpec> AttestationMap<E> {
     pub fn insert(&mut self, attestation: Attestation<E>, attesting_indices: Vec<u64>) {
         let SplitAttestation {
@@ -279,102 +317,85 @@ impl<E: EthSpec> AttestationMap<E> {
         }
     }
 
+    /// Aggregate Electra attestations for the same attestation data signed by different
+    /// committees.
+    ///
+    /// Non-Electra attestations are left as-is.
     pub fn aggregate_across_committees(&mut self, checkpoint_key: CheckpointKey) {
         let Some(attestation_map) = self.checkpoint_map.get_mut(&checkpoint_key) else {
             return;
         };
-        for (_, compact_indexed_attestations) in attestation_map.attestations.iter_mut() {
+        for compact_indexed_attestations in attestation_map.attestations.values_mut() {
             let unaggregated_attestations = std::mem::take(compact_indexed_attestations);
             let mut aggregated_attestations: Vec<CompactIndexedAttestation<E>> = vec![];
 
             // Aggregate the best attestations for each committee and leave the rest.
-            let mut best_attestations_by_committee: BTreeMap<u64, CompactIndexedAttestation<E>> =
+            let mut best_attestations_by_committee: BTreeMap<
-                BTreeMap::new();
+                u64,
+                CompactIndexedAttestationElectra<E>,
+            > = BTreeMap::new();
 
             for committee_attestation in unaggregated_attestations {
-                // TODO(electra)
+                let mut electra_attestation = match committee_attestation {
-                // compare to best attestations by committee
+                    CompactIndexedAttestation::Electra(att)
-                // could probably use `.entry` here
+                        if att.committee_bits.num_set_bits() == 1 =>
-                if let Some(existing_attestation) =
+                    {
-                    best_attestations_by_committee.get_mut(&committee_attestation.committee_index())
+                        att
-                {
-                    // compare and swap, put the discarded one straight into
-                    // `aggregated_attestations` in case we have room to pack it without
-                    // cross-committee aggregation
-                    if existing_attestation.should_aggregate(&committee_attestation) {
-                        existing_attestation.aggregate(&committee_attestation);
-
-                        best_attestations_by_committee.insert(
-                            committee_attestation.committee_index(),
-                            committee_attestation,
-                        );
-                    } else {
-                        aggregated_attestations.push(committee_attestation);
                     }
+                    CompactIndexedAttestation::Electra(att) => {
+                        // Aggregate already covers multiple committees, leave it as-is.
+                        aggregated_attestations.push(CompactIndexedAttestation::Electra(att));
+                        continue;
+                    }
+                    CompactIndexedAttestation::Base(att) => {
+                        // Leave as-is.
+                        aggregated_attestations.push(CompactIndexedAttestation::Base(att));
+                        continue;
+                    }
+                };
+                let committee_index = electra_attestation.committee_index();
+                if let Some(existing_attestation) =
+                    best_attestations_by_committee.get_mut(&committee_index)
+                {
+                    // Search for the best (most aggregation bits) attestation for this committee
+                    // index.
+                    if electra_attestation.aggregation_bits.num_set_bits()
+                        > existing_attestation.aggregation_bits.num_set_bits()
+                    {
+                        // New attestation is better than the previously known one for this
+                        // committee. Replace it.
+                        std::mem::swap(existing_attestation, &mut electra_attestation);
+                    }
+                    // Put the inferior attestation into the list of aggregated attestations
+                    // without performing any cross-committee aggregation.
+                    aggregated_attestations
+                        .push(CompactIndexedAttestation::Electra(electra_attestation));
                 } else {
-                    best_attestations_by_committee.insert(
+                    // First attestation seen for this committee. Place it in the map
-                        committee_attestation.committee_index(),
+                    // provisionally.
-                        committee_attestation,
+                    best_attestations_by_committee.insert(committee_index, electra_attestation);
-                    );
                 }
             }
 
-            // TODO(electra): aggregate all the best attestations by committee
+            if let Some(on_chain_aggregate) =
-            // (use btreemap sort order to get order by committee index)
+                Self::compute_on_chain_aggregate(best_attestations_by_committee)
-            aggregated_attestations.extend(Self::compute_on_chain_aggregate(
+            {
-                best_attestations_by_committee,
+                aggregated_attestations
-            ));
+                    .push(CompactIndexedAttestation::Electra(on_chain_aggregate));
+            }
 
             *compact_indexed_attestations = aggregated_attestations;
         }
     }
 
-    // TODO(electra) unwraps in this function should be cleaned up
-    // also in general this could be a bit more elegant
     pub fn compute_on_chain_aggregate(
-        mut attestations_by_committee: BTreeMap<u64, CompactIndexedAttestation<E>>,
+        mut attestations_by_committee: BTreeMap<u64, CompactIndexedAttestationElectra<E>>,
-    ) -> Vec<CompactIndexedAttestation<E>> {
+    ) -> Option<CompactIndexedAttestationElectra<E>> {
-        let mut aggregated_attestations = vec![];
+        let (_, mut on_chain_aggregate) = attestations_by_committee.pop_first()?;
-        if let Some((_, on_chain_aggregate)) = attestations_by_committee.pop_first() {
+        for (_, attestation) in attestations_by_committee {
-            match on_chain_aggregate {
+            on_chain_aggregate.aggregate_with_disjoint_committees(&attestation);
-                CompactIndexedAttestation::Base(a) => {
-                    aggregated_attestations.push(CompactIndexedAttestation::Base(a));
-                    aggregated_attestations.extend(
-                        attestations_by_committee
-                            .values()
-                            .map(|a| {
-                                CompactIndexedAttestation::Base(CompactIndexedAttestationBase {
-                                    attesting_indices: a.attesting_indices().clone(),
-                                    aggregation_bits: a.aggregation_bits_base().unwrap().clone(),
-                                    signature: a.signature().clone(),
-                                    index: *a.index(),
-                                })
-                            })
-                            .collect::<Vec<CompactIndexedAttestation<E>>>(),
-                    );
-                }
-                CompactIndexedAttestation::Electra(mut a) => {
-                    for (_, attestation) in attestations_by_committee.iter_mut() {
-                        let new_committee_bits = a
-                            .committee_bits
-                            .union(attestation.committee_bits().unwrap());
-                        a.aggregate(attestation.as_electra().unwrap());
-
-                        a = CompactIndexedAttestationElectra {
-                            attesting_indices: a.attesting_indices.clone(),
-                            aggregation_bits: a.aggregation_bits.clone(),
-                            signature: a.signature.clone(),
-                            index: a.index,
-                            committee_bits: new_committee_bits,
-                        };
-                    }
-
-                    aggregated_attestations.push(CompactIndexedAttestation::Electra(a));
-                }
-            }
         }
-
+        Some(on_chain_aggregate)
-        aggregated_attestations
     }
 
     /// Iterate all attestations matching the given `checkpoint_key`.

beacon_node/operation_pool/src/lib.rs

@@ -287,8 +287,10 @@ impl<E: EthSpec> OperationPool<E> {
         // TODO(electra): Work out how to do this more elegantly. This is a bit of a hack.
         let mut all_attestations = self.attestations.write();
 
-        all_attestations.aggregate_across_committees(prev_epoch_key);
+        if fork_name >= ForkName::Electra {
-        all_attestations.aggregate_across_committees(curr_epoch_key);
+            all_attestations.aggregate_across_committees(prev_epoch_key);
+            all_attestations.aggregate_across_committees(curr_epoch_key);
+        }
 
         let all_attestations = parking_lot::RwLockWriteGuard::downgrade(all_attestations);