Electra: Remaining Consensus Data Structures (#5712)

* Attestation superstruct changes for EIP 7549 (#5644) * update * experiment * superstruct changes * revert * superstruct changes * fix tests * indexed attestation * indexed attestation superstruct * updated TODOs * `superstruct` the `AttesterSlashing` (#5636) * `superstruct` Attester Fork Variants * Push a little further * Deal with Encode / Decode of AttesterSlashing * not so sure about this.. * Stop Encode/Decode Bounds from Propagating Out * Tons of Changes.. * More Conversions to AttestationRef * Add AsReference trait (#15) * Add AsReference trait * Fix some snafus * Got it Compiling! :D * Got Tests Building * Get beacon chain tests compiling --------- Co-authored-by: Michael Sproul <micsproul@gmail.com> * Merge remote-tracking branch 'upstream/unstable' into electra_attestation_changes * Make EF Tests Fork-Agnostic (#5713) * Finish EF Test Fork Agnostic (#5714) * Superstruct `AggregateAndProof` (#5715) * Upgrade `superstruct` to `0.8.0` * superstruct `AggregateAndProof` * Merge remote-tracking branch 'sigp/unstable' into electra_attestation_changes * cargo fmt * Merge pull request #5726 from realbigsean/electra_attestation_changes Merge unstable into Electra attestation changes * EIP7549 `get_attestation_indices` (#5657) * get attesting indices electra impl * fmt * get tests to pass * fmt * fix some beacon chain tests * fmt * fix slasher test * fmt got me again * fix more tests * fix tests * Some small changes (#5739) * cargo fmt (#5740) * Sketch op pool changes * fix get attesting indices (#5742) * fix get attesting indices * better errors * fix compile * only get committee index once * Ef test fixes (#5753) * attestation related ef test fixes * delete commented out stuff * Fix Aggregation Pool for Electra (#5754) * Fix Aggregation Pool for Electra * Remove Outdated Interface * fix ssz (#5755) * Get `electra_op_pool` up to date (#5756) * fix get attesting indices (#5742) * fix get attesting indices * better errors * fix compile * only get committee index once * Ef test fixes (#5753) * attestation related ef test fixes * delete commented out stuff * Fix Aggregation Pool for Electra (#5754) * Fix Aggregation Pool for Electra * Remove Outdated Interface * fix ssz (#5755) --------- Co-authored-by: realbigsean <sean@sigmaprime.io> * Revert "Get `electra_op_pool` up to date (#5756)" (#5757) This reverts commit ab9e58aa3d. * Merge branch 'electra_attestation_changes' of https://github.com/sigp/lighthouse into electra_op_pool * Compute on chain aggregate impl (#5752) * add compute_on_chain_agg impl to op pool changes * fmt * get op pool tests to pass * update the naive agg pool interface (#5760) * Fix bugs in cross-committee aggregation * Add comment to max cover optimisation * Fix assert * Merge pull request #5749 from sigp/electra_op_pool Optimise Electra op pool aggregation * update committee offset * Fix Electra Fork Choice Tests (#5764) * Subscribe to the correct subnets for electra attestations (#5782) * subscribe to the correct att subnets for electra * subscribe to the correct att subnets for electra * cargo fmt * fix slashing handling * Merge remote-tracking branch 'upstream/unstable' * Send unagg attestation based on fork * Publish all aggregates * just one more check bro plz.. * Merge pull request #5832 from ethDreamer/electra_attestation_changes_merge_unstable Merge `unstable` into `electra_attestation_changes` * Merge pull request #5835 from realbigsean/fix-validator-logic Fix validator logic * Merge pull request #5816 from realbigsean/electra-attestation-slashing-handling Electra slashing handling * Electra attestation changes rm decode impl (#5856) * Remove Crappy Decode impl for Attestation * Remove Inefficient Attestation Decode impl * Implement Schema Upgrade / Downgrade * Update beacon_node/beacon_chain/src/schema_change/migration_schema_v20.rs Co-authored-by: Michael Sproul <micsproul@gmail.com> --------- Co-authored-by: Michael Sproul <micsproul@gmail.com> * Fix failing attestation tests and misc electra attestation cleanup (#5810) * - get attestation related beacon chain tests to pass - observed attestations are now keyed off of data + committee index - rename op pool attestationref to compactattestationref - remove unwraps in agg pool and use options instead - cherry pick some changes from ef-tests-electra * cargo fmt * fix failing test * Revert dockerfile changes * make committee_index return option * function args shouldnt be a ref to attestation ref * fmt * fix dup imports --------- Co-authored-by: realbigsean <seananderson33@GMAIL.com> * fix some todos (#5817) * Merge branch 'unstable' of https://github.com/sigp/lighthouse into electra_attestation_changes * add consolidations to merkle calc for inclusion proof * Remove Duplicate KZG Commitment Merkle Proof Code (#5874) * Remove Duplicate KZG Commitment Merkle Proof Code * s/tree_lists/fields/ * Merge branch 'unstable' of https://github.com/sigp/lighthouse into electra_attestation_changes * fix compile * Fix slasher tests (#5906) * Fix electra tests * Add electra attestations to double vote tests * Update superstruct to 0.8 * Merge remote-tracking branch 'origin/unstable' into electra_attestation_changes * Small cleanup in slasher tests * Clean up Electra observed aggregates (#5929) * Use consistent key in observed_attestations * Remove unwraps from observed aggregates * Merge branch 'unstable' of https://github.com/sigp/lighthouse into electra_attestation_changes * De-dup attestation constructor logic * Remove unwraps in Attestation construction * Dedup match_attestation_data * Remove outdated TODO * Use ForkName Ord in fork-choice tests * Use ForkName Ord in BeaconBlockBody * Make to_electra not fallible * Remove TestRandom impl for IndexedAttestation * Remove IndexedAttestation faulty Decode impl * Drop TestRandom impl * Add PendingAttestationInElectra * Indexed att on disk (#35) * indexed att on disk * fix lints * Update slasher/src/migrate.rs Co-authored-by: ethDreamer <37123614+ethDreamer@users.noreply.github.com> --------- Co-authored-by: Lion - dapplion <35266934+dapplion@users.noreply.github.com> Co-authored-by: ethDreamer <37123614+ethDreamer@users.noreply.github.com> * add electra fork enabled fn to ForkName impl (#36) * add electra fork enabled fn to ForkName impl * remove inadvertent file * Update common/eth2/src/types.rs Co-authored-by: ethDreamer <37123614+ethDreamer@users.noreply.github.com> * Dedup attestation constructor logic in attester cache * Use if let Ok for committee_bits * Dedup Attestation constructor code * Diff reduction in tests * Fix beacon_chain tests * Diff reduction * Use Ord for ForkName in pubsub * Resolve into_attestation_and_indices todo * Remove stale TODO * Fix beacon_chain tests * Test spec invariant * Use electra_enabled in pubsub * Remove get_indexed_attestation_from_signed_aggregate * Use ok_or instead of if let else * committees are sorted * remove dup method `get_indexed_attestation_from_committees` * Merge pull request #5940 from dapplion/electra_attestation_changes_lionreview Electra attestations #5712 review * update default persisted op pool deserialization * ensure aggregate and proof uses serde untagged on ref * Fork aware ssz static attestation tests * Electra attestation changes from Lions review (#5971) * dedup/cleanup and remove unneeded hashset use * remove irrelevant TODOs * Merge branch 'unstable' of https://github.com/sigp/lighthouse into electra_attestation_changes * Electra attestation changes sean review (#5972) * instantiate empty bitlist in unreachable code * clean up error conversion * fork enabled bool cleanup * remove a couple todos * return bools instead of options in `aggregate` and use the result * delete commented out code * use map macros in simple transformations * remove signers_disjoint_from * get ef tests compiling * get ef tests compiling * update intentionally excluded files * Avoid changing slasher schema for Electra * Delete slasher schema v4 * Fix clippy * Fix compilation of beacon_chain tests * Update database.rs * Add electra lightclient types * Update slasher/src/database.rs * fix imports * Merge pull request #5980 from dapplion/electra-lightclient Add electra lightclient types * Merge pull request #5975 from michaelsproul/electra-slasher-no-migration Avoid changing slasher schema for Electra * Update beacon_node/beacon_chain/src/attestation_verification.rs * Update beacon_node/beacon_chain/src/attestation_verification.rs
2026-04-18 21:38:31 +00:00 · 2024-06-24 16:08:07 -05:00
parent 758b58c9e9
commit c52c598f69
118 changed files with 5076 additions and 1741 deletions
--- a/beacon_node/beacon_chain/src/naive_aggregation_pool.rs
+++ b/beacon_node/beacon_chain/src/naive_aggregation_pool.rs
@@ -1,14 +1,124 @@
 use crate::metrics;
+use crate::observed_aggregates::AsReference;
+use itertools::Itertools;
+use smallvec::SmallVec;
 use std::collections::HashMap;
-use tree_hash::TreeHash;
+use tree_hash::{MerkleHasher, TreeHash, TreeHashType};
 use types::consts::altair::SYNC_COMMITTEE_SUBNET_COUNT;
 use types::slot_data::SlotData;
 use types::sync_committee_contribution::SyncContributionData;
-use types::{Attestation, AttestationData, EthSpec, Hash256, Slot, SyncCommitteeContribution};
+use types::{
+    Attestation, AttestationData, AttestationRef, CommitteeIndex, EthSpec, Hash256, Slot,
+    SyncCommitteeContribution,
+};

-type AttestationDataRoot = Hash256;
+type AttestationKeyRoot = Hash256;
 type SyncDataRoot = Hash256;

+/// Post-Electra, we need a new key for Attestations that includes the committee index
+#[derive(Debug, Clone, PartialEq)]
+pub struct AttestationKey {
+    data_root: Hash256,
+    committee_index: Option<CommitteeIndex>,
+    slot: Slot,
+}
+
+// A custom implementation of `TreeHash` such that:
+//     AttestationKey(data, None).tree_hash_root() == data.tree_hash_root()
+//     AttestationKey(data, Some(index)).tree_hash_root() == (data, index).tree_hash_root()
+// This is necessary because pre-Electra, the validator will ask for the tree_hash_root()
+// of the `AttestationData`
+impl TreeHash for AttestationKey {
+    fn tree_hash_type() -> TreeHashType {
+        TreeHashType::Container
+    }
+
+    fn tree_hash_packed_encoding(&self) -> SmallVec<[u8; 32]> {
+        unreachable!("AttestationKey should never be packed.")
+    }
+
+    fn tree_hash_packing_factor() -> usize {
+        unreachable!("AttestationKey should never be packed.")
+    }
+
+    fn tree_hash_root(&self) -> Hash256 {
+        match self.committee_index {
+            None => self.data_root, // Return just the data root if no committee index is present
+            Some(index) => {
+                // Combine the hash of the data with the hash of the index
+                let mut hasher = MerkleHasher::with_leaves(2);
+                hasher
+                    .write(self.data_root.as_bytes())
+                    .expect("should write data hash");
+                hasher
+                    .write(&index.to_le_bytes())
+                    .expect("should write index");
+                hasher.finish().expect("should give tree hash")
+            }
+        }
+    }
+}
+
+impl AttestationKey {
+    pub fn from_attestation_ref<E: EthSpec>(attestation: AttestationRef<E>) -> Result<Self, Error> {
+        let slot = attestation.data().slot;
+        match attestation {
+            AttestationRef::Base(att) => Ok(Self {
+                data_root: att.data.tree_hash_root(),
+                committee_index: None,
+                slot,
+            }),
+            AttestationRef::Electra(att) => {
+                let committee_index = att
+                    .committee_bits
+                    .iter()
+                    .enumerate()
+                    .filter_map(|(i, bit)| if bit { Some(i) } else { None })
+                    .at_most_one()
+                    .map_err(|_| Error::MoreThanOneCommitteeBitSet)?
+                    .ok_or(Error::NoCommitteeBitSet)?;
+
+                Ok(Self {
+                    data_root: att.data.tree_hash_root(),
+                    committee_index: Some(committee_index as u64),
+                    slot,
+                })
+            }
+        }
+    }
+
+    pub fn new_base(data: &AttestationData) -> Self {
+        let slot = data.slot;
+        Self {
+            data_root: data.tree_hash_root(),
+            committee_index: None,
+            slot,
+        }
+    }
+
+    pub fn new_electra(slot: Slot, data_root: Hash256, committee_index: CommitteeIndex) -> Self {
+        Self {
+            data_root,
+            committee_index: Some(committee_index),
+            slot,
+        }
+    }
+
+    pub fn new_base_from_slot_and_root(slot: Slot, data_root: Hash256) -> Self {
+        Self {
+            data_root,
+            committee_index: None,
+            slot,
+        }
+    }
+}
+
+impl SlotData for AttestationKey {
+    fn get_slot(&self) -> Slot {
+        self.slot
+    }
+}
+
 /// The number of slots that will be stored in the pool.
 ///
 /// For example, if `SLOTS_RETAINED == 3` and the pool is pruned at slot `6`, then all items
@@ -46,6 +156,10 @@ pub enum Error {
    /// The given `aggregation_bits` field had more than one signature. The number of
    /// signatures found is included.
    MoreThanOneAggregationBitSet(usize),
+    /// The electra attestation has more than one committee bit set
+    MoreThanOneCommitteeBitSet,
+    /// The electra attestation has NO committee bit set
+    NoCommitteeBitSet,
    /// We have reached the maximum number of unique items that can be stored in a
    /// slot. This is a DoS protection function.
    ReachedMaxItemsPerSlot(usize),
@@ -59,12 +173,15 @@ pub enum Error {

 /// Implemented for items in the `NaiveAggregationPool`. Requires that items implement `SlotData`,
 /// which means they have an associated slot. This handles aggregation of items that are inserted.
-pub trait AggregateMap {
+pub trait AggregateMap
+where
+    for<'a> <Self::Value as AsReference>::Reference<'a>: SlotData,
+{
    /// `Key` should be a hash of `Data`.
    type Key;

    /// The item stored in the map
-    type Value: Clone + SlotData;
+    type Value: Clone + SlotData + AsReference;

    /// The unique fields of `Value`, hashed to create `Key`.
    type Data: SlotData;
@@ -73,7 +190,10 @@ pub trait AggregateMap {
    fn new(initial_capacity: usize) -> Self;

    /// Insert a `Value` into `Self`, returning a result.
-    fn insert(&mut self, value: &Self::Value) -> Result<InsertOutcome, Error>;
+    fn insert(
+        &mut self,
+        value: <Self::Value as AsReference>::Reference<'_>,
+    ) -> Result<InsertOutcome, Error>;

    /// Get a `Value` from `Self` based on `Data`.
    fn get(&self, data: &Self::Data) -> Option<Self::Value>;
@@ -81,9 +201,6 @@ pub trait AggregateMap {
    /// Get a reference to the inner `HashMap`.
    fn get_map(&self) -> &HashMap<Self::Key, Self::Value>;

-    /// Get a `Value` from `Self` based on `Key`, which is a hash of `Data`.
-    fn get_by_root(&self, root: &Self::Key) -> Option<&Self::Value>;
-
    /// The number of items store in `Self`.
    fn len(&self) -> usize;

@@ -103,13 +220,13 @@ pub trait AggregateMap {
 /// A collection of `Attestation` objects, keyed by their `attestation.data`. Enforces that all
 /// `attestation` are from the same slot.
 pub struct AggregatedAttestationMap<E: EthSpec> {
-    map: HashMap<AttestationDataRoot, Attestation<E>>,
+    map: HashMap<AttestationKeyRoot, Attestation<E>>,
 }

 impl<E: EthSpec> AggregateMap for AggregatedAttestationMap<E> {
-    type Key = AttestationDataRoot;
+    type Key = AttestationKeyRoot;
    type Value = Attestation<E>;
-    type Data = AttestationData;
+    type Data = AttestationKey;

    /// Create an empty collection with the given `initial_capacity`.
    fn new(initial_capacity: usize) -> Self {
@@ -121,48 +238,45 @@ impl<E: EthSpec> AggregateMap for AggregatedAttestationMap<E> {
    /// Insert an attestation into `self`, aggregating it into the pool.
    ///
    /// The given attestation (`a`) must only have one signature.
-    fn insert(&mut self, a: &Self::Value) -> Result<InsertOutcome, Error> {
+    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 = a
-            .aggregation_bits
+        let aggregation_bit = *a
+            .set_aggregation_bits()
            .iter()
-            .enumerate()
-            .filter(|(_i, bit)| *bit)
-            .map(|(i, _bit)| i)
-            .collect::<Vec<_>>();
-
-        let committee_index = set_bits
-            .first()
-            .copied()
+            .at_most_one()
+            .map_err(|iter| Error::MoreThanOneAggregationBitSet(iter.count()))?
            .ok_or(Error::NoAggregationBitsSet)?;

-        if set_bits.len() > 1 {
-            return Err(Error::MoreThanOneAggregationBitSet(set_bits.len()));
-        }
+        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();
-
-        if let Some(existing_attestation) = self.map.get_mut(&attestation_data_root) {
+        if let Some(existing_attestation) = self.map.get_mut(&attestation_key_root) {
            if existing_attestation
-                .aggregation_bits
-                .get(committee_index)
+                .get_aggregation_bit(aggregation_bit)
                .map_err(|_| Error::InconsistentBitfieldLengths)?
            {
-                Ok(InsertOutcome::SignatureAlreadyKnown { committee_index })
+                Ok(InsertOutcome::SignatureAlreadyKnown {
+                    committee_index: aggregation_bit,
+                })
            } else {
                let _timer =
                    metrics::start_timer(&metrics::ATTESTATION_PROCESSING_AGG_POOL_AGGREGATION);
                existing_attestation.aggregate(a);
-                Ok(InsertOutcome::SignatureAggregated { committee_index })
+                Ok(InsertOutcome::SignatureAggregated {
+                    committee_index: aggregation_bit,
+                })
            }
        } else {
            if self.map.len() >= MAX_ATTESTATIONS_PER_SLOT {
                return Err(Error::ReachedMaxItemsPerSlot(MAX_ATTESTATIONS_PER_SLOT));
            }

-            self.map.insert(attestation_data_root, a.clone());
-            Ok(InsertOutcome::NewItemInserted { committee_index })
+            self.map
+                .insert(attestation_key_root, a.clone_as_attestation());
+            Ok(InsertOutcome::NewItemInserted {
+                committee_index: aggregation_bit,
+            })
        }
    }

@@ -177,11 +291,6 @@ impl<E: EthSpec> AggregateMap for AggregatedAttestationMap<E> {
        &self.map
    }

-    /// Returns an aggregated `Attestation` with the given `root`, if any.
-    fn get_by_root(&self, root: &Self::Key) -> Option<&Self::Value> {
-        self.map.get(root)
-    }
-
    fn len(&self) -> usize {
        self.map.len()
    }
@@ -288,11 +397,6 @@ impl<E: EthSpec> AggregateMap for SyncContributionAggregateMap<E> {
        &self.map
    }

-    /// Returns an aggregated `SyncCommitteeContribution` with the given `root`, if any.
-    fn get_by_root(&self, root: &SyncDataRoot) -> Option<&SyncCommitteeContribution<E>> {
-        self.map.get(root)
-    }
-
    fn len(&self) -> usize {
        self.map.len()
    }
@@ -336,12 +440,20 @@ impl<E: EthSpec> AggregateMap for SyncContributionAggregateMap<E> {
 /// `current_slot - SLOTS_RETAINED` will be removed and any future item with a slot lower
 /// than that will also be refused. Pruning is done automatically based upon the items it
 /// receives and it can be triggered manually.
-pub struct NaiveAggregationPool<T: AggregateMap> {
+pub struct NaiveAggregationPool<T>
+where
+    T: AggregateMap,
+    for<'a> <T::Value as AsReference>::Reference<'a>: SlotData,
+{
    lowest_permissible_slot: Slot,
    maps: HashMap<Slot, T>,
 }

-impl<T: AggregateMap> Default for NaiveAggregationPool<T> {
+impl<T> Default for NaiveAggregationPool<T>
+where
+    T: AggregateMap,
+    for<'a> <T::Value as AsReference>::Reference<'a>: SlotData,
+{
    fn default() -> Self {
        Self {
            lowest_permissible_slot: Slot::new(0),
@@ -350,7 +462,11 @@ impl<T: AggregateMap> Default for NaiveAggregationPool<T> {
    }
 }

-impl<T: AggregateMap> NaiveAggregationPool<T> {
+impl<T> NaiveAggregationPool<T>
+where
+    T: AggregateMap,
+    for<'a> <T::Value as AsReference>::Reference<'a>: SlotData,
+{
    /// Insert an item into `self`, aggregating it into the pool.
    ///
    /// The given item must only have one signature and have an
@@ -358,7 +474,10 @@ impl<T: AggregateMap> NaiveAggregationPool<T> {
    ///
    /// The pool may be pruned if the given item has a slot higher than any
    /// previously seen.
-    pub fn insert(&mut self, item: &T::Value) -> Result<InsertOutcome, Error> {
+    pub fn insert(
+        &mut self,
+        item: <T::Value as AsReference>::Reference<'_>,
+    ) -> Result<InsertOutcome, Error> {
        let _timer = T::start_insert_timer();
        let slot = item.get_slot();
        let lowest_permissible_slot = self.lowest_permissible_slot;
@@ -412,13 +531,6 @@ impl<T: AggregateMap> NaiveAggregationPool<T> {
            .and_then(|map| map.get(data))
    }

-    /// Returns an aggregated `T::Value` with the given `slot` and `root`, if any.
-    pub fn get_by_slot_and_root(&self, slot: Slot, root: &T::Key) -> Option<T::Value> {
-        self.maps
-            .get(&slot)
-            .and_then(|map| map.get_by_root(root).cloned())
-    }
-
    /// Iterate all items in all slots of `self`.
    pub fn iter(&self) -> impl Iterator<Item = &T::Value> {
        self.maps.values().flat_map(|map| map.get_map().values())
@@ -467,18 +579,30 @@ mod tests {
    use super::*;
    use ssz_types::BitList;
    use store::BitVector;
+    use tree_hash::TreeHash;
    use types::{
        test_utils::{generate_deterministic_keypair, test_random_instance},
-        Fork, Hash256, SyncCommitteeMessage,
+        Attestation, AttestationBase, AttestationElectra, Fork, Hash256, SyncCommitteeMessage,
    };

    type E = types::MainnetEthSpec;

-    fn get_attestation(slot: Slot) -> Attestation<E> {
-        let mut a: Attestation<E> = test_random_instance();
+    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");
-        a
+        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> {
@@ -521,9 +645,16 @@ mod tests {
    }

    fn unset_attestation_bit(a: &mut Attestation<E>, i: usize) {
-        a.aggregation_bits
-            .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) {
@@ -533,19 +664,19 @@ mod tests {
    }

    fn mutate_attestation_block_root(a: &mut Attestation<E>, block_root: Hash256) {
-        a.data.beacon_block_root = block_root
+        a.data_mut().beacon_block_root = block_root
    }

    fn mutate_attestation_slot(a: &mut Attestation<E>, slot: Slot) {
-        a.data.slot = slot
+        a.data_mut().slot = slot
    }

    fn attestation_block_root_comparator(a: &Attestation<E>, block_root: Hash256) -> bool {
-        a.data.beacon_block_root == block_root
+        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(
@@ -570,6 +701,45 @@ mod tests {
        SyncContributionData::from_contribution(a)
    }

+    #[test]
+    fn attestation_key_tree_hash_tests() {
+        let attestation_base = get_attestation_base(Slot::new(42));
+        // for a base attestation, the tree_hash_root() of the key should be the same as the tree_hash_root() of the data
+        let attestation_key_base = AttestationKey::from_attestation_ref(attestation_base.to_ref())
+            .expect("should create attestation key");
+        assert_eq!(
+            attestation_key_base.tree_hash_root(),
+            attestation_base.data().tree_hash_root()
+        );
+        let mut attestation_electra = get_attestation_electra(Slot::new(42));
+        // for an electra attestation, the tree_hash_root() of the key should be different from the tree_hash_root() of the data
+        let attestation_key_electra =
+            AttestationKey::from_attestation_ref(attestation_electra.to_ref())
+                .expect("should create attestation key");
+        assert_ne!(
+            attestation_key_electra.tree_hash_root(),
+            attestation_electra.data().tree_hash_root()
+        );
+        // for an electra attestation, the tree_hash_root() of the key should be dependent on which committee bit is set
+        let committe_bits = attestation_electra
+            .committee_bits_mut()
+            .expect("should get committee bits");
+        committe_bits
+            .set(0, false)
+            .expect("should set committee bit");
+        committe_bits
+            .set(1, true)
+            .expect("should set committee bit");
+        let new_attestation_key_electra =
+            AttestationKey::from_attestation_ref(attestation_electra.to_ref())
+                .expect("should create attestation key");
+        // this new key should have a different tree_hash_root() than the previous key
+        assert_ne!(
+            attestation_key_electra.tree_hash_root(),
+            new_attestation_key_electra.tree_hash_root()
+        );
+    }
+
    macro_rules! test_suite {
        (
            $mod_name: ident,
@@ -592,10 +762,10 @@ mod tests {
                    let mut a = $get_method_name(Slot::new(0));

                    let mut pool: NaiveAggregationPool<$map_type<E>> =
-                        NaiveAggregationPool::default();
+                        NaiveAggregationPool::<$map_type<E>>::default();

                    assert_eq!(
-                        pool.insert(&a),
+                        pool.insert(a.as_reference()),
                        Err(Error::NoAggregationBitsSet),
                        "should not accept item without any signatures"
                    );
@@ -603,12 +773,12 @@ mod tests {
                    $sign_method_name(&mut a, 0, Hash256::random());

                    assert_eq!(
-                        pool.insert(&a),
+                        pool.insert(a.as_reference()),
                        Ok(InsertOutcome::NewItemInserted { committee_index: 0 }),
                        "should accept new item"
                    );
                    assert_eq!(
-                        pool.insert(&a),
+                        pool.insert(a.as_reference()),
                        Ok(InsertOutcome::SignatureAlreadyKnown { committee_index: 0 }),
                        "should acknowledge duplicate signature"
                    );
@@ -621,7 +791,7 @@ mod tests {
                    $sign_method_name(&mut a, 1, Hash256::random());

                    assert_eq!(
-                        pool.insert(&a),
+                        pool.insert(a.as_reference()),
                        Err(Error::MoreThanOneAggregationBitSet(2)),
                        "should not accept item with multiple signatures"
                    );
@@ -637,15 +807,15 @@ mod tests {
                    $sign_method_name(&mut a_1, 1, genesis_validators_root);

                    let mut pool: NaiveAggregationPool<$map_type<E>> =
-                        NaiveAggregationPool::default();
+                        NaiveAggregationPool::<$map_type<E>>::default();

                    assert_eq!(
-                        pool.insert(&a_0),
+                        pool.insert(a_0.as_reference()),
                        Ok(InsertOutcome::NewItemInserted { committee_index: 0 }),
                        "should accept a_0"
                    );
                    assert_eq!(
-                        pool.insert(&a_1),
+                        pool.insert(a_1.as_reference()),
                        Ok(InsertOutcome::SignatureAggregated { committee_index: 1 }),
                        "should accept a_1"
                    );
@@ -655,7 +825,7 @@ mod tests {
                        .expect("should not error while getting attestation");

                    let mut a_01 = a_0.clone();
-                    a_01.aggregate(&a_1);
+                    a_01.aggregate(a_1.as_reference());

                    assert_eq!(retrieved, a_01, "retrieved item should be aggregated");

@@ -671,7 +841,7 @@ mod tests {
                    $block_root_mutator(&mut a_different, different_root);

                    assert_eq!(
-                        pool.insert(&a_different),
+                        pool.insert(a_different.as_reference()),
                        Ok(InsertOutcome::NewItemInserted { committee_index: 2 }),
                        "should accept a_different"
                    );
@@ -690,7 +860,7 @@ mod tests {
                    $sign_method_name(&mut base, 0, Hash256::random());

                    let mut pool: NaiveAggregationPool<$map_type<E>> =
-                        NaiveAggregationPool::default();
+                        NaiveAggregationPool::<$map_type<E>>::default();

                    for i in 0..SLOTS_RETAINED * 2 {
                        let slot = Slot::from(i);
@@ -698,7 +868,7 @@ mod tests {
                        $slot_mutator(&mut a, slot);

                        assert_eq!(
-                            pool.insert(&a),
+                            pool.insert(a.as_reference()),
                            Ok(InsertOutcome::NewItemInserted { committee_index: 0 }),
                            "should accept new item"
                        );
@@ -739,7 +909,7 @@ mod tests {
                    $sign_method_name(&mut base, 0, Hash256::random());

                    let mut pool: NaiveAggregationPool<$map_type<E>> =
-                        NaiveAggregationPool::default();
+                        NaiveAggregationPool::<$map_type<E>>::default();

                    for i in 0..=$item_limit {
                        let mut a = base.clone();
@@ -747,13 +917,13 @@ mod tests {

                        if i < $item_limit {
                            assert_eq!(
-                                pool.insert(&a),
+                                pool.insert(a.as_reference()),
                                Ok(InsertOutcome::NewItemInserted { committee_index: 0 }),
                                "should accept item below limit"
                            );
                        } else {
                            assert_eq!(
-                                pool.insert(&a),
+                                pool.insert(a.as_reference()),
                                Err(Error::ReachedMaxItemsPerSlot($item_limit)),
                                "should not accept item above limit"
                            );
@@ -765,8 +935,21 @@ mod tests {
    }

    test_suite! {
-        attestation_tests,
-        get_attestation,
+        attestation_tests_base,
+        get_attestation_base,
+        sign_attestation,
+        unset_attestation_bit,
+        mutate_attestation_block_root,
+        mutate_attestation_slot,
+        attestation_block_root_comparator,
+        key_from_attestation,
+        AggregatedAttestationMap,
+        MAX_ATTESTATIONS_PER_SLOT
+    }
+
+    test_suite! {
+        attestation_tests_electra,
+        get_attestation_electra,
        sign_attestation,
        unset_attestation_bit,
        mutate_attestation_block_root,