Fix tests, probably

beacon_node/beacon_chain/tests/store_tests.rs

@@ -3535,13 +3535,12 @@ fn assert_chains_pretty_much_the_same<T: BeaconChainTypes>(a: &BeaconChain<T>, b
         a_head.beacon_block, b_head.beacon_block,
         "head blocks should be equal"
     );
-    // Clone with committee caches only to prevent other caches from messing with the equality
-    // check.
-    assert_eq!(
-        a_head.beacon_state.clone_with_only_committee_caches(),
-        b_head.beacon_state.clone_with_only_committee_caches(),
-        "head states should be equal"
-    );
+    // Drop all caches to prevent them messing with the equality check.
+    let mut a_head_state = a_head.beacon_state.clone();
+    a_head_state.drop_all_caches().unwrap();
+    let mut b_head_state = b_head.beacon_state.clone();
+    b_head_state.drop_all_caches().unwrap();
+    assert_eq!(a_head_state, b_head_state, "head states should be equal");
     assert_eq!(a.heads(), b.heads(), "heads() should be equal");
     assert_eq!(
         a.genesis_block_root, b.genesis_block_root,

consensus/state_processing/src/epoch_cache.rs

@@ -12,16 +12,21 @@ pub struct PreEpochCache {
 }
 
 impl PreEpochCache {
-    pub fn new_for_next_epoch<E: EthSpec>(state: &BeaconState<E>) -> Result<Self, EpochCacheError> {
+    pub fn new_for_next_epoch<E: EthSpec>(
+        state: &mut BeaconState<E>,
+    ) -> Result<Self, EpochCacheError> {
         // The decision block root for the next epoch is the latest block root from this epoch.
         let latest_block_header = state.latest_block_header();
 
-        // State root should already have been filled in by `process_slot`, except in the case
-        // of a `partial_state_advance`.
-        if latest_block_header.state_root.is_zero() {
-            return Err(EpochCacheError::ZeroStateRoot);
-        }
-        let decision_block_root = latest_block_header.canonical_root();
+        let decision_block_root = if !latest_block_header.state_root.is_zero() {
+            latest_block_header.canonical_root()
+        } else {
+            // State root should already have been filled in by `process_slot`, except in the case
+            // of a `partial_state_advance`. Once we have tree-states this can be an error, and
+            // `self` can be immutable.
+            let state_root = state.update_tree_hash_cache()?;
+            state.get_latest_block_root(state_root)
+        };
 
         let epoch_key = EpochCacheKey {
             epoch: state.next_epoch()?,

consensus/types/src/epoch_cache.rs

@@ -40,7 +40,6 @@ pub struct EpochCacheKey {
 pub enum EpochCacheError {
     IncorrectEpoch { cache: Epoch, state: Epoch },
     IncorrectDecisionBlock { cache: Hash256, state: Hash256 },
-    ZeroStateRoot,
     ValidatorIndexOutOfBounds { validator_index: usize },
     EffectiveBalanceOutOfBounds { effective_balance_eth: usize },
     InvalidSlot { slot: Slot },