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/types/src/beacon_state.rs

@@ -2,6 +2,7 @@ use self::committee_cache::get_active_validator_indices;
 use self::exit_cache::ExitCache;
 use crate::test_utils::TestRandom;
 use crate::*;
+use cached_tree_hash::{CachedTreeHash, MultiTreeHashCache, TreeHashCache};
 use compare_fields_derive::CompareFields;
 use eth2_hashing::hash;
 use int_to_bytes::{int_to_bytes32, int_to_bytes8};
@@ -12,7 +13,7 @@ use ssz_derive::{Decode, Encode};
 use ssz_types::{typenum::Unsigned, BitVector, FixedVector};
 use test_random_derive::TestRandom;
 use tree_hash::TreeHash;
-use tree_hash_derive::TreeHash;
+use tree_hash_derive::{CachedTreeHash, TreeHash};
 
 pub use self::committee_cache::CommitteeCache;
 pub use eth_spec::*;
@@ -57,6 +58,7 @@ pub enum Error {
     RelativeEpochError(RelativeEpochError),
     CommitteeCacheUninitialized(RelativeEpoch),
     SszTypesError(ssz_types::Error),
+    CachedTreeHashError(cached_tree_hash::Error),
 }
 
 /// Control whether an epoch-indexed field can be indexed at the next epoch or not.
@@ -75,6 +77,26 @@ impl AllowNextEpoch {
     }
 }
 
+#[derive(Debug, PartialEq, Clone, Default, Encode, Decode)]
+pub struct BeaconTreeHashCache {
+    initialized: bool,
+    block_roots: TreeHashCache,
+    state_roots: TreeHashCache,
+    historical_roots: TreeHashCache,
+    validators: MultiTreeHashCache,
+    balances: TreeHashCache,
+    randao_mixes: TreeHashCache,
+    active_index_roots: TreeHashCache,
+    compact_committees_roots: TreeHashCache,
+    slashings: TreeHashCache,
+}
+
+impl BeaconTreeHashCache {
+    pub fn is_initialized(&self) -> bool {
+        self.initialized
+    }
+}
+
 /// The state of the `BeaconChain` at some slot.
 ///
 /// Spec v0.8.0
@@ -88,9 +110,11 @@ impl AllowNextEpoch {
     Encode,
     Decode,
     TreeHash,
+    CachedTreeHash,
     CompareFields,
 )]
 #[serde(bound = "T: EthSpec")]
+#[cached_tree_hash(type = "BeaconTreeHashCache")]
 pub struct BeaconState<T>
 where
     T: EthSpec,
@@ -103,9 +127,12 @@ where
     // History
     pub latest_block_header: BeaconBlockHeader,
     #[compare_fields(as_slice)]
+    #[cached_tree_hash(block_roots)]
     pub block_roots: FixedVector<Hash256, T::SlotsPerHistoricalRoot>,
     #[compare_fields(as_slice)]
+    #[cached_tree_hash(state_roots)]
     pub state_roots: FixedVector<Hash256, T::SlotsPerHistoricalRoot>,
+    #[cached_tree_hash(historical_roots)]
     pub historical_roots: VariableList<Hash256, T::HistoricalRootsLimit>,
 
     // Ethereum 1.0 chain data
@@ -115,19 +142,25 @@ where
 
     // Registry
     #[compare_fields(as_slice)]
+    #[cached_tree_hash(validators)]
     pub validators: VariableList<Validator, T::ValidatorRegistryLimit>,
     #[compare_fields(as_slice)]
+    #[cached_tree_hash(balances)]
     pub balances: VariableList<u64, T::ValidatorRegistryLimit>,
 
     // Shuffling
     pub start_shard: u64,
+    #[cached_tree_hash(randao_mixes)]
     pub randao_mixes: FixedVector<Hash256, T::EpochsPerHistoricalVector>,
     #[compare_fields(as_slice)]
+    #[cached_tree_hash(active_index_roots)]
     pub active_index_roots: FixedVector<Hash256, T::EpochsPerHistoricalVector>,
     #[compare_fields(as_slice)]
+    #[cached_tree_hash(compact_committees_roots)]
     pub compact_committees_roots: FixedVector<Hash256, T::EpochsPerHistoricalVector>,
 
     // Slashings
+    #[cached_tree_hash(slashings)]
     pub slashings: FixedVector<u64, T::EpochsPerSlashingsVector>,
 
     // Attestations
@@ -164,6 +197,12 @@ where
     #[tree_hash(skip_hashing)]
     #[test_random(default)]
     pub exit_cache: ExitCache,
+    #[serde(skip_serializing, skip_deserializing)]
+    #[ssz(skip_serializing)]
+    #[ssz(skip_deserializing)]
+    #[tree_hash(skip_hashing)]
+    #[test_random(default)]
+    pub tree_hash_cache: BeaconTreeHashCache,
 }
 
 impl<T: EthSpec> BeaconState<T> {
@@ -225,6 +264,7 @@ impl<T: EthSpec> BeaconState<T> {
             ],
             pubkey_cache: PubkeyCache::default(),
             exit_cache: ExitCache::default(),
+            tree_hash_cache: BeaconTreeHashCache::default(),
         }
     }
 
@@ -825,7 +865,7 @@ impl<T: EthSpec> BeaconState<T> {
         self.build_committee_cache(RelativeEpoch::Current, spec)?;
         self.build_committee_cache(RelativeEpoch::Next, spec)?;
         self.update_pubkey_cache()?;
-        self.update_tree_hash_cache()?;
+        self.build_tree_hash_cache()?;
         self.exit_cache.build_from_registry(&self.validators, spec);
 
         Ok(())
@@ -936,41 +976,40 @@ impl<T: EthSpec> BeaconState<T> {
         self.pubkey_cache = PubkeyCache::default()
     }
 
-    /// Update the tree hash cache, building it for the first time if it is empty.
-    ///
-    /// Returns the `tree_hash_root` resulting from the update. This root can be considered the
-    /// canonical root of `self`.
-    ///
-    /// ## Note
-    ///
-    /// Cache not currently implemented, just performs a full tree hash.
-    pub fn update_tree_hash_cache(&mut self) -> Result<Hash256, Error> {
-        // TODO(#440): re-enable cached tree hash
-        Ok(Hash256::from_slice(&self.tree_hash_root()))
+    /// Initialize but don't fill the tree hash cache, if it isn't already initialized.
+    pub fn initialize_tree_hash_cache(&mut self) {
+        if !self.tree_hash_cache.initialized {
+            self.tree_hash_cache = Self::new_tree_hash_cache();
+        }
     }
 
-    /// Returns the tree hash root determined by the last execution of `self.update_tree_hash_cache(..)`.
+    /// Build and update the tree hash cache if it isn't already initialized.
+    pub fn build_tree_hash_cache(&mut self) -> Result<(), Error> {
+        self.update_tree_hash_cache().map(|_| ())
+    }
+
+    /// Build the tree hash cache, with blatant disregard for any existing cache.
+    pub fn force_build_tree_hash_cache(&mut self) -> Result<(), Error> {
+        self.tree_hash_cache.initialized = false;
+        self.build_tree_hash_cache()
+    }
+
+    /// Compute the tree hash root of the state using the tree hash cache.
     ///
-    /// Note: does _not_ update the cache and may return an outdated root.
-    ///
-    /// Returns an error if the cache is not initialized or if an error is encountered during the
-    /// cache update.
-    ///
-    /// ## Note
-    ///
-    /// Cache not currently implemented, just performs a full tree hash.
-    pub fn cached_tree_hash_root(&self) -> Result<Hash256, Error> {
-        // TODO(#440): re-enable cached tree hash
-        Ok(Hash256::from_slice(&self.tree_hash_root()))
+    /// Initialize the tree hash cache if it isn't already initialized.
+    pub fn update_tree_hash_cache(&mut self) -> Result<Hash256, Error> {
+        self.initialize_tree_hash_cache();
+
+        let mut cache = std::mem::replace(&mut self.tree_hash_cache, <_>::default());
+        let result = self.recalculate_tree_hash_root(&mut cache);
+        std::mem::replace(&mut self.tree_hash_cache, cache);
+
+        Ok(result?)
     }
 
     /// Completely drops the tree hash cache, replacing it with a new, empty cache.
-    ///
-    /// ## Note
-    ///
-    /// Cache not currently implemented, is a no-op.
     pub fn drop_tree_hash_cache(&mut self) {
-        // TODO(#440): re-enable cached tree hash
+        self.tree_hash_cache = BeaconTreeHashCache::default();
     }
 }
 
@@ -985,3 +1024,9 @@ impl From<ssz_types::Error> for Error {
         Error::SszTypesError(e)
     }
 }
+
+impl From<cached_tree_hash::Error> for Error {
+    fn from(e: cached_tree_hash::Error) -> Error {
+        Error::CachedTreeHashError(e)
+    }
+}

eth2/types/src/crosslink_committee.rs

@@ -1,7 +1,6 @@
 use crate::*;
-use tree_hash_derive::TreeHash;
 
-#[derive(Default, Clone, Debug, PartialEq, TreeHash)]
+#[derive(Default, Clone, Debug, PartialEq)]
 pub struct CrosslinkCommittee<'a> {
     pub slot: Slot,
     pub shard: Shard,
@@ -18,7 +17,7 @@ impl<'a> CrosslinkCommittee<'a> {
     }
 }
 
-#[derive(Default, Clone, Debug, PartialEq, TreeHash)]
+#[derive(Default, Clone, Debug, PartialEq)]
 pub struct OwnedCrosslinkCommittee {
     pub slot: Slot,
     pub shard: Shard,

eth2/types/src/lib.rs

@@ -38,6 +38,7 @@ pub mod slot_epoch_macros;
 pub mod relative_epoch;
 pub mod slot_epoch;
 pub mod slot_height;
+mod tree_hash_impls;
 pub mod validator;
 
 use ethereum_types::{H160, H256, U256};

eth2/types/src/tree_hash_impls.rs

@@ -0,0 +1,129 @@
+//! This module contains custom implementations of `CachedTreeHash` for ETH2-specific types.
+//!
+//! It makes some assumptions about the layouts and update patterns of other structs in this
+//! crate, and should be updated carefully whenever those structs are changed.
+use crate::{Hash256, Validator};
+use cached_tree_hash::{int_log, CachedTreeHash, Error, TreeHashCache};
+use tree_hash::TreeHash;
+
+/// Number of struct fields on `Validator`.
+const NUM_VALIDATOR_FIELDS: usize = 8;
+
+impl CachedTreeHash<TreeHashCache> for Validator {
+    fn new_tree_hash_cache() -> TreeHashCache {
+        TreeHashCache::new(int_log(NUM_VALIDATOR_FIELDS))
+    }
+
+    /// Efficiently tree hash a `Validator`, assuming it was updated by a valid state transition.
+    ///
+    /// Specifically, we assume that the `pubkey` and `withdrawal_credentials` fields are constant.
+    fn recalculate_tree_hash_root(&self, cache: &mut TreeHashCache) -> Result<Hash256, Error> {
+        // If the cache is empty, hash every field to fill it.
+        if cache.leaves().is_empty() {
+            return cache.recalculate_merkle_root(field_tree_hash_iter(self));
+        }
+
+        // Otherwise just check the fields which might have changed.
+        let dirty_indices = cache
+            .leaves()
+            .iter_mut()
+            .enumerate()
+            .flat_map(|(i, leaf)| {
+                // Fields pubkey and withdrawal_credentials are constant
+                if i == 0 || i == 1 {
+                    None
+                } else {
+                    let new_tree_hash = field_tree_hash_by_index(self, i);
+                    if leaf.as_bytes() != &new_tree_hash[..] {
+                        leaf.assign_from_slice(&new_tree_hash);
+                        Some(i)
+                    } else {
+                        None
+                    }
+                }
+            })
+            .collect();
+
+        cache.update_merkle_root(dirty_indices)
+    }
+}
+
+/// Get the tree hash root of a validator field by its position/index in the struct.
+fn field_tree_hash_by_index(v: &Validator, field_idx: usize) -> Vec<u8> {
+    match field_idx {
+        0 => v.pubkey.tree_hash_root(),
+        1 => v.withdrawal_credentials.tree_hash_root(),
+        2 => v.effective_balance.tree_hash_root(),
+        3 => v.slashed.tree_hash_root(),
+        4 => v.activation_eligibility_epoch.tree_hash_root(),
+        5 => v.activation_epoch.tree_hash_root(),
+        6 => v.exit_epoch.tree_hash_root(),
+        7 => v.withdrawable_epoch.tree_hash_root(),
+        _ => panic!(
+            "Validator type only has {} fields, {} out of bounds",
+            NUM_VALIDATOR_FIELDS, field_idx
+        ),
+    }
+}
+
+/// Iterator over the tree hash roots of `Validator` fields.
+fn field_tree_hash_iter<'a>(
+    v: &'a Validator,
+) -> impl Iterator<Item = [u8; 32]> + ExactSizeIterator + 'a {
+    (0..NUM_VALIDATOR_FIELDS)
+        .map(move |i| field_tree_hash_by_index(v, i))
+        .map(|tree_hash_root| {
+            let mut res = [0; 32];
+            res.copy_from_slice(&tree_hash_root[0..32]);
+            res
+        })
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+    use crate::test_utils::TestRandom;
+    use crate::Epoch;
+    use rand::SeedableRng;
+    use rand_xorshift::XorShiftRng;
+
+    fn test_validator_tree_hash(v: &Validator) {
+        let mut cache = Validator::new_tree_hash_cache();
+        // With a fresh cache
+        assert_eq!(
+            &v.tree_hash_root()[..],
+            v.recalculate_tree_hash_root(&mut cache).unwrap().as_bytes(),
+            "{:?}",
+            v
+        );
+        // With a completely up-to-date cache
+        assert_eq!(
+            &v.tree_hash_root()[..],
+            v.recalculate_tree_hash_root(&mut cache).unwrap().as_bytes(),
+            "{:?}",
+            v
+        );
+    }
+
+    #[test]
+    fn default_validator() {
+        test_validator_tree_hash(&Validator::default());
+    }
+
+    #[test]
+    fn zeroed_validator() {
+        let mut v = Validator::default();
+        v.activation_eligibility_epoch = Epoch::from(0u64);
+        v.activation_epoch = Epoch::from(0u64);
+        test_validator_tree_hash(&v);
+    }
+
+    #[test]
+    fn random_validators() {
+        let mut rng = XorShiftRng::from_seed([0xf1; 16]);
+        let num_validators = 1000;
+        (0..num_validators)
+            .map(|_| Validator::random_for_test(&mut rng))
+            .for_each(|v| test_validator_tree_hash(&v));
+    }
+}