From 282eefb91e250d6aa4ee0eaac7fe5b8ae80d63fc Mon Sep 17 00:00:00 2001
From: Michael Sproul <michael@sigmaprime.io>
Date: Fri, 20 Dec 2024 17:06:17 +1100
Subject: [PATCH] Implement "Use 16-bit random value in validator filter"

---
 consensus/types/src/beacon_state.rs | 46 +++++++++++++++++++++++++----
 1 file changed, 40 insertions(+), 6 deletions(-)

diff --git a/consensus/types/src/beacon_state.rs b/consensus/types/src/beacon_state.rs
index ad4484b86a..18b862f062 100644
--- a/consensus/types/src/beacon_state.rs
+++ b/consensus/types/src/beacon_state.rs
@@ -46,6 +46,7 @@ mod tests;
 
 pub const CACHED_EPOCHS: usize = 3;
 const MAX_RANDOM_BYTE: u64 = (1 << 8) - 1;
+const MAX_RANDOM_VALUE: u64 = (1 << 16) - 1;
 
 pub type Validators<E> = List<Validator, <E as EthSpec>::ValidatorRegistryLimit>;
 pub type Balances<E> = List<u64, <E as EthSpec>::ValidatorRegistryLimit>;
@@ -895,6 +896,11 @@ impl<E: EthSpec> BeaconState<E> {
         }
 
         let max_effective_balance = spec.max_effective_balance_for_fork(self.fork_name_unchecked());
+        let max_random_value = if self.fork_name_unchecked().electra_enabled() {
+            MAX_RANDOM_VALUE
+        } else {
+            MAX_RANDOM_BYTE
+        };
 
         let mut i = 0;
         loop {
@@ -908,10 +914,10 @@ impl<E: EthSpec> BeaconState<E> {
             let candidate_index = *indices
                 .get(shuffled_index)
                 .ok_or(Error::ShuffleIndexOutOfBounds(shuffled_index))?;
-            let random_byte = Self::shuffling_random_byte(i, seed)?;
+            let random_value = self.shuffling_random_value(i, seed)?;
             let effective_balance = self.get_effective_balance(candidate_index)?;
-            if effective_balance.safe_mul(MAX_RANDOM_BYTE)?
-                >= max_effective_balance.safe_mul(u64::from(random_byte))?
+            if effective_balance.safe_mul(max_random_value)?
+                >= max_effective_balance.safe_mul(random_value)?
             {
                 return Ok(candidate_index);
             }
@@ -919,6 +925,14 @@ impl<E: EthSpec> BeaconState<E> {
         }
     }
 
+    fn shuffling_random_value(&self, i: usize, seed: &[u8]) -> Result<u64, Error> {
+        if self.fork_name_unchecked().electra_enabled() {
+            Self::shuffling_random_u16_electra(i, seed).map(u64::from)
+        } else {
+            Self::shuffling_random_byte(i, seed).map(u64::from)
+        }
+    }
+
     /// Get a random byte from the given `seed`.
     ///
     /// Used by the proposer & sync committee selection functions.
@@ -932,6 +946,21 @@ impl<E: EthSpec> BeaconState<E> {
             .ok_or(Error::ShuffleIndexOutOfBounds(index))
     }
 
+    /// Get two random bytes from the given `seed`.
+    ///
+    /// This is used in place of the
+    fn shuffling_random_u16_electra(i: usize, seed: &[u8]) -> Result<u16, Error> {
+        let mut preimage = seed.to_vec();
+        preimage.append(&mut int_to_bytes8(i.safe_div(16)? as u64));
+        let offset = i.safe_rem(16)?.safe_mul(2)?;
+        hash(&preimage)
+            .get(offset..offset.safe_add(2)?)
+            .ok_or(Error::ShuffleIndexOutOfBounds(offset))?
+            .try_into()
+            .map(u16::from_le_bytes)
+            .map_err(|_| Error::ShuffleIndexOutOfBounds(offset))
+    }
+
     /// Convenience accessor for the `execution_payload_header` as an `ExecutionPayloadHeaderRef`.
     pub fn latest_execution_payload_header(&self) -> Result<ExecutionPayloadHeaderRef<E>, Error> {
         match self {
@@ -1093,6 +1122,11 @@ impl<E: EthSpec> BeaconState<E> {
 
         let seed = self.get_seed(epoch, Domain::SyncCommittee, spec)?;
         let max_effective_balance = spec.max_effective_balance_for_fork(self.fork_name_unchecked());
+        let max_random_value = if self.fork_name_unchecked().electra_enabled() {
+            MAX_RANDOM_VALUE
+        } else {
+            MAX_RANDOM_BYTE
+        };
 
         let mut i = 0;
         let mut sync_committee_indices = Vec::with_capacity(E::SyncCommitteeSize::to_usize());
@@ -1107,10 +1141,10 @@ impl<E: EthSpec> BeaconState<E> {
             let candidate_index = *active_validator_indices
                 .get(shuffled_index)
                 .ok_or(Error::ShuffleIndexOutOfBounds(shuffled_index))?;
-            let random_byte = Self::shuffling_random_byte(i, seed.as_slice())?;
+            let random_value = self.shuffling_random_value(i, seed.as_slice())?;
             let effective_balance = self.get_validator(candidate_index)?.effective_balance;
-            if effective_balance.safe_mul(MAX_RANDOM_BYTE)?
-                >= max_effective_balance.safe_mul(u64::from(random_byte))?
+            if effective_balance.safe_mul(max_random_value)?
+                >= max_effective_balance.safe_mul(random_value)?
             {
                 sync_committee_indices.push(candidate_index);
             }