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

tests/ef_tests/Cargo.toml

@@ -23,6 +23,7 @@ eth2_ssz = "0.1.2"
 eth2_ssz_derive = "0.1.0"
 tree_hash = "0.1.0"
 tree_hash_derive = "0.2"
+cached_tree_hash = { path = "../../eth2/utils/cached_tree_hash" }
 state_processing = { path = "../../eth2/state_processing" }
 swap_or_not_shuffle = { path = "../../eth2/utils/swap_or_not_shuffle" }
 types = { path = "../../eth2/types" }

tests/ef_tests/src/cases/ssz_generic.rs

@@ -218,7 +218,7 @@ fn ssz_generic_test<T: SszStaticType>(path: &Path) -> Result<(), Error> {
         check_serialization(&value, &serialized)?;
 
         if let Some(ref meta) = meta {
-            check_tree_hash(&meta.root, value.tree_hash_root())?;
+            check_tree_hash(&meta.root, &value.tree_hash_root())?;
         }
     }
     // Invalid

tests/ef_tests/src/cases/ssz_static.rs

@@ -2,8 +2,10 @@ use super::*;
 use crate::case_result::compare_result;
 use crate::cases::common::SszStaticType;
 use crate::decode::yaml_decode_file;
+use cached_tree_hash::CachedTreeHash;
 use serde_derive::Deserialize;
 use std::fs;
+use std::marker::PhantomData;
 use tree_hash::SignedRoot;
 use types::Hash256;
 
@@ -27,6 +29,14 @@ pub struct SszStaticSR<T> {
     value: T,
 }
 
+#[derive(Debug, Clone)]
+pub struct SszStaticTHC<T, C> {
+    roots: SszStaticRoots,
+    serialized: Vec<u8>,
+    value: T,
+    _phantom: PhantomData<C>,
+}
+
 fn load_from_dir<T: SszStaticType>(path: &Path) -> Result<(SszStaticRoots, Vec<u8>, T), Error> {
     let roots = yaml_decode_file(&path.join("roots.yaml"))?;
     let serialized = fs::read(&path.join("serialized.ssz")).expect("serialized.ssz exists");
@@ -55,6 +65,17 @@ impl<T: SszStaticType + SignedRoot> LoadCase for SszStaticSR<T> {
     }
 }
 
+impl<T: SszStaticType + CachedTreeHash<C>, C: Debug + Sync> LoadCase for SszStaticTHC<T, C> {
+    fn load_from_dir(path: &Path) -> Result<Self, Error> {
+        load_from_dir(path).map(|(roots, serialized, value)| Self {
+            roots,
+            serialized,
+            value,
+            _phantom: PhantomData,
+        })
+    }
+}
+
 pub fn check_serialization<T: SszStaticType>(value: &T, serialized: &[u8]) -> Result<(), Error> {
     // Check serialization
     let serialized_result = value.as_ssz_bytes();
@@ -68,18 +89,18 @@ pub fn check_serialization<T: SszStaticType>(value: &T, serialized: &[u8]) -> Re
     Ok(())
 }
 
-pub fn check_tree_hash(expected_str: &str, actual_root: Vec<u8>) -> Result<(), Error> {
+pub fn check_tree_hash(expected_str: &str, actual_root: &[u8]) -> Result<(), Error> {
     let expected_root = hex::decode(&expected_str[2..])
         .map_err(|e| Error::FailedToParseTest(format!("{:?}", e)))?;
     let expected_root = Hash256::from_slice(&expected_root);
-    let tree_hash_root = Hash256::from_slice(&actual_root);
+    let tree_hash_root = Hash256::from_slice(actual_root);
     compare_result::<Hash256, Error>(&Ok(tree_hash_root), &Some(expected_root))
 }
 
 impl<T: SszStaticType> Case for SszStatic<T> {
     fn result(&self, _case_index: usize) -> Result<(), Error> {
         check_serialization(&self.value, &self.serialized)?;
-        check_tree_hash(&self.roots.root, self.value.tree_hash_root())?;
+        check_tree_hash(&self.roots.root, &self.value.tree_hash_root())?;
         Ok(())
     }
 }
@@ -87,15 +108,28 @@ impl<T: SszStaticType> Case for SszStatic<T> {
 impl<T: SszStaticType + SignedRoot> Case for SszStaticSR<T> {
     fn result(&self, _case_index: usize) -> Result<(), Error> {
         check_serialization(&self.value, &self.serialized)?;
-        check_tree_hash(&self.roots.root, self.value.tree_hash_root())?;
+        check_tree_hash(&self.roots.root, &self.value.tree_hash_root())?;
         check_tree_hash(
             &self
                 .roots
                 .signing_root
                 .as_ref()
                 .expect("signed root exists"),
-            self.value.signed_root(),
+            &self.value.signed_root(),
         )?;
         Ok(())
     }
 }
+
+impl<T: SszStaticType + CachedTreeHash<C>, C: Debug + Sync> Case for SszStaticTHC<T, C> {
+    fn result(&self, _case_index: usize) -> Result<(), Error> {
+        check_serialization(&self.value, &self.serialized)?;
+        check_tree_hash(&self.roots.root, &self.value.tree_hash_root())?;
+
+        let mut cache = T::new_tree_hash_cache();
+        let cached_tree_hash_root = self.value.recalculate_tree_hash_root(&mut cache).unwrap();
+        check_tree_hash(&self.roots.root, cached_tree_hash_root.as_bytes())?;
+
+        Ok(())
+    }
+}

tests/ef_tests/src/handler.rs

@@ -1,6 +1,8 @@
 use crate::cases::{self, Case, Cases, EpochTransition, LoadCase, Operation};
 use crate::type_name;
 use crate::type_name::TypeName;
+use cached_tree_hash::CachedTreeHash;
+use std::fmt::Debug;
 use std::fs;
 use std::marker::PhantomData;
 use std::path::PathBuf;
@@ -93,6 +95,9 @@ pub struct SszStaticHandler<T, E>(PhantomData<(T, E)>);
 /// Handler for SSZ types that do implement `SignedRoot`.
 pub struct SszStaticSRHandler<T, E>(PhantomData<(T, E)>);
 
+/// Handler for SSZ types that implement `CachedTreeHash`.
+pub struct SszStaticTHCHandler<T, C, E>(PhantomData<(T, C, E)>);
+
 impl<T, E> Handler for SszStaticHandler<T, E>
 where
     T: cases::SszStaticType + TypeName,
@@ -133,6 +138,27 @@ where
     }
 }
 
+impl<T, C, E> Handler for SszStaticTHCHandler<T, C, E>
+where
+    T: cases::SszStaticType + CachedTreeHash<C> + TypeName,
+    C: Debug + Sync,
+    E: TypeName,
+{
+    type Case = cases::SszStaticTHC<T, C>;
+
+    fn config_name() -> &'static str {
+        E::name()
+    }
+
+    fn runner_name() -> &'static str {
+        "ssz_static"
+    }
+
+    fn handler_name() -> String {
+        T::name().into()
+    }
+}
+
 pub struct ShufflingHandler<E>(PhantomData<E>);
 
 impl<E: EthSpec + TypeName> Handler for ShufflingHandler<E> {

tests/ef_tests/tests/tests.rs

@@ -99,7 +99,7 @@ macro_rules! ssz_static_test {
     ($test_name:ident, $typ:ident$(<$generics:tt>)?, SR) => {
         ssz_static_test!($test_name, SszStaticSRHandler, $typ$(<$generics>)?);
     };
-    // Non-signed root
+    // Non-signed root, non-tree hash caching
     ($test_name:ident, $typ:ident$(<$generics:tt>)?) => {
         ssz_static_test!($test_name, SszStaticHandler, $typ$(<$generics>)?);
     };
@@ -122,11 +122,11 @@ macro_rules! ssz_static_test {
         );
     };
     // Base case
-    ($test_name:ident, $handler:ident, { $(($typ:ty, $spec:ident)),+ }) => {
+    ($test_name:ident, $handler:ident, { $(($($typ:ty),+)),+ }) => {
         #[test]
         fn $test_name() {
             $(
-                $handler::<$typ, $spec>::run();
+                $handler::<$($typ),+>::run();
             )+
         }
     };
@@ -134,7 +134,7 @@ macro_rules! ssz_static_test {
 
 #[cfg(feature = "fake_crypto")]
 mod ssz_static {
-    use ef_tests::{Handler, SszStaticHandler, SszStaticSRHandler};
+    use ef_tests::{Handler, SszStaticHandler, SszStaticSRHandler, SszStaticTHCHandler};
     use types::*;
 
     ssz_static_test!(attestation, Attestation<_>, SR);
@@ -147,7 +147,13 @@ mod ssz_static {
     ssz_static_test!(beacon_block, BeaconBlock<_>, SR);
     ssz_static_test!(beacon_block_body, BeaconBlockBody<_>);
     ssz_static_test!(beacon_block_header, BeaconBlockHeader, SR);
-    ssz_static_test!(beacon_state, BeaconState<_>);
+    ssz_static_test!(
+        beacon_state,
+        SszStaticTHCHandler, {
+            (BeaconState<MinimalEthSpec>, BeaconTreeHashCache, MinimalEthSpec),
+            (BeaconState<MainnetEthSpec>, BeaconTreeHashCache, MainnetEthSpec)
+        }
+    );
     ssz_static_test!(checkpoint, Checkpoint);
     ssz_static_test!(compact_committee, CompactCommittee<_>);
     ssz_static_test!(crosslink, Crosslink);