Apply CachedBeaconState to some functions

- Attestation aggregation
- Getting attestation duties

beacon_node/beacon_chain/src/attestation_aggregator.rs

@@ -1,3 +1,4 @@
+use crate::cached_beacon_state::CachedBeaconState;
 use state_processing::validate_attestation_without_signature;
 use std::collections::{HashMap, HashSet};
 use types::{
@@ -76,12 +77,12 @@ impl AttestationAggregator {
     ///  - The signature is verified against that of the validator at `validator_index`.
     pub fn process_free_attestation(
         &mut self,
-        state: &BeaconState,
+        cached_state: &CachedBeaconState,
         free_attestation: &FreeAttestation,
         spec: &ChainSpec,
     ) -> Result<Outcome, BeaconStateError> {
         let (slot, shard, committee_index) = some_or_invalid!(
-            state.attestation_slot_and_shard_for_validator(
+            cached_state.attestation_slot_and_shard_for_validator(
                 free_attestation.validator_index as usize,
                 spec,
             )?,
@@ -104,7 +105,8 @@ impl AttestationAggregator {
         let signable_message = free_attestation.data.signable_message(PHASE_0_CUSTODY_BIT);
 
         let validator_record = some_or_invalid!(
-            state
+            cached_state
+                .state
                 .validator_registry
                 .get(free_attestation.validator_index as usize),
             Message::BadValidatorIndex

beacon_node/beacon_chain/src/beacon_chain.rs

@@ -1,4 +1,5 @@
 use crate::attestation_aggregator::{AttestationAggregator, Outcome as AggregationOutcome};
+use crate::cached_beacon_state::CachedBeaconState;
 use crate::checkpoint::CheckPoint;
 use db::{
     stores::{BeaconBlockStore, BeaconStateStore},
@@ -69,6 +70,7 @@ pub struct BeaconChain<T: ClientDB + Sized, U: SlotClock, F: ForkChoice> {
     canonical_head: RwLock<CheckPoint>,
     finalized_head: RwLock<CheckPoint>,
     pub state: RwLock<BeaconState>,
+    pub cached_state: RwLock<CachedBeaconState>,
     pub spec: ChainSpec,
     pub fork_choice: RwLock<F>,
 }
@@ -107,6 +109,11 @@ where
         let block_root = genesis_block.canonical_root();
         block_store.put(&block_root, &ssz_encode(&genesis_block)[..])?;
 
+        let cached_state = RwLock::new(CachedBeaconState::from_beacon_state(
+            genesis_state.clone(),
+            spec.clone(),
+        )?);
+
         let finalized_head = RwLock::new(CheckPoint::new(
             genesis_block.clone(),
             block_root,
@@ -127,6 +134,7 @@ where
             slot_clock,
             attestation_aggregator,
             state: RwLock::new(genesis_state.clone()),
+            cached_state,
             finalized_head,
             canonical_head,
             spec,
@@ -280,7 +288,7 @@ where
             validator_index
         );
         if let Some((slot, shard, _committee)) = self
-            .state
+            .cached_state
             .read()
             .attestation_slot_and_shard_for_validator(validator_index, &self.spec)?
         {
@@ -338,9 +346,7 @@ where
         let aggregation_outcome = self
             .attestation_aggregator
             .write()
-            .process_free_attestation(&self.state.read(), &free_attestation, &self.spec)?;
-        // TODO: Check this comment
-        //.map_err(|e| e.into())?;
+            .process_free_attestation(&self.cached_state.read(), &free_attestation, &self.spec)?;
 
         // return if the attestation is invalid
         if !aggregation_outcome.valid {
@@ -495,6 +501,9 @@ where
             );
             // Update the local state variable.
             *self.state.write() = state.clone();
+            // Update the cached state variable.
+            *self.cached_state.write() =
+                CachedBeaconState::from_beacon_state(state.clone(), self.spec.clone())?;
         }
 
         Ok(BlockProcessingOutcome::ValidBlock(ValidBlock::Processed))
@@ -543,9 +552,15 @@ where
             },
         };
 
-        state
-            .per_block_processing_without_verifying_block_signature(&block, &self.spec)
-            .ok()?;
+        trace!("BeaconChain::produce_block: updating state for new block.",);
+
+        let result =
+            state.per_block_processing_without_verifying_block_signature(&block, &self.spec);
+        trace!(
+            "BeaconNode::produce_block: state processing result: {:?}",
+            result
+        );
+        result.ok()?;
 
         let state_root = state.canonical_root();
 

beacon_node/beacon_chain/src/cached_beacon_state.rs

@@ -1,49 +1,149 @@
+use log::debug;
+use std::collections::HashMap;
 use types::{beacon_state::BeaconStateError, BeaconState, ChainSpec, Epoch, Slot};
 
-pub const CACHED_EPOCHS: usize = 3; // previous, current, next.
+pub const CACHE_PREVIOUS: bool = false;
+pub const CACHE_CURRENT: bool = true;
+pub const CACHE_NEXT: bool = false;
 
 pub type CrosslinkCommittees = Vec<(Vec<usize>, u64)>;
+pub type Shard = u64;
+pub type CommitteeIndex = u64;
+pub type AttestationDuty = (Slot, Shard, CommitteeIndex);
+pub type AttestationDutyMap = HashMap<u64, AttestationDuty>;
 
-pub struct CachedBeaconState<'a> {
-    state: BeaconState,
-    crosslinks: Vec<Vec<CrosslinkCommittees>>,
-    spec: &'a ChainSpec,
+// TODO: CachedBeaconState is presently duplicating `BeaconState` and `ChainSpec`. This is a
+// massive memory waste, switch them to references.
+
+pub struct CachedBeaconState {
+    pub state: BeaconState,
+    committees: Vec<Vec<CrosslinkCommittees>>,
+    attestation_duties: Vec<AttestationDutyMap>,
+    next_epoch: Epoch,
+    current_epoch: Epoch,
+    previous_epoch: Epoch,
+    spec: ChainSpec,
 }
 
-impl<'a> CachedBeaconState<'a> {
+impl CachedBeaconState {
     pub fn from_beacon_state(
         state: BeaconState,
-        spec: &'a ChainSpec,
+        spec: ChainSpec,
     ) -> Result<Self, BeaconStateError> {
-        let current_epoch = state.current_epoch(spec);
+        let current_epoch = state.current_epoch(&spec);
         let previous_epoch = if current_epoch == spec.genesis_epoch {
             current_epoch
         } else {
             current_epoch.saturating_sub(1_u64)
         };
-        let next_epoch = state.next_epoch(spec);
+        let next_epoch = state.next_epoch(&spec);
 
-        let mut crosslinks: Vec<Vec<CrosslinkCommittees>> = Vec::with_capacity(3);
-        crosslinks.push(committees_for_all_slots(&state, previous_epoch, spec)?);
-        crosslinks.push(committees_for_all_slots(&state, current_epoch, spec)?);
-        crosslinks.push(committees_for_all_slots(&state, next_epoch, spec)?);
+        let mut committees: Vec<Vec<CrosslinkCommittees>> = Vec::with_capacity(3);
+        let mut attestation_duties: Vec<AttestationDutyMap> = Vec::with_capacity(3);
+
+        if CACHE_PREVIOUS {
+            debug!("CachedBeaconState::from_beacon_state: building previous epoch cache.");
+            let cache = build_epoch_cache(&state, previous_epoch, &spec)?;
+            committees.push(cache.committees);
+            attestation_duties.push(cache.attestation_duty_map);
+        } else {
+            committees.push(vec![]);
+            attestation_duties.push(HashMap::new());
+        }
+        if CACHE_CURRENT {
+            debug!("CachedBeaconState::from_beacon_state: building current epoch cache.");
+            let cache = build_epoch_cache(&state, current_epoch, &spec)?;
+            committees.push(cache.committees);
+            attestation_duties.push(cache.attestation_duty_map);
+        } else {
+            committees.push(vec![]);
+            attestation_duties.push(HashMap::new());
+        }
+        if CACHE_NEXT {
+            debug!("CachedBeaconState::from_beacon_state: building next epoch cache.");
+            let cache = build_epoch_cache(&state, next_epoch, &spec)?;
+            committees.push(cache.committees);
+            attestation_duties.push(cache.attestation_duty_map);
+        } else {
+            committees.push(vec![]);
+            attestation_duties.push(HashMap::new());
+        }
 
         Ok(Self {
             state,
-            crosslinks,
+            committees,
+            attestation_duties,
+            next_epoch,
+            current_epoch,
+            previous_epoch,
             spec,
         })
     }
+
+    fn slot_to_cache_index(&self, slot: Slot) -> Option<usize> {
+        match slot.epoch(self.spec.epoch_length) {
+            epoch if (epoch == self.previous_epoch) & CACHE_PREVIOUS => Some(0),
+            epoch if (epoch == self.current_epoch) & CACHE_CURRENT => Some(1),
+            epoch if (epoch == self.next_epoch) & CACHE_NEXT => Some(2),
+            _ => None,
+        }
+    }
+
+    /// Returns the `slot`, `shard` and `committee_index` for which a validator must produce an
+    /// attestation.
+    ///
+    /// Cached method.
+    ///
+    /// Spec v0.2.0
+    pub fn attestation_slot_and_shard_for_validator(
+        &self,
+        validator_index: usize,
+        _spec: &ChainSpec,
+    ) -> Result<Option<(Slot, u64, u64)>, BeaconStateError> {
+        // Get the result for this epoch.
+        let cache_index = self
+            .slot_to_cache_index(self.state.slot)
+            .expect("Current epoch should always have a cache index.");
+
+        let duties = self.attestation_duties[cache_index]
+            .get(&(validator_index as u64))
+            .and_then(|tuple| Some(*tuple));
+
+        Ok(duties)
+    }
 }
 
-fn committees_for_all_slots(
+struct EpochCacheResult {
+    committees: Vec<CrosslinkCommittees>,
+    attestation_duty_map: AttestationDutyMap,
+}
+
+fn build_epoch_cache(
     state: &BeaconState,
     epoch: Epoch,
     spec: &ChainSpec,
-) -> Result<Vec<CrosslinkCommittees>, BeaconStateError> {
-    let mut crosslinks: Vec<CrosslinkCommittees> = Vec::with_capacity(spec.epoch_length as usize);
+) -> Result<EpochCacheResult, BeaconStateError> {
+    let mut epoch_committees: Vec<CrosslinkCommittees> =
+        Vec::with_capacity(spec.epoch_length as usize);
+    let mut attestation_duty_map: AttestationDutyMap = HashMap::new();
+
     for slot in epoch.slot_iter(spec.epoch_length) {
-        crosslinks.push(state.get_crosslink_committees_at_slot(slot, false, spec)?)
+        let slot_committees = state.get_crosslink_committees_at_slot(slot, false, spec)?;
+
+        for (committee, shard) in slot_committees {
+            for (committee_index, validator_index) in committee.iter().enumerate() {
+                attestation_duty_map.insert(
+                    *validator_index as u64,
+                    (slot, shard, committee_index as u64),
+                );
+            }
+        }
+
+        epoch_committees.push(state.get_crosslink_committees_at_slot(slot, false, spec)?)
     }
-    Ok(crosslinks)
+
+    Ok(EpochCacheResult {
+        committees: epoch_committees,
+        attestation_duty_map,
+    })
 }

beacon_node/beacon_chain/test_harness/tests/chain.rs

@@ -6,8 +6,8 @@ use types::ChainSpec;
 #[test]
 fn it_can_build_on_genesis_block() {
     Builder::from_env(Env::default().default_filter_or("trace")).init();
-    let spec = ChainSpec::few_validators();
 
+    let spec = ChainSpec::few_validators();
     let validator_count = 8;
 
     let mut harness = BeaconChainHarness::new(spec, validator_count as usize);
@@ -18,13 +18,14 @@ fn it_can_build_on_genesis_block() {
 #[test]
 #[ignore]
 fn it_can_produce_past_first_epoch_boundary() {
-    Builder::from_env(Env::default().default_filter_or("trace")).init();
+    Builder::from_env(Env::default().default_filter_or("debug")).init();
 
-    let validator_count = 100;
+    let spec = ChainSpec::few_validators();
+    let validator_count = 8;
 
     debug!("Starting harness build...");
 
-    let mut harness = BeaconChainHarness::new(ChainSpec::foundation(), validator_count);
+    let mut harness = BeaconChainHarness::new(spec, validator_count);
 
     debug!("Harness built, tests starting..");