From 67fdb4a7fb42f210acb89d96c9d5f9b3d0213ec9 Mon Sep 17 00:00:00 2001
From: Paul Hauner <paul@paulhauner.com>
Date: Tue, 4 Jun 2019 13:13:58 +1000
Subject: [PATCH] Store beacon state committee cache in DB

---
 beacon_node/store/src/impls.rs                | 16 +----
 beacon_node/store/src/impls/beacon_state.rs   | 64 +++++++++++++++++++
 eth2/types/src/beacon_state.rs                |  3 +-
 .../types/src/beacon_state/committee_cache.rs |  3 +-
 4 files changed, 70 insertions(+), 16 deletions(-)
 create mode 100644 beacon_node/store/src/impls/beacon_state.rs

diff --git a/beacon_node/store/src/impls.rs b/beacon_node/store/src/impls.rs
index 91f8d52de6..418fcade1e 100644
--- a/beacon_node/store/src/impls.rs
+++ b/beacon_node/store/src/impls.rs
@@ -1,6 +1,8 @@
 use crate::*;
 use ssz::{Decode, Encode};
 
+mod beacon_state;
+
 impl StoreItem for BeaconBlock {
     fn db_column() -> DBColumn {
         DBColumn::BeaconBlock
@@ -14,17 +16,3 @@ impl StoreItem for BeaconBlock {
         Self::from_ssz_bytes(bytes).map_err(Into::into)
     }
 }
-
-impl<T: EthSpec> StoreItem for BeaconState<T> {
-    fn db_column() -> DBColumn {
-        DBColumn::BeaconState
-    }
-
-    fn as_store_bytes(&self) -> Vec<u8> {
-        self.as_ssz_bytes()
-    }
-
-    fn from_store_bytes(bytes: &mut [u8]) -> Result<Self, Error> {
-        Self::from_ssz_bytes(bytes).map_err(Into::into)
-    }
-}
diff --git a/beacon_node/store/src/impls/beacon_state.rs b/beacon_node/store/src/impls/beacon_state.rs
new file mode 100644
index 0000000000..40d2cacb80
--- /dev/null
+++ b/beacon_node/store/src/impls/beacon_state.rs
@@ -0,0 +1,64 @@
+use crate::*;
+use ssz::{Decode, DecodeError, Encode};
+use ssz_derive::{Decode, Encode};
+use std::convert::TryInto;
+use types::beacon_state::{CACHED_EPOCHS, CommitteeCache};
+
+/// A container for storing `BeaconState` components.
+#[derive(Encode, Decode)]
+struct StorageContainer {
+    state_bytes: Vec<u8>,
+    committee_caches_bytes: Vec<Vec<u8>>,
+}
+
+impl StorageContainer {
+    /// Create a new instance for storing a `BeaconState`.
+    pub fn new<T: EthSpec>(state: &BeaconState<T>) -> Self {
+        let mut committee_caches_bytes = vec![];
+
+        for cache in state.committee_caches[..].iter() {
+            committee_caches_bytes.push(cache.as_ssz_bytes());
+        }
+
+        Self {
+            state_bytes: state.as_ssz_bytes(),
+            committee_caches_bytes,
+        }
+    }
+}
+
+impl<T: EthSpec> TryInto<BeaconState<T>> for StorageContainer {
+    type Error = Error;
+
+    fn try_into(self) -> Result<BeaconState<T>, Error> {
+        let mut state: BeaconState<T> = BeaconState::from_ssz_bytes(&self.state_bytes)?;
+
+        for i in 0..CACHED_EPOCHS {
+            let bytes = &self.committee_caches_bytes.get(i).ok_or_else(|| {
+                Error::SszDecodeError(DecodeError::BytesInvalid(
+                    "Insufficient committees for BeaconState".to_string(),
+                ))
+            })?;
+
+            state.committee_caches[i] = CommitteeCache::from_ssz_bytes(bytes)?;
+        }
+
+        Ok(state)
+    }
+}
+
+impl<T: EthSpec> StoreItem for BeaconState<T> {
+    fn db_column() -> DBColumn {
+        DBColumn::BeaconState
+    }
+
+    fn as_store_bytes(&self) -> Vec<u8> {
+        let container = StorageContainer::new(self);
+        container.as_ssz_bytes()
+    }
+
+    fn from_store_bytes(bytes: &mut [u8]) -> Result<Self, Error> {
+        let container = StorageContainer::from_ssz_bytes(bytes)?;
+        container.try_into()
+    }
+}
diff --git a/eth2/types/src/beacon_state.rs b/eth2/types/src/beacon_state.rs
index 4e510940f0..471b82b20b 100644
--- a/eth2/types/src/beacon_state.rs
+++ b/eth2/types/src/beacon_state.rs
@@ -1,4 +1,4 @@
-use self::committee_cache::{get_active_validator_indices, CommitteeCache};
+use self::committee_cache::get_active_validator_indices;
 use self::exit_cache::ExitCache;
 use crate::test_utils::TestRandom;
 use crate::*;
@@ -15,6 +15,7 @@ use test_random_derive::TestRandom;
 use tree_hash::TreeHash;
 use tree_hash_derive::{CachedTreeHash, TreeHash};
 
+pub use self::committee_cache::CommitteeCache;
 pub use beacon_state_types::*;
 
 mod beacon_state_types;
diff --git a/eth2/types/src/beacon_state/committee_cache.rs b/eth2/types/src/beacon_state/committee_cache.rs
index 27374c3391..d1e0e70307 100644
--- a/eth2/types/src/beacon_state/committee_cache.rs
+++ b/eth2/types/src/beacon_state/committee_cache.rs
@@ -2,6 +2,7 @@ use super::BeaconState;
 use crate::*;
 use core::num::NonZeroUsize;
 use serde_derive::{Deserialize, Serialize};
+use ssz_derive::{Decode, Encode};
 use std::ops::Range;
 use swap_or_not_shuffle::shuffle_list;
 
@@ -9,7 +10,7 @@ mod tests;
 
 /// Computes and stores the shuffling for an epoch. Provides various getters to allow callers to
 /// read the committees for the given epoch.
-#[derive(Debug, Default, PartialEq, Clone, Serialize, Deserialize)]
+#[derive(Debug, Default, PartialEq, Clone, Serialize, Deserialize, Encode, Decode)]
 pub struct CommitteeCache {
     initialized_epoch: Option<Epoch>,
     shuffling: Vec<usize>,