From 2f64ebd435b57e2abf86f592c32161f9273cb1ec Mon Sep 17 00:00:00 2001
From: Leo Eichhorn <leo.eichhorn@dfinity.org>
Date: Mon, 15 Jan 2024 17:09:32 +0000
Subject: [PATCH] feat(ecdsa): CON-1207 Refill matched quadruples

---
 rs/consensus/src/ecdsa/payload_builder.rs     |  7 +++-
 .../src/ecdsa/payload_builder/quadruples.rs   | 39 ++++++++++++++-----
 2 files changed, 36 insertions(+), 10 deletions(-)

diff --git a/rs/consensus/src/ecdsa/payload_builder.rs b/rs/consensus/src/ecdsa/payload_builder.rs
index e0db12fb6ee..16dbc59cf29 100644
--- a/rs/consensus/src/ecdsa/payload_builder.rs
+++ b/rs/consensus/src/ecdsa/payload_builder.rs
@@ -632,7 +632,12 @@ pub(crate) fn create_data_payload_helper_2(
         quadruples::purge_old_key_quadruples(ecdsa_payload, all_signing_requests);
     }
 
-    quadruples::make_new_quadruples_if_needed(ecdsa_config, ecdsa_payload);
+    let matched_quadruples = all_signing_requests
+        .values()
+        .filter_map(|context| context.matched_quadruple.as_ref())
+        .filter(|(qid, _)| ecdsa_payload.available_quadruples.contains_key(qid))
+        .count();
+    quadruples::make_new_quadruples_if_needed(ecdsa_config, ecdsa_payload, matched_quadruples);
 
     let mut new_transcripts =
         quadruples::update_quadruples_in_creation(ecdsa_payload, transcript_builder, height, log)?;
diff --git a/rs/consensus/src/ecdsa/payload_builder/quadruples.rs b/rs/consensus/src/ecdsa/payload_builder/quadruples.rs
index 17786c46f08..4a3afef2047 100644
--- a/rs/consensus/src/ecdsa/payload_builder/quadruples.rs
+++ b/rs/consensus/src/ecdsa/payload_builder/quadruples.rs
@@ -240,6 +240,8 @@ pub(super) fn purge_old_key_quadruples(
 pub(super) fn make_new_quadruples_if_needed(
     ecdsa_config: &EcdsaConfig,
     ecdsa_payload: &mut ecdsa::EcdsaPayload,
+    //TODO(CON-1149): Pass value to helper
+    _matched_quadruples: usize,
 ) {
     if let Some(key_transcript) = &ecdsa_payload.key_transcript.current {
         let node_ids: Vec<_> = key_transcript.receivers().iter().copied().collect();
@@ -248,6 +250,7 @@ pub(super) fn make_new_quadruples_if_needed(
             key_transcript.registry_version(),
             ecdsa_config,
             ecdsa_payload,
+            0,
         )
     }
 }
@@ -257,8 +260,12 @@ fn make_new_quadruples_if_needed_helper(
     registry_version: RegistryVersion,
     ecdsa_config: &EcdsaConfig,
     ecdsa_payload: &mut ecdsa::EcdsaPayload,
+    matched_quadruples: usize,
 ) {
-    let unassigned_quadruples = ecdsa_payload.unassigned_quadruple_ids().count();
+    let unassigned_quadruples = ecdsa_payload
+        .unassigned_quadruple_ids()
+        .count()
+        .saturating_sub(matched_quadruples);
     let quadruples_to_create = ecdsa_config.quadruples_to_create_in_advance as usize;
     if quadruples_to_create > unassigned_quadruples {
         let quadruples_in_creation = &mut ecdsa_payload.quadruples_in_creation;
@@ -393,24 +400,41 @@ pub(super) mod tests {
     #[test]
     fn test_ecdsa_make_new_quadruples_if_needed() {
         let mut rng = reproducible_rng();
-        let quadruples_to_create_in_advance = 3;
         let subnet_id = subnet_test_id(1);
         let height = Height::new(10);
         let (mut ecdsa_payload, env, _block_reader) = set_up(&mut rng, subnet_id, height);
+
+        // 4 Quadruples should be created in advance (in creation + unmatched available = 4)
+        let quadruples_to_create_in_advance = 4;
         let ecdsa_config = EcdsaConfig {
             quadruples_to_create_in_advance,
             ..EcdsaConfig::default()
         };
 
+        // Add 3 available quadruples
+        for i in 0..3 {
+            create_available_quadruple(&mut ecdsa_payload, i);
+        }
+
+        // 2 available quadruples are already matched
+        let quadruples_already_matched = 2;
+
+        // We expect 3 quadruples in creation to be added
+        let expected_quadruples_in_creation = quadruples_to_create_in_advance as usize
+            - (ecdsa_payload.available_quadruples.len() - quadruples_already_matched);
+        assert_eq!(expected_quadruples_in_creation, 3);
+
         make_new_quadruples_if_needed_helper(
             &env.nodes.ids::<Vec<_>>(),
             env.newest_registry_version,
             &ecdsa_config,
             &mut ecdsa_payload,
+            quadruples_already_matched,
         );
 
         assert_eq!(
-            ecdsa_payload.quadruples_in_creation.len(),
+            ecdsa_payload.quadruples_in_creation.len() + ecdsa_payload.available_quadruples.len()
+                - quadruples_already_matched,
             quadruples_to_create_in_advance as usize
         );
         // Verify the generated transcript ids.
@@ -419,10 +443,7 @@ pub(super) mod tests {
             transcript_ids.insert(quadruple.1.kappa_config.as_ref().transcript_id);
             transcript_ids.insert(quadruple.1.lambda_config.as_ref().transcript_id);
         }
-        assert_eq!(
-            transcript_ids.len(),
-            2 * quadruples_to_create_in_advance as usize
-        );
+        assert_eq!(transcript_ids.len(), 2 * expected_quadruples_in_creation);
         assert_eq!(
             transcript_ids,
             BTreeSet::from([
@@ -435,8 +456,8 @@ pub(super) mod tests {
             ])
         );
         assert_eq!(
-            ecdsa_payload.peek_next_transcript_id().id() as u32,
-            2 * quadruples_to_create_in_advance,
+            ecdsa_payload.peek_next_transcript_id().id() as usize,
+            2 * expected_quadruples_in_creation,
         );
     }