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);