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/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(())
     }