Implement tree hash caching (#584)

* Implement basic tree hash caching

* Use spaces to indent top-level Cargo.toml

* Optimize BLS tree hash by hashing bytes directly

* Implement tree hash caching for validator registry

* Persist BeaconState tree hash cache to disk

* Address Paul's review comments

eth2/utils/tree_hash/Cargo.toml

@@ -15,8 +15,8 @@ criterion = "0.3.0"
 rand = "0.7.2"
 tree_hash_derive = "0.2"
 types = { path = "../../types" }
+lazy_static = "1.4.0"
 
 [dependencies]
 ethereum-types = "0.8.0"
 eth2_hashing = "0.1.0"
-lazy_static = "1.4.0"

eth2/utils/tree_hash/benches/benches.rs

@@ -1,8 +1,6 @@
-#[macro_use]
-extern crate lazy_static;
-
 use criterion::Criterion;
 use criterion::{black_box, criterion_group, criterion_main, Benchmark};
+use lazy_static::lazy_static;
 use types::test_utils::{generate_deterministic_keypairs, TestingBeaconStateBuilder};
 use types::{BeaconState, EthSpec, Keypair, MainnetEthSpec, MinimalEthSpec};
 
@@ -27,25 +25,61 @@ fn build_state<T: EthSpec>(validator_count: usize) -> BeaconState<T> {
     state
 }
 
+// Note: `state.canonical_root()` uses whatever `tree_hash` that the `types` crate
+// uses, which is not necessarily this crate. If you want to ensure that types is
+// using this local version of `tree_hash`, ensure you add a workspace-level
+// [dependency
+// patch](https://doc.rust-lang.org/cargo/reference/manifest.html#the-patch-section).
 fn bench_suite<T: EthSpec>(c: &mut Criterion, spec_desc: &str, validator_count: usize) {
-    let state = build_state::<T>(validator_count);
+    let state1 = build_state::<T>(validator_count);
+    let state2 = state1.clone();
+    let mut state3 = state1.clone();
+    state3.build_tree_hash_cache().unwrap();
 
     c.bench(
-        &format!("{}/{}_validators", spec_desc, validator_count),
+        &format!("{}/{}_validators/no_cache", spec_desc, validator_count),
         Benchmark::new("genesis_state", move |b| {
             b.iter_batched_ref(
-                || state.clone(),
-                // Note: `state.canonical_root()` uses whatever `tree_hash` that the `types` crate
-                // uses, which is not necessarily this crate. If you want to ensure that types is
-                // using this local version of `tree_hash`, ensure you add a workspace-level
-                // [dependency
-                // patch](https://doc.rust-lang.org/cargo/reference/manifest.html#the-patch-section).
+                || state1.clone(),
                 |state| black_box(state.canonical_root()),
                 criterion::BatchSize::SmallInput,
             )
         })
         .sample_size(10),
     );
+
+    c.bench(
+        &format!("{}/{}_validators/empty_cache", spec_desc, validator_count),
+        Benchmark::new("genesis_state", move |b| {
+            b.iter_batched_ref(
+                || state2.clone(),
+                |state| {
+                    assert!(!state.tree_hash_cache.is_initialized());
+                    black_box(state.update_tree_hash_cache().unwrap())
+                },
+                criterion::BatchSize::SmallInput,
+            )
+        })
+        .sample_size(10),
+    );
+
+    c.bench(
+        &format!(
+            "{}/{}_validators/up_to_date_cache",
+            spec_desc, validator_count
+        ),
+        Benchmark::new("genesis_state", move |b| {
+            b.iter_batched_ref(
+                || state3.clone(),
+                |state| {
+                    assert!(state.tree_hash_cache.is_initialized());
+                    black_box(state.update_tree_hash_cache().unwrap())
+                },
+                criterion::BatchSize::SmallInput,
+            )
+        })
+        .sample_size(10),
+    );
 }
 
 fn all_benches(c: &mut Criterion) {

eth2/utils/tree_hash/src/impls.rs

@@ -131,36 +131,6 @@ impl TreeHash for H256 {
     }
 }
 
-// TODO: this implementation always panics, it only exists to allow us to compile whilst
-// refactoring tree hash. Should be removed.
-macro_rules! impl_for_list {
-    ($type: ty) => {
-        impl<T> TreeHash for $type
-        where
-            T: TreeHash,
-        {
-            fn tree_hash_type() -> TreeHashType {
-                unimplemented!("TreeHash is not implemented for Vec or slice")
-            }
-
-            fn tree_hash_packed_encoding(&self) -> Vec<u8> {
-                unimplemented!("TreeHash is not implemented for Vec or slice")
-            }
-
-            fn tree_hash_packing_factor() -> usize {
-                unimplemented!("TreeHash is not implemented for Vec or slice")
-            }
-
-            fn tree_hash_root(&self) -> Vec<u8> {
-                unimplemented!("TreeHash is not implemented for Vec or slice")
-            }
-        }
-    };
-}
-
-impl_for_list!(Vec<T>);
-impl_for_list!(&[T]);
-
 /// Returns `int` as little-endian bytes with a length of 32.
 fn int_to_bytes32(int: u64) -> Vec<u8> {
     let mut vec = int.to_le_bytes().to_vec();

eth2/utils/tree_hash/src/lib.rs

@@ -1,6 +1,3 @@
-#[macro_use]
-extern crate lazy_static;
-
 pub mod impls;
 mod merkleize_padded;
 mod merkleize_standard;
@@ -27,7 +24,7 @@ pub fn mix_in_length(root: &[u8], length: usize) -> Vec<u8> {
     let mut length_bytes = length.to_le_bytes().to_vec();
     length_bytes.resize(BYTES_PER_CHUNK, 0);
 
-    merkleize_padded::hash_concat(root, &length_bytes)
+    eth2_hashing::hash_concat(root, &length_bytes)
 }
 
 #[derive(Debug, PartialEq, Clone)]

eth2/utils/tree_hash/src/merkleize_padded.rs

@@ -1,25 +1,10 @@
 use super::BYTES_PER_CHUNK;
-use eth2_hashing::hash;
+use eth2_hashing::{hash, hash_concat, ZERO_HASHES, ZERO_HASHES_MAX_INDEX};
 
 /// The size of the cache that stores padding nodes for a given height.
 ///
 /// Currently, we panic if we encounter a tree with a height larger than `MAX_TREE_DEPTH`.
-///
-/// It is set to 48 as we expect it to be sufficiently high that we won't exceed it.
-pub const MAX_TREE_DEPTH: usize = 48;
-
-lazy_static! {
-    /// Cached zero hashes where `ZERO_HASHES[i]` is the hash of a Merkle tree with 2^i zero leaves.
-    static ref ZERO_HASHES: Vec<Vec<u8>> = {
-        let mut hashes = vec![vec![0; 32]; MAX_TREE_DEPTH + 1];
-
-        for i in 0..MAX_TREE_DEPTH {
-            hashes[i + 1] = hash_concat(&hashes[i], &hashes[i]);
-        }
-
-        hashes
-    };
-}
+pub const MAX_TREE_DEPTH: usize = ZERO_HASHES_MAX_INDEX;
 
 /// Merkleize `bytes` and return the root, optionally padding the tree out to `min_leaves` number of
 /// leaves.
@@ -236,17 +221,6 @@ fn get_zero_hash(height: usize) -> &'static [u8] {
     }
 }
 
-/// Concatenate two vectors.
-fn concat(mut vec1: Vec<u8>, mut vec2: Vec<u8>) -> Vec<u8> {
-    vec1.append(&mut vec2);
-    vec1
-}
-
-/// Compute the hash of two other hashes concatenated.
-pub fn hash_concat(h1: &[u8], h2: &[u8]) -> Vec<u8> {
-    hash(&concat(h1.to_vec(), h2.to_vec()))
-}
-
 /// Returns the next even number following `n`. If `n` is even, `n` is returned.
 fn next_even_number(n: usize) -> usize {
     n + n % 2