Unimplement TreeHash for BeaconState (sigp#6083)

* Unimplement `TreeHash` for `BeaconState`

beacon_node/beacon_chain/src/block_verification.rs

@@ -93,7 +93,6 @@ use std::io::Write;
 use std::sync::Arc;
 use store::{Error as DBError, HotStateSummary, KeyValueStore, StoreOp};
 use task_executor::JoinHandle;
-use tree_hash::TreeHash;
 use types::{
     BeaconBlockRef, BeaconState, BeaconStateError, ChainSpec, Epoch, EthSpec, ExecutionBlockHash,
     Hash256, InconsistentFork, PublicKey, PublicKeyBytes, RelativeEpoch, SignedBeaconBlock,
@@ -2107,7 +2106,14 @@ pub fn verify_header_signature<T: BeaconChainTypes, Err: BlockBlobError>(
 
 fn write_state<E: EthSpec>(prefix: &str, state: &BeaconState<E>, log: &Logger) {
     if WRITE_BLOCK_PROCESSING_SSZ {
-        let root = state.tree_hash_root();
+        let mut state = state.clone();
+        let Ok(root) = state.canonical_root() else {
+            error!(
+                log,
+                "Unable to hash state for writing";
+            );
+            return;
+        };
         let filename = format!("{}_slot_{}_root_{}.ssz", prefix, state.slot(), root);
         let mut path = std::env::temp_dir().join("lighthouse");
         let _ = fs::create_dir_all(path.clone());

beacon_node/beacon_chain/src/builder.rs

@@ -1195,7 +1195,7 @@ mod test {
 
         let head = chain.head_snapshot();
 
-        let state = &head.beacon_state;
+        let mut state = head.beacon_state.clone();
         let block = &head.beacon_block;
 
         assert_eq!(state.slot(), Slot::new(0), "should start from genesis");
@@ -1206,7 +1206,7 @@ mod test {
         );
         assert_eq!(
             block.state_root(),
-            state.canonical_root(),
+            state.canonical_root().unwrap(),
             "block should have correct state root"
         );
         assert_eq!(

beacon_node/beacon_chain/src/data_availability_checker/overflow_lru_cache.rs

@@ -720,12 +720,12 @@ mod test {
         let mut state_roots = Vec::new();
         // Get enough blocks to fill the cache to capacity, ensuring all blocks have blobs
         while pending_blocks.len() < capacity {
-            let (pending_block, _) = availability_pending_block(&harness).await;
+            let (mut pending_block, _) = availability_pending_block(&harness).await;
             if pending_block.num_blobs_expected() == 0 {
                 // we need blocks with blobs
                 continue;
             }
-            let state_root = pending_block.import_data.state.canonical_root();
+            let state_root = pending_block.import_data.state.canonical_root().unwrap();
             states.push(pending_block.import_data.state.clone());
             pending_blocks.push_back(pending_block);
             state_roots.push(state_root);

beacon_node/beacon_chain/src/data_availability_checker/state_lru_cache.rs

@@ -209,10 +209,10 @@ impl<T: BeaconChainTypes> StateLRUCache<T> {
 impl<E: EthSpec> From<AvailabilityPendingExecutedBlock<E>>
     for DietAvailabilityPendingExecutedBlock<E>
 {
-    fn from(value: AvailabilityPendingExecutedBlock<E>) -> Self {
+    fn from(mut value: AvailabilityPendingExecutedBlock<E>) -> Self {
         Self {
             block: value.block,
-            state_root: value.import_data.state.canonical_root(),
+            state_root: value.import_data.state.canonical_root().unwrap(),
             parent_block: value.import_data.parent_block,
             parent_eth1_finalization_data: value.import_data.parent_eth1_finalization_data,
             confirmed_state_roots: value.import_data.confirmed_state_roots,

beacon_node/beacon_chain/src/test_utils.rs

@@ -2246,7 +2246,7 @@ where
                 .unwrap();
             state = new_state;
             block_hash_from_slot.insert(*slot, block_hash);
-            state_hash_from_slot.insert(*slot, state.tree_hash_root().into());
+            state_hash_from_slot.insert(*slot, state.canonical_root().unwrap().into());
             latest_block_hash = Some(block_hash);
         }
         (

beacon_node/beacon_chain/tests/store_tests.rs

@@ -35,7 +35,6 @@ use store::{
 };
 use tempfile::{tempdir, TempDir};
 use tokio::time::sleep;
-use tree_hash::TreeHash;
 use types::test_utils::{SeedableRng, XorShiftRng};
 use types::*;
 
@@ -199,8 +198,8 @@ async fn heal_freezer_block_roots_with_skip_slots() {
     );
     let harness = get_harness(store.clone(), LOW_VALIDATOR_COUNT);
 
-    let current_state = harness.get_current_state();
-    let state_root = harness.get_current_state().tree_hash_root();
+    let mut current_state = harness.get_current_state();
+    let state_root = current_state.canonical_root().unwrap();
     let all_validators = &harness.get_all_validators();
     harness
         .add_attested_blocks_at_slots(
@@ -611,12 +610,13 @@ async fn epoch_boundary_state_attestation_processing() {
             .get_blinded_block(&block_root)
             .unwrap()
             .expect("block exists");
-        let epoch_boundary_state = store
+        let mut epoch_boundary_state = store
             .load_epoch_boundary_state(&block.state_root())
             .expect("no error")
             .expect("epoch boundary state exists");
+        let ebs_state_root = epoch_boundary_state.canonical_root().unwrap();
         let ebs_of_ebs = store
-            .load_epoch_boundary_state(&epoch_boundary_state.canonical_root())
+            .load_epoch_boundary_state(&ebs_state_root)
             .expect("no error")
             .expect("ebs of ebs exists");
         assert_eq!(epoch_boundary_state, ebs_of_ebs);
@@ -2604,9 +2604,9 @@ async fn weak_subjectivity_sync_test(slots: Vec<Slot>, checkpoint_slot: Slot) {
         .unwrap()
         .map(Result::unwrap)
     {
-        let state = store.get_state(&state_root, Some(slot)).unwrap().unwrap();
+        let mut state = store.get_state(&state_root, Some(slot)).unwrap().unwrap();
         assert_eq!(state.slot(), slot);
-        assert_eq!(state.canonical_root(), state_root);
+        assert_eq!(state.canonical_root().unwrap(), state_root);
     }
 
     // Anchor slot is still set to the slot of the checkpoint block.

beacon_node/http_api/tests/broadcast_validation_tests.rs

@@ -7,7 +7,6 @@ use eth2::types::{BroadcastValidation, PublishBlockRequest};
 use http_api::test_utils::InteractiveTester;
 use http_api::{publish_blinded_block, publish_block, reconstruct_block, ProvenancedBlock};
 use std::sync::Arc;
-use tree_hash::TreeHash;
 use types::{Epoch, EthSpec, ForkName, Hash256, MainnetEthSpec, Slot};
 use warp::Rejection;
 use warp_utils::reject::CustomBadRequest;
@@ -353,13 +352,20 @@ pub async fn consensus_partial_pass_only_consensus() {
     let slot_b = slot_a + 1;
 
     let state_a = tester.harness.get_current_state();
-    let ((block_a, _), state_after_a) = tester.harness.make_block(state_a.clone(), slot_b).await;
-    let ((block_b, blobs_b), state_after_b) = tester.harness.make_block(state_a, slot_b).await;
+    let ((block_a, _), mut state_after_a) =
+        tester.harness.make_block(state_a.clone(), slot_b).await;
+    let ((block_b, blobs_b), mut state_after_b) = tester.harness.make_block(state_a, slot_b).await;
     let block_b_root = block_b.canonical_root();
 
     /* check for `make_block` curios */
-    assert_eq!(block_a.state_root(), state_after_a.tree_hash_root());
-    assert_eq!(block_b.state_root(), state_after_b.tree_hash_root());
+    assert_eq!(
+        block_a.state_root(),
+        state_after_a.canonical_root().unwrap()
+    );
+    assert_eq!(
+        block_b.state_root(),
+        state_after_b.canonical_root().unwrap()
+    );
     assert_ne!(block_a.state_root(), block_b.state_root());
 
     let gossip_block_contents_b = PublishBlockRequest::new(block_b, blobs_b)
@@ -516,13 +522,19 @@ pub async fn equivocation_consensus_early_equivocation() {
     let slot_b = slot_a + 1;
 
     let state_a = tester.harness.get_current_state();
-    let ((block_a, blobs_a), state_after_a) =
+    let ((block_a, blobs_a), mut state_after_a) =
         tester.harness.make_block(state_a.clone(), slot_b).await;
-    let ((block_b, blobs_b), state_after_b) = tester.harness.make_block(state_a, slot_b).await;
+    let ((block_b, blobs_b), mut state_after_b) = tester.harness.make_block(state_a, slot_b).await;
 
     /* check for `make_block` curios */
-    assert_eq!(block_a.state_root(), state_after_a.tree_hash_root());
-    assert_eq!(block_b.state_root(), state_after_b.tree_hash_root());
+    assert_eq!(
+        block_a.state_root(),
+        state_after_a.canonical_root().unwrap()
+    );
+    assert_eq!(
+        block_b.state_root(),
+        state_after_b.canonical_root().unwrap()
+    );
     assert_ne!(block_a.state_root(), block_b.state_root());
 
     /* submit `block_a` as valid */
@@ -642,13 +654,19 @@ pub async fn equivocation_consensus_late_equivocation() {
     let slot_b = slot_a + 1;
 
     let state_a = tester.harness.get_current_state();
-    let ((block_a, blobs_a), state_after_a) =
+    let ((block_a, blobs_a), mut state_after_a) =
         tester.harness.make_block(state_a.clone(), slot_b).await;
-    let ((block_b, blobs_b), state_after_b) = tester.harness.make_block(state_a, slot_b).await;
+    let ((block_b, blobs_b), mut state_after_b) = tester.harness.make_block(state_a, slot_b).await;
 
     /* check for `make_block` curios */
-    assert_eq!(block_a.state_root(), state_after_a.tree_hash_root());
-    assert_eq!(block_b.state_root(), state_after_b.tree_hash_root());
+    assert_eq!(
+        block_a.state_root(),
+        state_after_a.canonical_root().unwrap()
+    );
+    assert_eq!(
+        block_b.state_root(),
+        state_after_b.canonical_root().unwrap()
+    );
     assert_ne!(block_a.state_root(), block_b.state_root());
 
     let gossip_block_contents_b = PublishBlockRequest::new(block_b, blobs_b)
@@ -1135,15 +1153,21 @@ pub async fn blinded_equivocation_consensus_early_equivocation() {
     let slot_b = slot_a + 1;
 
     let state_a = tester.harness.get_current_state();
-    let (block_a, state_after_a) = tester
+    let (block_a, mut state_after_a) = tester
         .harness
         .make_blinded_block(state_a.clone(), slot_b)
         .await;
-    let (block_b, state_after_b) = tester.harness.make_blinded_block(state_a, slot_b).await;
+    let (block_b, mut state_after_b) = tester.harness.make_blinded_block(state_a, slot_b).await;
 
     /* check for `make_blinded_block` curios */
-    assert_eq!(block_a.state_root(), state_after_a.tree_hash_root());
-    assert_eq!(block_b.state_root(), state_after_b.tree_hash_root());
+    assert_eq!(
+        block_a.state_root(),
+        state_after_a.canonical_root().unwrap()
+    );
+    assert_eq!(
+        block_b.state_root(),
+        state_after_b.canonical_root().unwrap()
+    );
     assert_ne!(block_a.state_root(), block_b.state_root());
 
     /* submit `block_a` as valid */
@@ -1259,16 +1283,22 @@ pub async fn blinded_equivocation_consensus_late_equivocation() {
     let slot_b = slot_a + 1;
 
     let state_a = tester.harness.get_current_state();
-    let (block_a, state_after_a) = tester
+    let (block_a, mut state_after_a) = tester
         .harness
         .make_blinded_block(state_a.clone(), slot_b)
         .await;
-    let (block_b, state_after_b) = tester.harness.make_blinded_block(state_a, slot_b).await;
+    let (block_b, mut state_after_b) = tester.harness.make_blinded_block(state_a, slot_b).await;
     let block_b = Arc::new(block_b);
 
     /* check for `make_blinded_block` curios */
-    assert_eq!(block_a.state_root(), state_after_a.tree_hash_root());
-    assert_eq!(block_b.state_root(), state_after_b.tree_hash_root());
+    assert_eq!(
+        block_a.state_root(),
+        state_after_a.canonical_root().unwrap()
+    );
+    assert_eq!(
+        block_b.state_root(),
+        state_after_b.canonical_root().unwrap()
+    );
     assert_ne!(block_a.state_root(), block_b.state_root());
 
     let unblinded_block_a = reconstruct_block(

beacon_node/http_api/tests/fork_tests.rs

@@ -55,7 +55,7 @@ async fn sync_committee_duties_across_fork() {
     // though the head state hasn't transitioned yet.
     let fork_slot = fork_epoch.start_slot(E::slots_per_epoch());
     let (genesis_state, genesis_state_root) = harness.get_current_state_and_root();
-    let (_, state) = harness
+    let (_, mut state) = harness
         .add_attested_block_at_slot(
             fork_slot - 1,
             genesis_state,
@@ -76,7 +76,7 @@ async fn sync_committee_duties_across_fork() {
     assert_eq!(sync_duties.len(), E::sync_committee_size());
 
     // After applying a block at the fork slot the duties should remain unchanged.
-    let state_root = state.canonical_root();
+    let state_root = state.canonical_root().unwrap();
     harness
         .add_attested_block_at_slot(fork_slot, state, state_root, &all_validators)
         .await
@@ -257,7 +257,7 @@ async fn sync_committee_indices_across_fork() {
     // applied.
     let fork_slot = fork_epoch.start_slot(E::slots_per_epoch());
     let (genesis_state, genesis_state_root) = harness.get_current_state_and_root();
-    let (_, state) = harness
+    let (_, mut state) = harness
         .add_attested_block_at_slot(
             fork_slot - 1,
             genesis_state,
@@ -295,7 +295,7 @@ async fn sync_committee_indices_across_fork() {
 
     // Once the head is updated it should be useable for requests, including in the next sync
     // committee period.
-    let state_root = state.canonical_root();
+    let state_root = state.canonical_root().unwrap();
     harness
         .add_attested_block_at_slot(fork_slot + 1, state, state_root, &all_validators)
         .await

beacon_node/http_api/tests/interactive_tests.rs

@@ -17,7 +17,6 @@ use state_processing::{
 use std::collections::HashMap;
 use std::sync::Arc;
 use std::time::Duration;
-use tree_hash::TreeHash;
 use types::{
     Address, Epoch, EthSpec, ExecPayload, ExecutionBlockHash, ForkName, MainnetEthSpec,
     MinimalEthSpec, ProposerPreparationData, Slot,
@@ -515,16 +514,17 @@ pub async fn proposer_boost_re_org_test(
     }
 
     harness.advance_slot();
-    let (block_a_root, block_a, state_a) = harness
+    let (block_a_root, block_a, mut state_a) = harness
         .add_block_at_slot(slot_a, harness.get_current_state())
         .await
         .unwrap();
+    let state_a_root = state_a.canonical_root().unwrap();
 
     // Attest to block A during slot A.
     let (block_a_parent_votes, _) = harness.make_attestations_with_limit(
         &all_validators,
         &state_a,
-        state_a.canonical_root(),
+        state_a_root,
         block_a_root,
         slot_a,
         num_parent_votes,
@@ -538,7 +538,7 @@ pub async fn proposer_boost_re_org_test(
     let (block_a_empty_votes, block_a_attesters) = harness.make_attestations_with_limit(
         &all_validators,
         &state_a,
-        state_a.canonical_root(),
+        state_a_root,
         block_a_root,
         slot_b,
         num_empty_votes,
@@ -553,6 +553,7 @@ pub async fn proposer_boost_re_org_test(
 
     // Produce block B and process it halfway through the slot.
     let (block_b, mut state_b) = harness.make_block(state_a.clone(), slot_b).await;
+    let state_b_root = state_b.canonical_root().unwrap();
     let block_b_root = block_b.0.canonical_root();
 
     let obs_time = slot_clock.start_of(slot_b).unwrap() + slot_clock.slot_duration() / 2;
@@ -570,7 +571,7 @@ pub async fn proposer_boost_re_org_test(
     let (block_b_head_votes, _) = harness.make_attestations_with_limit(
         &remaining_attesters,
         &state_b,
-        state_b.canonical_root(),
+        state_b_root,
         block_b_root.into(),
         slot_b,
         num_head_votes,
@@ -774,32 +775,34 @@ pub async fn fork_choice_before_proposal() {
     let slot_d = slot_a + 3;
 
     let state_a = harness.get_current_state();
-    let (block_b, state_b) = harness.make_block(state_a.clone(), slot_b).await;
+    let (block_b, mut state_b) = harness.make_block(state_a.clone(), slot_b).await;
     let block_root_b = harness
         .process_block(slot_b, block_b.0.canonical_root(), block_b)
         .await
         .unwrap();
+    let state_root_b = state_b.canonical_root().unwrap();
 
     // Create attestations to B but keep them in reserve until after C has been processed.
     let attestations_b = harness.make_attestations(
         &all_validators,
         &state_b,
-        state_b.tree_hash_root(),
+        state_root_b,
         block_root_b,
         slot_b,
     );
 
-    let (block_c, state_c) = harness.make_block(state_a, slot_c).await;
+    let (block_c, mut state_c) = harness.make_block(state_a, slot_c).await;
     let block_root_c = harness
         .process_block(slot_c, block_c.0.canonical_root(), block_c.clone())
         .await
         .unwrap();
+    let state_root_c = state_c.canonical_root().unwrap();
 
     // Create attestations to C from a small number of validators and process them immediately.
     let attestations_c = harness.make_attestations(
         &all_validators[..validator_count / 2],
         &state_c,
-        state_c.tree_hash_root(),
+        state_root_c,
         block_root_c,
         slot_c,
     );