Make BeaconState.balances a tree list!

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)]
-    #[serde(with = "ssz_types::serde_utils::quoted_u64_var_list")]
-    pub balances: VariableList<u64, T::ValidatorRegistryLimit>,
+    // FIXME(sproul): serde quoting
+    // #[serde(with = "ssz_types::serde_utils::quoted_u64_var_list")]
+    #[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.
-            balances: VariableList::empty(), // Set later.
+            validators: VList::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::Altair(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(), 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))
+    }
+}