From af7ba6ff70241589df3159e35e6fa36ce297ac50 Mon Sep 17 00:00:00 2001
From: ethDreamer <37123614+ethDreamer@users.noreply.github.com>
Date: Sun, 12 May 2024 04:24:19 -0500
Subject: [PATCH] Fix Electra Fork Choice Tests (#5764)

---
 beacon_node/store/src/consensus_context.rs    |  5 +-
 .../state_processing/src/consensus_context.rs | 61 ++++++-------------
 2 files changed, 22 insertions(+), 44 deletions(-)

diff --git a/beacon_node/store/src/consensus_context.rs b/beacon_node/store/src/consensus_context.rs
index 409b364992d..281106d9aaa 100644
--- a/beacon_node/store/src/consensus_context.rs
+++ b/beacon_node/store/src/consensus_context.rs
@@ -1,7 +1,7 @@
 use ssz_derive::{Decode, Encode};
 use state_processing::ConsensusContext;
 use std::collections::HashMap;
-use types::{AttestationData, BitList, EthSpec, Hash256, IndexedAttestation, Slot};
+use types::{EthSpec, Hash256, IndexedAttestation, Slot};
 
 /// The consensus context is stored on disk as part of the data availability overflow cache.
 ///
@@ -21,8 +21,7 @@ pub struct OnDiskConsensusContext<E: EthSpec> {
     ///
     /// They are not part of the on-disk format.
     #[ssz(skip_serializing, skip_deserializing)]
-    indexed_attestations:
-        HashMap<(AttestationData, BitList<E::MaxValidatorsPerSlot>), IndexedAttestation<E>>,
+    indexed_attestations: HashMap<Hash256, IndexedAttestation<E>>,
 }
 
 impl<E: EthSpec> OnDiskConsensusContext<E> {
diff --git a/consensus/state_processing/src/consensus_context.rs b/consensus/state_processing/src/consensus_context.rs
index 6b0f2eb8fee..b0eaf3422d3 100644
--- a/consensus/state_processing/src/consensus_context.rs
+++ b/consensus/state_processing/src/consensus_context.rs
@@ -4,9 +4,8 @@ use crate::EpochCacheError;
 use std::collections::{hash_map::Entry, HashMap};
 use tree_hash::TreeHash;
 use types::{
-    AbstractExecPayload, AttestationData, AttestationRef, BeaconState, BeaconStateError, BitList,
-    ChainSpec, Epoch, EthSpec, Hash256, IndexedAttestation, IndexedAttestationRef,
-    SignedBeaconBlock, Slot,
+    AbstractExecPayload, AttestationRef, BeaconState, BeaconStateError, ChainSpec, Epoch, EthSpec,
+    Hash256, IndexedAttestation, IndexedAttestationRef, SignedBeaconBlock, Slot,
 };
 
 #[derive(Debug, PartialEq, Clone)]
@@ -22,8 +21,7 @@ pub struct ConsensusContext<E: EthSpec> {
     /// Block root of the block at `slot`.
     pub current_block_root: Option<Hash256>,
     /// Cache of indexed attestations constructed during block processing.
-    pub indexed_attestations:
-        HashMap<(AttestationData, BitList<E::MaxValidatorsPerSlot>), IndexedAttestation<E>>,
+    pub indexed_attestations: HashMap<Hash256, IndexedAttestation<E>>,
 }
 
 #[derive(Debug, PartialEq, Clone)]
@@ -154,41 +152,25 @@ impl<E: EthSpec> ConsensusContext<E> {
         state: &BeaconState<E>,
         attestation: AttestationRef<'a, E>,
     ) -> Result<IndexedAttestationRef<E>, BlockOperationError<AttestationInvalid>> {
+        let key = attestation.tree_hash_root();
         match attestation {
-            AttestationRef::Base(attn) => {
-                let extended_aggregation_bits = attn
-                    .extend_aggregation_bits()
-                    .map_err(BeaconStateError::from)?;
-
-                let key = (attn.data.clone(), extended_aggregation_bits);
-
-                match self.indexed_attestations.entry(key) {
-                    Entry::Occupied(occupied) => Ok(occupied.into_mut()),
-                    Entry::Vacant(vacant) => {
-                        let committee =
-                            state.get_beacon_committee(attn.data.slot, attn.data.index)?;
-                        let indexed_attestation = attesting_indices_base::get_indexed_attestation(
-                            committee.committee,
-                            attn,
-                        )?;
-                        Ok(vacant.insert(indexed_attestation))
-                    }
+            AttestationRef::Base(attn) => match self.indexed_attestations.entry(key) {
+                Entry::Occupied(occupied) => Ok(occupied.into_mut()),
+                Entry::Vacant(vacant) => {
+                    let committee = state.get_beacon_committee(attn.data.slot, attn.data.index)?;
+                    let indexed_attestation =
+                        attesting_indices_base::get_indexed_attestation(committee.committee, attn)?;
+                    Ok(vacant.insert(indexed_attestation))
                 }
-            }
-            AttestationRef::Electra(attn) => {
-                let key = (attn.data.clone(), attn.aggregation_bits.clone());
-
-                match self.indexed_attestations.entry(key) {
-                    Entry::Occupied(occupied) => Ok(occupied.into_mut()),
-                    Entry::Vacant(vacant) => {
-                        let indexed_attestation =
-                            attesting_indices_electra::get_indexed_attestation_from_state(
-                                state, attn,
-                            )?;
-                        Ok(vacant.insert(indexed_attestation))
-                    }
+            },
+            AttestationRef::Electra(attn) => match self.indexed_attestations.entry(key) {
+                Entry::Occupied(occupied) => Ok(occupied.into_mut()),
+                Entry::Vacant(vacant) => {
+                    let indexed_attestation =
+                        attesting_indices_electra::get_indexed_attestation_from_state(state, attn)?;
+                    Ok(vacant.insert(indexed_attestation))
                 }
-            }
+            },
         }
         .map(|indexed_attestation| (*indexed_attestation).to_ref())
     }
@@ -200,10 +182,7 @@ impl<E: EthSpec> ConsensusContext<E> {
     #[must_use]
     pub fn set_indexed_attestations(
         mut self,
-        attestations: HashMap<
-            (AttestationData, BitList<E::MaxValidatorsPerSlot>),
-            IndexedAttestation<E>,
-        >,
+        attestations: HashMap<Hash256, IndexedAttestation<E>>,
     ) -> Self {
         self.indexed_attestations = attestations;
         self