From 416517a5d462798d6ed131ea2cb2c0067f8920d6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Esteban=20Dimitroff=20H=C3=B3di?=
 <esteban.dimitroff@lambdaclass.com>
Date: Fri, 15 Mar 2024 20:15:40 -0300
Subject: [PATCH] Simplified weights representation

---
 .../roles/src/validator/messages/consensus.rs | 20 +++++++++++++------
 1 file changed, 14 insertions(+), 6 deletions(-)

diff --git a/node/libs/roles/src/validator/messages/consensus.rs b/node/libs/roles/src/validator/messages/consensus.rs
index 3865fcd0..3ce32b70 100644
--- a/node/libs/roles/src/validator/messages/consensus.rs
+++ b/node/libs/roles/src/validator/messages/consensus.rs
@@ -8,6 +8,7 @@ use bit_vec::BitVec;
 use std::{
     collections::{BTreeMap, BTreeSet},
     fmt,
+    iter::zip,
 };
 use zksync_consensus_crypto::{keccak256::Keccak256, ByteFmt, Text, TextFmt};
 use zksync_consensus_utils::enum_util::{BadVariantError, Variant};
@@ -83,7 +84,7 @@ impl Default for Fork {
 pub struct ValidatorCommittee {
     vec: Vec<validator::PublicKey>,
     indexes: BTreeMap<validator::PublicKey, usize>,
-    weights: BTreeMap<validator::PublicKey, usize>,
+    weights: Vec<usize>,
 }
 
 impl ValidatorCommittee {
@@ -107,8 +108,13 @@ impl ValidatorCommittee {
         );
         Ok(Self {
             vec: set.iter().cloned().collect(),
-            indexes: set.into_iter().enumerate().map(|(i, pk)| (pk, i)).collect(),
-            weights,
+            indexes: set
+                .clone()
+                .into_iter()
+                .enumerate()
+                .map(|(i, pk)| (pk, i))
+                .collect(),
+            weights: set.iter().map(|pk| weights[pk]).collect(),
         })
     }
 
@@ -119,7 +125,7 @@ impl ValidatorCommittee {
 
     /// Iterates over validators.
     pub fn weighted_validators_iter(&self) -> impl Iterator<Item = WeightedValidator> + '_ {
-        self.weights.iter().map(|(key, weight)| WeightedValidator {
+        zip(&self.vec, &self.weights).map(|(key, weight)| WeightedValidator {
             key: key.clone(),
             weight: *weight as u32,
         })
@@ -164,10 +170,12 @@ impl ValidatorCommittee {
 
     /// Compute the sum of signers weights.
     pub fn weight(&self, signers: Signers) -> usize {
-        self.iter()
+        self.weights
+            .iter()
             .enumerate()
             .filter(|(i, _)| signers.0[*i])
-            .fold(0, |acc, (_, pk)| acc + self.weights[pk])
+            .map(|(_, weight)| weight)
+            .sum()
     }
 }