Make BeaconState.balances a tree list!

Cargo.lock

@@ -3481,6 +3481,7 @@ dependencies = [
  "eth2_ssz 0.4.0 (registry+https://github.com/rust-lang/crates.io-index)",
  "parking_lot",
  "serde",
+ "smallvec",
  "tree_hash 0.4.0 (registry+https://github.com/rust-lang/crates.io-index)",
  "typenum",
 ]
@@ -5611,6 +5612,7 @@ dependencies = [
  "sloggers",
  "state_processing",
  "tempfile",
+ "tree_hash 0.4.0 (registry+https://github.com/rust-lang/crates.io-index)",
  "types",
 ]
 

beacon_node/beacon_chain/src/beacon_fork_choice_store.rs

@@ -217,12 +217,13 @@ where
         };
         let finalized_checkpoint = justified_checkpoint;
 
+        // FIXME(sproul): avoid `to_vec` perf penalty
         Self {
             store,
             balances_cache: <_>::default(),
             time: anchor_state.slot(),
             justified_checkpoint,
-            justified_balances: anchor_state.balances().clone().into(),
+            justified_balances: anchor_state.balances().to_vec(),
             finalized_checkpoint,
             best_justified_checkpoint: justified_checkpoint,
             _phantom: PhantomData,
@@ -327,8 +328,7 @@ where
                 .map_err(Error::FailedToReadState)?
                 .ok_or_else(|| Error::MissingState(justified_block.state_root()))?
                 .balances()
-                .clone()
+                .to_vec();
-                .into();
         }
 
         Ok(())

beacon_node/store/Cargo.toml

@@ -25,6 +25,7 @@ lighthouse_metrics = { path = "../../common/lighthouse_metrics" }
 lru = "0.6.0"
 sloggers = "2.0.2"
 directory = { path = "../../common/directory" }
+tree_hash = "0.4.0"
 
 [features]
 milhouse = ["state_processing/milhouse"]

beacon_node/store/src/chunked_vector.rs

@@ -17,6 +17,7 @@
 use self::UpdatePattern::*;
 use crate::*;
 use ssz::{Decode, Encode};
+use tree_hash::TreeHash;
 use typenum::Unsigned;
 use types::VList;
 
@@ -56,7 +57,7 @@ pub trait Field<E: EthSpec>: Copy {
     /// The type of value stored in this field: the `T` from `FixedVector<T, N>`.
     ///
     /// The `Default` impl will be used to fill extra vector entries.
-    type Value: Decode + Encode + Default + Clone + PartialEq + std::fmt::Debug;
+    type Value: TreeHash + Decode + Encode + Default + Clone + PartialEq + std::fmt::Debug;
 
     /// The length of this field: the `N` from `FixedVector<T, N>`.
     type Length: Unsigned;

beacon_node/store/src/partial_beacon_state.rs

@@ -48,7 +48,7 @@ where
 
     // Registry
     pub validators: VList<Validator, T::ValidatorRegistryLimit>,
-    pub balances: VariableList<u64, T::ValidatorRegistryLimit>,
+    pub balances: VList<u64, T::ValidatorRegistryLimit>,
 
     // Shuffling
     /// Randao value from the current slot, for patching into the per-epoch randao vector.

consensus/fork_choice/src/fork_choice.rs

@@ -352,12 +352,15 @@ where
 
         let store = &mut self.fc_store;
 
+        // FIXME(sproul): plumb VList through fork choice
+        let justified_balances = store.justified_balances().to_vec();
+
         self.proto_array
             .find_head(
                 store.justified_checkpoint().epoch,
                 store.justified_checkpoint().root,
                 store.finalized_checkpoint().epoch,
-                store.justified_balances(),
+                &justified_balances,
             )
             .map_err(Into::into)
     }

consensus/state_processing/src/common/mod.rs

@@ -16,7 +16,7 @@ pub use initiate_validator_exit::initiate_validator_exit;
 pub use slash_validator::slash_validator;
 
 use safe_arith::SafeArith;
-use types::{BeaconState, BeaconStateError, EthSpec};
+use types::{BeaconState, BeaconStateError, EthSpec, GetBalanceMut};
 
 /// Increase the balance of a validator, erroring upon overflow, as per the spec.
 pub fn increase_balance<E: EthSpec>(
@@ -24,7 +24,10 @@ pub fn increase_balance<E: EthSpec>(
     index: usize,
     delta: u64,
 ) -> Result<(), BeaconStateError> {
-    state.get_balance_mut(index)?.safe_add_assign(delta)?;
+    state
+        .balances_mut()
+        .get_balance_mut(index)?
+        .safe_add_assign(delta)?;
     Ok(())
 }
 
@@ -34,7 +37,8 @@ pub fn decrease_balance<E: EthSpec>(
     index: usize,
     delta: u64,
 ) -> Result<(), BeaconStateError> {
-    let balance = state.get_balance_mut(index)?;
+    let mut balances = state.balances_mut();
+    let balance = balances.get_balance_mut(index)?;
     *balance = balance.saturating_sub(delta);
     Ok(())
 }

consensus/state_processing/src/per_epoch_processing/slashings.rs

@@ -1,6 +1,6 @@
 use crate::per_epoch_processing::Error;
 use safe_arith::{SafeArith, SafeArithIter};
-use types::{BeaconState, BeaconStateError, ChainSpec, EthSpec, GetValidatorMut, Unsigned};
+use types::{BeaconState, ChainSpec, EthSpec, GetBalanceMut, GetValidatorMut, Unsigned};
 
 /// Process slashings.
 pub fn process_slashings<T: EthSpec>(
@@ -15,7 +15,7 @@ pub fn process_slashings<T: EthSpec>(
     let adjusted_total_slashing_balance =
         std::cmp::min(sum_slashings.safe_mul(slashing_multiplier)?, total_balance);
 
-    let (validators, balances) = state.validators_and_balances_mut();
+    let (validators, mut balances) = state.validators_and_balances_mut();
     for index in 0..validators.len() {
         let validator = validators.get_validator(index)?;
         if validator.slashed
@@ -32,9 +32,7 @@ pub fn process_slashings<T: EthSpec>(
                 .safe_mul(increment)?;
 
             // Equivalent to `decrease_balance(state, index, penalty)`, but avoids borrowing `state`.
-            let balance = balances
+            let balance = balances.get_balance_mut(index)?;
-                .get_mut(index)
-                .ok_or(BeaconStateError::BalancesOutOfBounds(index))?;
             *balance = balance.saturating_sub(penalty);
         }
     }

consensus/types/src/beacon_state.rs

@@ -53,6 +53,11 @@ pub type ValidatorsMut<'a, N> = ListMut<'a, Validator, N>;
 #[cfg(not(feature = "milhouse"))]
 pub type ValidatorsMut<'a, N> = &'a mut VList<Validator, N>;
 
+#[cfg(feature = "milhouse")]
+pub type BalancesMut<'a, N> = ListMut<'a, u64, N>;
+#[cfg(not(feature = "milhouse"))]
+pub type BalancesMut<'a, N> = ListMut<'a, u64, N>;
+
 pub const CACHED_EPOCHS: usize = 3;
 const MAX_RANDOM_BYTE: u64 = (1 << 8) - 1;
 
@@ -251,9 +256,11 @@ where
     #[superstruct(getter(rename = "validators_raw"))]
     #[test_random(default)]
     pub validators: VList<Validator, T::ValidatorRegistryLimit>,
-    #[compare_fields(as_slice)]
+    // FIXME(sproul): serde quoting
-    #[serde(with = "ssz_types::serde_utils::quoted_u64_var_list")]
+    // #[serde(with = "ssz_types::serde_utils::quoted_u64_var_list")]
-    pub balances: VariableList<u64, T::ValidatorRegistryLimit>,
+    #[superstruct(getter(rename = "balances_raw"))]
+    #[test_random(default)]
+    pub balances: VList<u64, T::ValidatorRegistryLimit>,
 
     // Randomness
     pub randao_mixes: FixedVector<Hash256, T::EpochsPerHistoricalVector>,
@@ -365,8 +372,8 @@ impl<T: EthSpec> BeaconState<T> {
             eth1_deposit_index: 0,
 
             // Validator registry
-            validators: VList::empty(),      // Set later.
+            validators: VList::empty(), // Set later.
-            balances: VariableList::empty(), // Set later.
+            balances: VList::empty(),   // Set later.
 
             // Randomness
             randao_mixes: FixedVector::from_elem(Hash256::zero()),
@@ -1122,10 +1129,28 @@ impl<T: EthSpec> BeaconState<T> {
         }
     }
 
+    pub fn balances(&self) -> &VList<u64, T::ValidatorRegistryLimit> {
+        self.balances_raw()
+    }
+
+    pub fn balances_mut(&mut self) -> BalancesMut<T::ValidatorRegistryLimit> {
+        #[cfg(not(feature = "milhouse"))]
+        {
+            self.balances_raw_mut()
+        }
+        #[cfg(feature = "milhouse")]
+        {
+            self.balances_raw_mut().as_mut()
+        }
+    }
+
     /// Convenience accessor for validators and balances simultaneously.
     pub fn validators_and_balances_mut(
         &mut self,
-    ) -> (ValidatorsMut<T::ValidatorRegistryLimit>, &mut [u64]) {
+    ) -> (
+        ValidatorsMut<T::ValidatorRegistryLimit>,
+        BalancesMut<T::ValidatorRegistryLimit>,
+    ) {
         #[cfg(not(feature = "milhouse"))]
         match self {
             BeaconState::Base(state) => (&mut state.validators, &mut state.balances),
@@ -1134,8 +1159,8 @@ impl<T: EthSpec> BeaconState<T> {
 
         #[cfg(feature = "milhouse")]
         match self {
-            BeaconState::Base(state) => (state.validators.as_mut(), &mut state.balances),
+            BeaconState::Base(state) => (state.validators.as_mut(), state.balances.as_mut()),
-            BeaconState::Altair(state) => (state.validators.as_mut(), &mut state.balances),
+            BeaconState::Altair(state) => (state.validators.as_mut(), state.balances.as_mut()),
         }
     }
 
@@ -1215,13 +1240,6 @@ impl<T: EthSpec> BeaconState<T> {
             .ok_or(Error::InactivityScoresOutOfBounds(validator_index))
     }
 
-    /// Get a mutable reference to the balance of a single validator.
-    pub fn get_balance_mut(&mut self, validator_index: usize) -> Result<&mut u64, Error> {
-        self.balances_mut()
-            .get_mut(validator_index)
-            .ok_or(Error::BalancesOutOfBounds(validator_index))
-    }
-
     ///  Return the epoch at which an activation or exit triggered in ``epoch`` takes effect.
     ///
     ///  Spec v0.12.1

consensus/types/src/lib.rs

@@ -118,7 +118,7 @@ pub use crate::free_attestation::FreeAttestation;
 pub use crate::graffiti::{Graffiti, GRAFFITI_BYTES_LEN};
 pub use crate::historical_batch::HistoricalBatch;
 pub use crate::indexed_attestation::IndexedAttestation;
-pub use crate::mixin::GetValidatorMut;
+pub use crate::mixin::{GetBalanceMut, GetValidatorMut};
 pub use crate::participation_flags::ParticipationFlags;
 pub use crate::participation_list::ParticipationList;
 pub use crate::pending_attestation::PendingAttestation;

consensus/types/src/mixin.rs

@@ -1,4 +1,4 @@
-use crate::beacon_state::{Error, ValidatorsMut};
+use crate::beacon_state::{BalancesMut, Error, ValidatorsMut};
 use crate::{Unsigned, Validator};
 
 pub trait GetValidatorMut {
@@ -16,3 +16,21 @@ impl<'a, N: Unsigned> GetValidatorMut for ValidatorsMut<'a, N> {
         self.get_mut(index).ok_or(Error::UnknownValidator(index))
     }
 }
+
+pub trait GetBalanceMut {
+    fn get_balance(&self, index: usize) -> Result<u64, Error>;
+
+    fn get_balance_mut(&mut self, index: usize) -> Result<&mut u64, Error>;
+}
+
+impl<'a, N: Unsigned> GetBalanceMut for BalancesMut<'a, N> {
+    fn get_balance(&self, index: usize) -> Result<u64, Error> {
+        self.get(index)
+            .copied()
+            .ok_or(Error::BalancesOutOfBounds(index))
+    }
+
+    fn get_balance_mut(&mut self, index: usize) -> Result<&mut u64, Error> {
+        self.get_mut(index).ok_or(Error::BalancesOutOfBounds(index))
+    }
+}