From 921f450fe1c36bee94819895a526523d2ff2ca13 Mon Sep 17 00:00:00 2001
From: realbigsean <sean@sigmaprime.io>
Date: Thu, 9 May 2024 21:29:31 -0400
Subject: [PATCH] update the naive agg pool interface (#5760)

---
 beacon_node/beacon_chain/src/beacon_chain.rs          | 11 +++++++----
 .../beacon_chain/src/naive_aggregation_pool.rs        | 10 +++++-----
 2 files changed, 12 insertions(+), 9 deletions(-)

diff --git a/beacon_node/beacon_chain/src/beacon_chain.rs b/beacon_node/beacon_chain/src/beacon_chain.rs
index 719b5df76be..ca0dcbb80bb 100644
--- a/beacon_node/beacon_chain/src/beacon_chain.rs
+++ b/beacon_node/beacon_chain/src/beacon_chain.rs
@@ -1677,14 +1677,17 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
         }
     }
 
-    // TODO(electra): call this function from the new beacon API method
     pub fn get_aggregated_attestation_electra(
         &self,
-        data: &AttestationData,
+        slot: Slot,
+        attestation_data_root: &Hash256,
         committee_index: CommitteeIndex,
     ) -> Result<Option<Attestation<T::EthSpec>>, Error> {
-        let attestation_key =
-            crate::naive_aggregation_pool::AttestationKey::new_electra(data, committee_index);
+        let attestation_key = crate::naive_aggregation_pool::AttestationKey::new_electra(
+            slot,
+            *attestation_data_root,
+            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 6c4f7cdae72..a1c736cd0ef 100644
--- a/beacon_node/beacon_chain/src/naive_aggregation_pool.rs
+++ b/beacon_node/beacon_chain/src/naive_aggregation_pool.rs
@@ -8,7 +8,8 @@ use types::consts::altair::SYNC_COMMITTEE_SUBNET_COUNT;
 use types::slot_data::SlotData;
 use types::sync_committee_contribution::SyncContributionData;
 use types::{
-    Attestation, AttestationData, AttestationRef, EthSpec, Hash256, Slot, SyncCommitteeContribution,
+    Attestation, AttestationData, AttestationRef, CommitteeIndex, EthSpec, Hash256, Slot,
+    SyncCommitteeContribution,
 };
 
 type AttestationKeyRoot = Hash256;
@@ -18,7 +19,7 @@ type SyncDataRoot = Hash256;
 #[derive(Debug, Clone, PartialEq)]
 pub struct AttestationKey {
     data_root: Hash256,
-    committee_index: Option<u64>,
+    committee_index: Option<CommitteeIndex>,
     slot: Slot,
 }
 
@@ -95,10 +96,9 @@ impl AttestationKey {
         }
     }
 
-    pub fn new_electra(data: &AttestationData, committee_index: u64) -> Self {
-        let slot = data.slot;
+    pub fn new_electra(slot: Slot, data_root: Hash256, committee_index: CommitteeIndex) -> Self {
         Self {
-            data_root: data.tree_hash_root(),
+            data_root,
             committee_index: Some(committee_index),
             slot,
         }