From 9de8dec7502f987f93ecf9e3a95532bdb05ad5ff Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tom=C3=A1s=20Gr=C3=BCner?=
 <47506558+MegaRedHand@users.noreply.github.com>
Date: Tue, 12 Dec 2023 14:22:05 -0300
Subject: [PATCH] perf: reduce `get_beacon_committee` calls (#518)

---
 .../state_transition/accessors.ex             | 29 +++++++++++++++++--
 .../state_transition/operations.ex            | 24 +++++++--------
 2 files changed, 37 insertions(+), 16 deletions(-)

diff --git a/lib/lambda_ethereum_consensus/state_transition/accessors.ex b/lib/lambda_ethereum_consensus/state_transition/accessors.ex
index 588641ba9..5e0131e88 100644
--- a/lib/lambda_ethereum_consensus/state_transition/accessors.ex
+++ b/lib/lambda_ethereum_consensus/state_transition/accessors.ex
@@ -251,8 +251,9 @@ defmodule LambdaEthereumConsensus.StateTransition.Accessors do
   """
   @spec get_committee_count_per_slot(BeaconState.t(), SszTypes.epoch()) :: SszTypes.uint64()
   def get_committee_count_per_slot(%BeaconState{} = state, epoch) do
-    get_active_validator_indices(state, epoch)
-    |> length()
+    state.validators
+    |> Stream.filter(&Predicates.is_active_validator(&1, epoch))
+    |> Enum.count()
     |> div(ChainSpec.get("SLOTS_PER_EPOCH"))
     |> div(ChainSpec.get("TARGET_COMMITTEE_SIZE"))
     |> min(ChainSpec.get("MAX_COMMITTEES_PER_SLOT"))
@@ -444,6 +445,20 @@ defmodule LambdaEthereumConsensus.StateTransition.Accessors do
     end
   end
 
+  @spec get_committee_indexed_attestation([SszTypes.validator_index()], Attestation.t()) ::
+          IndexedAttestation.t()
+  def get_committee_indexed_attestation(beacon_committee, attestation) do
+    get_committee_attesting_indices(beacon_committee, attestation.aggregation_bits)
+    |> Enum.sort()
+    |> then(
+      &%IndexedAttestation{
+        attesting_indices: &1,
+        data: attestation.data,
+        signature: attestation.signature
+      }
+    )
+  end
+
   @doc """
   Return the set of attesting indices corresponding to ``data`` and ``bits``.
   """
@@ -460,6 +475,16 @@ defmodule LambdaEthereumConsensus.StateTransition.Accessors do
     end
   end
 
+  @spec get_committee_attesting_indices([SszTypes.validator_index()], SszTypes.bitlist()) ::
+          MapSet.t()
+  def get_committee_attesting_indices(committee, bits) do
+    committee
+    |> Stream.with_index()
+    |> Stream.filter(fn {_value, index} -> participated?(bits, index) end)
+    |> Stream.map(fn {value, _index} -> value end)
+    |> MapSet.new()
+  end
+
   defp participated?(bits, index) do
     # The bit order inside the byte is reversed (e.g. bits[0] is the 8th bit).
     # Here we keep the byte index the same, but reverse the bit index.
diff --git a/lib/lambda_ethereum_consensus/state_transition/operations.ex b/lib/lambda_ethereum_consensus/state_transition/operations.ex
index 64bd832d7..01ddec148 100644
--- a/lib/lambda_ethereum_consensus/state_transition/operations.ex
+++ b/lib/lambda_ethereum_consensus/state_transition/operations.ex
@@ -645,29 +645,27 @@ defmodule LambdaEthereumConsensus.StateTransition.Operations do
   @spec process_attestation(BeaconState.t(), Attestation.t()) ::
           {:ok, BeaconState.t()} | {:error, binary()}
   def process_attestation(state, %Attestation{data: data} = attestation) do
-    # TODO: optimize (takes ~3s)
     with :ok <- check_valid_target_epoch(data, state),
          :ok <- check_epoch_matches(data),
          :ok <- check_valid_slot_range(data, state),
          :ok <- check_committee_count(data, state),
          {:ok, beacon_committee} <- Accessors.get_beacon_committee(state, data.slot, data.index),
          :ok <- check_matching_aggregation_bits_length(attestation, beacon_committee),
-         {:ok, indexed_attestation} <- Accessors.get_indexed_attestation(state, attestation),
+         indexed_attestation =
+           Accessors.get_committee_indexed_attestation(beacon_committee, attestation),
          :ok <- check_valid_signature(state, indexed_attestation) do
-      # TODO: optimize (takes ~1s)
-      process_attestation(state, data, attestation.aggregation_bits)
+      inner_process_attestation(state, data, attestation.aggregation_bits, beacon_committee)
     end
   end
 
-  defp process_attestation(state, data, aggregation_bits) do
+  defp inner_process_attestation(state, data, aggregation_bits, committee) do
+    slot = state.slot - data.slot
+
     with {:ok, participation_flag_indices} <-
-           Accessors.get_attestation_participation_flag_indices(
-             state,
-             data,
-             state.slot - data.slot
-           ),
-         {:ok, attesting_indices} <-
-           Accessors.get_attesting_indices(state, data, aggregation_bits) do
+           Accessors.get_attestation_participation_flag_indices(state, data, slot) do
+      attesting_indices =
+        Accessors.get_committee_attesting_indices(committee, aggregation_bits)
+
       is_current_epoch = data.target.epoch == Accessors.get_current_epoch(state)
       initial_epoch_participation = get_initial_epoch_participation(state, is_current_epoch)
 
@@ -687,8 +685,6 @@ defmodule LambdaEthereumConsensus.StateTransition.Operations do
       |> Mutators.increase_balance(proposer_index, proposer_reward)
       |> update_state(is_current_epoch, updated_epoch_participation)
       |> then(&{:ok, &1})
-    else
-      {:error, reason} -> {:error, reason}
     end
   end