Optimization: avoid recomputing known state roots (#762)

* Start adding optimization

* Add temp fix for protobuf issue

* Fix compile errors

* Fix protobuf import

beacon_node/beacon_chain/src/beacon_chain.rs

@@ -515,7 +515,9 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
                     });
                 }
 
-                match per_slot_processing(&mut state, &self.spec) {
+                // Note: supplying some `state_root` when it is known would be a cheap and easy
+                // optimization.
+                match per_slot_processing(&mut state, None, &self.spec) {
                     Ok(()) => (),
                     Err(e) => {
                         warn!(
@@ -887,7 +889,7 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
                     .start_slot(T::EthSpec::slots_per_epoch())
                     .as_u64()
             {
-                per_slot_processing(&mut state, &self.spec)?;
+                per_slot_processing(&mut state, None, &self.spec)?;
             }
 
             state.build_committee_cache(RelativeEpoch::Current, &self.spec)?;
@@ -1258,7 +1260,14 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
             if i > 0 {
                 intermediate_states.push(state.clone());
             }
-            per_slot_processing(&mut state, &self.spec)?;
+
+            let state_root = if i == 0 {
+                Some(parent_block.state_root)
+            } else {
+                None
+            };
+
+            per_slot_processing(&mut state, state_root, &self.spec)?;
         }
 
         metrics::stop_timer(catchup_timer);
@@ -1426,8 +1435,11 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
             .ok_or_else(|| BlockProductionError::NoEth1ChainConnection)?;
 
         // If required, transition the new state to the present slot.
+        //
+        // Note: supplying some `state_root` when it it is known would be a cheap and easy
+        // optimization.
         while state.slot < produce_at_slot {
-            per_slot_processing(&mut state, &self.spec)?;
+            per_slot_processing(&mut state, None, &self.spec)?;
         }
 
         state.build_committee_cache(RelativeEpoch::Current, &self.spec)?;

beacon_node/beacon_chain/src/fork_choice.rs

@@ -155,7 +155,7 @@ impl<T: BeaconChainTypes> ForkChoice<T> {
 
             // Fast-forward the state to the start slot of the epoch where it was justified.
             for _ in block.slot.as_u64()..block_justified_slot.as_u64() {
-                per_slot_processing(&mut state, &chain.spec)
+                per_slot_processing(&mut state, None, &chain.spec)
                     .map_err(BeaconChainError::SlotProcessingError)?
             }
 

beacon_node/beacon_chain/src/test_utils.rs

@@ -287,7 +287,7 @@ where
         }
 
         while state.slot < slot {
-            per_slot_processing(&mut state, &self.spec)
+            per_slot_processing(&mut state, None, &self.spec)
                 .expect("should be able to advance state to slot");
         }
 

beacon_node/beacon_chain/tests/tests.rs

@@ -41,7 +41,7 @@ fn massive_skips() {
 
     // Run per_slot_processing until it returns an error.
     let error = loop {
-        match per_slot_processing(&mut state, spec) {
+        match per_slot_processing(&mut state, None, spec) {
             Ok(_) => continue,
             Err(e) => break e,
         }

beacon_node/rest_api/src/helpers.rs

@@ -216,7 +216,7 @@ pub fn state_root_at_slot<T: BeaconChainTypes>(
             // Ensure the next epoch state caches are built in case of an epoch transition.
             state.build_committee_cache(RelativeEpoch::Next, spec)?;
 
-            state_processing::per_slot_processing(&mut state, spec)?;
+            state_processing::per_slot_processing(&mut state, None, spec)?;
         }
 
         // Note: this is an expensive operation. Once the tree hash cache is implement it may be

beacon_node/store/src/hot_cold_store.rs

@@ -424,7 +424,7 @@ impl<E: EthSpec> HotColdDB<E> {
 
         for block in blocks {
             while state.slot < block.slot {
-                per_slot_processing(&mut state, &self.spec)
+                per_slot_processing(&mut state, None, &self.spec)
                     .map_err(HotColdDbError::BlockReplaySlotError)?;
             }
             per_block_processing(
@@ -438,7 +438,7 @@ impl<E: EthSpec> HotColdDB<E> {
         }
 
         while state.slot < target_slot {
-            per_slot_processing(&mut state, &self.spec)
+            per_slot_processing(&mut state, None, &self.spec)
                 .map_err(HotColdDbError::BlockReplaySlotError)?;
         }
 

eth2/state_processing/src/per_slot_processing.rs

@@ -9,12 +9,17 @@ pub enum Error {
 
 /// Advances a state forward by one slot, performing per-epoch processing if required.
 ///
+/// If the root of the supplied `state` is known, then it can be passed as `state_root`. If
+/// `state_root` is `None`, the root of `state` will be computed using a cached tree hash.
+/// Providing the `state_root` makes this function several orders of magniude faster.
+///
 /// Spec v0.9.1
 pub fn per_slot_processing<T: EthSpec>(
     state: &mut BeaconState<T>,
+    state_root: Option<Hash256>,
     spec: &ChainSpec,
 ) -> Result<(), Error> {
-    cache_state(state)?;
+    cache_state(state, state_root)?;
 
     if state.slot > spec.genesis_slot && (state.slot + 1) % T::slots_per_epoch() == 0 {
         per_epoch_processing(state, spec)?;
@@ -25,8 +30,15 @@ pub fn per_slot_processing<T: EthSpec>(
     Ok(())
 }
 
-fn cache_state<T: EthSpec>(state: &mut BeaconState<T>) -> Result<(), Error> {
+fn cache_state<T: EthSpec>(
-    let previous_state_root = state.update_tree_hash_cache()?;
+    state: &mut BeaconState<T>,
+    state_root: Option<Hash256>,
+) -> Result<(), Error> {
+    let previous_state_root = if let Some(root) = state_root {
+        root
+    } else {
+        state.update_tree_hash_cache()?
+    };
 
     // Note: increment the state slot here to allow use of our `state_root` and `block_root`
     // getter/setter functions.

lcli/src/transition_blocks.rs

@@ -56,7 +56,7 @@ fn do_transition<T: EthSpec>(
 
     // Transition the parent state to the block slot.
     for i in pre_state.slot.as_u64()..block.slot.as_u64() {
-        per_slot_processing(&mut pre_state, spec)
+        per_slot_processing(&mut pre_state, None, spec)
             .map_err(|e| format!("Failed to advance slot on iteration {}: {:?}", i, e))?;
     }
 

lighthouse/Cargo.toml

@@ -19,3 +19,5 @@ environment = { path = "./environment" }
 futures = "0.1.25"
 validator_client = { "path" = "../validator_client" }
 account_manager = { "path" = "../account_manager" }
+# TODO: remove this once @agemanning adds a permanent fix.
+protobuf = "=2.8.1"

tests/ef_tests/src/cases/sanity_blocks.rs

@@ -73,7 +73,7 @@ impl<E: EthSpec> Case for SanityBlocks<E> {
             .iter()
             .try_for_each(|block| {
                 while state.slot < block.slot {
-                    per_slot_processing(&mut state, spec).unwrap();
+                    per_slot_processing(&mut state, None, spec).unwrap();
                 }
 
                 state

tests/ef_tests/src/cases/sanity_slots.rs

@@ -66,7 +66,7 @@ impl<E: EthSpec> Case for SanitySlots<E> {
         state.build_all_caches(spec).unwrap();
 
         let mut result = (0..self.slots)
-            .try_for_each(|_| per_slot_processing(&mut state, spec))
+            .try_for_each(|_| per_slot_processing(&mut state, None, spec))
             .map(|_| state);
 
         compare_beacon_state_results_without_caches(&mut result, &mut expected)