Simplify exit cache (#5280)

consensus/types/src/beacon_state.rs

@@ -105,6 +105,10 @@ pub enum Error {
     },
     RelativeEpochError(RelativeEpochError),
     ExitCacheUninitialized,
+    ExitCacheInvalidEpoch {
+        max_exit_epoch: Epoch,
+        request_epoch: Epoch,
+    },
     SlashingsCacheUninitialized {
         initialized_slot: Option<Slot>,
         latest_block_slot: Slot,

consensus/types/src/beacon_state/exit_cache.rs

@@ -1,21 +1,25 @@
 use super::{BeaconStateError, ChainSpec, Epoch, Validator};
-use safe_arith::SafeArith;
 use serde::{Deserialize, Serialize};
-use std::collections::HashMap;
+use std::cmp::Ordering;
 
 /// Map from exit epoch to the number of validators with that exit epoch.
 #[derive(Debug, Default, Clone, PartialEq, Serialize, Deserialize)]
 pub struct ExitCache {
+    /// True if the cache has been initialized.
     initialized: bool,
-    exit_epoch_counts: HashMap<Epoch, u64>,
+    /// Maximum `exit_epoch` of any validator.
+    max_exit_epoch: Epoch,
+    /// Number of validators known to be exiting at `max_exit_epoch`.
+    max_exit_epoch_churn: u64,
 }
 
 impl ExitCache {
     /// Initialize a new cache for the given list of validators.
     pub fn new(validators: &[Validator], spec: &ChainSpec) -> Result<Self, BeaconStateError> {
         let mut exit_cache = ExitCache {
             initialized: true,
-            ..ExitCache::default()
+            max_exit_epoch: Epoch::new(0),
+            max_exit_epoch_churn: 0,
         };
         // Add all validators with a non-default exit epoch to the cache.
         validators
@@ -37,27 +41,44 @@ impl ExitCache {
     /// Record the exit epoch of a validator. Must be called only once per exiting validator.
     pub fn record_validator_exit(&mut self, exit_epoch: Epoch) -> Result<(), BeaconStateError> {
         self.check_initialized()?;
-        self.exit_epoch_counts
-            .entry(exit_epoch)
-            .or_insert(0)
-            .safe_add_assign(1)?;
+        match exit_epoch.cmp(&self.max_exit_epoch) {
+            // Update churn for the current maximum epoch.
+            Ordering::Equal => {
+                self.max_exit_epoch_churn += 1;
+            }
+            // Increase the max exit epoch, reset the churn to 1.
+            Ordering::Greater => {
+                self.max_exit_epoch = exit_epoch;
+                self.max_exit_epoch_churn = 1;
+            }
+            // Older exit epochs are not relevant.
+            Ordering::Less => (),
+        }
         Ok(())
     }
 
     /// Get the largest exit epoch with a non-zero exit epoch count.
     pub fn max_epoch(&self) -> Result<Option<Epoch>, BeaconStateError> {
         self.check_initialized()?;
-        Ok(self.exit_epoch_counts.keys().max().cloned())
+        Ok((self.max_exit_epoch_churn > 0).then_some(self.max_exit_epoch))
     }
 
     /// Get number of validators with the given exit epoch. (Return 0 for the default exit epoch.)
     pub fn get_churn_at(&self, exit_epoch: Epoch) -> Result<u64, BeaconStateError> {
         self.check_initialized()?;
-        Ok(self
-            .exit_epoch_counts
-            .get(&exit_epoch)
-            .cloned()
-            .unwrap_or(0))
+        match exit_epoch.cmp(&self.max_exit_epoch) {
+            // Epochs are equal, we know the churn exactly.
+            Ordering::Equal => Ok(self.max_exit_epoch_churn),
+            // If exiting at an epoch later than the cached epoch then the churn is 0. This is a
+            // common case which happens when there are no exits for an epoch.
+            Ordering::Greater => Ok(0),
+            // Consensus code should never require the churn at an epoch prior to the cached epoch.
+            // That's a bug.
+            Ordering::Less => Err(BeaconStateError::ExitCacheInvalidEpoch {
+                max_exit_epoch: self.max_exit_epoch,
+                request_epoch: exit_epoch,
+            }),
+        }
     }
 }