Update network and rpc to BeaconStateTypes

beacon_node/network/src/beacon_chain.rs

@@ -8,19 +8,22 @@ use beacon_chain::{
     AttestationValidationError, CheckPoint,
 };
 use eth2_libp2p::rpc::HelloMessage;
-use types::{Attestation, BeaconBlock, BeaconBlockBody, BeaconBlockHeader, Epoch, Hash256, Slot};
+use types::{
+    Attestation, BeaconBlock, BeaconBlockBody, BeaconBlockHeader, BeaconStateTypes, Epoch, Hash256,
+    Slot,
+};
 
 pub use beacon_chain::{BeaconChainError, BlockProcessingOutcome, InvalidBlock};
 
 /// The network'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 slot(&self) -> Slot;
 
-    fn head(&self) -> RwLockReadGuard<CheckPoint>;
+    fn head(&self) -> RwLockReadGuard<CheckPoint<B>>;
 
     fn get_block(&self, block_root: &Hash256) -> Result<Option<BeaconBlock>, BeaconChainError>;
 
@@ -28,7 +31,7 @@ pub trait BeaconChain: Send + Sync {
 
     fn best_block_root(&self) -> Hash256;
 
-    fn finalized_head(&self) -> RwLockReadGuard<CheckPoint>;
+    fn finalized_head(&self) -> RwLockReadGuard<CheckPoint<B>>;
 
     fn finalized_epoch(&self) -> Epoch;
 
@@ -62,17 +65,18 @@ pub trait BeaconChain: Send + Sync {
     fn is_new_block_root(&self, beacon_block_root: &Hash256) -> Result<bool, BeaconChainError>;
 }
 
-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()
     }
 
@@ -80,7 +84,7 @@ where
         self.get_state().slot
     }
 
-    fn head(&self) -> RwLockReadGuard<CheckPoint> {
+    fn head(&self) -> RwLockReadGuard<CheckPoint<B>> {
         self.head()
     }
 
@@ -92,7 +96,7 @@ where
         self.get_state().finalized_epoch
     }
 
-    fn finalized_head(&self) -> RwLockReadGuard<CheckPoint> {
+    fn finalized_head(&self) -> RwLockReadGuard<CheckPoint<B>> {
         self.finalized_head()
     }
 

beacon_node/network/src/message_handler.rs

@@ -13,6 +13,7 @@ use slog::{debug, warn};
 use std::collections::HashMap;
 use std::sync::Arc;
 use std::time::Instant;
+use types::BeaconStateTypes;
 
 /// Timeout for RPC requests.
 // const REQUEST_TIMEOUT: Duration = Duration::from_secs(30);
@@ -20,11 +21,11 @@ use std::time::Instant;
 // const HELLO_TIMEOUT: Duration = Duration::from_secs(30);
 
 /// Handles messages received from the network and client and organises syncing.
-pub struct MessageHandler {
+pub struct MessageHandler<B: BeaconStateTypes> {
     /// Currently loaded and initialised beacon chain.
-    _chain: Arc<BeaconChain>,
+    _chain: Arc<BeaconChain<B>>,
     /// The syncing framework.
-    sync: SimpleSync,
+    sync: SimpleSync<B>,
     /// The context required to send messages to, and process messages from peers.
     network_context: NetworkContext,
     /// The `MessageHandler` logger.
@@ -44,10 +45,10 @@ pub enum HandlerMessage {
     PubsubMessage(PeerId, Box<PubsubMessage>),
 }
 
-impl MessageHandler {
+impl<B: BeaconStateTypes> MessageHandler<B> {
     /// Initializes and runs the MessageHandler.
     pub fn spawn(
-        beacon_chain: Arc<BeaconChain>,
+        beacon_chain: Arc<BeaconChain<B>>,
         network_send: crossbeam_channel::Sender<NetworkMessage>,
         executor: &tokio::runtime::TaskExecutor,
         log: slog::Logger,

beacon_node/network/src/service.rs

@@ -10,22 +10,23 @@ use futures::prelude::*;
 use futures::sync::oneshot;
 use futures::Stream;
 use slog::{debug, info, o, trace};
+use std::marker::PhantomData;
 use std::sync::Arc;
 use tokio::runtime::TaskExecutor;
-use types::Topic;
+use types::{BeaconStateTypes, Topic};
 
 /// Service that handles communication between internal services and the eth2_libp2p network service.
-pub struct Service {
+pub struct Service<B: BeaconStateTypes> {
     //libp2p_service: Arc<Mutex<LibP2PService>>,
     _libp2p_exit: oneshot::Sender<()>,
     network_send: crossbeam_channel::Sender<NetworkMessage>,
-    //message_handler: MessageHandler,
-    //message_handler_send: Sender<HandlerMessage>,
+    _phantom: PhantomData<B>, //message_handler: MessageHandler,
+                              //message_handler_send: Sender<HandlerMessage>
 }
 
-impl Service {
+impl<B: BeaconStateTypes> Service<B> {
     pub fn new(
-        beacon_chain: Arc<BeaconChain>,
+        beacon_chain: Arc<BeaconChain<B>>,
         config: &NetworkConfig,
         executor: &TaskExecutor,
         log: slog::Logger,
@@ -56,6 +57,7 @@ impl Service {
         let network_service = Service {
             _libp2p_exit: libp2p_exit,
             network_send: network_send.clone(),
+            _phantom: PhantomData,
         };
 
         Ok((Arc::new(network_service), network_send))

beacon_node/network/src/sync/import_queue.rs

@@ -5,7 +5,7 @@ use slog::{debug, error};
 use std::sync::Arc;
 use std::time::{Duration, Instant};
 use tree_hash::TreeHash;
-use types::{BeaconBlock, BeaconBlockBody, BeaconBlockHeader, Hash256, Slot};
+use types::{BeaconBlock, BeaconBlockBody, BeaconBlockHeader, BeaconStateTypes, Hash256, Slot};
 
 /// Provides a queue for fully and partially built `BeaconBlock`s.
 ///
@@ -19,8 +19,8 @@ use types::{BeaconBlock, BeaconBlockBody, BeaconBlockHeader, Hash256, Slot};
 /// `BeaconBlockBody` as the key.
 /// - It is possible for multiple distinct blocks to have identical `BeaconBlockBodies`. Therefore
 /// we cannot use a `HashMap` keyed by the root of `BeaconBlockBody`.
-pub struct ImportQueue {
-    pub chain: Arc<BeaconChain>,
+pub struct ImportQueue<B: BeaconStateTypes> {
+    pub chain: Arc<BeaconChain<B>>,
     /// Partially imported blocks, keyed by the root of `BeaconBlockBody`.
     pub partials: Vec<PartialBeaconBlock>,
     /// Time before a queue entry is considered state.
@@ -29,9 +29,9 @@ pub struct ImportQueue {
     log: slog::Logger,
 }
 
-impl ImportQueue {
+impl<B: BeaconStateTypes> ImportQueue<B> {
     /// Return a new, empty queue.
-    pub fn new(chain: Arc<BeaconChain>, stale_time: Duration, log: slog::Logger) -> Self {
+    pub fn new(chain: Arc<BeaconChain<B>>, stale_time: Duration, log: slog::Logger) -> Self {
         Self {
             chain,
             partials: vec![],

beacon_node/network/src/sync/simple_sync.rs

@@ -9,7 +9,7 @@ use std::collections::HashMap;
 use std::sync::Arc;
 use std::time::Duration;
 use tree_hash::TreeHash;
-use types::{Attestation, BeaconBlock, Epoch, Hash256, Slot};
+use types::{Attestation, BeaconBlock, BeaconStateTypes, Epoch, Hash256, Slot};
 
 /// The number of slots that we can import blocks ahead of us, before going into full Sync mode.
 const SLOT_IMPORT_TOLERANCE: u64 = 100;
@@ -88,8 +88,8 @@ impl From<HelloMessage> for PeerSyncInfo {
     }
 }
 
-impl From<&Arc<BeaconChain>> for PeerSyncInfo {
-    fn from(chain: &Arc<BeaconChain>) -> PeerSyncInfo {
+impl<B: BeaconStateTypes> From<&Arc<BeaconChain<B>>> for PeerSyncInfo {
+    fn from(chain: &Arc<BeaconChain<B>>) -> PeerSyncInfo {
         Self::from(chain.hello_message())
     }
 }
@@ -103,22 +103,22 @@ pub enum SyncState {
 }
 
 /// Simple Syncing protocol.
-pub struct SimpleSync {
+pub struct SimpleSync<B: BeaconStateTypes> {
     /// A reference to the underlying beacon chain.
-    chain: Arc<BeaconChain>,
+    chain: Arc<BeaconChain<B>>,
     /// A mapping of Peers to their respective PeerSyncInfo.
     known_peers: HashMap<PeerId, PeerSyncInfo>,
     /// A queue to allow importing of blocks
-    import_queue: ImportQueue,
+    import_queue: ImportQueue<B>,
     /// The current state of the syncing protocol.
     state: SyncState,
     /// Sync logger.
     log: slog::Logger,
 }
 
-impl SimpleSync {
+impl<B: BeaconStateTypes> SimpleSync<B> {
     /// Instantiate a `SimpleSync` instance, with no peers and an empty queue.
-    pub fn new(beacon_chain: Arc<BeaconChain>, log: &slog::Logger) -> Self {
+    pub fn new(beacon_chain: Arc<BeaconChain<B>>, log: &slog::Logger) -> Self {
         let sync_logger = log.new(o!("Service"=> "Sync"));
 
         let queue_item_stale_time = Duration::from_secs(QUEUE_STALE_SECS);

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.

eth2/types/src/beacon_state/beacon_state_types.rs

@@ -2,7 +2,7 @@ use crate::*;
 use fixed_len_vec::typenum::{Unsigned, U1024, U8, U8192};
 use std::fmt::Debug;
 
-pub trait BeaconStateTypes: Default + Sync + Send + Clone + Debug + PartialEq {
+pub trait BeaconStateTypes: 'static + Default + Sync + Send + Clone + Debug + PartialEq {
     type ShardCount: Unsigned + Clone + Sync + Send + Debug + PartialEq;
     type SlotsPerHistoricalRoot: Unsigned + Clone + Sync + Send + Debug + PartialEq;
     type LatestRandaoMixesLength: Unsigned + Clone + Sync + Send + Debug + PartialEq;