Add a cache for public keys to BeaconState

This allows for a fast lookup of "is this public key already in the
validator registry".

eth2/state_processing/benches/benches.rs

@@ -1,14 +1,11 @@
-use criterion::Benchmark;
 use criterion::Criterion;
 use criterion::{criterion_group, criterion_main};
 use env_logger::{Builder, Env};
-use types::test_utils::TestingBeaconStateBuilder;
-use types::*;
 
 mod bench_block_processing;
 mod bench_epoch_processing;
 
-pub const VALIDATOR_COUNT: usize = 300_032;
+pub const VALIDATOR_COUNT: usize = 16_384;
 
 // `LOG_LEVEL == "debug"` gives logs, but they're very noisy and slow down benching.
 pub const LOG_LEVEL: &str = "";

eth2/state_processing/src/per_block_processing.rs

@@ -373,19 +373,20 @@ pub fn process_deposits(
                 .map_err(|e| e.into_with_index(i))
         })?;
 
-    let public_key_to_index_hashmap = build_public_key_hashmap(&state);
-
     // Check `state.deposit_index` and update the state in series.
     for (i, deposit) in deposits.iter().enumerate() {
         verify_deposit_index(state, deposit).map_err(|e| e.into_with_index(i))?;
 
+        // Ensure the state's pubkey cache is fully up-to-date, it will be used to check to see if the
+        // depositing validator already exists in the registry.
+        state.update_pubkey_cache()?;
+
         // Get an `Option<u64>` where `u64` is the validator index if this deposit public key
         // already exists in the beacon_state.
         //
         // This function also verifies the withdrawal credentials.
         let validator_index =
-            get_existing_validator_index(state, deposit, &public_key_to_index_hashmap)
-                .map_err(|e| e.into_with_index(i))?;
+            get_existing_validator_index(state, deposit).map_err(|e| e.into_with_index(i))?;
 
         let deposit_data = &deposit.deposit_data;
         let deposit_input = &deposit.deposit_data.deposit_input;

eth2/state_processing/src/per_block_processing/errors.rs

@@ -294,6 +294,8 @@ impl_into_with_index_without_beacon_error!(
 pub enum DepositValidationError {
     /// Validation completed successfully and the object is invalid.
     Invalid(DepositInvalid),
+    /// Encountered a `BeaconStateError` whilst attempting to determine validity.
+    BeaconStateError(BeaconStateError),
 }
 
 /// Describes why an object is invalid.
@@ -313,7 +315,8 @@ pub enum DepositInvalid {
     BadMerkleProof,
 }
 
-impl_into_with_index_without_beacon_error!(DepositValidationError, DepositInvalid);
+impl_from_beacon_state_error!(DepositValidationError);
+impl_into_with_index_with_beacon_error!(DepositValidationError, DepositInvalid);
 
 /*
  * `Exit` Validation

eth2/state_processing/src/per_block_processing/verify_deposit.rs

@@ -72,11 +72,12 @@ pub fn build_public_key_hashmap(state: &BeaconState) -> PublicKeyValidatorIndexH
 pub fn get_existing_validator_index(
     state: &BeaconState,
     deposit: &Deposit,
-    pubkey_map: &HashMap<PublicKey, u64>,
 ) -> Result<Option<u64>, Error> {
     let deposit_input = &deposit.deposit_data.deposit_input;
 
-    let validator_index = pubkey_map.get(&deposit_input.pubkey).and_then(|i| Some(*i));
+    let validator_index = state
+        .get_validator_index(&deposit_input.pubkey)?
+        .and_then(|i| Some(i));
 
     match validator_index {
         None => Ok(None),
@@ -86,7 +87,7 @@ pub fn get_existing_validator_index(
                     == state.validator_registry[index as usize].withdrawal_credentials,
                 Invalid::BadWithdrawalCredentials
             );
-            Ok(Some(index))
+            Ok(Some(index as u64))
         }
     }
 }

eth2/types/src/beacon_state.rs

@@ -5,6 +5,7 @@ use helpers::*;
 use honey_badger_split::SplitExt;
 use int_to_bytes::int_to_bytes32;
 use log::{debug, error, trace};
+use pubkey_cache::PubkeyCache;
 use rand::RngCore;
 use serde_derive::Serialize;
 use ssz::{hash, Decodable, DecodeError, Encodable, SignedRoot, SszStream, TreeHash};
@@ -16,6 +17,7 @@ pub use builder::BeaconStateBuilder;
 mod builder;
 mod epoch_cache;
 pub mod helpers;
+mod pubkey_cache;
 mod tests;
 
 pub type Committee = Vec<usize>;
@@ -52,6 +54,11 @@ pub enum Error {
     InsufficientAttestations,
     InsufficientCommittees,
     EpochCacheUninitialized(RelativeEpoch),
+    PubkeyCacheInconsistent,
+    PubkeyCacheIncomplete {
+        cache_len: usize,
+        registry_len: usize,
+    },
 }
 
 macro_rules! safe_add_assign {
@@ -108,6 +115,7 @@ pub struct BeaconState {
     // Caching (not in the spec)
     pub cache_index_offset: usize,
     pub caches: Vec<EpochCache>,
+    pub pubkey_cache: PubkeyCache,
 }
 
 impl BeaconState {
@@ -186,6 +194,7 @@ impl BeaconState {
              */
             cache_index_offset: 0,
             caches: vec![EpochCache::empty(); CACHED_EPOCHS],
+            pubkey_cache: PubkeyCache::empty(),
         }
     }
 
@@ -293,6 +302,41 @@ impl BeaconState {
         }
     }
 
+    /// Updates the pubkey cache, if required.
+    ///
+    /// Adds all `pubkeys` from the `validator_registry` which are not already in the cache. Will
+    /// never re-add a pubkey.
+    pub fn update_pubkey_cache(&mut self) -> Result<(), Error> {
+        for (i, validator) in self
+            .validator_registry
+            .iter()
+            .enumerate()
+            .skip(self.pubkey_cache.len())
+        {
+            let success = self.pubkey_cache.insert(validator.pubkey.clone(), i);
+            if !success {
+                return Err(Error::PubkeyCacheInconsistent);
+            }
+        }
+
+        Ok(())
+    }
+
+    /// If a validator pubkey exists in the validator registry, returns `Some(i)`, otherwise
+    /// returns `None`.
+    ///
+    /// Requires a fully up-to-date `pubkey_cache`, returns an error if this is not the case.
+    pub fn get_validator_index(&self, pubkey: &PublicKey) -> Result<Option<usize>, Error> {
+        if self.pubkey_cache.len() == self.validator_registry.len() {
+            Ok(self.pubkey_cache.get(pubkey))
+        } else {
+            Err(Error::PubkeyCacheIncomplete {
+                cache_len: self.pubkey_cache.len(),
+                registry_len: self.validator_registry.len(),
+            })
+        }
+    }
+
     /// The epoch corresponding to `self.slot`.
     ///
     /// Spec v0.4.0
@@ -1188,6 +1232,7 @@ impl Decodable for BeaconState {
                 deposit_index,
                 cache_index_offset: 0,
                 caches: vec![EpochCache::empty(); CACHED_EPOCHS],
+                pubkey_cache: PubkeyCache::empty(),
             },
             i,
         ))
@@ -1258,6 +1303,7 @@ impl<T: RngCore> TestRandom<T> for BeaconState {
             deposit_index: <_>::random_for_test(rng),
             cache_index_offset: 0,
             caches: vec![EpochCache::empty(); CACHED_EPOCHS],
+            pubkey_cache: PubkeyCache::empty(),
         }
     }
 }

eth2/types/src/beacon_state/pubkey_cache.rs

@@ -0,0 +1,45 @@
+use crate::*;
+use serde_derive::Serialize;
+use std::collections::HashMap;
+
+type ValidatorIndex = usize;
+
+#[derive(Debug, PartialEq, Clone, Default, Serialize)]
+pub struct PubkeyCache {
+    map: HashMap<PublicKey, ValidatorIndex>,
+}
+
+impl PubkeyCache {
+    /// Instantiates a new, empty cache.
+    pub fn empty() -> Self {
+        Self {
+            map: HashMap::new(),
+        }
+    }
+
+    /// Returns the number of validator indices already in the map.
+    pub fn len(&self) -> ValidatorIndex {
+        self.map.len()
+    }
+
+    /// Inserts a validator index into the map.
+    ///
+    /// 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: PublicKey, index: ValidatorIndex) -> bool {
+        if index == self.map.len() {
+            self.map.insert(pubkey, index);
+            true
+        } else {
+            false
+        }
+    }
+
+    /// Inserts a validator index into the map.
+    ///
+    /// The added index must equal the number of validators already added to the map. This ensures
+    /// that an index is never skipped.
+    pub fn get(&self, pubkey: &PublicKey) -> Option<ValidatorIndex> {
+        self.map.get(pubkey).cloned()
+    }
+}

eth2/types/src/lib.rs

@@ -1,4 +1,4 @@
-pub mod test_utils;
+//! Ethereum 2.0 types
 
 pub mod attestation;
 pub mod attestation_data;
@@ -22,6 +22,7 @@ pub mod proposer_slashing;
 pub mod readers;
 pub mod shard_reassignment_record;
 pub mod slashable_attestation;
+pub mod test_utils;
 pub mod transfer;
 pub mod voluntary_exit;
 #[macro_use]

eth2/types/src/test_utils/testing_beacon_state_builder.rs

@@ -144,6 +144,8 @@ impl TestingBeaconStateBuilder {
         state.build_epoch_cache(RelativeEpoch::Current, &spec)?;
         state.build_epoch_cache(RelativeEpoch::Next, &spec)?;
 
+        state.update_pubkey_cache()?;
+
         Ok(())
     }