From 2a9bcb5f59d6e1a528e9f2df35eedaaaf2c62aa8 Mon Sep 17 00:00:00 2001
From: Michal Nazarewicz <mina86@mina86.com>
Date: Tue, 9 Apr 2024 12:26:23 +0200
Subject: [PATCH] light-client-verifier: optimise validator lookup in
 voting_power_in
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Rather than performing a linear search for a validator by its address,
construct an index for the validator set sorted by address.  While
building the index takes O(N log N) time, it cuts lookup to O(log n).

On Solana, with ~50 validators, this change causes reduction of the
cost of a single lookup from 20k compute units to 700.

Furthermore, merge the index with the seen_validators map so that only
a single lookup is now performed which allows us to find the validator
and check whether it’s not a duplicate.

Co-authored-by: Dhruv D Jain <dhruv@iamsizzling.com>
---
 .../1406-optimise-voting_power_in.md          |  2 +
 .../src/operations/voting_power.rs            | 38 +++++++++++++------
 2 files changed, 29 insertions(+), 11 deletions(-)
 create mode 100644 .changelog/unreleased/improvements/1406-optimise-voting_power_in.md

diff --git a/.changelog/unreleased/improvements/1406-optimise-voting_power_in.md b/.changelog/unreleased/improvements/1406-optimise-voting_power_in.md
new file mode 100644
index 000000000..fbe14dcf4
--- /dev/null
+++ b/.changelog/unreleased/improvements/1406-optimise-voting_power_in.md
@@ -0,0 +1,2 @@
+- [light-client-verifier]: Optimise validators lookup in
+  ProvidedVotingPowerCalculator::voting_power_in method.
diff --git a/light-client-verifier/src/operations/voting_power.rs b/light-client-verifier/src/operations/voting_power.rs
index 8fb8ecd85..cc6899f06 100644
--- a/light-client-verifier/src/operations/voting_power.rs
+++ b/light-client-verifier/src/operations/voting_power.rs
@@ -1,6 +1,6 @@
 //! Provides an interface and default implementation for the `VotingPower` operation
 
-use alloc::collections::BTreeSet as HashSet;
+use alloc::vec::Vec;
 use core::{fmt, marker::PhantomData};
 
 use serde::{Deserialize, Serialize};
@@ -134,7 +134,6 @@ impl<V: signature::Verifier> VotingPowerCalculator for ProvidedVotingPowerCalcul
         let total_voting_power = self.total_power_of(validator_set);
 
         let mut tallied_voting_power = 0_u64;
-        let mut seen_validators = HashSet::new();
 
         // Get non-absent votes from the signatures
         let non_absent_votes = signatures.iter().enumerate().flat_map(|(idx, signature)| {
@@ -146,20 +145,37 @@ impl<V: signature::Verifier> VotingPowerCalculator for ProvidedVotingPowerCalcul
             .map(|vote| (signature, vote))
         });
 
+        // Create index of validators sorted by address.  The index stores
+        // reference to the validaotr’s Info object held by the validator_set
+        // and boolean flag indicating whether we've already seen that
+        // validator.
+        let mut validators: Vec<(&_, bool)> = validator_set
+            .validators()
+            .iter()
+            .map(|v| (v, false))
+            .collect();
+        validators.sort_unstable_by_key(|item| &item.0.address);
+
         for (signature, vote) in non_absent_votes {
+            // Find the validator by address.
+            let index = validators
+                .binary_search_by_key(&&vote.validator_address, |item| &item.0.address)
+                .map(|index| {
+                    let item = &mut validators[index];
+                    (item.0, &mut item.1)
+                });
+            let (validator, seen) = match index {
+                Ok(it) => it,
+                Err(_) => continue, // Cannot find matching validator, so we skip the vote
+            };
+
             // Ensure we only count a validator's power once
-            if seen_validators.contains(&vote.validator_address) {
+            if *seen {
                 return Err(VerificationError::duplicate_validator(
                     vote.validator_address,
                 ));
-            } else {
-                seen_validators.insert(vote.validator_address);
             }
-
-            let validator = match validator_set.validator(vote.validator_address) {
-                Some(validator) => validator,
-                None => continue, // Cannot find matching validator, so we skip the vote
-            };
+            *seen = true;
 
             let signed_vote =
                 SignedVote::from_vote(vote.clone(), signed_header.header.chain_id.clone())
@@ -173,7 +189,7 @@ impl<V: signature::Verifier> VotingPowerCalculator for ProvidedVotingPowerCalcul
             {
                 return Err(VerificationError::invalid_signature(
                     signed_vote.signature().as_bytes().to_vec(),
-                    Box::new(validator),
+                    Box::new(validator.clone()),
                     sign_bytes,
                 ));
             }