From f1584dada470213d2477832ffc2c400d0e1694b5 Mon Sep 17 00:00:00 2001
From: Paul Hauner <paul@paulhauner.com>
Date: Tue, 21 May 2019 13:36:14 +1000
Subject: [PATCH] Update BeaconChain struct to use new store

---
 beacon_node/beacon_chain/src/beacon_chain.rs  | 22 +++++++------------
 .../testing_beacon_chain_builder.rs           | 14 ++++--------
 2 files changed, 12 insertions(+), 24 deletions(-)

diff --git a/beacon_node/beacon_chain/src/beacon_chain.rs b/beacon_node/beacon_chain/src/beacon_chain.rs
index db5ea1cdb..57b697019 100644
--- a/beacon_node/beacon_chain/src/beacon_chain.rs
+++ b/beacon_node/beacon_chain/src/beacon_chain.rs
@@ -1,9 +1,6 @@
 use crate::checkpoint::CheckPoint;
 use crate::errors::{BeaconChainError as Error, BlockProductionError};
-use db::{
-    stores::{BeaconBlockStore, BeaconStateStore},
-    ClientDB, DBError,
-};
+use db::{Error as DBError, Store};
 use fork_choice::{ForkChoice, ForkChoiceError};
 use log::{debug, trace};
 use operation_pool::DepositInsertStatus;
@@ -83,9 +80,8 @@ impl BlockProcessingOutcome {
     }
 }
 
-pub struct BeaconChain<T: ClientDB + Sized, U: SlotClock, F: ForkChoice, E: EthSpec> {
-    pub block_store: Arc<BeaconBlockStore<T>>,
-    pub state_store: Arc<BeaconStateStore<T>>,
+pub struct BeaconChain<T, U, F, E: EthSpec> {
+    pub store: Arc<T>,
     pub slot_clock: U,
     pub op_pool: OperationPool<E>,
     canonical_head: RwLock<CheckPoint<E>>,
@@ -97,15 +93,14 @@ pub struct BeaconChain<T: ClientDB + Sized, U: SlotClock, F: ForkChoice, E: EthS
 
 impl<T, U, F, E> BeaconChain<T, U, F, E>
 where
-    T: ClientDB,
+    T: Store,
     U: SlotClock,
     F: ForkChoice,
     E: EthSpec,
 {
     /// Instantiate a new Beacon Chain, from genesis.
     pub fn from_genesis(
-        state_store: Arc<BeaconStateStore<T>>,
-        block_store: Arc<BeaconBlockStore<T>>,
+        store: Arc<T>,
         slot_clock: U,
         mut genesis_state: BeaconState<E>,
         genesis_block: BeaconBlock,
@@ -113,10 +108,10 @@ where
         fork_choice: F,
     ) -> Result<Self, Error> {
         let state_root = genesis_state.canonical_root();
-        state_store.put(&state_root, &ssz_encode(&genesis_state)[..])?;
+        store.put(&state_root, &genesis_state)?;
 
         let block_root = genesis_block.block_header().canonical_root();
-        block_store.put(&block_root, &ssz_encode(&genesis_block)[..])?;
+        store.put(&block_root, &genesis_block)?;
 
         let finalized_head = RwLock::new(CheckPoint::new(
             genesis_block.clone(),
@@ -134,8 +129,7 @@ where
         genesis_state.build_all_caches(&spec)?;
 
         Ok(Self {
-            block_store,
-            state_store,
+            store,
             slot_clock,
             op_pool: OperationPool::new(),
             state: RwLock::new(genesis_state),
diff --git a/beacon_node/beacon_chain/src/test_utils/testing_beacon_chain_builder.rs b/beacon_node/beacon_chain/src/test_utils/testing_beacon_chain_builder.rs
index f7ff3cdae..ce3588674 100644
--- a/beacon_node/beacon_chain/src/test_utils/testing_beacon_chain_builder.rs
+++ b/beacon_node/beacon_chain/src/test_utils/testing_beacon_chain_builder.rs
@@ -1,8 +1,5 @@
 pub use crate::{BeaconChain, BeaconChainError, CheckPoint};
-use db::{
-    stores::{BeaconBlockStore, BeaconStateStore},
-    MemoryDB,
-};
+use db::MemoryDB;
 use fork_choice::BitwiseLMDGhost;
 use slot_clock::TestingSlotClock;
 use std::sync::Arc;
@@ -19,11 +16,9 @@ pub struct TestingBeaconChainBuilder<E: EthSpec> {
 
 impl<E: EthSpec> TestingBeaconChainBuilder<E> {
     pub fn build(self, spec: &ChainSpec) -> TestingBeaconChain<E> {
-        let db = Arc::new(MemoryDB::open());
-        let block_store = Arc::new(BeaconBlockStore::new(db.clone()));
-        let state_store = Arc::new(BeaconStateStore::new(db.clone()));
+        let store = Arc::new(MemoryDB::open());
         let slot_clock = TestingSlotClock::new(spec.genesis_slot.as_u64());
-        let fork_choice = BitwiseLMDGhost::new(block_store.clone(), state_store.clone());
+        let fork_choice = BitwiseLMDGhost::new(store.clone());
 
         let (genesis_state, _keypairs) = self.state_builder.build();
 
@@ -32,8 +27,7 @@ impl<E: EthSpec> TestingBeaconChainBuilder<E> {
 
         // Create the Beacon Chain
         BeaconChain::from_genesis(
-            state_store.clone(),
-            block_store.clone(),
+            store,
             slot_clock,
             genesis_state,
             genesis_block,