add missing benchmarks

Author: Mr-Leshiy

rust/catalyst-voting/benches/vote_protocol.rs

@@ -6,12 +6,20 @@
     clippy::similar_names
 )]
 
-use catalyst_voting::vote_protocol::{
-    committee::{ElectionPublicKey, ElectionSecretKey},
-    voter::{
-        encrypt_vote_with_default_rng,
-        proof::{generate_voter_proof_with_default_rng, verify_voter_proof, VoterProofCommitment},
-        Vote,
+use catalyst_voting::{
+    crypto::default_rng,
+    vote_protocol::{
+        committee::{ElectionPublicKey, ElectionSecretKey},
+        tally::{
+            decrypt_tally,
+            proof::{generate_tally_proof, verify_tally_proof},
+            tally, DecryptionTallySetup,
+        },
+        voter::{
+            encrypt_vote,
+            proof::{generate_voter_proof, verify_voter_proof, VoterProofCommitment},
+            Vote,
+        },
     },
 };
 use criterion::{criterion_group, criterion_main, Criterion};
@@ -27,20 +35,27 @@ const VOTERS_NUMBER: usize = 100;
 
 #[derive(Arbitrary, Debug)]
 struct Voter {
-    _voting_power: u32,
+    voting_power: u32,
     #[strategy(0..VOTING_OPTIONS)]
     choice: usize,
 }
 
 fn initial_setup() -> (
-    [Voter; VOTERS_NUMBER],
+    Vec<usize>,
+    Vec<u64>,
     ElectionSecretKey,
     ElectionPublicKey,
     VoterProofCommitment,
 ) {
     let mut runner = TestRunner::default();
 
-    let voters = any::<[Voter; VOTERS_NUMBER]>()
+    let (choices, voting_powers) = any::<[Voter; VOTERS_NUMBER]>()
+        .prop_map(|voter| {
+            (
+                voter.iter().map(|v| v.choice).collect(),
+                voter.iter().map(|v| v.voting_power.into()).collect(),
+            )
+        })
         .new_tree(&mut runner)
         .unwrap()
         .current();
@@ -50,31 +65,34 @@ fn initial_setup() -> (
     let election_public_key = election_secret_key.public_key();
 
     (
-        voters,
+        choices,
+        voting_powers,
         election_secret_key,
         election_public_key,
         voter_proof_commitment,
     )
 }
 
+#[allow(clippy::too_many_lines)]
 fn vote_protocol_benches(c: &mut Criterion) {
-    let (voters, _election_secret_key, election_public_key, voter_proof_commitment) =
-        initial_setup();
+    let mut group = c.benchmark_group("vote protocol benchmark");
+    group.sample_size(10);
 
-    let votes: Vec<_> = voters
+    let (choices, voting_powers, election_secret_key, election_public_key, voter_proof_commitment) =
+        initial_setup();
+    let votes: Vec<_> = choices
         .iter()
-        .map(|voter| Vote::new(voter.choice, VOTING_OPTIONS).unwrap())
+        .map(|choice| Vote::new(*choice, VOTING_OPTIONS).unwrap())
         .collect();
-
-    let mut group = c.benchmark_group("vote protocol benchmark");
+    let mut rng = default_rng();
 
     let mut encrypted_votes = Vec::new();
     let mut randomness = Vec::new();
     group.bench_function("vote encryption", |b| {
         b.iter(|| {
             (encrypted_votes, randomness) = votes
                 .iter()
-                .map(|vote| encrypt_vote_with_default_rng(vote, &election_public_key))
+                .map(|vote| encrypt_vote(vote, &election_public_key, &mut rng))
                 .unzip();
         });
     });
@@ -87,12 +105,13 @@ fn vote_protocol_benches(c: &mut Criterion) {
                 .zip(encrypted_votes.iter())
                 .zip(randomness.iter())
                 .map(|((v, enc_v), r)| {
-                    generate_voter_proof_with_default_rng(
+                    generate_voter_proof(
                         v,
                         enc_v.clone(),
                         r.clone(),
                         &election_public_key,
                         &voter_proof_commitment,
+                        &mut rng,
                     )
                     .unwrap()
                 })
@@ -117,6 +136,57 @@ fn vote_protocol_benches(c: &mut Criterion) {
         });
     });
 
+    let mut encrypted_tallies = Vec::new();
+    group.bench_function("tally", |b| {
+        b.iter(|| {
+            encrypted_tallies = (0..VOTING_OPTIONS)
+                .map(|voting_option| {
+                    tally(voting_option, &encrypted_votes, &voting_powers).unwrap()
+                })
+                .collect();
+        });
+    });
+
+    let total_voting_power = voting_powers.iter().sum();
+    let mut decryption_tally_setup = None;
+    group.bench_function("decryption tally setup initialization", |b| {
+        b.iter(|| {
+            decryption_tally_setup = Some(DecryptionTallySetup::new(total_voting_power).unwrap());
+        });
+    });
+    let decryption_tally_setup = decryption_tally_setup.unwrap();
+
+    let mut decrypted_tallies = Vec::new();
+    group.bench_function("decrypt tally", |b| {
+        b.iter(|| {
+            decrypted_tallies = encrypted_tallies
+                .iter()
+                .map(|t| decrypt_tally(t, &election_secret_key, &decryption_tally_setup).unwrap())
+                .collect();
+        });
+    });
+
+    let mut tally_proofs = Vec::new();
+    group.bench_function("tally proof generation", |b| {
+        b.iter(|| {
+            tally_proofs = encrypted_tallies
+                .iter()
+                .map(|t| generate_tally_proof(t, &election_secret_key, &mut rng))
+                .collect();
+        });
+    });
+
+    group.bench_function("tally proof verification", |b| {
+        b.iter(|| {
+            let is_ok = tally_proofs
+                .iter()
+                .zip(encrypted_tallies.iter())
+                .zip(decrypted_tallies.iter())
+                .all(|((p, enc_t), t)| verify_tally_proof(enc_t, *t, &election_public_key, p));
+            assert!(is_ok);
+        });
+    });
+
     group.finish();
 }