From e741caee33bd1fc89c11ed3b71c0a0a2e14ab903 Mon Sep 17 00:00:00 2001
From: dapplion <35266934+dapplion@users.noreply.github.com>
Date: Mon, 26 Feb 2024 14:05:38 +0800
Subject: [PATCH 1/5] Drop beacon_chain pubkey to index map cache

---
 beacon_node/beacon_chain/src/beacon_chain.rs  | 44 ----------
 .../beacon_chain/src/block_verification.rs    |  2 +-
 .../src/sync_committee_verification.rs        | 39 ++++++---
 beacon_node/beacon_chain/src/test_utils.rs    | 16 ++--
 .../src/validator_pubkey_cache.rs             | 37 +--------
 beacon_node/http_api/src/lib.rs               | 82 +++++++++----------
 beacon_node/http_api/src/validator.rs         | 21 -----
 .../per_block_processing/signature_sets.rs    |  8 +-
 consensus/types/src/beacon_state.rs           |  4 +
 lcli/src/transition_blocks.rs                 |  4 +-
 10 files changed, 87 insertions(+), 170 deletions(-)
 delete mode 100644 beacon_node/http_api/src/validator.rs

diff --git a/beacon_node/beacon_chain/src/beacon_chain.rs b/beacon_node/beacon_chain/src/beacon_chain.rs
index 20a93e31e8d..21db349505b 100644
--- a/beacon_node/beacon_chain/src/beacon_chain.rs
+++ b/beacon_node/beacon_chain/src/beacon_chain.rs
@@ -1462,50 +1462,6 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
         self.state_at_slot(self.slot()?, StateSkipConfig::WithStateRoots)
     }
 
-    /// Returns the validator index (if any) for the given public key.
-    ///
-    /// ## Notes
-    ///
-    /// This query uses the `validator_pubkey_cache` which contains _all_ validators ever seen,
-    /// even if those validators aren't included in the head state. It is important to remember
-    /// that just because a validator exists here, it doesn't necessarily exist in all
-    /// `BeaconStates`.
-    ///
-    /// ## Errors
-    ///
-    /// May return an error if acquiring a read-lock on the `validator_pubkey_cache` times out.
-    pub fn validator_index(&self, pubkey: &PublicKeyBytes) -> Result<Option<usize>, Error> {
-        let pubkey_cache = self
-            .validator_pubkey_cache
-            .try_read_for(VALIDATOR_PUBKEY_CACHE_LOCK_TIMEOUT)
-            .ok_or(Error::ValidatorPubkeyCacheLockTimeout)?;
-
-        Ok(pubkey_cache.get_index(pubkey))
-    }
-
-    /// Return the validator indices of all public keys fetched from an iterator.
-    ///
-    /// If any public key doesn't belong to a known validator then an error will be returned.
-    /// We could consider relaxing this by returning `Vec<Option<usize>>` in future.
-    pub fn validator_indices<'a>(
-        &self,
-        validator_pubkeys: impl Iterator<Item = &'a PublicKeyBytes>,
-    ) -> Result<Vec<u64>, Error> {
-        let pubkey_cache = self
-            .validator_pubkey_cache
-            .try_read_for(VALIDATOR_PUBKEY_CACHE_LOCK_TIMEOUT)
-            .ok_or(Error::ValidatorPubkeyCacheLockTimeout)?;
-
-        validator_pubkeys
-            .map(|pubkey| {
-                pubkey_cache
-                    .get_index(pubkey)
-                    .map(|id| id as u64)
-                    .ok_or(Error::ValidatorPubkeyUnknown(*pubkey))
-            })
-            .collect()
-    }
-
     /// Returns the validator pubkey (if any) for the given validator index.
     ///
     /// ## Notes
diff --git a/beacon_node/beacon_chain/src/block_verification.rs b/beacon_node/beacon_chain/src/block_verification.rs
index ac3d3e3ab80..9d2679f5d58 100644
--- a/beacon_node/beacon_chain/src/block_verification.rs
+++ b/beacon_node/beacon_chain/src/block_verification.rs
@@ -2083,7 +2083,7 @@ fn get_signature_verifier<'a, T: BeaconChainTypes>(
 
     let decompressor = move |pk_bytes| {
         // Map compressed pubkey to validator index.
-        let validator_index = validator_pubkey_cache.get_index(pk_bytes)?;
+        let validator_index = state.get_validator_index_readonly(pk_bytes)?;
         // Map validator index to pubkey (respecting guard on unknown validators).
         get_pubkey(validator_index)
     };
diff --git a/beacon_node/beacon_chain/src/sync_committee_verification.rs b/beacon_node/beacon_chain/src/sync_committee_verification.rs
index 5c6710bfd6c..4eec3848cb2 100644
--- a/beacon_node/beacon_chain/src/sync_committee_verification.rs
+++ b/beacon_node/beacon_chain/src/sync_committee_verification.rs
@@ -376,12 +376,22 @@ impl<T: BeaconChainTypes> VerifiedSyncContribution<T> {
             .filter_map(|(pubkey, bit)| bit.then_some(pubkey))
             .collect::<Vec<_>>();
 
+        let participant_indices = {
+            // TODO(lion): okay to get head state here?
+            let head_state = &chain.head_snapshot().beacon_state;
+
+            participant_pubkeys
+                .iter()
+                .map(|pubkey| {
+                    head_state
+                        .get_validator_index_readonly(pubkey)
+                        .ok_or(Error::UnknownValidatorPubkey(*pubkey))
+                })
+                .collect::<Result<Vec<_>, _>>()?
+        };
+
         // Ensure that all signatures are valid.
-        if !verify_signed_aggregate_signatures(
-            chain,
-            &signed_aggregate,
-            participant_pubkeys.as_slice(),
-        )? {
+        if !verify_signed_aggregate_signatures(chain, &signed_aggregate, &participant_indices)? {
             return Err(Error::InvalidSignature);
         }
 
@@ -617,7 +627,7 @@ pub fn verify_propagation_slot_range<S: SlotClock, U: SlotData>(
 pub fn verify_signed_aggregate_signatures<T: BeaconChainTypes>(
     chain: &BeaconChain<T>,
     signed_aggregate: &SignedContributionAndProof<T::EthSpec>,
-    participant_pubkeys: &[PublicKeyBytes],
+    participant_indexes: &[usize],
 ) -> Result<bool, Error> {
     let pubkey_cache = chain
         .validator_pubkey_cache
@@ -651,12 +661,8 @@ pub fn verify_signed_aggregate_signatures<T: BeaconChainTypes>(
         )
         .map_err(BeaconChainError::SignatureSetError)?,
         sync_committee_contribution_signature_set_from_pubkeys::<T::EthSpec, _>(
-            |validator_index| {
-                pubkey_cache
-                    .get_pubkey_from_pubkey_bytes(validator_index)
-                    .map(Cow::Borrowed)
-            },
-            participant_pubkeys,
+            |validator_index| pubkey_cache.get(validator_index).map(Cow::Borrowed),
+            participant_indexes,
             &signed_aggregate.message.contribution.signature,
             signed_aggregate
                 .message
@@ -683,13 +689,20 @@ pub fn verify_sync_committee_message<T: BeaconChainTypes>(
     let signature_setup_timer =
         metrics::start_timer(&metrics::SYNC_MESSAGE_PROCESSING_SIGNATURE_SETUP_TIMES);
 
+    let index = chain
+        .head_snapshot()
+        .beacon_state
+        .pubkey_cache()
+        .get(pubkey_bytes)
+        .ok_or(Error::UnknownValidatorPubkey(*pubkey_bytes))?;
+
     let pubkey_cache = chain
         .validator_pubkey_cache
         .try_read_for(VALIDATOR_PUBKEY_CACHE_LOCK_TIMEOUT)
         .ok_or(BeaconChainError::ValidatorPubkeyCacheLockTimeout)?;
 
     let pubkey = pubkey_cache
-        .get_pubkey_from_pubkey_bytes(pubkey_bytes)
+        .get(index)
         .map(Cow::Borrowed)
         .ok_or(Error::UnknownValidatorPubkey(*pubkey_bytes))?;
 
diff --git a/beacon_node/beacon_chain/src/test_utils.rs b/beacon_node/beacon_chain/src/test_utils.rs
index 6b85c8e4931..fe8d722bb9f 100644
--- a/beacon_node/beacon_chain/src/test_utils.rs
+++ b/beacon_node/beacon_chain/src/test_utils.rs
@@ -1191,11 +1191,9 @@ where
                     .iter()
                     .enumerate()
                     .map(|(subcommittee_position, pubkey)| {
-                        let validator_index = self
-                            .chain
-                            .validator_index(pubkey)
-                            .expect("should find validator index")
-                            .expect("pubkey should exist in the beacon chain");
+                        let validator_index = state
+                            .get_validator_index_readonly(pubkey)
+                            .expect("should find validator index");
 
                         let sync_message = SyncCommitteeMessage::new::<E>(
                             message_slot,
@@ -1403,11 +1401,9 @@ where
                         .unwrap()
                         .iter()
                         .find_map(|pubkey| {
-                            let validator_index = self
-                                .chain
-                                .validator_index(pubkey)
-                                .expect("should find validator index")
-                                .expect("pubkey should exist in the beacon chain");
+                            let validator_index = state
+                                .get_validator_index_readonly(pubkey)
+                                .expect("should find validator index");
 
                             let selection_proof = SyncSelectionProof::new::<E>(
                                 slot,
diff --git a/beacon_node/beacon_chain/src/validator_pubkey_cache.rs b/beacon_node/beacon_chain/src/validator_pubkey_cache.rs
index 00140dd6ec0..51b871b8d31 100644
--- a/beacon_node/beacon_chain/src/validator_pubkey_cache.rs
+++ b/beacon_node/beacon_chain/src/validator_pubkey_cache.rs
@@ -1,7 +1,6 @@
 use crate::errors::BeaconChainError;
 use crate::{BeaconChainTypes, BeaconStore};
 use ssz::{Decode, Encode};
-use std::collections::HashMap;
 use std::convert::TryInto;
 use std::marker::PhantomData;
 use store::{DBColumn, Error as StoreError, StoreItem, StoreOp};
@@ -17,7 +16,6 @@ use types::{BeaconState, Hash256, PublicKey, PublicKeyBytes};
 ///    Decompression is expensive when many keys are involved.
 pub struct ValidatorPubkeyCache<T: BeaconChainTypes> {
     pubkeys: Vec<PublicKey>,
-    indices: HashMap<PublicKeyBytes, usize>,
     pubkey_bytes: Vec<PublicKeyBytes>,
     _phantom: PhantomData<T>,
 }
@@ -32,7 +30,6 @@ impl<T: BeaconChainTypes> ValidatorPubkeyCache<T> {
     ) -> Result<Self, BeaconChainError> {
         let mut cache = Self {
             pubkeys: vec![],
-            indices: HashMap::new(),
             pubkey_bytes: vec![],
             _phantom: PhantomData,
         };
@@ -46,7 +43,6 @@ impl<T: BeaconChainTypes> ValidatorPubkeyCache<T> {
     /// Load the pubkey cache from the given on-disk database.
     pub fn load_from_store(store: BeaconStore<T>) -> Result<Self, BeaconChainError> {
         let mut pubkeys = vec![];
-        let mut indices = HashMap::new();
         let mut pubkey_bytes = vec![];
 
         for validator_index in 0.. {
@@ -57,7 +53,6 @@ impl<T: BeaconChainTypes> ValidatorPubkeyCache<T> {
                     BeaconChainError::ValidatorPubkeyCacheError(format!("{:?}", e))
                 })?);
                 pubkey_bytes.push(pubkey);
-                indices.insert(pubkey, validator_index);
             } else {
                 break;
             }
@@ -65,7 +60,6 @@ impl<T: BeaconChainTypes> ValidatorPubkeyCache<T> {
 
         Ok(ValidatorPubkeyCache {
             pubkeys,
-            indices,
             pubkey_bytes,
             _phantom: PhantomData,
         })
@@ -101,16 +95,11 @@ impl<T: BeaconChainTypes> ValidatorPubkeyCache<T> {
     {
         self.pubkey_bytes.reserve(validator_keys.len());
         self.pubkeys.reserve(validator_keys.len());
-        self.indices.reserve(validator_keys.len());
 
         let mut store_ops = Vec::with_capacity(validator_keys.len());
         for pubkey in validator_keys {
             let i = self.pubkeys.len();
 
-            if self.indices.contains_key(&pubkey) {
-                return Err(BeaconChainError::DuplicateValidatorPublicKey);
-            }
-
             // Stage the new validator key for writing to disk.
             // It will be committed atomically when the block that introduced it is written to disk.
             // Notably it is NOT written while the write lock on the cache is held.
@@ -125,8 +114,6 @@ impl<T: BeaconChainTypes> ValidatorPubkeyCache<T> {
                     .map_err(BeaconChainError::InvalidValidatorPubkeyBytes)?,
             );
             self.pubkey_bytes.push(pubkey);
-
-            self.indices.insert(pubkey, i);
         }
 
         Ok(store_ops)
@@ -137,29 +124,19 @@ impl<T: BeaconChainTypes> ValidatorPubkeyCache<T> {
         self.pubkeys.get(i)
     }
 
-    /// Get the `PublicKey` for a validator with `PublicKeyBytes`.
-    pub fn get_pubkey_from_pubkey_bytes(&self, pubkey: &PublicKeyBytes) -> Option<&PublicKey> {
-        self.get_index(pubkey).and_then(|index| self.get(index))
-    }
-
     /// Get the public key (in bytes form) for a validator with index `i`.
     pub fn get_pubkey_bytes(&self, i: usize) -> Option<&PublicKeyBytes> {
         self.pubkey_bytes.get(i)
     }
 
-    /// Get the index of a validator with `pubkey`.
-    pub fn get_index(&self, pubkey: &PublicKeyBytes) -> Option<usize> {
-        self.indices.get(pubkey).copied()
-    }
-
     /// Returns the number of validators in the cache.
     pub fn len(&self) -> usize {
-        self.indices.len()
+        self.pubkeys.len()
     }
 
     /// Returns `true` if there are no validators in the cache.
     pub fn is_empty(&self) -> bool {
-        self.indices.is_empty()
+        self.pubkeys.is_empty()
     }
 }
 
@@ -226,16 +203,6 @@ mod test {
             if i < validator_count {
                 let pubkey = cache.get(i).expect("pubkey should be present");
                 assert_eq!(pubkey, &keypairs[i].pk, "pubkey should match cache");
-
-                let pubkey_bytes: PublicKeyBytes = pubkey.clone().into();
-
-                assert_eq!(
-                    i,
-                    cache
-                        .get_index(&pubkey_bytes)
-                        .expect("should resolve index"),
-                    "index should match cache"
-                );
             } else {
                 assert_eq!(
                     cache.get(i),
diff --git a/beacon_node/http_api/src/lib.rs b/beacon_node/http_api/src/lib.rs
index a9b245e7987..38dd9f290a4 100644
--- a/beacon_node/http_api/src/lib.rs
+++ b/beacon_node/http_api/src/lib.rs
@@ -25,7 +25,6 @@ mod sync_committees;
 mod task_spawner;
 pub mod test_utils;
 mod ui;
-mod validator;
 mod validator_inclusion;
 mod validators;
 mod version;
@@ -86,7 +85,6 @@ use types::{
     SignedContributionAndProof, SignedValidatorRegistrationData, SignedVoluntaryExit, Slot,
     SyncCommitteeMessage, SyncContributionData,
 };
-use validator::pubkey_to_validator_index;
 use version::{
     add_consensus_version_header, add_ssz_content_type_header,
     execution_optimistic_finalized_fork_versioned_response, inconsistent_fork_rejection,
@@ -771,14 +769,9 @@ pub fn serve<T: BeaconChainTypes>(
                             &chain,
                             |state, execution_optimistic, finalized| {
                                 let index_opt = match &validator_id {
-                                    ValidatorId::PublicKey(pubkey) => pubkey_to_validator_index(
-                                        &chain, state, pubkey,
-                                    )
-                                    .map_err(|e| {
-                                        warp_utils::reject::custom_not_found(format!(
-                                            "unable to access pubkey cache: {e:?}",
-                                        ))
-                                    })?,
+                                    ValidatorId::PublicKey(pubkey) => {
+                                        state.get_validator_index_readonly(pubkey)
+                                    }
                                     ValidatorId::Index(index) => Some(*index as usize),
                                 };
 
@@ -1017,42 +1010,50 @@ pub fn serve<T: BeaconChainTypes>(
              chain: Arc<BeaconChain<T>>,
              query: api_types::SyncCommitteesQuery| {
                 task_spawner.blocking_json_task(Priority::P1, move || {
-                    let (sync_committee, execution_optimistic, finalized) = state_id
+                    let (validators, execution_optimistic, finalized) = state_id
                         .map_state_and_execution_optimistic_and_finalized(
                             &chain,
                             |state, execution_optimistic, finalized| {
                                 let current_epoch = state.current_epoch();
                                 let epoch = query.epoch.unwrap_or(current_epoch);
-                                Ok((
-                                    state
-                                        .get_built_sync_committee(epoch, &chain.spec)
-                                        .cloned()
-                                        .map_err(|e| match e {
-                                            BeaconStateError::SyncCommitteeNotKnown { .. } => {
-                                                warp_utils::reject::custom_bad_request(format!(
-                                                    "state at epoch {} has no \
+                                let sync_committee = state
+                                    .get_built_sync_committee(epoch, &chain.spec)
+                                    .cloned()
+                                    .map_err(|e| match e {
+                                        BeaconStateError::SyncCommitteeNotKnown { .. } => {
+                                            warp_utils::reject::custom_bad_request(format!(
+                                                "state at epoch {} has no \
                                                      sync committee for epoch {}",
-                                                    current_epoch, epoch
-                                                ))
-                                            }
-                                            BeaconStateError::IncorrectStateVariant => {
-                                                warp_utils::reject::custom_bad_request(format!(
-                                                    "state at epoch {} is not activated for Altair",
-                                                    current_epoch,
-                                                ))
-                                            }
-                                            e => warp_utils::reject::beacon_state_error(e),
-                                        })?,
-                                    execution_optimistic,
-                                    finalized,
-                                ))
+                                                current_epoch, epoch
+                                            ))
+                                        }
+                                        BeaconStateError::IncorrectStateVariant => {
+                                            warp_utils::reject::custom_bad_request(format!(
+                                                "state at epoch {} is not activated for Altair",
+                                                current_epoch,
+                                            ))
+                                        }
+                                        e => warp_utils::reject::beacon_state_error(e),
+                                    })?;
+
+                                let validators = sync_committee
+                                    .pubkeys
+                                    .iter()
+                                    .map(|pubkey| {
+                                        state
+                                            .get_validator_index_readonly(pubkey)
+                                            .map(|index| index as u64)
+                                            .ok_or(BeaconChainError::ValidatorPubkeyUnknown(
+                                                *pubkey,
+                                            ))
+                                    })
+                                    .collect::<Result<Vec<_>, _>>()
+                                    .map_err(warp_utils::reject::beacon_chain_error)?;
+
+                                Ok((validators, execution_optimistic, finalized))
                             },
                         )?;
 
-                    let validators = chain
-                        .validator_indices(sync_committee.pubkeys.iter())
-                        .map_err(warp_utils::reject::beacon_chain_error)?;
-
                     let validator_aggregates = validators
                         .chunks_exact(T::EthSpec::sync_subcommittee_size())
                         .map(|indices| api_types::SyncSubcommittee {
@@ -3580,10 +3581,9 @@ pub fn serve<T: BeaconChainTypes>(
                         ) = register_val_data
                             .into_iter()
                             .filter_map(|register_data| {
-                                chain
-                                    .validator_index(&register_data.message.pubkey)
-                                    .ok()
-                                    .flatten()
+                                head_snapshot
+                                    .beacon_state
+                                    .get_validator_index_readonly(&register_data.message.pubkey)
                                     .and_then(|validator_index| {
                                         let validator = head_snapshot
                                             .beacon_state
diff --git a/beacon_node/http_api/src/validator.rs b/beacon_node/http_api/src/validator.rs
deleted file mode 100644
index 7f11ddd8f43..00000000000
--- a/beacon_node/http_api/src/validator.rs
+++ /dev/null
@@ -1,21 +0,0 @@
-use beacon_chain::{BeaconChain, BeaconChainError, BeaconChainTypes};
-use types::{BeaconState, PublicKeyBytes};
-
-/// Uses the `chain.validator_pubkey_cache` to resolve a pubkey to a validator
-/// index and then ensures that the validator exists in the given `state`.
-pub fn pubkey_to_validator_index<T: BeaconChainTypes>(
-    chain: &BeaconChain<T>,
-    state: &BeaconState<T::EthSpec>,
-    pubkey: &PublicKeyBytes,
-) -> Result<Option<usize>, BeaconChainError> {
-    chain
-        .validator_index(pubkey)?
-        .filter(|&index| {
-            state
-                .validators()
-                .get(index)
-                .map_or(false, |v| v.pubkey == *pubkey)
-        })
-        .map(Result::Ok)
-        .transpose()
-}
diff --git a/consensus/state_processing/src/per_block_processing/signature_sets.rs b/consensus/state_processing/src/per_block_processing/signature_sets.rs
index fcd324e9eb1..6b054333dd5 100644
--- a/consensus/state_processing/src/per_block_processing/signature_sets.rs
+++ b/consensus/state_processing/src/per_block_processing/signature_sets.rs
@@ -553,7 +553,7 @@ where
 #[allow(clippy::too_many_arguments)]
 pub fn sync_committee_contribution_signature_set_from_pubkeys<'a, T, F>(
     get_pubkey: F,
-    pubkey_bytes: &[PublicKeyBytes],
+    indices: &[usize],
     signature: &'a AggregateSignature,
     epoch: Epoch,
     beacon_block_root: Hash256,
@@ -563,11 +563,11 @@ pub fn sync_committee_contribution_signature_set_from_pubkeys<'a, T, F>(
 ) -> Result<SignatureSet<'a>>
 where
     T: EthSpec,
-    F: Fn(&PublicKeyBytes) -> Option<Cow<'a, PublicKey>>,
+    F: Fn(usize) -> Option<Cow<'a, PublicKey>>,
 {
     let mut pubkeys = Vec::with_capacity(T::SyncSubcommitteeSize::to_usize());
-    for pubkey in pubkey_bytes {
-        pubkeys.push(get_pubkey(pubkey).ok_or(Error::ValidatorPubkeyUnknown(*pubkey))?);
+    for index in indices {
+        pubkeys.push(get_pubkey(*index).ok_or(Error::ValidatorUnknown(*index as u64))?);
     }
 
     let domain = spec.get_domain(epoch, Domain::SyncCommittee, fork, genesis_validators_root);
diff --git a/consensus/types/src/beacon_state.rs b/consensus/types/src/beacon_state.rs
index 4c0ee1bfa20..bcfc2f687bd 100644
--- a/consensus/types/src/beacon_state.rs
+++ b/consensus/types/src/beacon_state.rs
@@ -500,6 +500,10 @@ impl<T: EthSpec> BeaconState<T> {
         Ok(self.pubkey_cache().get(pubkey))
     }
 
+    pub fn get_validator_index_readonly(&self, pubkey: &PublicKeyBytes) -> Option<usize> {
+        self.pubkey_cache().get(pubkey)
+    }
+
     /// The epoch corresponding to `self.slot()`.
     pub fn current_epoch(&self) -> Epoch {
         self.slot().epoch(T::slots_per_epoch())
diff --git a/lcli/src/transition_blocks.rs b/lcli/src/transition_blocks.rs
index 23b0ae26206..434e6c06118 100644
--- a/lcli/src/transition_blocks.rs
+++ b/lcli/src/transition_blocks.rs
@@ -354,9 +354,11 @@ fn do_transition<T: EthSpec>(
                 .map(Cow::Borrowed)
         };
 
+        // TODO(lion): shit hack, revisit latter
+        let state = pre_state.clone();
         let decompressor = move |pk_bytes| {
             // Map compressed pubkey to validator index.
-            let validator_index = validator_pubkey_cache.get_index(pk_bytes)?;
+            let validator_index = state.get_validator_index_readonly(pk_bytes)?;
             // Map validator index to pubkey (respecting guard on unknown validators).
             get_pubkey(validator_index)
         };

From fc3e22a0b17ceb223c418a1639f75868fbed46a5 Mon Sep 17 00:00:00 2001
From: dapplion <35266934+dapplion@users.noreply.github.com>
Date: Mon, 26 Feb 2024 15:01:46 +0800
Subject: [PATCH 2/5] De-duplicate pubkey_cache between states

---
 Cargo.lock                                    | 26 +++++++++++++++++++
 beacon_node/store/src/hot_cold_store.rs       | 18 ++++++++++---
 .../process_operations.rs                     |  3 +++
 consensus/types/Cargo.toml                    |  3 ++-
 consensus/types/src/beacon_state.rs           | 23 ++++++++++++++--
 .../types/src/beacon_state/pubkey_cache.rs    | 16 ++++++------
 6 files changed, 75 insertions(+), 14 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index aa17e2f4ceb..8e16a45f918 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -318,6 +318,16 @@ version = "1.6.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "bddcadddf5e9015d310179a59bb28c4d4b9920ad0f11e8e14dbadf654890c9a6"
 
+[[package]]
+name = "archery"
+version = "1.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8967cd1cc9e9e1954f644e14fbd6042fe9a37da96c52a67e44a2ac18261f8561"
+dependencies = [
+ "static_assertions",
+ "triomphe",
+]
+
 [[package]]
 name = "ark-ff"
 version = "0.3.0"
@@ -7173,6 +7183,15 @@ dependencies = [
  "winapi",
 ]
 
+[[package]]
+name = "rpds"
+version = "1.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a0e15515d3ce3313324d842629ea4905c25a13f81953eadb88f85516f59290a4"
+dependencies = [
+ "archery",
+]
+
 [[package]]
 name = "rtnetlink"
 version = "0.10.1"
@@ -8960,6 +8979,12 @@ dependencies = [
  "rlp",
 ]
 
+[[package]]
+name = "triomphe"
+version = "0.1.11"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "859eb650cfee7434994602c3a68b25d77ad9e68c8a6cd491616ef86661382eb3"
+
 [[package]]
 name = "try-lock"
 version = "0.2.5"
@@ -9005,6 +9030,7 @@ dependencies = [
  "rand_xorshift",
  "rayon",
  "regex",
+ "rpds",
  "rusqlite",
  "safe_arith",
  "serde",
diff --git a/beacon_node/store/src/hot_cold_store.rs b/beacon_node/store/src/hot_cold_store.rs
index 4bdb0deca33..fe1e61114d6 100644
--- a/beacon_node/store/src/hot_cold_store.rs
+++ b/beacon_node/store/src/hot_cold_store.rs
@@ -21,7 +21,7 @@ use crate::{
     get_key_for_col, ChunkWriter, DBColumn, DatabaseBlock, Error, ItemStore, KeyValueStoreOp,
     PartialBeaconState, StoreItem, StoreOp,
 };
-use itertools::process_results;
+use itertools::{process_results, Itertools};
 use leveldb::iterator::LevelDBIterator;
 use lru::LruCache;
 use parking_lot::{Mutex, RwLock};
@@ -1115,11 +1115,13 @@ impl<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>> HotColdDB<E, Hot, Cold>
             epoch_boundary_state_root,
         }) = self.load_hot_state_summary(state_root)?
         {
-            let boundary_state =
+            let mut boundary_state =
                 get_full_state(&self.hot_db, &epoch_boundary_state_root, &self.spec)?.ok_or(
                     HotColdDBError::MissingEpochBoundaryState(epoch_boundary_state_root),
                 )?;
 
+            self.rebase_caches(&mut boundary_state);
+
             // Optimization to avoid even *thinking* about replaying blocks if we're already
             // on an epoch boundary.
             let state = if slot % E::slots_per_epoch() == 0 {
@@ -1142,6 +1144,14 @@ impl<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>> HotColdDB<E, Hot, Cold>
         }
     }
 
+    /// Rebase state pubkey_cache on some existing state to prevent re-populating the cache on the
+    /// first process_deposit call.
+    fn rebase_caches(&self, state: &mut BeaconState<E>) {
+        if let Ok((_, first_state)) = self.state_cache.lock().iter().exactly_one() {
+            state.rebase_caches_on(first_state);
+        }
+    }
+
     /// Store a pre-finalization state in the freezer database.
     ///
     /// If the state doesn't lie on a restore point boundary then just its summary will be stored.
@@ -1274,7 +1284,7 @@ impl<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>> HotColdDB<E, Hot, Cold>
         }
 
         // If low_state is still None, use load_restore_point_by_index to load the state.
-        let low_state = match low_state {
+        let mut low_state = match low_state {
             Some(state) => state,
             None => self.load_restore_point_by_index(low_restore_point_idx)?,
         };
@@ -1301,6 +1311,8 @@ impl<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>> HotColdDB<E, Hot, Cold>
             &self.spec,
         )?;
 
+        self.rebase_caches(&mut low_state);
+
         let state = self.replay_blocks(
             low_state,
             blocks,
diff --git a/consensus/state_processing/src/per_block_processing/process_operations.rs b/consensus/state_processing/src/per_block_processing/process_operations.rs
index cb24a7ba7ec..7305d9eb490 100644
--- a/consensus/state_processing/src/per_block_processing/process_operations.rs
+++ b/consensus/state_processing/src/per_block_processing/process_operations.rs
@@ -419,6 +419,9 @@ pub fn process_deposit<T: EthSpec>(
         state.validators_mut().push(validator)?;
         state.balances_mut().push(deposit.data.amount)?;
 
+        // Register newly added key to the pubkey_cache
+        state.update_pubkey_cache()?;
+
         // Altair or later initializations.
         if let Ok(previous_epoch_participation) = state.previous_epoch_participation_mut() {
             previous_epoch_participation.push(ParticipationFlags::default())?;
diff --git a/consensus/types/Cargo.toml b/consensus/types/Cargo.toml
index db15f53537e..68e9240e98e 100644
--- a/consensus/types/Cargo.toml
+++ b/consensus/types/Cargo.toml
@@ -52,6 +52,7 @@ serde_json = { workspace = true }
 smallvec = { workspace = true }
 maplit = { workspace = true }
 strum = { workspace = true }
+rpds = "1.1.0"
 
 [dev-dependencies]
 criterion = { workspace = true }
@@ -68,4 +69,4 @@ sqlite = []
 # The `arbitrary-fuzz` feature is a no-op provided for backwards compatibility.
 # For simplicity `Arbitrary` is now derived regardless of the feature's presence.
 arbitrary-fuzz = []
-portable = ["bls/supranational-portable"]
\ No newline at end of file
+portable = ["bls/supranational-portable"]
diff --git a/consensus/types/src/beacon_state.rs b/consensus/types/src/beacon_state.rs
index bcfc2f687bd..0f542f2d2a0 100644
--- a/consensus/types/src/beacon_state.rs
+++ b/consensus/types/src/beacon_state.rs
@@ -497,11 +497,19 @@ impl<T: EthSpec> BeaconState<T> {
     /// otherwise returns `None`.
     pub fn get_validator_index(&mut self, pubkey: &PublicKeyBytes) -> Result<Option<usize>, Error> {
         self.update_pubkey_cache()?;
-        Ok(self.pubkey_cache().get(pubkey))
+        Ok(self.get_validator_index_readonly(pubkey))
     }
 
+    /// pubkey_cache may contain pubkeys from future blocks not yet known to this state. Ignore
+    /// pubkeys that resolve to indexes beyond the current validators list.
     pub fn get_validator_index_readonly(&self, pubkey: &PublicKeyBytes) -> Option<usize> {
-        self.pubkey_cache().get(pubkey)
+        self.pubkey_cache().get(pubkey).and_then(|index| {
+            if index >= self.validators().len() {
+                Some(index)
+            } else {
+                None
+            }
+        })
     }
 
     /// The epoch corresponding to `self.slot()`.
@@ -1840,6 +1848,17 @@ impl<T: EthSpec> BeaconState<T> {
         Ok(sync_committee)
     }
 
+    pub fn rebase_caches_on(&mut self, base: &Self) {
+        // Use pubkey cache from `base` if it contains superior information (likely if our cache is
+        // uninitialized).
+        let pubkey_cache = self.pubkey_cache_mut();
+        let base_pubkey_cache = base.pubkey_cache();
+        // Note: It's okay to clone a cache with pubkeys from future blocks.
+        if pubkey_cache.len() < base_pubkey_cache.len() {
+            *pubkey_cache = base_pubkey_cache.clone();
+        }
+    }
+
     pub fn compute_merkle_proof(
         &mut self,
         generalized_index: usize,
diff --git a/consensus/types/src/beacon_state/pubkey_cache.rs b/consensus/types/src/beacon_state/pubkey_cache.rs
index c56c9077e1a..27b7b86a623 100644
--- a/consensus/types/src/beacon_state/pubkey_cache.rs
+++ b/consensus/types/src/beacon_state/pubkey_cache.rs
@@ -1,20 +1,20 @@
 use crate::*;
-use serde::{Deserialize, Serialize};
-use std::collections::HashMap;
+use rpds::HashTrieMapSync;
 
 type ValidatorIndex = usize;
 
-#[derive(Debug, PartialEq, Clone, Default, Serialize, Deserialize)]
+#[derive(Debug, PartialEq, Clone, Default)]
 pub struct PubkeyCache {
-    /// Maintain the number of keys added to the map. It is not sufficient to just use the HashMap
-    /// len, as it does not increase when duplicate keys are added. Duplicate keys are used during
-    /// testing.
+    /// Maintain the number of keys added to the map. It is not sufficient to just use the
+    /// HashTrieMap len, as it does not increase when duplicate keys are added. Duplicate keys are
+    /// used during testing.
     len: usize,
-    map: HashMap<PublicKeyBytes, ValidatorIndex>,
+    map: HashTrieMapSync<PublicKeyBytes, ValidatorIndex>,
 }
 
 impl PubkeyCache {
     /// Returns the number of validator indices added to the map so far.
+    #[allow(clippy::len_without_is_empty)]
     pub fn len(&self) -> ValidatorIndex {
         self.len
     }
@@ -25,7 +25,7 @@ impl PubkeyCache {
     /// that an index is never skipped.
     pub fn insert(&mut self, pubkey: PublicKeyBytes, index: ValidatorIndex) -> bool {
         if index == self.len {
-            self.map.insert(pubkey, index);
+            self.map.insert_mut(pubkey, index);
             self.len = self
                 .len
                 .checked_add(1)

From bd682be4ce5c0d4ac111853e4e0fbea2d07c5d82 Mon Sep 17 00:00:00 2001
From: dapplion <35266934+dapplion@users.noreply.github.com>
Date: Mon, 26 Feb 2024 15:45:13 +0800
Subject: [PATCH 3/5] Handle key re-additions

---
 .../beacon_chain/src/block_verification.rs    |  3 +-
 .../src/sync_committee_verification.rs        |  8 ++--
 beacon_node/beacon_chain/src/test_utils.rs    | 13 +++++-
 beacon_node/http_api/src/lib.rs               | 18 ++++----
 beacon_node/operation_pool/src/lib.rs         | 12 +++--
 consensus/types/src/beacon_state.rs           | 45 ++++++++++++++-----
 .../types/src/beacon_state/pubkey_cache.rs    | 22 +++++----
 lcli/src/transition_blocks.rs                 |  2 +-
 8 files changed, 84 insertions(+), 39 deletions(-)

diff --git a/beacon_node/beacon_chain/src/block_verification.rs b/beacon_node/beacon_chain/src/block_verification.rs
index 9d2679f5d58..56b212db403 100644
--- a/beacon_node/beacon_chain/src/block_verification.rs
+++ b/beacon_node/beacon_chain/src/block_verification.rs
@@ -2083,7 +2083,8 @@ fn get_signature_verifier<'a, T: BeaconChainTypes>(
 
     let decompressor = move |pk_bytes| {
         // Map compressed pubkey to validator index.
-        let validator_index = state.get_validator_index_readonly(pk_bytes)?;
+        // TODO(lion): How to ensure the cache is updated here?
+        let validator_index = state.get_validator_index_readonly_unchecked(pk_bytes)?;
         // Map validator index to pubkey (respecting guard on unknown validators).
         get_pubkey(validator_index)
     };
diff --git a/beacon_node/beacon_chain/src/sync_committee_verification.rs b/beacon_node/beacon_chain/src/sync_committee_verification.rs
index 4eec3848cb2..9b3c3603f75 100644
--- a/beacon_node/beacon_chain/src/sync_committee_verification.rs
+++ b/beacon_node/beacon_chain/src/sync_committee_verification.rs
@@ -383,11 +383,11 @@ impl<T: BeaconChainTypes> VerifiedSyncContribution<T> {
             participant_pubkeys
                 .iter()
                 .map(|pubkey| {
-                    head_state
-                        .get_validator_index_readonly(pubkey)
-                        .ok_or(Error::UnknownValidatorPubkey(*pubkey))
+                    Ok(head_state
+                        .get_validator_index_readonly(pubkey)?
+                        .ok_or(Error::UnknownValidatorPubkey(*pubkey))?)
                 })
-                .collect::<Result<Vec<_>, _>>()?
+                .collect::<Result<Vec<_>, Error>>()?
         };
 
         // Ensure that all signatures are valid.
diff --git a/beacon_node/beacon_chain/src/test_utils.rs b/beacon_node/beacon_chain/src/test_utils.rs
index fe8d722bb9f..c3bfd838726 100644
--- a/beacon_node/beacon_chain/src/test_utils.rs
+++ b/beacon_node/beacon_chain/src/test_utils.rs
@@ -736,7 +736,16 @@ where
     }
 
     pub fn get_current_state(&self) -> BeaconState<E> {
-        self.chain.head_beacon_state_cloned()
+        self.chain
+            .head_snapshot()
+            .beacon_state
+            .clone_with(CloneConfig {
+                committee_caches: true,
+                pubkey_cache: true,
+                exit_cache: false,
+                tree_hash_cache: false,
+                progressive_balances_cache: false,
+            })
     }
 
     pub fn get_timestamp_at_slot(&self) -> u64 {
@@ -1193,6 +1202,7 @@ where
                     .map(|(subcommittee_position, pubkey)| {
                         let validator_index = state
                             .get_validator_index_readonly(pubkey)
+                            .expect("pubkey cache not updated")
                             .expect("should find validator index");
 
                         let sync_message = SyncCommitteeMessage::new::<E>(
@@ -1403,6 +1413,7 @@ where
                         .find_map(|pubkey| {
                             let validator_index = state
                                 .get_validator_index_readonly(pubkey)
+                                .expect("pubkey cache not updated")
                                 .expect("should find validator index");
 
                             let selection_proof = SyncSelectionProof::new::<E>(
diff --git a/beacon_node/http_api/src/lib.rs b/beacon_node/http_api/src/lib.rs
index 38dd9f290a4..a6dcc8f3b83 100644
--- a/beacon_node/http_api/src/lib.rs
+++ b/beacon_node/http_api/src/lib.rs
@@ -769,9 +769,9 @@ pub fn serve<T: BeaconChainTypes>(
                             &chain,
                             |state, execution_optimistic, finalized| {
                                 let index_opt = match &validator_id {
-                                    ValidatorId::PublicKey(pubkey) => {
-                                        state.get_validator_index_readonly(pubkey)
-                                    }
+                                    ValidatorId::PublicKey(pubkey) => state
+                                        .get_validator_index_readonly(pubkey)
+                                        .map_err(|e| warp_utils::reject::beacon_state_error(e))?,
                                     ValidatorId::Index(index) => Some(*index as usize),
                                 };
 
@@ -1040,14 +1040,14 @@ pub fn serve<T: BeaconChainTypes>(
                                     .pubkeys
                                     .iter()
                                     .map(|pubkey| {
-                                        state
-                                            .get_validator_index_readonly(pubkey)
+                                        Ok(state
+                                            .get_validator_index_readonly(pubkey)?
                                             .map(|index| index as u64)
                                             .ok_or(BeaconChainError::ValidatorPubkeyUnknown(
                                                 *pubkey,
-                                            ))
+                                            ))?)
                                     })
-                                    .collect::<Result<Vec<_>, _>>()
+                                    .collect::<Result<Vec<_>, BeaconChainError>>()
                                     .map_err(warp_utils::reject::beacon_chain_error)?;
 
                                 Ok((validators, execution_optimistic, finalized))
@@ -3583,7 +3583,9 @@ pub fn serve<T: BeaconChainTypes>(
                             .filter_map(|register_data| {
                                 head_snapshot
                                     .beacon_state
-                                    .get_validator_index_readonly(&register_data.message.pubkey)
+                                    .get_validator_index_readonly_unchecked(
+                                        &register_data.message.pubkey,
+                                    )
                                     .and_then(|validator_index| {
                                         let validator = head_snapshot
                                             .beacon_state
diff --git a/beacon_node/operation_pool/src/lib.rs b/beacon_node/operation_pool/src/lib.rs
index 7e1ddb1fd2f..6f4ac374b2a 100644
--- a/beacon_node/operation_pool/src/lib.rs
+++ b/beacon_node/operation_pool/src/lib.rs
@@ -1529,7 +1529,8 @@ mod release_tests {
         let (harness, _) = sync_contribution_test_state::<MainnetEthSpec>(1).await;
 
         let op_pool = OperationPool::<MainnetEthSpec>::new();
-        let state = harness.get_current_state();
+        let mut state = harness.get_current_state();
+        state.update_pubkey_cache().unwrap();
 
         let block_root = *state
             .get_block_root(state.slot() - Slot::new(1))
@@ -1589,7 +1590,8 @@ mod release_tests {
         let (harness, _) = sync_contribution_test_state::<MainnetEthSpec>(1).await;
 
         let op_pool = OperationPool::<MainnetEthSpec>::new();
-        let state = harness.get_current_state();
+        let mut state = harness.get_current_state();
+        state.update_pubkey_cache().unwrap();
         let block_root = *state
             .get_block_root(state.slot() - Slot::new(1))
             .ok()
@@ -1626,7 +1628,8 @@ mod release_tests {
         let (harness, _) = sync_contribution_test_state::<MainnetEthSpec>(1).await;
 
         let op_pool = OperationPool::<MainnetEthSpec>::new();
-        let state = harness.get_current_state();
+        let mut state = harness.get_current_state();
+        state.update_pubkey_cache().unwrap();
         let block_root = *state
             .get_block_root(state.slot() - Slot::new(1))
             .ok()
@@ -1706,7 +1709,8 @@ mod release_tests {
         let (harness, _) = sync_contribution_test_state::<MainnetEthSpec>(1).await;
 
         let op_pool = OperationPool::<MainnetEthSpec>::new();
-        let state = harness.get_current_state();
+        let mut state = harness.get_current_state();
+        state.update_pubkey_cache().unwrap();
         let block_root = *state
             .get_block_root(state.slot() - Slot::new(1))
             .ok()
diff --git a/consensus/types/src/beacon_state.rs b/consensus/types/src/beacon_state.rs
index 0f542f2d2a0..73b243ee7d2 100644
--- a/consensus/types/src/beacon_state.rs
+++ b/consensus/types/src/beacon_state.rs
@@ -85,7 +85,7 @@ pub enum Error {
         index: CommitteeIndex,
     },
     ZeroSlotsPerEpoch,
-    PubkeyCacheInconsistent,
+    PubkeyCacheMissingKey(PublicKeyBytes),
     PubkeyCacheIncomplete {
         cache_len: usize,
         registry_len: usize,
@@ -497,14 +497,40 @@ impl<T: EthSpec> BeaconState<T> {
     /// otherwise returns `None`.
     pub fn get_validator_index(&mut self, pubkey: &PublicKeyBytes) -> Result<Option<usize>, Error> {
         self.update_pubkey_cache()?;
-        Ok(self.get_validator_index_readonly(pubkey))
+        Ok(self.get_validator_index_readonly_unchecked(pubkey))
     }
 
     /// pubkey_cache may contain pubkeys from future blocks not yet known to this state. Ignore
     /// pubkeys that resolve to indexes beyond the current validators list.
-    pub fn get_validator_index_readonly(&self, pubkey: &PublicKeyBytes) -> Option<usize> {
+    pub fn get_validator_index_readonly(
+        &self,
+        pubkey: &PublicKeyBytes,
+    ) -> Result<Option<usize>, Error> {
+        if self.pubkey_cache().len() < self.validators().len() {
+            return Err(Error::PubkeyCacheIncomplete {
+                cache_len: self.pubkey_cache().len(),
+                registry_len: self.validators().len(),
+            });
+        }
+
+        Ok(self.get_validator_index_readonly_unchecked(pubkey))
+    }
+
+    pub fn assert_pubkey_cache(&self) -> Result<(), Error> {
+        if self.pubkey_cache().len() < self.validators().len() {
+            return Err(Error::PubkeyCacheIncomplete {
+                cache_len: self.pubkey_cache().len(),
+                registry_len: self.validators().len(),
+            });
+        }
+        Ok(())
+    }
+
+    /// pubkey_cache may contain pubkeys from future blocks not yet known to this state. Ignore
+    /// pubkeys that resolve to indexes beyond the current validators list.
+    pub fn get_validator_index_readonly_unchecked(&self, pubkey: &PublicKeyBytes) -> Option<usize> {
         self.pubkey_cache().get(pubkey).and_then(|index| {
-            if index >= self.validators().len() {
+            if index < self.validators().len() {
                 Some(index)
             } else {
                 None
@@ -898,7 +924,7 @@ impl<T: EthSpec> BeaconState<T> {
         for pubkey in sync_committee.pubkeys.iter() {
             indices.push(
                 self.get_validator_index(pubkey)?
-                    .ok_or(Error::PubkeyCacheInconsistent)?,
+                    .ok_or(Error::PubkeyCacheMissingKey(*pubkey))?,
             )
         }
         Ok(indices)
@@ -1689,8 +1715,7 @@ impl<T: EthSpec> BeaconState<T> {
 
     /// Updates the pubkey cache, if required.
     ///
-    /// Adds all `pubkeys` from the `validators` which are not already in the cache. Will
-    /// never re-add a pubkey.
+    /// Adds all `pubkeys` from the `validators`. May re-add the same pubkey
     pub fn update_pubkey_cache(&mut self) -> Result<(), Error> {
         let mut pubkey_cache = mem::take(self.pubkey_cache_mut());
         for (i, validator) in self
@@ -1699,10 +1724,8 @@ impl<T: EthSpec> BeaconState<T> {
             .enumerate()
             .skip(pubkey_cache.len())
         {
-            let success = pubkey_cache.insert(validator.pubkey, i);
-            if !success {
-                return Err(Error::PubkeyCacheInconsistent);
-            }
+            // Ok to read pubkeys before EIP-6110
+            pubkey_cache.insert(validator.pubkey, i);
         }
         *self.pubkey_cache_mut() = pubkey_cache;
 
diff --git a/consensus/types/src/beacon_state/pubkey_cache.rs b/consensus/types/src/beacon_state/pubkey_cache.rs
index 27b7b86a623..dfe3feee885 100644
--- a/consensus/types/src/beacon_state/pubkey_cache.rs
+++ b/consensus/types/src/beacon_state/pubkey_cache.rs
@@ -1,5 +1,6 @@
 use crate::*;
 use rpds::HashTrieMapSync;
+use std::cmp::Ordering;
 
 type ValidatorIndex = usize;
 
@@ -24,15 +25,18 @@ impl PubkeyCache {
     /// The added index must equal the number of validators already added to the map. This ensures
     /// that an index is never skipped.
     pub fn insert(&mut self, pubkey: PublicKeyBytes, index: ValidatorIndex) -> bool {
-        if index == self.len {
-            self.map.insert_mut(pubkey, index);
-            self.len = self
-                .len
-                .checked_add(1)
-                .expect("map length cannot exceed usize");
-            true
-        } else {
-            false
+        match index.cmp(&self.len) {
+            Ordering::Greater => false,
+            Ordering::Equal => {
+                self.map.insert_mut(pubkey, index);
+                self.len = self
+                    .len
+                    .checked_add(1)
+                    .expect("map length cannot exceed usize");
+                true
+            }
+            // Ignore inserts for already known keys
+            Ordering::Less => true,
         }
     }
 
diff --git a/lcli/src/transition_blocks.rs b/lcli/src/transition_blocks.rs
index 434e6c06118..6bd53e879aa 100644
--- a/lcli/src/transition_blocks.rs
+++ b/lcli/src/transition_blocks.rs
@@ -358,7 +358,7 @@ fn do_transition<T: EthSpec>(
         let state = pre_state.clone();
         let decompressor = move |pk_bytes| {
             // Map compressed pubkey to validator index.
-            let validator_index = state.get_validator_index_readonly(pk_bytes)?;
+            let validator_index = state.get_validator_index_readonly_unchecked(pk_bytes)?;
             // Map validator index to pubkey (respecting guard on unknown validators).
             get_pubkey(validator_index)
         };

From 5e4c129e271562c1b348d2f29758fbfe558ad257 Mon Sep 17 00:00:00 2001
From: dapplion <35266934+dapplion@users.noreply.github.com>
Date: Mon, 26 Feb 2024 17:39:03 +0800
Subject: [PATCH 4/5] fix ci

---
 .../beacon_chain/src/sync_committee_verification.rs      | 4 ++--
 .../beacon_chain/tests/sync_committee_verification.rs    | 9 ++++-----
 2 files changed, 6 insertions(+), 7 deletions(-)

diff --git a/beacon_node/beacon_chain/src/sync_committee_verification.rs b/beacon_node/beacon_chain/src/sync_committee_verification.rs
index 9b3c3603f75..f3dad4edff7 100644
--- a/beacon_node/beacon_chain/src/sync_committee_verification.rs
+++ b/beacon_node/beacon_chain/src/sync_committee_verification.rs
@@ -383,9 +383,9 @@ impl<T: BeaconChainTypes> VerifiedSyncContribution<T> {
             participant_pubkeys
                 .iter()
                 .map(|pubkey| {
-                    Ok(head_state
+                    head_state
                         .get_validator_index_readonly(pubkey)?
-                        .ok_or(Error::UnknownValidatorPubkey(*pubkey))?)
+                        .ok_or(Error::UnknownValidatorPubkey(*pubkey))
                 })
                 .collect::<Result<Vec<_>, Error>>()?
         };
diff --git a/beacon_node/beacon_chain/tests/sync_committee_verification.rs b/beacon_node/beacon_chain/tests/sync_committee_verification.rs
index 0e4745ff6b8..64d4c17f120 100644
--- a/beacon_node/beacon_chain/tests/sync_committee_verification.rs
+++ b/beacon_node/beacon_chain/tests/sync_committee_verification.rs
@@ -140,11 +140,10 @@ fn get_non_aggregator(
         .enumerate()
         .find_map(|(subcommittee_index, subcommittee)| {
             subcommittee.iter().find_map(|pubkey| {
-                let validator_index = harness
-                    .chain
-                    .validator_index(pubkey)
-                    .expect("should get validator index")
-                    .expect("pubkey should exist in beacon chain");
+                let validator_index = state
+                    .get_validator_index_readonly(pubkey)
+                    .expect("pubkey cache not updated")
+                    .expect("should find validator index");
 
                 let selection_proof = SyncSelectionProof::new::<E>(
                     slot,

From 0ad7641ed15b170e37b6106784c5dcacb893e1fb Mon Sep 17 00:00:00 2001
From: dapplion <35266934+dapplion@users.noreply.github.com>
Date: Mon, 26 Feb 2024 22:14:19 +0800
Subject: [PATCH 5/5] fix CI

---
 beacon_node/beacon_chain/src/test_utils.rs         | 14 +++++++++-----
 beacon_node/beacon_chain/tests/rewards.rs          |  2 +-
 .../tests/sync_committee_verification.rs           | 11 ++++++-----
 beacon_node/http_api/src/lib.rs                    |  6 +++---
 4 files changed, 19 insertions(+), 14 deletions(-)

diff --git a/beacon_node/beacon_chain/src/test_utils.rs b/beacon_node/beacon_chain/src/test_utils.rs
index c3bfd838726..56c93cc0dcd 100644
--- a/beacon_node/beacon_chain/src/test_utils.rs
+++ b/beacon_node/beacon_chain/src/test_utils.rs
@@ -1172,7 +1172,7 @@ where
     /// A list of sync messages for the given state.
     pub fn make_sync_committee_messages(
         &self,
-        state: &BeaconState<E>,
+        state: &mut BeaconState<E>,
         head_block_root: Hash256,
         message_slot: Slot,
         relative_sync_committee: RelativeSyncCommittee,
@@ -1191,6 +1191,8 @@ where
             .spec
             .fork_at_epoch(message_slot.epoch(E::slots_per_epoch()));
 
+        state.update_pubkey_cache().unwrap();
+
         sync_committee
             .pubkeys
             .as_ref()
@@ -1387,7 +1389,7 @@ where
 
     pub fn make_sync_contributions(
         &self,
-        state: &BeaconState<E>,
+        state: &mut BeaconState<E>,
         block_hash: Hash256,
         slot: Slot,
         relative_sync_committee: RelativeSyncCommittee,
@@ -1395,6 +1397,8 @@ where
         let sync_messages =
             self.make_sync_committee_messages(state, block_hash, slot, relative_sync_committee);
 
+        state.update_pubkey_cache().unwrap();
+
         let sync_contributions: Vec<Option<SignedContributionAndProof<E>>> = sync_messages
             .iter()
             .enumerate()
@@ -2008,7 +2012,7 @@ where
 
     pub fn sync_committee_sign_block(
         &self,
-        state: &BeaconState<E>,
+        state: &mut BeaconState<E>,
         block_hash: Hash256,
         slot: Slot,
         relative_sync_committee: RelativeSyncCommittee,
@@ -2043,14 +2047,14 @@ where
         validators: &[usize],
         sync_committee_strategy: SyncCommitteeStrategy,
     ) -> Result<(SignedBeaconBlockHash, BeaconState<E>), BlockError<E>> {
-        let (block_hash, block, state) = self.add_block_at_slot(slot, state).await?;
+        let (block_hash, block, mut state) = self.add_block_at_slot(slot, state).await?;
         self.attest_block(&state, state_root, block_hash, &block.0, validators);
 
         if sync_committee_strategy == SyncCommitteeStrategy::AllValidators
             && state.current_sync_committee().is_ok()
         {
             self.sync_committee_sign_block(
-                &state,
+                &mut state,
                 block_hash.into(),
                 slot,
                 if (slot + 1).epoch(E::slots_per_epoch())
diff --git a/beacon_node/beacon_chain/tests/rewards.rs b/beacon_node/beacon_chain/tests/rewards.rs
index a78463ef5d7..e7ce47800f4 100644
--- a/beacon_node/beacon_chain/tests/rewards.rs
+++ b/beacon_node/beacon_chain/tests/rewards.rs
@@ -54,7 +54,7 @@ async fn test_sync_committee_rewards() {
 
     // Create and add sync committee message to op_pool
     let sync_contributions = harness.make_sync_contributions(
-        &harness.get_current_state(),
+        &mut harness.get_current_state(),
         latest_block_root,
         harness.get_current_slot(),
         RelativeSyncCommittee::Current,
diff --git a/beacon_node/beacon_chain/tests/sync_committee_verification.rs b/beacon_node/beacon_chain/tests/sync_committee_verification.rs
index 64d4c17f120..fc96cc4a1a5 100644
--- a/beacon_node/beacon_chain/tests/sync_committee_verification.rs
+++ b/beacon_node/beacon_chain/tests/sync_committee_verification.rs
@@ -71,9 +71,9 @@ fn get_valid_sync_committee_message_for_block(
     message_index: usize,
     block_root: Hash256,
 ) -> (SyncCommitteeMessage, usize, SecretKey, SyncSubnetId) {
-    let head_state = harness.chain.head_beacon_state_cloned();
+    let mut head_state = harness.chain.head_beacon_state_cloned();
     let (signature, _) = harness
-        .make_sync_committee_messages(&head_state, block_root, slot, relative_sync_committee)
+        .make_sync_committee_messages(&mut head_state, block_root, slot, relative_sync_committee)
         .get(0)
         .expect("sync messages should exist")
         .get(message_index)
@@ -94,13 +94,14 @@ fn get_valid_sync_contribution(
     harness: &BeaconChainHarness<EphemeralHarnessType<E>>,
     relative_sync_committee: RelativeSyncCommittee,
 ) -> (SignedContributionAndProof<E>, usize, SecretKey) {
-    let head_state = harness.chain.head_beacon_state_cloned();
+    let mut head_state = harness.chain.head_beacon_state_cloned();
 
     let head_block_root = harness.chain.head_snapshot().beacon_block_root;
+    let head_state_slot = head_state.slot();
     let sync_contributions = harness.make_sync_contributions(
-        &head_state,
+        &mut head_state,
         head_block_root,
-        head_state.slot(),
+        head_state_slot,
         relative_sync_committee,
     );
 
diff --git a/beacon_node/http_api/src/lib.rs b/beacon_node/http_api/src/lib.rs
index a6dcc8f3b83..60215361604 100644
--- a/beacon_node/http_api/src/lib.rs
+++ b/beacon_node/http_api/src/lib.rs
@@ -771,7 +771,7 @@ pub fn serve<T: BeaconChainTypes>(
                                 let index_opt = match &validator_id {
                                     ValidatorId::PublicKey(pubkey) => state
                                         .get_validator_index_readonly(pubkey)
-                                        .map_err(|e| warp_utils::reject::beacon_state_error(e))?,
+                                        .map_err(warp_utils::reject::beacon_state_error)?,
                                     ValidatorId::Index(index) => Some(*index as usize),
                                 };
 
@@ -1040,12 +1040,12 @@ pub fn serve<T: BeaconChainTypes>(
                                     .pubkeys
                                     .iter()
                                     .map(|pubkey| {
-                                        Ok(state
+                                        state
                                             .get_validator_index_readonly(pubkey)?
                                             .map(|index| index as u64)
                                             .ok_or(BeaconChainError::ValidatorPubkeyUnknown(
                                                 *pubkey,
-                                            ))?)
+                                            ))
                                     })
                                     .collect::<Result<Vec<_>, BeaconChainError>>()
                                     .map_err(warp_utils::reject::beacon_chain_error)?;