From 2350a955e86bc4584e65a69e788b55ddd9f2cd2d Mon Sep 17 00:00:00 2001
From: Michael Sproul <michael@sigmaprime.io>
Date: Thu, 20 Oct 2022 23:05:07 +1100
Subject: [PATCH] Store pubkey cache uncompressed

---
 beacon_node/store/Cargo.toml                  |  1 +
 .../store/src/validator_pubkey_cache.rs       | 58 ++++++++++++++-----
 consensus/ssz/src/decode/impls.rs             |  1 +
 consensus/ssz/src/encode/impls.rs             |  1 +
 4 files changed, 45 insertions(+), 16 deletions(-)

diff --git a/beacon_node/store/Cargo.toml b/beacon_node/store/Cargo.toml
index 5b1a6b7846..53e2a983aa 100644
--- a/beacon_node/store/Cargo.toml
+++ b/beacon_node/store/Cargo.toml
@@ -31,3 +31,4 @@ take-until = "0.1.0"
 zstd = "0.11.0"
 strum = { version = "0.24.0", features = ["derive"] }
 bls = { path = "../../crypto/bls" }
+smallvec = "1.0.0"
diff --git a/beacon_node/store/src/validator_pubkey_cache.rs b/beacon_node/store/src/validator_pubkey_cache.rs
index 4f130a91ab..1d8d365da2 100644
--- a/beacon_node/store/src/validator_pubkey_cache.rs
+++ b/beacon_node/store/src/validator_pubkey_cache.rs
@@ -1,5 +1,8 @@
 use crate::{DBColumn, Error, HotColdDB, ItemStore, StoreItem};
+use bls::PUBLIC_KEY_UNCOMPRESSED_BYTES_LEN;
+use smallvec::SmallVec;
 use ssz::{Decode, Encode};
+use ssz_derive::{Decode, Encode};
 use std::collections::HashMap;
 use std::convert::TryInto;
 use std::marker::PhantomData;
@@ -64,16 +67,14 @@ impl<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>> ValidatorPubkeyCache<E,
         let mut validators = vec![];
 
         for validator_index in 0.. {
-            if let Some(DatabaseValidator(validator)) =
+            if let Some(db_validator) =
                 store.get_item(&DatabaseValidator::key_for_index(validator_index))?
             {
-                pubkeys.push(
-                    (&validator.pubkey)
-                        .try_into()
-                        .map_err(|e| Error::ValidatorPubkeyCacheError(format!("{:?}", e)))?,
-                );
+                let (pubkey, validator) =
+                    DatabaseValidator::into_immutable_validator(&db_validator)?;
+                pubkeys.push(pubkey);
                 indices.insert(validator.pubkey, validator_index);
-                validators.push(validator);
+                validators.push(Arc::new(validator));
             } else {
                 break;
             }
@@ -125,6 +126,10 @@ impl<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>> ValidatorPubkeyCache<E,
                 return Err(Error::DuplicateValidatorPublicKey);
             }
 
+            let pubkey = (&validator.pubkey)
+                .try_into()
+                .map_err(Error::InvalidValidatorPubkeyBytes)?;
+
             // The item is written to disk _before_ it is written into
             // the local struct.
             //
@@ -136,14 +141,10 @@ impl<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>> ValidatorPubkeyCache<E,
             // that are never referenced in a state.
             store.put_item(
                 &DatabaseValidator::key_for_index(i),
-                &DatabaseValidator(validator.clone()),
+                &DatabaseValidator::from_immutable_validator(&pubkey, &validator),
             )?;
 
-            self.pubkeys.push(
-                (&validator.pubkey)
-                    .try_into()
-                    .map_err(Error::InvalidValidatorPubkeyBytes)?,
-            );
+            self.pubkeys.push(pubkey);
             self.indices.insert(validator.pubkey, i);
             self.validators.push(validator);
         }
@@ -190,7 +191,11 @@ impl<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>> ValidatorPubkeyCache<E,
 /// Wrapper for a public key stored in the database.
 ///
 /// Keyed by the validator index as `Hash256::from_low_u64_be(index)`.
-struct DatabaseValidator(Arc<ValidatorImmutable>);
+#[derive(Encode, Decode)]
+struct DatabaseValidator {
+    pubkey: SmallVec<[u8; PUBLIC_KEY_UNCOMPRESSED_BYTES_LEN]>,
+    withdrawal_credentials: Hash256,
+}
 
 impl StoreItem for DatabaseValidator {
     fn db_column() -> DBColumn {
@@ -198,11 +203,11 @@ impl StoreItem for DatabaseValidator {
     }
 
     fn as_store_bytes(&self) -> Result<Vec<u8>, Error> {
-        Ok(self.0.as_ssz_bytes())
+        Ok(self.as_ssz_bytes())
     }
 
     fn from_store_bytes(bytes: &[u8]) -> Result<Self, Error> {
-        Ok(Self(Arc::new(ValidatorImmutable::from_ssz_bytes(bytes)?)))
+        Ok(Self::from_ssz_bytes(bytes)?)
     }
 }
 
@@ -210,6 +215,27 @@ impl DatabaseValidator {
     fn key_for_index(index: usize) -> Hash256 {
         Hash256::from_low_u64_be(index as u64)
     }
+
+    fn from_immutable_validator(pubkey: &PublicKey, validator: &ValidatorImmutable) -> Self {
+        DatabaseValidator {
+            pubkey: pubkey.serialize_uncompressed().into(),
+            withdrawal_credentials: validator.withdrawal_credentials,
+        }
+    }
+
+    fn into_immutable_validator(&self) -> Result<(PublicKey, ValidatorImmutable), Error> {
+        let pubkey = PublicKey::deserialize_uncompressed(&self.pubkey)
+            .map_err(Error::InvalidValidatorPubkeyBytes)?;
+        let pubkey_bytes = pubkey.compress();
+        let withdrawal_credentials = self.withdrawal_credentials;
+        Ok((
+            pubkey,
+            ValidatorImmutable {
+                pubkey: pubkey_bytes,
+                withdrawal_credentials,
+            },
+        ))
+    }
 }
 
 #[cfg(test)]
diff --git a/consensus/ssz/src/decode/impls.rs b/consensus/ssz/src/decode/impls.rs
index d91ddabe02..0e7c3a040b 100644
--- a/consensus/ssz/src/decode/impls.rs
+++ b/consensus/ssz/src/decode/impls.rs
@@ -407,6 +407,7 @@ impl_for_vec!(SmallVec<[T; 5]>, None);
 impl_for_vec!(SmallVec<[T; 6]>, None);
 impl_for_vec!(SmallVec<[T; 7]>, None);
 impl_for_vec!(SmallVec<[T; 8]>, None);
+impl_for_vec!(SmallVec<[T; 96]>, None);
 
 impl<K, V> Decode for BTreeMap<K, V>
 where
diff --git a/consensus/ssz/src/encode/impls.rs b/consensus/ssz/src/encode/impls.rs
index cfd95ba40d..5ac5bcd413 100644
--- a/consensus/ssz/src/encode/impls.rs
+++ b/consensus/ssz/src/encode/impls.rs
@@ -307,6 +307,7 @@ impl_for_vec!(SmallVec<[T; 5]>);
 impl_for_vec!(SmallVec<[T; 6]>);
 impl_for_vec!(SmallVec<[T; 7]>);
 impl_for_vec!(SmallVec<[T; 8]>);
+impl_for_vec!(SmallVec<[T; 96]>);
 
 impl<K, V> Encode for BTreeMap<K, V>
 where