Update network and rpc to BeaconStateTypes

beacon_node/rpc/src/attestation.rs

@@ -9,15 +9,15 @@ use protos::services_grpc::AttestationService;
 use slog::{error, info, trace, warn};
 use ssz::{ssz_encode, Decodable};
 use std::sync::Arc;
-use types::Attestation;
+use types::{Attestation, BeaconStateTypes};
 
 #[derive(Clone)]
-pub struct AttestationServiceInstance {
-    pub chain: Arc<BeaconChain>,
+pub struct AttestationServiceInstance<B: BeaconStateTypes> {
+    pub chain: Arc<BeaconChain<B>>,
     pub log: slog::Logger,
 }
 
-impl AttestationService for AttestationServiceInstance {
+impl<B: BeaconStateTypes> AttestationService for AttestationServiceInstance<B> {
     /// Produce the `AttestationData` for signing by a validator.
     fn produce_attestation_data(
         &mut self,

beacon_node/rpc/src/beacon_block.rs

@@ -13,16 +13,16 @@ use slog::Logger;
 use slog::{error, info, trace, warn};
 use ssz::{ssz_encode, Decodable};
 use std::sync::Arc;
-use types::{BeaconBlock, Signature, Slot};
+use types::{BeaconBlock, BeaconStateTypes, Signature, Slot};
 
 #[derive(Clone)]
-pub struct BeaconBlockServiceInstance {
-    pub chain: Arc<BeaconChain>,
+pub struct BeaconBlockServiceInstance<B: BeaconStateTypes> {
+    pub chain: Arc<BeaconChain<B>>,
     pub network_chan: crossbeam_channel::Sender<NetworkMessage>,
     pub log: Logger,
 }
 
-impl BeaconBlockService for BeaconBlockServiceInstance {
+impl<B: BeaconStateTypes> BeaconBlockService for BeaconBlockServiceInstance<B> {
     /// Produce a `BeaconBlock` for signing by a validator.
     fn produce_beacon_block(
         &mut self,

beacon_node/rpc/src/beacon_chain.rs

@@ -8,15 +8,15 @@ use beacon_chain::{
     AttestationValidationError, BlockProductionError,
 };
 pub use beacon_chain::{BeaconChainError, BlockProcessingOutcome};
-use types::{Attestation, AttestationData, BeaconBlock};
+use types::{Attestation, AttestationData, BeaconBlock, BeaconStateTypes};
 
 /// The RPC's API to the beacon chain.
-pub trait BeaconChain: Send + Sync {
+pub trait BeaconChain<B: BeaconStateTypes>: Send + Sync {
     fn get_spec(&self) -> &ChainSpec;
 
-    fn get_state(&self) -> RwLockReadGuard<BeaconState>;
+    fn get_state(&self) -> RwLockReadGuard<BeaconState<B>>;
 
-    fn get_mut_state(&self) -> RwLockWriteGuard<BeaconState>;
+    fn get_mut_state(&self) -> RwLockWriteGuard<BeaconState<B>>;
 
     fn process_block(&self, block: BeaconBlock)
         -> Result<BlockProcessingOutcome, BeaconChainError>;
@@ -24,7 +24,7 @@ pub trait BeaconChain: Send + Sync {
     fn produce_block(
         &self,
         randao_reveal: Signature,
-    ) -> Result<(BeaconBlock, BeaconState), BlockProductionError>;
+    ) -> Result<(BeaconBlock, BeaconState<B>), BlockProductionError>;
 
     fn produce_attestation_data(&self, shard: u64) -> Result<AttestationData, BeaconChainError>;
 
@@ -34,21 +34,22 @@ pub trait BeaconChain: Send + Sync {
     ) -> Result<(), AttestationValidationError>;
 }
 
-impl<T, U, F> BeaconChain for RawBeaconChain<T, U, F>
+impl<T, U, F, B> BeaconChain<B> for RawBeaconChain<T, U, F, B>
 where
     T: ClientDB + Sized,
     U: SlotClock,
     F: ForkChoice,
+    B: BeaconStateTypes,
 {
     fn get_spec(&self) -> &ChainSpec {
         &self.spec
     }
 
-    fn get_state(&self) -> RwLockReadGuard<BeaconState> {
+    fn get_state(&self) -> RwLockReadGuard<BeaconState<B>> {
         self.state.read()
     }
 
-    fn get_mut_state(&self) -> RwLockWriteGuard<BeaconState> {
+    fn get_mut_state(&self) -> RwLockWriteGuard<BeaconState<B>> {
         self.state.write()
     }
 
@@ -62,7 +63,7 @@ where
     fn produce_block(
         &self,
         randao_reveal: Signature,
-    ) -> Result<(BeaconBlock, BeaconState), BlockProductionError> {
+    ) -> Result<(BeaconBlock, BeaconState<B>), BlockProductionError> {
         self.produce_block(randao_reveal)
     }
 

beacon_node/rpc/src/beacon_node.rs

@@ -5,14 +5,15 @@ use protos::services::{Empty, Fork, NodeInfoResponse};
 use protos::services_grpc::BeaconNodeService;
 use slog::{trace, warn};
 use std::sync::Arc;
+use types::BeaconStateTypes;
 
 #[derive(Clone)]
-pub struct BeaconNodeServiceInstance {
-    pub chain: Arc<BeaconChain>,
+pub struct BeaconNodeServiceInstance<B: BeaconStateTypes> {
+    pub chain: Arc<BeaconChain<B>>,
     pub log: slog::Logger,
 }
 
-impl BeaconNodeService for BeaconNodeServiceInstance {
+impl<B: BeaconStateTypes> BeaconNodeService for BeaconNodeServiceInstance<B> {
     /// Provides basic node information.
     fn info(&mut self, ctx: RpcContext, _req: Empty, sink: UnarySink<NodeInfoResponse>) {
         trace!(self.log, "Node info requested via RPC");

beacon_node/rpc/src/lib.rs

@@ -21,12 +21,13 @@ use protos::services_grpc::{
 use slog::{info, o, warn};
 use std::sync::Arc;
 use tokio::runtime::TaskExecutor;
+use types::BeaconStateTypes;
 
-pub fn start_server(
+pub fn start_server<B: BeaconStateTypes>(
     config: &RPCConfig,
     executor: &TaskExecutor,
     network_chan: crossbeam_channel::Sender<NetworkMessage>,
-    beacon_chain: Arc<BeaconChain>,
+    beacon_chain: Arc<BeaconChain<B>>,
     log: &slog::Logger,
 ) -> exit_future::Signal {
     let log = log.new(o!("Service"=>"RPC"));

beacon_node/rpc/src/validator.rs

@@ -7,16 +7,16 @@ use protos::services_grpc::ValidatorService;
 use slog::{trace, warn};
 use ssz::decode;
 use std::sync::Arc;
-use types::{Epoch, RelativeEpoch};
+use types::{BeaconStateTypes, Epoch, RelativeEpoch};
 
 #[derive(Clone)]
-pub struct ValidatorServiceInstance {
-    pub chain: Arc<BeaconChain>,
+pub struct ValidatorServiceInstance<B: BeaconStateTypes> {
+    pub chain: Arc<BeaconChain<B>>,
     pub log: slog::Logger,
 }
 //TODO: Refactor Errors
 
-impl ValidatorService for ValidatorServiceInstance {
+impl<B: BeaconStateTypes> ValidatorService for ValidatorServiceInstance<B> {
     /// For a list of validator public keys, this function returns the slot at which each
     /// validator must propose a block, attest to a shard, their shard committee and the shard they
     /// need to attest to.