Use THC for state.inactivity_scores (#2504)

## Issue Addressed

- Resolves #2502

## Proposed Changes

Adds tree-hash caching (THC 🍁) for `state.inactivity_scores`, as per #2502.

Since the `inactivity_scores` field is introduced during Altair, the cache must be optional (i.e., not present pre-Altair). The mechanism for optional caches was already implemented via the `ParticipationTreeHashCache`, albeit not quite generically enough. To this end, I made the `ParticipationTreeHashCache` more generic and renamed it to `OptionalTreeHashCache`. This made the code a little more verbose around the previous/current epoch participation fields, but overall less verbose when the needs of `inactivity_scores` are considered.

All changes to `ParticipationTreeHashCache` should be *non-substantial*.

## Additional Info

NA

Author: paulhauner

consensus/types/src/beacon_state/tree_hash_cache.rs

@@ -3,9 +3,7 @@
 #![allow(clippy::indexing_slicing)]
 
 use super::Error;
-use crate::{
-    BeaconState, EthSpec, Hash256, ParticipationFlags, ParticipationList, Slot, Unsigned, Validator,
-};
+use crate::{BeaconState, EthSpec, Hash256, ParticipationList, Slot, Unsigned, Validator};
 use cached_tree_hash::{int_log, CacheArena, CachedTreeHash, TreeHashCache};
 use rayon::prelude::*;
 use ssz_derive::{Decode, Encode};
@@ -141,9 +139,10 @@ pub struct BeaconTreeHashCacheInner<T: EthSpec> {
     randao_mixes: TreeHashCache,
     slashings: TreeHashCache,
     eth1_data_votes: Eth1DataVotesTreeHashCache<T>,
+    inactivity_scores: OptionalTreeHashCache,
     // Participation caches
-    previous_epoch_participation: ParticipationTreeHashCache,
-    current_epoch_participation: ParticipationTreeHashCache,
+    previous_epoch_participation: OptionalTreeHashCache,
+    current_epoch_participation: OptionalTreeHashCache,
 }
 
 impl<T: EthSpec> BeaconTreeHashCacheInner<T> {
@@ -168,10 +167,22 @@ impl<T: EthSpec> BeaconTreeHashCacheInner<T> {
         let mut slashings_arena = CacheArena::default();
         let slashings = state.slashings().new_tree_hash_cache(&mut slashings_arena);
 
-        let previous_epoch_participation =
-            ParticipationTreeHashCache::new(state, BeaconState::previous_epoch_participation);
-        let current_epoch_participation =
-            ParticipationTreeHashCache::new(state, BeaconState::current_epoch_participation);
+        let inactivity_scores = OptionalTreeHashCache::new(state.inactivity_scores().ok());
+
+        let previous_epoch_participation = OptionalTreeHashCache::new(
+            state
+                .previous_epoch_participation()
+                .ok()
+                .map(ParticipationList::new)
+                .as_ref(),
+        );
+        let current_epoch_participation = OptionalTreeHashCache::new(
+            state
+                .current_epoch_participation()
+                .ok()
+                .map(ParticipationList::new)
+                .as_ref(),
+        );
 
         Self {
             previous_state: None,
@@ -185,6 +196,7 @@ impl<T: EthSpec> BeaconTreeHashCacheInner<T> {
             balances,
             randao_mixes,
             slashings,
+            inactivity_scores,
             eth1_data_votes: Eth1DataVotesTreeHashCache::new(state),
             previous_epoch_participation,
             current_epoch_participation,
@@ -287,12 +299,16 @@ impl<T: EthSpec> BeaconTreeHashCacheInner<T> {
         } else {
             hasher.write(
                 self.previous_epoch_participation
-                    .recalculate_tree_hash_root(state.previous_epoch_participation()?)?
+                    .recalculate_tree_hash_root(&ParticipationList::new(
+                        state.previous_epoch_participation()?,
+                    ))?
                     .as_bytes(),
             )?;
             hasher.write(
                 self.current_epoch_participation
-                    .recalculate_tree_hash_root(state.current_epoch_participation()?)?
+                    .recalculate_tree_hash_root(&ParticipationList::new(
+                        state.current_epoch_participation()?,
+                    ))?
                     .as_bytes(),
             )?;
         }
@@ -314,8 +330,11 @@ impl<T: EthSpec> BeaconTreeHashCacheInner<T> {
 
         // Inactivity & light-client sync committees
         if let BeaconState::Altair(ref state) = state {
-            // FIXME(altair): add cache for this field
-            hasher.write(state.inactivity_scores.tree_hash_root().as_bytes())?;
+            hasher.write(
+                self.inactivity_scores
+                    .recalculate_tree_hash_root(&state.inactivity_scores)?
+                    .as_bytes(),
+            )?;
 
             hasher.write(state.current_sync_committee.tree_hash_root().as_bytes())?;
             hasher.write(state.next_sync_committee.tree_hash_root().as_bytes())?;
@@ -513,53 +532,43 @@ impl ParallelValidatorTreeHash {
 }
 
 #[derive(Debug, PartialEq, Clone)]
-pub struct ParticipationTreeHashCache {
-    inner: Option<ParticipationTreeHashCacheInner>,
+pub struct OptionalTreeHashCache {
+    inner: Option<OptionalTreeHashCacheInner>,
 }
 
 #[derive(Debug, PartialEq, Clone)]
-pub struct ParticipationTreeHashCacheInner {
+pub struct OptionalTreeHashCacheInner {
     arena: CacheArena,
     tree_hash_cache: TreeHashCache,
 }
 
-impl ParticipationTreeHashCache {
-    /// Initialize a new cache for the participation list returned by `field` (if any).
-    fn new<T: EthSpec>(
-        state: &BeaconState<T>,
-        field: impl FnOnce(
-            &BeaconState<T>,
-        ) -> Result<
-            &VariableList<ParticipationFlags, T::ValidatorRegistryLimit>,
-            Error,
-        >,
-    ) -> Self {
-        let inner = field(state).map(ParticipationTreeHashCacheInner::new).ok();
+impl OptionalTreeHashCache {
+    /// Initialize a new cache if `item.is_some()`.
+    fn new<C: CachedTreeHash<TreeHashCache>>(item: Option<&C>) -> Self {
+        let inner = item.map(OptionalTreeHashCacheInner::new);
         Self { inner }
     }
 
-    /// Compute the tree hash root for the given `epoch_participation`.
+    /// Compute the tree hash root for the given `item`.
     ///
     /// This function will initialize the inner cache if necessary (e.g. when crossing the fork).
-    fn recalculate_tree_hash_root<N: Unsigned>(
+    fn recalculate_tree_hash_root<C: CachedTreeHash<TreeHashCache>>(
         &mut self,
-        epoch_participation: &VariableList<ParticipationFlags, N>,
+        item: &C,
     ) -> Result<Hash256, Error> {
         let cache = self
             .inner
-            .get_or_insert_with(|| ParticipationTreeHashCacheInner::new(epoch_participation));
-        ParticipationList::new(epoch_participation)
-            .recalculate_tree_hash_root(&mut cache.arena, &mut cache.tree_hash_cache)
+            .get_or_insert_with(|| OptionalTreeHashCacheInner::new(item));
+        item.recalculate_tree_hash_root(&mut cache.arena, &mut cache.tree_hash_cache)
             .map_err(Into::into)
     }
 }
 
-impl ParticipationTreeHashCacheInner {
-    fn new<N: Unsigned>(epoch_participation: &VariableList<ParticipationFlags, N>) -> Self {
+impl OptionalTreeHashCacheInner {
+    fn new<C: CachedTreeHash<TreeHashCache>>(item: &C) -> Self {
         let mut arena = CacheArena::default();
-        let tree_hash_cache =
-            ParticipationList::new(epoch_participation).new_tree_hash_cache(&mut arena);
-        ParticipationTreeHashCacheInner {
+        let tree_hash_cache = item.new_tree_hash_cache(&mut arena);
+        OptionalTreeHashCacheInner {
             arena,
             tree_hash_cache,
         }
@@ -576,7 +585,7 @@ impl<T: EthSpec> arbitrary::Arbitrary for BeaconTreeHashCache<T> {
 #[cfg(test)]
 mod test {
     use super::*;
-    use crate::MainnetEthSpec;
+    use crate::{MainnetEthSpec, ParticipationFlags};
 
     #[test]
     fn validator_node_count() {
@@ -594,13 +603,13 @@ mod test {
         test_flag.add_flag(0).unwrap();
         let epoch_participation = VariableList::<_, N>::new(vec![test_flag; len]).unwrap();
 
-        let mut cache = ParticipationTreeHashCache { inner: None };
+        let mut cache = OptionalTreeHashCache { inner: None };
 
         let cache_root = cache
-            .recalculate_tree_hash_root(&epoch_participation)
+            .recalculate_tree_hash_root(&ParticipationList::new(&epoch_participation))
             .unwrap();
         let recalc_root = cache
-            .recalculate_tree_hash_root(&epoch_participation)
+            .recalculate_tree_hash_root(&ParticipationList::new(&epoch_participation))
             .unwrap();
 
         assert_eq!(cache_root, recalc_root, "recalculated root should match");