From d4a867231b794f9c7ecb40dfcb62ea4212d03071 Mon Sep 17 00:00:00 2001
From: Michael Sproul <michael@sigmaprime.io>
Date: Fri, 23 Feb 2024 13:52:43 +1100
Subject: [PATCH] Simplify exit cache

---
 consensus/types/src/beacon_state.rs           |  4 ++
 .../types/src/beacon_state/exit_cache.rs      | 49 +++++++++++++------
 2 files changed, 39 insertions(+), 14 deletions(-)

diff --git a/consensus/types/src/beacon_state.rs b/consensus/types/src/beacon_state.rs
index 77f15c0ed1f..36d26aa3f2b 100644
--- a/consensus/types/src/beacon_state.rs
+++ b/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,
diff --git a/consensus/types/src/beacon_state/exit_cache.rs b/consensus/types/src/beacon_state/exit_cache.rs
index cb96fba6917..1ab7a8e3119 100644
--- a/consensus/types/src/beacon_state/exit_cache.rs
+++ b/consensus/types/src/beacon_state/exit_cache.rs
@@ -1,13 +1,16 @@
 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 {
@@ -15,7 +18,8 @@ impl ExitCache {
     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,
+            }),
+        }
     }
 }