feat(rust/catalyst-voting): Vote protocol benchmarks (#65)

Mr-Leshiy · minikin · web-flow · commit 2f9277fa21ed · 2024-10-29T09:50:48.000Z
* add vote protocol benchmarks

* add missing benchmarks

* refactor

* cleanup

* wip

---------

Co-authored-by: Oleksandr Prokhorenko &lt;djminikin@gmail.com&gt;
diff --git a/rust/.cargo/config.toml b/rust/.cargo/config.toml
@@ -77,8 +77,8 @@ lint-vscode = "clippy --message-format=json-diagnostic-rendered-ansi --all-targe
 
 docs = "doc --release --no-deps --document-private-items --bins --lib --examples"
 # nightly docs build broken... when they are'nt we can enable these docs... --unit-graph --timings=html,json -Z unstable-options"
-testunit = "nextest run --release --bins --lib --tests --benches --no-fail-fast -P ci"
-testcov = "llvm-cov nextest --release --bins --lib --tests --benches --no-fail-fast -P ci"
+testunit = "nextest run --release --bins --lib --tests --no-fail-fast -P ci"
+testcov = "llvm-cov nextest --release --bins --lib --tests --no-fail-fast -P ci"
 testdocs = "test --doc --release"
 
 # Rust formatting, MUST be run with +nightly
diff --git a/rust/Earthfile b/rust/Earthfile
@@ -1,6 +1,6 @@
 VERSION 0.8
 
-IMPORT github.com/input-output-hk/catalyst-ci/earthly/rust:v3.2.15 AS rust-ci
+IMPORT github.com/input-output-hk/catalyst-ci/earthly/rust:v3.2.19 AS rust-ci
 
 COPY_SRC:
     FUNCTION
diff --git a/rust/catalyst-voting/Cargo.toml b/rust/catalyst-voting/Cargo.toml
@@ -12,6 +12,10 @@ license.workspace = true
 [lints]
 workspace = true
 
+[[bench]]
+name = "vote_protocol"
+harness = false
+
 [dependencies]
 anyhow = "1.0.89"
 rand_core = { version = "0.6.4", features = ["getrandom"] }
@@ -21,7 +25,8 @@ ed25519-dalek = { version = "2.1.1", features = ["rand_core"] }
 blake2b_simd = "1.0.2"
 
 [dev-dependencies]
-proptest = {version = "1.5.0" }
+criterion = "0.5.1"
+proptest = { version = "1.5.0" }
 # Potentially it could be replaced with using `proptest::property_test` attribute macro,
 # after this PR will be merged https://github.com/proptest-rs/proptest/pull/523
 test-strategy = "0.4.0"
diff --git a/rust/catalyst-voting/benches/vote_protocol.rs b/rust/catalyst-voting/benches/vote_protocol.rs
@@ -0,0 +1,219 @@
+//! `catalyst_voting::vote_protocol` benchmark
+//!
+//! To run these benchmarks use
+//! ```shell
+//! SAMPLE_SIZE=<sample size> VOTERS_NUMBER=<voters number> cargo bench -p catalyst-voting vote_protocol
+//! ```
+#![allow(
+    missing_docs,
+    clippy::missing_docs_in_private_items,
+    clippy::unwrap_used,
+    clippy::similar_names
+)]
+
+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};
+use proptest::{
+    prelude::{any_with, Strategy},
+    sample::size_range,
+    strategy::ValueTree,
+    test_runner::TestRunner,
+};
+use test_strategy::Arbitrary;
+
+const VOTERS_NUMBER_ENV: &str = "VOTERS_NUMBER";
+const SAMPLE_SIZE_ENV: &str = "SAMPLE_SIZE";
+const DEFAULT_SAMPLE_SIZE: usize = 10;
+const DEFAULT_VOTERS_NUMBER: usize = 1;
+
+const VOTING_OPTIONS: usize = 3;
+
+#[derive(Arbitrary, Debug)]
+struct Voter {
+    voting_power: u32,
+    #[strategy(0..VOTING_OPTIONS)]
+    choice: usize,
+}
+
+struct Choices(Vec<usize>);
+struct VotingPowers(Vec<u64>);
+
+fn rand_generate_vote_data(
+    voters_number: usize,
+) -> (
+    Choices,
+    VotingPowers,
+    ElectionSecretKey,
+    ElectionPublicKey,
+    VoterProofCommitment,
+) {
+    let mut runner = TestRunner::default();
+
+    let (choices, voting_powers) = any_with::<Vec<Voter>>((size_range(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();
+
+    let election_secret_key = ElectionSecretKey::random_with_default_rng();
+    let voter_proof_commitment = VoterProofCommitment::random_with_default_rng();
+    let election_public_key = election_secret_key.public_key();
+
+    (
+        Choices(choices),
+        VotingPowers(voting_powers),
+        election_secret_key,
+        election_public_key,
+        voter_proof_commitment,
+    )
+}
+
+#[allow(clippy::too_many_lines)]
+fn vote_protocol_benches(c: &mut Criterion) {
+    let sample_size = std::env::var(SAMPLE_SIZE_ENV)
+        .map(|s| s.parse().unwrap())
+        .unwrap_or(DEFAULT_SAMPLE_SIZE);
+    let voters_number = std::env::var(VOTERS_NUMBER_ENV)
+        .map(|s| s.parse().unwrap())
+        .unwrap_or(DEFAULT_VOTERS_NUMBER);
+
+    let mut group = c.benchmark_group("vote protocol benchmark");
+    group.sample_size(sample_size);
+
+    let (choices, voting_powers, election_secret_key, election_public_key, voter_proof_commitment) =
+        rand_generate_vote_data(voters_number);
+
+    let votes: Vec<_> = choices
+        .0
+        .iter()
+        .map(|choice| Vote::new(*choice, VOTING_OPTIONS).unwrap())
+        .collect();
+    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(vote, &election_public_key, &mut rng))
+                .unzip();
+        });
+    });
+
+    let mut voter_proofs = Vec::new();
+    group.bench_function("voter proof generation", |b| {
+        b.iter(|| {
+            voter_proofs = votes
+                .iter()
+                .zip(encrypted_votes.iter())
+                .zip(randomness.iter())
+                .map(|((v, enc_v), r)| {
+                    generate_voter_proof(
+                        v,
+                        enc_v.clone(),
+                        r.clone(),
+                        &election_public_key,
+                        &voter_proof_commitment,
+                        &mut rng,
+                    )
+                    .unwrap()
+                })
+                .collect();
+        });
+    });
+
+    group.bench_function("voter proof verification", |b| {
+        b.iter(|| {
+            let is_ok = voter_proofs
+                .iter()
+                .zip(encrypted_votes.iter())
+                .all(|(p, enc_v)| {
+                    verify_voter_proof(
+                        enc_v.clone(),
+                        &election_public_key,
+                        &voter_proof_commitment,
+                        p,
+                    )
+                });
+            assert!(is_ok);
+        });
+    });
+
+    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.0).unwrap()
+                })
+                .collect();
+        });
+    });
+
+    let total_voting_power = voting_powers.0.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();
+}
+
+criterion_group!(benches, vote_protocol_benches);
+
+criterion_main!(benches);
diff --git a/rust/catalyst-voting/src/crypto/zk_unit_vector/decoding.rs b/rust/catalyst-voting/src/crypto/zk_unit_vector/decoding.rs
@@ -27,22 +27,34 @@ impl UnitVectorProof {
         let ann = (0..len)
             .map(|i| {
                 let bytes = read_array(reader)?;
-                Announcement::from_bytes(&bytes)
-                    .map_err(|e| anyhow!("Cannot decode announcement at {i}, error: {e}."))
+                Announcement::from_bytes(&bytes).map_err(|e| {
+                    anyhow!(
+                        "Cannot decode announcement at {i}, \
+                        error: {e}."
+                    )
+                })
             })
             .collect::<anyhow::Result<_>>()?;
         let dl = (0..len)
             .map(|i| {
                 let bytes = read_array(reader)?;
-                Ciphertext::from_bytes(&bytes)
-                    .map_err(|e| anyhow!("Cannot decode ciphertext at {i}, error: {e}."))
+                Ciphertext::from_bytes(&bytes).map_err(|e| {
+                    anyhow!(
+                        "Cannot decode ciphertext at {i}, \
+                        error: {e}."
+                    )
+                })
             })
             .collect::<anyhow::Result<_>>()?;
         let rr = (0..len)
             .map(|i| {
                 let bytes = read_array(reader)?;
-                ResponseRandomness::from_bytes(&bytes)
-                    .map_err(|e| anyhow!("Cannot decode response randomness at {i}, error: {e}."))
+                ResponseRandomness::from_bytes(&bytes).map_err(|e| {
+                    anyhow!(
+                        "Cannot decode response randomness at {i}, \
+                        error: {e}."
+                    )
+                })
             })
             .collect::<anyhow::Result<_>>()?;
 
diff --git a/rust/catalyst-voting/src/txs/v1/decoding.rs b/rust/catalyst-voting/src/txs/v1/decoding.rs
@@ -112,14 +112,16 @@ impl Tx {
         let padding_tag = read_be_u8(reader).map_err(|_| anyhow!("Missing padding tag field."))?;
         ensure!(
             padding_tag == PADDING_TAG,
-            "Invalid padding tag field value, must be equals to {PADDING_TAG}, provided: {padding_tag}.",
+            "Invalid padding tag field value, must be equals to {PADDING_TAG}, \
+            provided: {padding_tag}.",
         );
 
         let fragment_tag =
             read_be_u8(reader).map_err(|_| anyhow!("Missing fragment tag field."))?;
         ensure!(
             fragment_tag == FRAGMENT_TAG,
-            "Invalid fragment tag field value, must be equals to {FRAGMENT_TAG}, provided: {fragment_tag}.",
+            "Invalid fragment tag field value, must be equals to {FRAGMENT_TAG}, \
+            provided: {fragment_tag}.",
         );
 
         let vote_plan_id =
@@ -148,7 +150,8 @@ impl Tx {
             },
             tag => {
                 bail!(
-                    "Invalid vote tag value, must be equals to {PUBLIC_VOTE_TAG} or {PRIVATE_VOTE_TAG}, provided: {tag}"
+                    "Invalid vote tag value, \
+                    must be equals to {PUBLIC_VOTE_TAG} or {PRIVATE_VOTE_TAG}, provided: {tag}"
                 )
             },
         };
@@ -160,20 +163,23 @@ impl Tx {
             read_be_u8(reader).map_err(|_| anyhow!("Missing inputs amount field."))?;
         ensure!(
             inputs_amount == NUMBER_OF_INPUTS,
-            "Invalid number of inputs, expected: {NUMBER_OF_INPUTS}, provided: {inputs_amount}",
+            "Invalid number of inputs, expected: {NUMBER_OF_INPUTS}, \
+            provided: {inputs_amount}",
         );
 
         let outputs_amount =
             read_be_u8(reader).map_err(|_| anyhow!("Missing outputs amount field."))?;
         ensure!(
             outputs_amount == NUMBER_OF_OUTPUTS,
-            "Invalid number of outputs, expected: {NUMBER_OF_OUTPUTS}, provided: {outputs_amount}",
+            "Invalid number of outputs, expected: {NUMBER_OF_OUTPUTS}, \
+            provided: {outputs_amount}",
         );
 
         let input_tag = read_be_u8(reader).map_err(|_| anyhow!("Missing input tag field."))?;
         ensure!(
             input_tag == INPUT_TAG,
-            "Invalid input tag, expected: {INPUT_TAG}, provided: {input_tag}",
+            "Invalid input tag, expected: {INPUT_TAG}, \
+            provided: {input_tag}",
         );
 
         // skip value
@@ -187,7 +193,8 @@ impl Tx {
         let witness_tag = read_be_u8(reader).map_err(|_| anyhow!("Missing witness tag field."))?;
         ensure!(
             witness_tag == WITNESS_TAG,
-            "Invalid witness tag, expected: {WITNESS_TAG}, provided: {witness_tag}",
+            "Invalid witness tag, expected: {WITNESS_TAG}, \
+            provided: {witness_tag}",
         );
 
         // Skip nonce field
diff --git a/rust/catalyst-voting/src/vote_protocol/tally/mod.rs b/rust/catalyst-voting/src/vote_protocol/tally/mod.rs
diff --git a/rust/catalyst-voting/tests/voting_test.rs b/rust/catalyst-voting/tests/voting_test.rs
diff --git a/rust/deny.toml b/rust/deny.toml
