diff --git a/beacon_node/beacon_chain/src/beacon_chain.rs b/beacon_node/beacon_chain/src/beacon_chain.rs index ca0dcbb80bb..719b5df76be 100644 --- a/beacon_node/beacon_chain/src/beacon_chain.rs +++ b/beacon_node/beacon_chain/src/beacon_chain.rs @@ -1677,17 +1677,14 @@ impl BeaconChain { } } + // 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>, 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) diff --git a/beacon_node/beacon_chain/src/naive_aggregation_pool.rs b/beacon_node/beacon_chain/src/naive_aggregation_pool.rs index 13b22df2fed..6c4f7cdae72 100644 --- a/beacon_node/beacon_chain/src/naive_aggregation_pool.rs +++ b/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, + committee_index: Option, 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 AggregateMap for AggregatedAttestationMap { fn insert(&mut self, a: AttestationRef) -> Result { 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::>(), + .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::>(), + .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 AggregateMap for AggregatedAttestationMap { 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 { - let mut a: Attestation = 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 { + let mut a: AttestationBase = 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 { + let mut a: AttestationElectra = 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 { @@ -647,10 +657,16 @@ mod tests { } fn unset_attestation_bit(a: &mut Attestation, 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, i: usize) { @@ -671,8 +687,8 @@ mod tests { a.data().beacon_block_root == block_root } - fn key_from_attestation(a: &Attestation) -> AttestationData { - a.data().clone() + fn key_from_attestation(a: &Attestation) -> AttestationKey { + AttestationKey::from_attestation_ref(a.to_ref()).expect("should create attestation key") } fn mutate_sync_contribution_block_root( diff --git a/beacon_node/beacon_chain/src/test_utils.rs b/beacon_node/beacon_chain/src/test_utils.rs index 7122d1837bf..5d9e959cfbb 100644 --- a/beacon_node/beacon_chain/src/test_utils.rs +++ b/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( diff --git a/beacon_node/beacon_chain/tests/payload_invalidation.rs b/beacon_node/beacon_chain/tests/payload_invalidation.rs index c4c1ee19915..30af35c5703 100644 --- a/beacon_node/beacon_chain/tests/payload_invalidation.rs +++ b/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()) }; /* diff --git a/consensus/types/src/attestation.rs b/consensus/types/src/attestation.rs index 21512bb01e6..1fc838400e8 100644 --- a/consensus/types/src/attestation.rs +++ b/consensus/types/src/attestation.rs @@ -97,6 +97,7 @@ impl Decode for Attestation { } } +// TODO(electra): think about how to handle fork variants here impl TestRandom for Attestation { fn random_for_test(rng: &mut impl RngCore) -> Self { let aggregation_bits: BitList = BitList::random_for_test(rng);