From a50a2486da4b791ae976807c47914ceac800a624 Mon Sep 17 00:00:00 2001
From: EchoAlice <echoalice0@gmail.com>
Date: Mon, 2 Dec 2024 13:13:34 -0700
Subject: [PATCH] feat: optimize phase0 attestation processing

---
 ethereum-consensus/Cargo.toml                 |  3 +-
 .../src/phase0/epoch_processing.rs            | 85 +++++++++++++++----
 2 files changed, 72 insertions(+), 16 deletions(-)

diff --git a/ethereum-consensus/Cargo.toml b/ethereum-consensus/Cargo.toml
index fee91112b..220f9abaf 100644
--- a/ethereum-consensus/Cargo.toml
+++ b/ethereum-consensus/Cargo.toml
@@ -9,8 +9,9 @@ license = "MIT OR Apache-2.0"
 default = ["serde", "async"]
 serde = ["hex", "serde_json", "serde_yaml"]
 async = ["tokio", "tokio-stream"]
-optimized = ["shuffling"]
+optimized = ["shuffling", "attestation-processing"]
 shuffling = [] # supports optimized shuffling routines
+attestation-processing = [] # supports optimized attestation processing
 secret-key-debug = [
 ] # enable if you want to be able to print `crypto::SecretKey`
 spec-tests = [] # enable extra features for testing
diff --git a/ethereum-consensus/src/phase0/epoch_processing.rs b/ethereum-consensus/src/phase0/epoch_processing.rs
index 74b5e6a06..19cb14b3a 100644
--- a/ethereum-consensus/src/phase0/epoch_processing.rs
+++ b/ethereum-consensus/src/phase0/epoch_processing.rs
@@ -923,23 +923,78 @@ pub fn get_inclusion_delay_deltas<
     let previous_epoch = get_previous_epoch(state, context);
     let validator_count = state.validators.len();
     let mut rewards = vec![0; validator_count];
-    let matching_source_attestations =
-        get_matching_source_attestations(state, previous_epoch, context)?;
-    for i in get_unslashed_attesting_indices(state, matching_source_attestations.iter(), context)? {
-        let mut attestations = Vec::new();
-        for a in matching_source_attestations.iter() {
-            if get_attesting_indices(state, &a.data, &a.aggregation_bits, context)?.contains(&i) {
-                attestations.push(a)
+
+    #[cfg(feature = "attestation-processing")]
+    {
+        #[derive(Default, Clone)]
+        struct AttesterStatus {
+            min_inclusion_delay: u64,
+            proposer_index: usize,
+        }
+
+        let eligible_validators: Vec<bool> = state.validators.iter().map(|v| !v.slashed).collect();
+        let matching_source_attestations =
+            get_matching_source_attestations(state, previous_epoch, context)?;
+        let mut attester_statuses: Vec<AttesterStatus> =
+            vec![AttesterStatus::default(); validator_count];
+
+        // Process all attestations once and store minimum inclusion delays
+        for attestation in matching_source_attestations.iter() {
+            let attesting_indices = get_attesting_indices(
+                state,
+                &attestation.data,
+                &attestation.aggregation_bits,
+                context,
+            )?;
+            for &validator_index in &attesting_indices {
+                if !eligible_validators[validator_index] {
+                    continue;
+                }
+                let current_status = &mut attester_statuses[validator_index];
+                if current_status.min_inclusion_delay == 0 ||
+                    attestation.inclusion_delay < current_status.min_inclusion_delay
+                {
+                    current_status.min_inclusion_delay = attestation.inclusion_delay;
+                    current_status.proposer_index = attestation.proposer_index;
+                }
             }
         }
-        let attestation = attestations
-            .iter()
-            .min_by(|&a, &b| a.inclusion_delay.cmp(&b.inclusion_delay))
-            .expect("at least one attestation in collection");
-        rewards[attestation.proposer_index] += get_proposer_reward(state, i, context)?;
-        let max_attester_reward =
-            get_base_reward(state, i, context)? - get_proposer_reward(state, i, context)?;
-        rewards[i] += max_attester_reward / attestation.inclusion_delay;
+        // Calculate rewards based on pre-computed data
+        for (validator_index, status) in attester_statuses.iter().enumerate() {
+            if status.min_inclusion_delay == 0 || !eligible_validators[validator_index] {
+                continue;
+            }
+            let proposer_reward = get_proposer_reward(state, validator_index, context)?;
+            rewards[status.proposer_index] += proposer_reward;
+            let max_attester_reward =
+                get_base_reward(state, validator_index, context)? - proposer_reward;
+            rewards[validator_index] += max_attester_reward / status.min_inclusion_delay;
+        }
+    }
+
+    #[cfg(not(feature = "attestation-processing"))]
+    {
+        let matching_source_attestations =
+            get_matching_source_attestations(state, previous_epoch, context)?;
+        for i in
+            get_unslashed_attesting_indices(state, matching_source_attestations.iter(), context)?
+        {
+            let mut attestations = Vec::new();
+            for a in matching_source_attestations.iter() {
+                if get_attesting_indices(state, &a.data, &a.aggregation_bits, context)?.contains(&i)
+                {
+                    attestations.push(a)
+                }
+            }
+            let attestation = attestations
+                .iter()
+                .min_by(|&a, &b| a.inclusion_delay.cmp(&b.inclusion_delay))
+                .expect("at least one attestation in collection");
+            rewards[attestation.proposer_index] += get_proposer_reward(state, i, context)?;
+            let max_attester_reward =
+                get_base_reward(state, i, context)? - get_proposer_reward(state, i, context)?;
+            rewards[i] += max_attester_reward / attestation.inclusion_delay;
+        }
     }
     Ok((rewards, vec![0; validator_count]))
 }