refactor(consensus): Refactor code in NiDkg to prepare VetKD implementation

rs/consensus/dkg/src/dkg_key_manager.rs

@@ -484,7 +484,7 @@ impl DkgKeyManager {
             .set(summary.registry_version.get() as i64);
 
         for tag in [NiDkgTag::LowThreshold, NiDkgTag::HighThreshold].iter() {
-            let current_transcript = summary.current_transcript(tag);
+            let current_transcript = summary.current_transcript(tag).unwrap();
             let metric_label = &format!("{:?}", tag);
 
             self.metrics

rs/consensus/dkg/src/lib.rs

@@ -1646,7 +1646,7 @@ mod tests {
                 RegistryVersion::from(5)
             );
             for tag in TAGS.iter() {
-                let current_transcript = dkg_summary.current_transcript(tag);
+                let current_transcript = dkg_summary.current_transcript(tag).unwrap();
                 assert_eq!(
                     current_transcript.dkg_id.start_block_height,
                     Height::from(0)
@@ -1700,7 +1700,7 @@ mod tests {
             );
             for tag in TAGS.iter() {
                 // We reused the transcript.
-                let current_transcript = dkg_summary.current_transcript(tag);
+                let current_transcript = dkg_summary.current_transcript(tag).unwrap();
                 assert_eq!(
                     current_transcript.dkg_id.start_block_height,
                     Height::from(0)
@@ -1763,7 +1763,7 @@ mod tests {
                     conf.receivers().get(),
                     &committee3.clone().into_iter().collect::<BTreeSet<_>>()
                 );
-                let current_transcript = dkg_summary.current_transcript(tag);
+                let current_transcript = dkg_summary.current_transcript(tag).unwrap();
                 assert_eq!(
                     current_transcript.dkg_id.start_block_height,
                     Height::from(0)
@@ -1803,7 +1803,7 @@ mod tests {
                     conf.receivers().get(),
                     &committee3.clone().into_iter().collect::<BTreeSet<_>>()
                 );
-                let current_transcript = dkg_summary.current_transcript(tag);
+                let current_transcript = dkg_summary.current_transcript(tag).unwrap();
                 assert_eq!(
                     current_transcript.dkg_id.start_block_height,
                     Height::from(5)
@@ -1841,7 +1841,7 @@ mod tests {
                     conf.receivers().get(),
                     &committee3.clone().into_iter().collect::<BTreeSet<_>>()
                 );
-                let current_transcript = dkg_summary.current_transcript(tag);
+                let current_transcript = dkg_summary.current_transcript(tag).unwrap();
                 assert_eq!(
                     current_transcript.dkg_id.start_block_height,
                     Height::from(10)

rs/consensus/dkg/src/payload_builder.rs

@@ -212,7 +212,7 @@ pub(super) fn create_summary_payload(
         subnet_id,
     )?;
     // Current transcripts come from next transcripts of the last_summary.
-    let current_transcripts = last_summary.clone().into_next_transcripts();
+    let current_transcripts = as_next_transcripts(last_summary, &logger);
 
     // If the config for the currently computed DKG intervals requires a transcript
     // resharing (currently for high-threshold DKG only), we are going to re-share
@@ -266,6 +266,24 @@ fn create_transcript(
     ic_interfaces::crypto::NiDkgAlgorithm::create_transcript(crypto, config, dealings)
 }
 
+/// Return the set of next transcripts for all tags. If for some tag
+/// the next transcript is not available, the current transcript is used.
+fn as_next_transcripts(
+    summary: &Summary,
+    logger: &ReplicaLogger,
+) -> BTreeMap<NiDkgTag, NiDkgTranscript> {
+    let mut next_transcripts = summary.next_transcripts().clone();
+
+    for (tag, transcript) in summary.current_transcripts().iter() {
+        if !next_transcripts.contains_key(tag) {
+            warn!(logger, "Reusing current transcript for tag {:?}", tag);
+            next_transcripts.insert(tag.clone(), transcript.clone());
+        }
+    }
+
+    next_transcripts
+}
+
 #[allow(clippy::type_complexity)]
 #[allow(clippy::too_many_arguments)]
 fn compute_remote_dkg_data(
@@ -1135,7 +1153,8 @@ mod tests {
                     assert_eq!(
                         next_summary
                             .clone()
-                            .current_transcript(&NiDkgTag::HighThreshold),
+                            .current_transcript(&NiDkgTag::HighThreshold)
+                            .unwrap(),
                         &conf.resharing_transcript().clone().unwrap()
                     )
                 }

rs/consensus/dkg/src/payload_validator.rs

@@ -1,6 +1,4 @@
-use std::collections::HashSet;
-
-use super::{payload_builder, utils, PayloadCreationError};
+use crate::{payload_builder, utils, PayloadCreationError};
 use ic_consensus_utils::{crypto::ConsensusCrypto, pool_reader::PoolReader};
 use ic_interfaces::{
     dkg::DkgPool,
@@ -24,6 +22,7 @@ use ic_types::{
     Height, NodeId, SubnetId,
 };
 use prometheus::IntCounterVec;
+use std::collections::HashSet;
 
 /// Reasons for why a dkg payload might be invalid.
 // The `Debug` implementation is ignored during the dead code analysis and we are getting a `field

rs/consensus/src/cup_utils.rs

@@ -74,10 +74,12 @@ pub fn make_registry_cup_from_cup_contents(
 
     let low_dkg_id = dkg_summary
         .current_transcript(&NiDkgTag::LowThreshold)
+        .expect("No current low threshold transcript available")
         .dkg_id
         .clone();
     let high_dkg_id = dkg_summary
         .current_transcript(&NiDkgTag::HighThreshold)
+        .expect("No current high threshold transcript available")
         .dkg_id
         .clone();
 

rs/consensus/utils/src/lib.rs

@@ -263,7 +263,10 @@ pub fn active_low_threshold_nidkg_id(
 ) -> Option<NiDkgId> {
     get_active_data_at(reader, height, |block, height| {
         get_transcript_data_at_given_summary(block, height, NiDkgTag::LowThreshold, |transcript| {
-            transcript.dkg_id.clone()
+            transcript
+                .expect("No active low threshold transcript available for tag {:?}")
+                .dkg_id
+                .clone()
         })
     })
 }
@@ -275,7 +278,10 @@ pub fn active_high_threshold_nidkg_id(
 ) -> Option<NiDkgId> {
     get_active_data_at(reader, height, |block, height| {
         get_transcript_data_at_given_summary(block, height, NiDkgTag::HighThreshold, |transcript| {
-            transcript.dkg_id.clone()
+            transcript
+                .expect("No active high threshold transcript available for tag {:?}")
+                .dkg_id
+                .clone()
         })
     })
 }
@@ -287,6 +293,7 @@ pub fn active_low_threshold_committee(
 ) -> Option<(Threshold, NiDkgReceivers)> {
     get_active_data_at(reader, height, |block, height| {
         get_transcript_data_at_given_summary(block, height, NiDkgTag::LowThreshold, |transcript| {
+            let transcript = transcript.expect("No active low threshold transcript available");
             (
                 transcript.threshold.get().get() as usize,
                 transcript.committee.clone(),
@@ -302,6 +309,7 @@ pub fn active_high_threshold_committee(
 ) -> Option<(Threshold, NiDkgReceivers)> {
     get_active_data_at(reader, height, |block, height| {
         get_transcript_data_at_given_summary(block, height, NiDkgTag::HighThreshold, |transcript| {
+            let transcript = transcript.expect("No active high threshold transcript available");
             (
                 transcript.threshold.get().get() as usize,
                 transcript.committee.clone(),
@@ -357,15 +365,15 @@ fn get_transcript_data_at_given_summary<T>(
     summary_block: &Block,
     height: Height,
     tag: NiDkgTag,
-    getter: impl Fn(&NiDkgTranscript) -> T,
+    getter: impl Fn(Option<&NiDkgTranscript>) -> T,
 ) -> Option<T> {
     let dkg_summary = &summary_block.payload.as_ref().as_summary().dkg;
     if dkg_summary.current_interval_includes(height) {
         Some(getter(dkg_summary.current_transcript(&tag)))
     } else if dkg_summary.next_interval_includes(height) {
         let transcript = dkg_summary
             .next_transcript(&tag)
-            .unwrap_or_else(|| dkg_summary.current_transcript(&tag));
+            .or(dkg_summary.current_transcript(&tag));
         Some(getter(transcript))
     } else {
         None

rs/replay/src/lib.rs

@@ -240,10 +240,12 @@ fn cmd_get_recovery_cup(
     let low_threshold_transcript = summary
         .dkg
         .current_transcript(&NiDkgTag::LowThreshold)
+        .expect("No current low threshold transcript available")
         .clone();
     let high_threshold_transcript = summary
         .dkg
         .current_transcript(&NiDkgTag::HighThreshold)
+        .expect("No current high threshold transcript available")
         .clone();
     let initial_ni_dkg_transcript_low_threshold =
         Some(InitialNiDkgTranscriptRecord::from(low_threshold_transcript));

rs/test_utilities/consensus/src/lib.rs

@@ -140,10 +140,12 @@ pub fn make_genesis(summary: dkg::Summary) -> CatchUpPackage {
     let height = summary.height;
     let low_dkg_id = summary
         .current_transcript(&NiDkgTag::LowThreshold)
+        .unwrap()
         .dkg_id
         .clone();
     let high_dkg_id = summary
         .current_transcript(&NiDkgTag::HighThreshold)
+        .unwrap()
         .dkg_id
         .clone();
     let block = Block::new(

rs/types/types/src/consensus/dkg.rs

@@ -209,10 +209,8 @@ impl Summary {
     /// Returns a reference to the current transcript for the given tag. Note
     /// that currently we expect that a valid summary contains the current
     /// transcript for any DKG tag.
-    pub fn current_transcript(&self, tag: &NiDkgTag) -> &NiDkgTranscript {
-        self.current_transcripts
-            .get(tag)
-            .unwrap_or_else(|| panic!("No current transcript available for tag {:?}", tag))
+    pub fn current_transcript(&self, tag: &NiDkgTag) -> Option<&NiDkgTranscript> {
+        self.current_transcripts.get(tag)
     }
 
     /// Returns a reference to the current transcripts.
@@ -241,20 +239,6 @@ impl Summary {
             .collect()
     }
 
-    /// Return the set of next transcripts for all tags. If for some tag
-    /// the next transcript is not available, the current transcript is used.
-    /// This function avoids expensive copying when transcripts are large.
-    pub fn into_next_transcripts(self) -> BTreeMap<NiDkgTag, NiDkgTranscript> {
-        let mut next_transcripts = self.next_transcripts;
-        self.current_transcripts
-            .into_iter()
-            .map(|(tag, current)| {
-                let new_next_transcripts = next_transcripts.remove(&tag).unwrap_or(current);
-                (tag, new_next_transcripts)
-            })
-            .collect()
-    }
-
     /// Returns `true` if the provided height is included in the DKG interval
     /// corresponding to the current summary. Note that the summary block is
     /// considered to be part of the interval. For example, if the start height