Implement tree states & hierarchical state DB

consensus/state_processing/src/common/altair.rs

@@ -24,14 +24,12 @@ impl BaseRewardPerIncrement {
 /// shown to be a significant optimisation.
 ///
 /// Spec v1.1.0
-pub fn get_base_reward<T: EthSpec>(
-    state: &BeaconState<T>,
-    index: usize,
+pub fn get_base_reward(
+    validator_effective_balance: u64,
     base_reward_per_increment: BaseRewardPerIncrement,
     spec: &ChainSpec,
 ) -> Result<u64, Error> {
-    state
-        .get_effective_balance(index)?
+    validator_effective_balance
         .safe_div(spec.effective_balance_increment)?
         .safe_mul(base_reward_per_increment.as_u64())
         .map_err(Into::into)

consensus/state_processing/src/common/base.rs

@@ -1,19 +1,30 @@
 use integer_sqrt::IntegerSquareRoot;
-use safe_arith::SafeArith;
+use safe_arith::{ArithError, SafeArith};
 use types::*;
 
-/// Returns the base reward for some validator.
-pub fn get_base_reward<T: EthSpec>(
-    state: &BeaconState<T>,
-    index: usize,
-    // Should be == get_total_active_balance(state, spec)
-    total_active_balance: u64,
-    spec: &ChainSpec,
-) -> Result<u64, BeaconStateError> {
-    state
-        .get_effective_balance(index)?
-        .safe_mul(spec.base_reward_factor)?
-        .safe_div(total_active_balance.integer_sqrt())?
-        .safe_div(spec.base_rewards_per_epoch)
-        .map_err(Into::into)
+/// This type exists to avoid confusing `total_active_balance` with `sqrt_total_active_balance`,
+/// since they are used in close proximity and the same type (`u64`).
+#[derive(Copy, Clone)]
+pub struct SqrtTotalActiveBalance(u64);
+
+impl SqrtTotalActiveBalance {
+    pub fn new(total_active_balance: u64) -> Self {
+        Self(total_active_balance.integer_sqrt())
+    }
+
+    pub fn as_u64(&self) -> u64 {
+        self.0
+    }
+}
+
+/// Returns the base reward for some validator.
+pub fn get_base_reward(
+    validator_effective_balance: u64,
+    sqrt_total_active_balance: SqrtTotalActiveBalance,
+    spec: &ChainSpec,
+) -> Result<u64, ArithError> {
+    validator_effective_balance
+        .safe_mul(spec.base_reward_factor)?
+        .safe_div(sqrt_total_active_balance.as_u64())?
+        .safe_div(spec.base_rewards_per_epoch)
 }

consensus/state_processing/src/common/initiate_validator_exit.rs

@@ -8,10 +8,10 @@ pub fn initiate_validator_exit<T: EthSpec>(
     index: usize,
     spec: &ChainSpec,
 ) -> Result<(), Error> {
-    // Return if the validator already initiated exit
-    if state.get_validator(index)?.exit_epoch != spec.far_future_epoch {
-        return Ok(());
-    }
+    // We do things in a slightly different order to the spec here. Instead of immediately checking
+    // whether the validator has already exited, we instead prepare the exit cache and compute the
+    // cheap-to-calculate values from that. *Then* we look up the validator a single time in the
+    // validator tree (expensive), make the check and mutate as appropriate.
 
     // Ensure the exit cache is built.
     state.build_exit_cache(spec)?;
@@ -28,12 +28,21 @@ pub fn initiate_validator_exit<T: EthSpec>(
         exit_queue_epoch.safe_add_assign(1)?;
     }
 
+    let validator = state.get_validator_cow(index)?;
+
+    // Return if the validator already initiated exit
+    if validator.exit_epoch() != spec.far_future_epoch {
+        return Ok(());
+    }
+
+    let validator = validator.to_mut();
+    validator.mutable.exit_epoch = exit_queue_epoch;
+    validator.mutable.withdrawable_epoch =
+        exit_queue_epoch.safe_add(spec.min_validator_withdrawability_delay)?;
+
     state
         .exit_cache_mut()
         .record_validator_exit(exit_queue_epoch)?;
-    state.get_validator_mut(index)?.exit_epoch = exit_queue_epoch;
-    state.get_validator_mut(index)?.withdrawable_epoch =
-        exit_queue_epoch.safe_add(spec.min_validator_withdrawability_delay)?;
 
     Ok(())
 }

consensus/state_processing/src/common/mod.rs

@@ -24,8 +24,7 @@ pub fn increase_balance<E: EthSpec>(
     index: usize,
     delta: u64,
 ) -> Result<(), BeaconStateError> {
-    state.get_balance_mut(index)?.safe_add_assign(delta)?;
-    Ok(())
+    increase_balance_directly(state.get_balance_mut(index)?, delta)
 }
 
 /// Decrease the balance of a validator, saturating upon overflow, as per the spec.
@@ -34,7 +33,17 @@ pub fn decrease_balance<E: EthSpec>(
     index: usize,
     delta: u64,
 ) -> Result<(), BeaconStateError> {
-    let balance = state.get_balance_mut(index)?;
+    decrease_balance_directly(state.get_balance_mut(index)?, delta)
+}
+
+/// Increase the balance of a validator, erroring upon overflow, as per the spec.
+pub fn increase_balance_directly(balance: &mut u64, delta: u64) -> Result<(), BeaconStateError> {
+    balance.safe_add_assign(delta)?;
+    Ok(())
+}
+
+/// Decrease the balance of a validator, saturating upon overflow, as per the spec.
+pub fn decrease_balance_directly(balance: &mut u64, delta: u64) -> Result<(), BeaconStateError> {
     *balance = balance.saturating_sub(delta);
     Ok(())
 }

consensus/state_processing/src/common/slash_validator.rs

@@ -23,12 +23,12 @@ pub fn slash_validator<T: EthSpec>(
     initiate_validator_exit(state, slashed_index, spec)?;
 
     let validator = state.get_validator_mut(slashed_index)?;
-    validator.slashed = true;
-    validator.withdrawable_epoch = cmp::max(
-        validator.withdrawable_epoch,
+    validator.mutable.slashed = true;
+    validator.mutable.withdrawable_epoch = cmp::max(
+        validator.withdrawable_epoch(),
         epoch.safe_add(T::EpochsPerSlashingsVector::to_u64())?,
     );
-    let validator_effective_balance = validator.effective_balance;
+    let validator_effective_balance = validator.effective_balance();
     state.set_slashings(
         epoch,
         state