From 5ed8ad557de8026aa8592aed38ffafd69466e891 Mon Sep 17 00:00:00 2001
From: Satya Vusirikala <satyasatya123456@gmail.com>
Date: Wed, 18 Sep 2024 23:31:45 -0700
Subject: [PATCH] Addressing PR comments

---
 types/src/aggregate_signature.rs |  6 +++---
 types/src/ledger_info.rs         | 21 ++++++++-------------
 2 files changed, 11 insertions(+), 16 deletions(-)

diff --git a/types/src/aggregate_signature.rs b/types/src/aggregate_signature.rs
index 3202583b6252d..24ac789da671e 100644
--- a/types/src/aggregate_signature.rs
+++ b/types/src/aggregate_signature.rs
@@ -86,12 +86,12 @@ impl PartialSignatures {
         self.signatures.is_empty()
     }
 
-    pub fn remove_signature(&mut self, validator: AccountAddress) {
-        self.signatures.remove(&validator);
+    pub fn remove_signature(&mut self, validator: AccountAddress) -> Option<bls12381::Signature> {
+        self.signatures.remove(&validator)
     }
 
     pub fn add_signature(&mut self, validator: AccountAddress, signature: bls12381::Signature) {
-        self.signatures.entry(validator).or_insert(signature);
+        self.signatures.insert(validator, signature);
     }
 
     pub fn signatures(&self) -> &BTreeMap<AccountAddress, bls12381::Signature> {
diff --git a/types/src/ledger_info.rs b/types/src/ledger_info.rs
index 7e0da229864e0..9f3aead311f81 100644
--- a/types/src/ledger_info.rs
+++ b/types/src/ledger_info.rs
@@ -318,11 +318,6 @@ impl LedgerInfoWithV0 {
     }
 }
 
-pub enum VerificationStatus {
-    Verified,
-    Unverified,
-}
-
 /// Contains the ledger info and partially aggregated signature from a set of validators, this data
 /// is only used during the aggregating the votes from different validators and is not persisted in DB.
 #[derive(Clone, Debug, Eq, PartialEq)]
@@ -437,9 +432,6 @@ impl LedgerInfoWithUnverifiedSignatures {
         if self.verified_signatures.contains_voter(&validator) {
             return;
         }
-        if self.unverified_signatures.contains_voter(&validator) {
-            self.unverified_signatures.remove_signature(validator);
-        }
         self.unverified_signatures
             .add_signature(validator, signature);
     }
@@ -527,15 +519,18 @@ impl LedgerInfoWithUnverifiedSignatures {
                             .verify(*account_address, self.ledger_info(), signature)
                             .is_ok()
                         {
-                            return Some((*account_address, signature.clone()));
+                            return Some(*account_address);
                         }
                         None
                     })
                     .collect::<Vec<_>>();
-                for (account_address, signature) in verified {
-                    self.verified_signatures
-                        .add_signature(account_address, signature);
-                    self.unverified_signatures.remove_signature(account_address);
+                for account_address in verified {
+                    if let Some(signature) =
+                        self.unverified_signatures.remove_signature(account_address)
+                    {
+                        self.verified_signatures
+                            .add_signature(account_address, signature);
+                    }
                 }
 
                 // For these authors, we will not use optimistic signature verification in the future.