Add benchmarks for AVSS + optimise complete_dkg (#894)

* Add avss benchmarks

* Benchmark complete_dkg + optimise

* clippy

Author: jonas-lj

fastcrypto-tbls/Cargo.toml

@@ -46,6 +46,10 @@ required-features = ["experimental"]
 name = "tbls"
 harness = false
 
+[[bench]]
+name = "avss"
+harness = false
+
 [features]
 default = []
 experimental = []

fastcrypto-tbls/benches/avss.rs

@@ -0,0 +1,178 @@
+// Copyright (c) 2022, Mysten Labs, Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+use criterion::{criterion_group, criterion_main, BenchmarkGroup, Criterion};
+use fastcrypto::groups::ristretto255;
+use fastcrypto_tbls::ecies_v1;
+use fastcrypto_tbls::nodes::{Node, Nodes, PartyId};
+use fastcrypto_tbls::threshold_schnorr::avss;
+use itertools::iproduct;
+use rand::thread_rng;
+
+type EG = ristretto255::RistrettoPoint;
+
+fn generate_ecies_keys(
+    n: u16,
+) -> Vec<(PartyId, ecies_v1::PrivateKey<EG>, ecies_v1::PublicKey<EG>)> {
+    (0..n)
+        .map(|id| {
+            let sk = ecies_v1::PrivateKey::<EG>::new(&mut thread_rng());
+            let pk = ecies_v1::PublicKey::<EG>::from_private_key(&sk);
+            (id, sk, pk)
+        })
+        .collect()
+}
+
+pub fn setup_receiver(
+    id: PartyId,
+    threshold: u16,
+    weight: u16, // Per node
+    keys: &[(PartyId, ecies_v1::PrivateKey<EG>, ecies_v1::PublicKey<EG>)],
+) -> avss::Receiver {
+    let nodes = keys
+        .iter()
+        .map(|(id, _sk, pk)| Node::<EG> {
+            id: *id,
+            pk: pk.clone(),
+            weight,
+        })
+        .collect();
+    avss::Receiver::new(
+        Nodes::new(nodes).unwrap(),
+        id,
+        threshold,
+        b"avss".to_vec(),
+        None,
+        keys.get(id as usize).unwrap().1.clone(),
+    )
+}
+
+pub fn setup_dealer(
+    threshold: u16,
+    f: u16,
+    weight: u16, // Per node
+    keys: &[(PartyId, ecies_v1::PrivateKey<EG>, ecies_v1::PublicKey<EG>)],
+) -> avss::Dealer {
+    let nodes = keys
+        .iter()
+        .map(|(id, _sk, pk)| Node::<EG> {
+            id: *id,
+            pk: pk.clone(),
+            weight,
+        })
+        .collect();
+    avss::Dealer::new(
+        None,
+        Nodes::new(nodes).unwrap(),
+        threshold,
+        f,
+        b"avss".to_vec(),
+    )
+    .unwrap()
+}
+
+mod avss_benches {
+    use super::*;
+    use fastcrypto_tbls::threshold_schnorr::avss::ProcessedMessage::Valid;
+    use fastcrypto_tbls::threshold_schnorr::avss::{PartialOutput, ReceiverOutput};
+    use itertools::Itertools;
+    use std::collections::HashMap;
+
+    fn dkg(c: &mut Criterion) {
+        const SIZES: [u16; 1] = [100];
+        const TOTAL_WEIGHTS: [u16; 4] = [500, 1000, 2000, 2500];
+
+        {
+            let mut create: BenchmarkGroup<_> = c.benchmark_group("AVSS create_message");
+            for (n, total_w) in iproduct!(SIZES.iter(), TOTAL_WEIGHTS.iter()) {
+                let w = total_w / n;
+                let total_w = w * n;
+                let t = total_w / 3 - 1;
+                let keys = generate_ecies_keys(*n);
+                let d0 = setup_dealer(t, t - 1, w, &keys);
+                create.bench_function(
+                    format!("n={}, total_weight={}, t={}, w={}", n, total_w, t, w).as_str(),
+                    |b| b.iter(|| d0.create_message(&mut thread_rng())),
+                );
+            }
+        }
+
+        {
+            let mut verify: BenchmarkGroup<_> = c.benchmark_group("AVSS process_message");
+            for (n, total_w) in iproduct!(SIZES.iter(), TOTAL_WEIGHTS.iter()) {
+                let w = total_w / n;
+                let total_w = w * n;
+                let t = total_w / 3 - 1;
+                let keys = generate_ecies_keys(*n);
+                let d0 = setup_dealer(t, t - 1, w, &keys);
+                let r1 = setup_receiver(1, t, w, &keys);
+                let message = d0.create_message(&mut thread_rng()).unwrap();
+
+                verify.bench_function(
+                    format!("n={}, total_weight={}, t={}, w={}", n, total_w, t, w).as_str(),
+                    |b| b.iter(|| r1.process_message(&message).unwrap()),
+                );
+            }
+        }
+
+        {
+            let mut verify: BenchmarkGroup<_> = c.benchmark_group("AVSS complete_*");
+            for (n, total_w) in iproduct!(SIZES.iter(), TOTAL_WEIGHTS.iter()) {
+                let w = total_w / n;
+                let total_w = w * n;
+                let t = total_w / 3 - 1;
+                let keys = generate_ecies_keys(*n);
+                let nodes = Nodes::new(
+                    keys.iter()
+                        .map(|(id, _sk, pk)| Node::<EG> {
+                            id: *id,
+                            pk: pk.clone(),
+                            weight: w,
+                        })
+                        .collect(),
+                )
+                .unwrap();
+                let dealers = (0..*n)
+                    .map(|_| setup_dealer(t, t - 1, w, &keys))
+                    .collect_vec();
+                let r1 = setup_receiver(1, t, w, &keys);
+                let messages: HashMap<PartyId, avss::Message> = dealers
+                    .iter()
+                    .enumerate()
+                    .map(|(i, d)| {
+                        (
+                            PartyId::from(i as u16),
+                            d.create_message(&mut thread_rng()).unwrap(),
+                        )
+                    })
+                    .collect();
+
+                let outputs: HashMap<PartyId, PartialOutput> = messages
+                    .iter()
+                    .map(|(i, m)| {
+                        let output = r1.process_message(m).unwrap();
+                        if let Valid(o) = output {
+                            return (*i, o);
+                        }
+                        panic!()
+                    })
+                    .collect();
+
+                verify.bench_function(
+                    format!("DKG n={}, total_weight={}, t={}, w={}", n, total_w, t, w).as_str(),
+                    |b| {
+                        b.iter(|| ReceiverOutput::complete_dkg(t, &nodes, outputs.clone()).unwrap())
+                    },
+                );
+            }
+        }
+    }
+
+    criterion_group! {
+        name = avss_benches;
+        config = Criterion::default().sample_size(10);
+        targets = dkg,
+    }
+}
+
+criterion_main!(avss_benches::avss_benches);

fastcrypto-tbls/src/threshold_schnorr/avss.rs

@@ -28,6 +28,7 @@ use fastcrypto::traits::AllowedRng;
 use itertools::Itertools;
 use serde::{Deserialize, Serialize};
 use std::collections::HashMap;
+use std::ops::Add;
 
 /// This represents a Dealer in the AVSS. There is exactly one dealer, who creates the shares and broadcasts the encrypted shares.
 #[allow(dead_code)]
@@ -402,34 +403,15 @@ impl ReceiverOutput {
         let outputs = outputs.into_values().collect_vec();
 
         // Sanity check: Outputs cannot be empty and all outputs must have the same weight.
-        if outputs.is_empty() || !outputs.iter().map(|output| output.weight()).all_equal() {
+        if !outputs.iter().map(|output| output.weight()).all_equal() {
             return Err(InvalidInput);
         }
 
-        Ok(outputs
+        outputs
             .into_iter()
-            .map(|output| output.into_receiver_output(nodes))
-            .reduce(|acc, output| {
-                let shares = acc
-                    .my_shares
-                    .shares
-                    .iter()
-                    .zip_eq(&output.my_shares.shares)
-                    .map(types::sum)
-                    .collect_vec();
-                let commitments = acc
-                    .commitments
-                    .iter()
-                    .zip_eq(&output.commitments)
-                    .map(types::sum)
-                    .collect_vec();
-                ReceiverOutput {
-                    my_shares: SharesForNode { shares },
-                    commitments,
-                    vk: acc.vk + output.vk,
-                }
-            })
-            .expect("Should not be empty"))
+            .reduce(|acc, output| acc + output)
+            .ok_or(InvalidInput)
+            .map(|o| o.into_receiver_output(nodes))
     }
 
     /// Interpolate shares from multiple outputs to create new shares for the given indices.
@@ -537,6 +519,24 @@ impl PartialOutput {
     }
 }
 
+impl Add<Self> for PartialOutput {
+    type Output = Self;
+
+    fn add(self, rhs: Self) -> Self::Output {
+        let shares = self
+            .my_shares
+            .shares
+            .iter()
+            .zip_eq(&rhs.my_shares.shares)
+            .map(types::sum)
+            .collect_vec();
+        Self {
+            my_shares: SharesForNode { shares },
+            feldman_commitment: self.feldman_commitment + &rhs.feldman_commitment,
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use crate::ecies_v1;