From 8824c0b71f7e01c604d7db35c9e7309aca4497de Mon Sep 17 00:00:00 2001
From: Michael Sproul <michael@sigmaprime.io>
Date: Thu, 4 Apr 2024 19:09:50 +1100
Subject: [PATCH] Fix tests, probably

---
 beacon_node/beacon_chain/tests/store_tests.rs | 13 ++++++-------
 consensus/state_processing/src/epoch_cache.rs | 19 ++++++++++++-------
 consensus/types/src/epoch_cache.rs            |  1 -
 3 files changed, 18 insertions(+), 15 deletions(-)

diff --git a/beacon_node/beacon_chain/tests/store_tests.rs b/beacon_node/beacon_chain/tests/store_tests.rs
index ff201729821..66f4138afb4 100644
--- a/beacon_node/beacon_chain/tests/store_tests.rs
+++ b/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,
diff --git a/consensus/state_processing/src/epoch_cache.rs b/consensus/state_processing/src/epoch_cache.rs
index 8f8006e2bfa..b2f2d85407e 100644
--- a/consensus/state_processing/src/epoch_cache.rs
+++ b/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()?,
diff --git a/consensus/types/src/epoch_cache.rs b/consensus/types/src/epoch_cache.rs
index 86ac7b45cca..b447e9b71e0 100644
--- a/consensus/types/src/epoch_cache.rs
+++ b/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 },