From 5f0a015fd680c3ef1f49b10423492c361cf06ef1 Mon Sep 17 00:00:00 2001 From: Dam Date: Fri, 10 Oct 2025 11:55:37 +0200 Subject: [PATCH 1/3] Initial commit for reworked select k indices function. --- .../src/aggregate_signature/basic_verifier.rs | 116 ++++++++++++++++++ 1 file changed, 116 insertions(+) diff --git a/mithril-stm/src/aggregate_signature/basic_verifier.rs b/mithril-stm/src/aggregate_signature/basic_verifier.rs index ab59658650f..d443846de57 100644 --- a/mithril-stm/src/aggregate_signature/basic_verifier.rs +++ b/mithril-stm/src/aggregate_signature/basic_verifier.rs @@ -172,6 +172,122 @@ impl BasicVerifier { Err(AggregationError::NotEnoughSignatures(count, params.k)) } + /// Modification of the function select_valid_signatures_for_k_indices to try to improve readability. + /// This function follows the same structure of iterating first over the indices inside all the signatures to select k indices, + /// then over the selected signatures to modify the list of indices to fit the selected ones. + pub fn reworked_select_valid_signatures_for_k_indices( + total_stake: &Stake, + params: &Parameters, + msg: &[u8], + sigs: &[SingleSignatureWithRegisteredParty], + ) -> Result, AggregationError> { + let (idx_by_mtidx, btm) = Self::get_k_indices(total_stake, params, msg, sigs); + + Self::valid_signatures_from_k_indices(params, idx_by_mtidx, btm) + } + + /// Iterates over the indices list in each signature to select the tuple (index, signature) with the smallest corresponding sigma. + /// Uses a BTreeMap to keep track of the correspondance between merkle tree index and signature + /// and a HashMap to connect the selected indices to the merkle tree index of the signature. + /// TODO: Merge those two as the MT index is already in the signature + pub fn get_k_indices<'a>( + total_stake: &Stake, + params: &Parameters, + msg: &[u8], + sigs: &'a [SingleSignatureWithRegisteredParty], + ) -> ( + HashMap, + BTreeMap, + ) { + let mut sig_by_mt_index: BTreeMap = + BTreeMap::new(); + let mut indices_by_mt_idx: HashMap = HashMap::new(); + + for sig_reg in sigs.iter() { + if sig_reg + .sig + .basic_verify( + params, + &sig_reg.reg_party.0, + &sig_reg.reg_party.1, + msg, + total_stake, + ) + .is_err() + { + continue; + } + for index in sig_reg.sig.indexes.iter() { + if let Some(mt_idx) = indices_by_mt_idx.get(index) { + if let Some(prev_sig) = sig_by_mt_index.get(mt_idx) { + if prev_sig.sig.sigma < sig_reg.sig.sigma { + continue; + } else { + indices_by_mt_idx.insert(*index, sig_reg.sig.signer_index); + sig_by_mt_index.insert(sig_reg.sig.signer_index, sig_reg); + } + } + } else { + // Should we test for k indices here? + indices_by_mt_idx.insert(*index, sig_reg.sig.signer_index); + sig_by_mt_index.insert(sig_reg.sig.signer_index, sig_reg); + } + } + } + (indices_by_mt_idx, sig_by_mt_index) + } + + /// Takes a list of tuples (selected index, MT index) and outputs a list of signatures with a total of k valid indices + /// First iterates over the tuples to create a vector of selected indices for each signature. + /// Then modify and aggregates the signatures starting with the signature with the most valid indices. + /// This has the effect of minimizing the number of different signature used. + pub fn valid_signatures_from_k_indices( + params: &Parameters, + list_k_valid_indices: HashMap, + sig_by_mt_index: BTreeMap, + ) -> Result, AggregationError> { + let mut valid_idx_for_mt_idx: HashMap> = HashMap::new(); + + // Uses the HashMap of tuples (index, MT index) to create a vector of indices connected to one signature + for (&valid_idx, &mt_idx) in list_k_valid_indices.iter() { + if let Some(val) = valid_idx_for_mt_idx.get_mut(&mt_idx) { + val.push(valid_idx); + } else { + valid_idx_for_mt_idx.insert(mt_idx, vec![valid_idx]); + } + } + + // sort the HashMap according to the number of selected indices of each signature + let mut vec_valid_idx_per_mt_idx = valid_idx_for_mt_idx.into_iter().collect::>(); + vec_valid_idx_per_mt_idx.sort_by(|(_, v1), (_, v2)| v2.len().cmp(&v1.len())); + + let mut uniques_sig: Vec = Vec::new(); + let mut count_idx = 0; + // Modify and aggregate the signatures + for (mt_idx, indices) in vec_valid_idx_per_mt_idx.into_iter() { + let mut single_sig = if let Some(sig) = sig_by_mt_index.get(&mt_idx) { + (*sig).clone() + } else { + // Change the error + return Err(AggregationError::NotEnoughSignatures(0, params.k)); + }; + if indices.len() > (params.k - count_idx) as usize { + single_sig.sig.indexes = indices[0..(params.k - count_idx) as usize].to_vec(); + count_idx += params.k - count_idx; + } else { + count_idx += indices.len() as u64; + single_sig.sig.indexes = indices; + } + uniques_sig.push(single_sig); + + if count_idx >= params.k { + return Ok(uniques_sig); + } + } + + Err(AggregationError::NotEnoughSignatures(count_idx, params.k)) + } + /// Given a slice of `sig_reg_list`, this function returns a new list of `sig_reg_list` with only valid indices. /// In case of conflict (having several signatures for the same index) /// it selects the smallest signature (i.e. takes the signature with the smallest scalar). From e0d81aa13339bb1e0eff5f1b6d96a1fcdd03e6ce Mon Sep 17 00:00:00 2001 From: Dam Date: Tue, 4 Nov 2025 09:52:01 +0100 Subject: [PATCH 2/3] Added small test function for comparison. --- .../test_extensions/basic_verifier_tests.rs | 150 ++++++++++++++++++ mithril-stm/tests/test_extensions/mod.rs | 1 + 2 files changed, 151 insertions(+) create mode 100644 mithril-stm/tests/test_extensions/basic_verifier_tests.rs diff --git a/mithril-stm/tests/test_extensions/basic_verifier_tests.rs b/mithril-stm/tests/test_extensions/basic_verifier_tests.rs new file mode 100644 index 00000000000..254db8bdb57 --- /dev/null +++ b/mithril-stm/tests/test_extensions/basic_verifier_tests.rs @@ -0,0 +1,150 @@ +use std::{collections::BTreeMap, process::id, time::Instant}; + +use blake2::{Blake2b, digest::consts::U32}; +use mithril_stm::{ + AggregateSignature, AggregateVerificationKey, AggregationError, BasicVerifier, Clerk, + Initializer, KeyRegistration, Parameters, Signer, SingleSignature, + SingleSignatureWithRegisteredParty, Stake, VerificationKey, +}; +use rand_chacha::ChaCha20Rng; +use rand_core::{RngCore, SeedableRng}; +use rayon::{prelude::*, vec}; + +type H = Blake2b; + +#[test] +pub fn select_sig_test() { + let nparties = 3000; // Use a small number of parties for this example + type D = Blake2b; // Setting the hash function for convenience + + let mut rng = ChaCha20Rng::from_seed([0u8; 32]); // create and initialize rng + let mut rng = ChaCha20Rng::from_entropy(); + let mut msg = [0u8; 16]; // setting an arbitrary message + rng.fill_bytes(&mut msg); + + let params = Parameters { + m: 2728, // Security parameter XXX: not for production + k: 445, // Quorum parameter XXX: not for production + phi_f: 0.2, // Lottery parameter XXX: not for production + }; + + // Generate some arbitrary stake for each party + // Stake is an integer. + // Total stake of all parties is total stake in the system. + let stakes = (0..nparties) + .into_iter() + .map(|_| 1 + (rng.next_u64() % 9999)) + .collect::>(); + + // Create a new key registry from the parties and their stake + let mut key_reg = KeyRegistration::init(); + + // For each party, crate a Initializer. + // This struct can create keys for the party. + let mut ps: Vec = Vec::with_capacity(nparties); + for stake in stakes { + // Create keys for this party + let p = Initializer::new(params, stake, &mut rng); + // Register keys with the KeyRegistration service + key_reg + .register(p.stake, p.get_verification_key_proof_of_possession()) + .unwrap(); + ps.push(p); + } + + let closed_reg = key_reg.close(); + + // println!("{:?}", closed_reg.merkle_tree); + + let ps = ps + .into_par_iter() + .map(|p| p.create_signer(closed_reg.clone()).unwrap()) + .collect::>>(); + + // println!("Number of Signers: {:?}", ps.len()); + + let sigs = ps + .par_iter() + .filter_map(|p| { + return p.sign(&msg); + }) + .collect::>(); + + // Clerk can aggregate and verify signatures. + let clerk = Clerk::new_clerk_from_closed_key_registration(¶ms, &closed_reg); + + // Aggregate and verify the signatures + let sig_reg_list = sigs + .iter() + .map(|sig| SingleSignatureWithRegisteredParty { + sig: sig.clone(), + reg_party: closed_reg.reg_parties[sig.signer_index as usize], + }) + .collect::>(); + + let avk = AggregateVerificationKey::from(&closed_reg); + let msgp = avk.get_mt_commitment().concat_with_msg(&msg); + + // println!("NEW function!"); + let now = Instant::now(); + let (idx_by_mtidx, btm) = + BasicVerifier::get_k_indices(&closed_reg.total_stake, ¶ms, &msgp, &sig_reg_list); + // println!("list idx: {:?}", idx_by_mtidx); + // println!("nb of idx: {:?}", idx_by_mtidx.len()); + let mut vec_single_sig = + BasicVerifier::valid_signatures_from_k_indices(¶ms, idx_by_mtidx, btm).unwrap(); + vec_single_sig.sort_unstable(); + println!( + "Time to get the indices NEW: {:?} ms.", + (Instant::now() - now).as_millis() + ); + + let now = Instant::now(); + let mut unique_sigs = BasicVerifier::select_valid_signatures_for_k_indices( + &closed_reg.total_stake, + ¶ms, + &msgp, + &sig_reg_list, + ) + .unwrap(); + unique_sigs.sort_unstable(); + println!( + "Time to get the indices OLD: {:?} ms.", + (Instant::now() - now).as_millis() + ); + + // println!("bytes sigma: {:?}", vec_single_sig[0].sig.sigma.to_bytes().len()); + + for sig_new in vec_single_sig.iter_mut() { + sig_new.sig.indexes.sort(); + // println!("sig new: {:?}", sig_new.sig.indexes); + } + println!(); + for sig_old in unique_sigs.iter_mut() { + sig_old.sig.indexes.sort(); + // println!("sig old: {:?}", sig_old.sig.indexes); + } + + let nb_bytes_new: usize = vec_single_sig.iter().map(|s| s.to_bytes().len()).sum(); + let nb_bytes_old: usize = unique_sigs.iter().map(|s| s.to_bytes().len()).sum(); + + println!("nb_bytes_new: {:?}", nb_bytes_new); + println!("nb_bytes_old: {:?}", nb_bytes_old); + + let mt_index_list = unique_sigs + .iter() + .map(|sig_reg| sig_reg.sig.signer_index as usize) + .collect::>(); + + println!("batch_proof_old: {:?}", mt_index_list); + let batch_proof_old = closed_reg.merkle_tree.get_batched_path(mt_index_list); + + let mt_index_list = vec_single_sig + .iter() + .map(|sig_reg| sig_reg.sig.signer_index as usize) + .collect::>(); + + println!("batch_proof_new: {:?}", mt_index_list); + let batch_proof_new = closed_reg.merkle_tree.get_batched_path(mt_index_list); +} + diff --git a/mithril-stm/tests/test_extensions/mod.rs b/mithril-stm/tests/test_extensions/mod.rs index 651995e6909..5bb6ebde7d7 100644 --- a/mithril-stm/tests/test_extensions/mod.rs +++ b/mithril-stm/tests/test_extensions/mod.rs @@ -2,4 +2,5 @@ // is a different compilation target). #![allow(dead_code)] +pub mod basic_verifier_tests; pub mod protocol_phase; From e3c2b196006832b2687dbd638c48209e5519bd11 Mon Sep 17 00:00:00 2001 From: Dam Date: Tue, 4 Nov 2025 18:46:41 +0100 Subject: [PATCH 3/3] Added bench function for select function and try to optimize new version. --- mithril-stm/Cargo.toml | 4 + mithril-stm/benches/select_function.rs | 156 ++++++++++++++++++ mithril-stm/benches/stm.rs | 11 +- .../src/aggregate_signature/basic_verifier.rs | 132 ++++++++++++++- .../test_extensions/basic_verifier_tests.rs | 59 +++++-- 5 files changed, 339 insertions(+), 23 deletions(-) create mode 100644 mithril-stm/benches/select_function.rs diff --git a/mithril-stm/Cargo.toml b/mithril-stm/Cargo.toml index 3f660b78bf7..0851e6ba916 100644 --- a/mithril-stm/Cargo.toml +++ b/mithril-stm/Cargo.toml @@ -63,3 +63,7 @@ harness = false [[bench]] name = "size_benches" harness = false + +[[bench]] +name = "select_function" +harness = false diff --git a/mithril-stm/benches/select_function.rs b/mithril-stm/benches/select_function.rs new file mode 100644 index 00000000000..15272c6cc25 --- /dev/null +++ b/mithril-stm/benches/select_function.rs @@ -0,0 +1,156 @@ +use std::{collections::BTreeMap, process::id, time::Instant}; + +use blake2::{Blake2b, digest::consts::U32}; +use mithril_stm::{ + AggregateSignature, AggregateVerificationKey, AggregationError, BasicVerifier, Clerk, + Initializer, KeyRegistration, Parameters, Signer, SingleSignature, + SingleSignatureWithRegisteredParty, Stake, VerificationKey, +}; +use rand_chacha::ChaCha20Rng; +use rand_core::{RngCore, SeedableRng}; +use rayon::{prelude::*, vec}; + +type H = Blake2b; + +pub fn select_sig_test() { + let nparties = 3000; // Use a small number of parties for this example + type D = Blake2b; // Setting the hash function for convenience + + // let mut rng = ChaCha20Rng::from_seed([0u8; 32]); // create and initialize rng + let mut rng = ChaCha20Rng::from_entropy(); + let mut msg = [0u8; 16]; // setting an arbitrary message + rng.fill_bytes(&mut msg); + + let params = Parameters { + m: 2728, // Security parameter XXX: not for production + k: 445, // Quorum parameter XXX: not for production + phi_f: 0.2, // Lottery parameter XXX: not for production + }; + + // Generate some arbitrary stake for each party + // Stake is an integer. + // Total stake of all parties is total stake in the system. + let stakes = (0..nparties) + .into_iter() + .map(|_| 1 + (rng.next_u64() % 9999)) + .collect::>(); + + // Create a new key registry from the parties and their stake + let mut key_reg = KeyRegistration::init(); + + // For each party, crate a Initializer. + // This struct can create keys for the party. + let mut ps: Vec = Vec::with_capacity(nparties); + for stake in stakes { + // Create keys for this party + let p = Initializer::new(params, stake, &mut rng); + // Register keys with the KeyRegistration service + key_reg + .register(p.stake, p.get_verification_key_proof_of_possession()) + .unwrap(); + ps.push(p); + } + + let closed_reg = key_reg.close(); + + // println!("{:?}", closed_reg.merkle_tree); + + let ps = ps + .into_par_iter() + .map(|p| p.create_signer(closed_reg.clone()).unwrap()) + .collect::>>(); + + // println!("Number of Signers: {:?}", ps.len()); + + let sigs = ps + .par_iter() + .filter_map(|p| { + return p.sign(&msg); + }) + .collect::>(); + + // Clerk can aggregate and verify signatures. + let clerk = Clerk::new_clerk_from_closed_key_registration(¶ms, &closed_reg); + + // Aggregate and verify the signatures + let sig_reg_list = sigs + .iter() + .map(|sig| SingleSignatureWithRegisteredParty { + sig: sig.clone(), + reg_party: closed_reg.reg_parties[sig.signer_index as usize], + }) + .collect::>(); + + let avk = AggregateVerificationKey::from(&closed_reg); + let msgp = avk.get_mt_commitment().concat_with_msg(&msg); + + // println!("NEW function!"); + // let (idx_by_mtidx, btm) = + // BasicVerifier::get_k_indices(&closed_reg.total_stake, ¶ms, &msgp, &sig_reg_list); + // // println!("list idx: {:?}", idx_by_mtidx); + // // println!("nb of idx: {:?}", idx_by_mtidx.len()); + // let mut vec_single_sig = + // BasicVerifier::valid_signatures_from_k_indices(¶ms, idx_by_mtidx, btm).unwrap(); + // vec_single_sig.sort_unstable(); + let nb_loop = 20; + let now = Instant::now(); + let mut time_rework = 0; + let mut time_rework_opti = 0; + for _ in 0..nb_loop { + let now = Instant::now(); + let _ = BasicVerifier::reworked_select_valid_signatures_for_k_indices( + &closed_reg.total_stake, + ¶ms, + &msgp, + &sig_reg_list, + ) + .unwrap(); + time_rework += now.elapsed().as_millis(); + let now = Instant::now(); + let _ = BasicVerifier::reworked_select_valid_signatures_for_k_indices_opti( + &closed_reg.total_stake, + ¶ms, + &msgp, + &sig_reg_list, + ) + .unwrap(); + time_rework_opti += now.elapsed().as_millis(); + } + println!("Rework loop took: {:?}ms.", time_rework); + println!("Rework opti loop took: {:?}ms.", time_rework_opti); + + + println!( + "Time to get the indices NEW: {:?} ms.", + ((Instant::now() - now)/nb_loop).as_millis() + ); + let mut new_unique_sigs = BasicVerifier::reworked_select_valid_signatures_for_k_indices( + &closed_reg.total_stake, + ¶ms, + &msgp, + &sig_reg_list, + ) + .unwrap(); + new_unique_sigs.sort_unstable(); + + let now = Instant::now(); + for _ in 0..nb_loop { + let _ = BasicVerifier::select_valid_signatures_for_k_indices( + &closed_reg.total_stake, + ¶ms, + &msgp, + &sig_reg_list, + ) + .unwrap(); + } + println!( + "Time to get the indices OLD: {:?} ms.", + ((Instant::now() - now)/nb_loop).as_millis() + ); + + +} + +fn main() { + select_sig_test(); +} \ No newline at end of file diff --git a/mithril-stm/benches/stm.rs b/mithril-stm/benches/stm.rs index d4274d170e0..fa51d7996c9 100644 --- a/mithril-stm/benches/stm.rs +++ b/mithril-stm/benches/stm.rs @@ -53,6 +53,7 @@ where }) }); + let closed_reg = key_reg.close(); let signers = initializers @@ -287,13 +288,13 @@ fn core_verifier_benches_blake_2000(c: &mut Criterion) { } criterion_group!(name = benches; - config = Criterion::default().nresamples(1000); + config = Criterion::default().nresamples(10).sample_size(10); targets = - core_verifier_benches_blake_300, - core_verifier_benches_blake_2000, + // core_verifier_benches_blake_300, + // core_verifier_benches_blake_2000, stm_benches_blake_300, stm_benches_blake_2000, - batch_stm_benches_blake_300, - batch_stm_benches_blake_2000, + // batch_stm_benches_blake_300, + // batch_stm_benches_blake_2000, ); criterion_main!(benches); diff --git a/mithril-stm/src/aggregate_signature/basic_verifier.rs b/mithril-stm/src/aggregate_signature/basic_verifier.rs index d443846de57..12055a82d37 100644 --- a/mithril-stm/src/aggregate_signature/basic_verifier.rs +++ b/mithril-stm/src/aggregate_signature/basic_verifier.rs @@ -1,4 +1,7 @@ use std::collections::{BTreeMap, HashMap, HashSet}; +use std::time::Instant; + +use rug::integer::IsPrime; use crate::bls_multi_signature::{BlsSignature, BlsVerificationKey}; use crate::key_registration::RegisteredParty; @@ -99,7 +102,10 @@ impl BasicVerifier { let mut removal_idx_by_vk: HashMap<&SingleSignatureWithRegisteredParty, Vec> = HashMap::new(); + let mut basic_verif_time = 0; + let now = Instant::now(); for sig_reg in sigs.iter() { + let now = Instant::now(); if sig_reg .sig .basic_verify( @@ -113,6 +119,7 @@ impl BasicVerifier { { continue; } + basic_verif_time += now.elapsed().as_millis(); for index in sig_reg.sig.indexes.iter() { let mut insert_this_sig = false; if let Some(&previous_sig) = sig_by_index.get(index) { @@ -137,10 +144,13 @@ impl BasicVerifier { } } } + // println!("First loop took: {:?}ms.", now.elapsed().as_millis()); + // println!("Basic verif took: {:?}ms.", basic_verif_time); + let mut dedup_sigs: HashSet = HashSet::new(); let mut count: u64 = 0; - + let now = Instant::now(); for (_, &sig_reg) in sig_by_index.iter() { if dedup_sigs.contains(sig_reg) { continue; @@ -165,10 +175,12 @@ impl BasicVerifier { count += size; if count >= params.k { + // println!("Second loop took: {:?}micros.", now.elapsed().as_micros()); return Ok(dedup_sigs.into_iter().collect()); } } } + // println!("Second loop took: {:?}micros.", now.elapsed().as_micros()); Err(AggregationError::NotEnoughSignatures(count, params.k)) } @@ -181,9 +193,32 @@ impl BasicVerifier { msg: &[u8], sigs: &[SingleSignatureWithRegisteredParty], ) -> Result, AggregationError> { + let now = Instant::now(); let (idx_by_mtidx, btm) = Self::get_k_indices(total_stake, params, msg, sigs); + // println!("First function took: {:?}ms.", now.elapsed().as_millis()); - Self::valid_signatures_from_k_indices(params, idx_by_mtidx, btm) + let now = Instant::now(); + let result = Self::valid_signatures_from_k_indices(params, idx_by_mtidx, btm); + // println!("Second function took: {:?}micros.", now.elapsed().as_micros()); + + result + } + + pub fn reworked_select_valid_signatures_for_k_indices_opti( + total_stake: &Stake, + params: &Parameters, + msg: &[u8], + sigs: &[SingleSignatureWithRegisteredParty], + ) -> Result, AggregationError> { + let now = Instant::now(); + let btm = Self::get_k_indices_opti(total_stake, params, msg, sigs); + // println!("First function OPTI? took: {:?}ms.\n", now.elapsed().as_millis()); + + let now = Instant::now(); + let result = Self::valid_signatures_from_k_indices_opti(params, btm); + // println!("Second function took: {:?}micros.", now.elapsed().as_micros()); + + result } /// Iterates over the indices list in each signature to select the tuple (index, signature) with the smallest corresponding sigma. @@ -237,6 +272,54 @@ impl BasicVerifier { (indices_by_mt_idx, sig_by_mt_index) } + /// Iterates over the indices list in each signature to select the tuple (index, signature) with the smallest corresponding sigma. + /// Uses a BTreeMap to keep track of the correspondance between merkle tree index and signature + /// and a HashMap to connect the selected indices to the merkle tree index of the signature. + /// TODO: Merge those two as the MT index is already in the signature + pub fn get_k_indices_opti<'a>( + total_stake: &Stake, + params: &Parameters, + msg: &[u8], + sigs: &'a [SingleSignatureWithRegisteredParty], + // ) -> BTreeMap + ) -> HashMap + { + // let mut indices_by_sig: BTreeMap = + // BTreeMap::new(); + let mut indices_by_sig: HashMap = + HashMap::new(); + + for sig_reg in sigs.iter() { + if sig_reg + .sig + .basic_verify( + params, + &sig_reg.reg_party.0, + &sig_reg.reg_party.1, + msg, + total_stake, + ) + .is_err() + { + continue; + } + for index in sig_reg.sig.indexes.iter() { + if let Some(&prev_sig) = indices_by_sig.get(index) { + // if let Some(prev_sig) = sig_by_mt_index.get(mt_idx) { + if prev_sig.sig.sigma < sig_reg.sig.sigma { + continue; + } else { + indices_by_sig.insert(*index, sig_reg); + } + } else { + // Should we test for k indices here? + indices_by_sig.insert(*index, sig_reg); + } + } + } + indices_by_sig + } + /// Takes a list of tuples (selected index, MT index) and outputs a list of signatures with a total of k valid indices /// First iterates over the tuples to create a vector of selected indices for each signature. /// Then modify and aggregates the signatures starting with the signature with the most valid indices. @@ -247,6 +330,7 @@ impl BasicVerifier { sig_by_mt_index: BTreeMap, ) -> Result, AggregationError> { let mut valid_idx_for_mt_idx: HashMap> = HashMap::new(); + let mut valid_idx_for_sig: HashMap<&SingleSignatureWithRegisteredParty, Vec> = HashMap::new(); // Uses the HashMap of tuples (index, MT index) to create a vector of indices connected to one signature for (&valid_idx, &mt_idx) in list_k_valid_indices.iter() { @@ -288,6 +372,50 @@ impl BasicVerifier { Err(AggregationError::NotEnoughSignatures(count_idx, params.k)) } + pub fn valid_signatures_from_k_indices_opti( + params: &Parameters, + sig_by_mt_index: HashMap, + ) -> Result, AggregationError> { + // let mut valid_idx_for_mt_idx: HashMap> = HashMap::new(); + let mut valid_idx_for_sig: HashMap<&SingleSignatureWithRegisteredParty, Vec> = HashMap::new(); + + // Uses the HashMap of tuples (index, MT index) to create a vector of indices connected to one signature + for (&valid_idx, &sig) in sig_by_mt_index.iter() { + if let Some(val) = valid_idx_for_sig.get_mut(sig) { + val.push(valid_idx); + } else { + // valid_idx_for_mt_idx.insert(mt_idx, vec![valid_idx]); + valid_idx_for_sig.insert(sig, vec![valid_idx]); + + } + } + + // sort the HashMap according to the number of selected indices of each signature + let mut vec_valid_idx_per_mt_idx = valid_idx_for_sig.into_iter().collect::>(); + vec_valid_idx_per_mt_idx.sort_by(|(_, v1), (_, v2)| v2.len().cmp(&v1.len())); + + let mut uniques_sig: Vec = Vec::new(); + let mut count_idx = 0; + // Modify and aggregate the signatures + for (sig, indices) in vec_valid_idx_per_mt_idx.into_iter() { + let mut single_sig = (*sig).clone(); + if indices.len() > (params.k - count_idx) as usize { + single_sig.sig.indexes = indices[0..(params.k - count_idx) as usize].to_vec(); + count_idx += params.k - count_idx; + } else { + count_idx += indices.len() as u64; + single_sig.sig.indexes = indices; + } + uniques_sig.push(single_sig); + + if count_idx >= params.k { + return Ok(uniques_sig); + } + } + + Err(AggregationError::NotEnoughSignatures(count_idx, params.k)) + } + /// Given a slice of `sig_reg_list`, this function returns a new list of `sig_reg_list` with only valid indices. /// In case of conflict (having several signatures for the same index) /// it selects the smallest signature (i.e. takes the signature with the smallest scalar). diff --git a/mithril-stm/tests/test_extensions/basic_verifier_tests.rs b/mithril-stm/tests/test_extensions/basic_verifier_tests.rs index 254db8bdb57..51c4478c1d5 100644 --- a/mithril-stm/tests/test_extensions/basic_verifier_tests.rs +++ b/mithril-stm/tests/test_extensions/basic_verifier_tests.rs @@ -17,7 +17,7 @@ pub fn select_sig_test() { let nparties = 3000; // Use a small number of parties for this example type D = Blake2b; // Setting the hash function for convenience - let mut rng = ChaCha20Rng::from_seed([0u8; 32]); // create and initialize rng + // let mut rng = ChaCha20Rng::from_seed([0u8; 32]); // create and initialize rng let mut rng = ChaCha20Rng::from_entropy(); let mut msg = [0u8; 16]; // setting an arbitrary message rng.fill_bytes(&mut msg); @@ -86,20 +86,51 @@ pub fn select_sig_test() { let msgp = avk.get_mt_commitment().concat_with_msg(&msg); // println!("NEW function!"); + // let (idx_by_mtidx, btm) = + // BasicVerifier::get_k_indices(&closed_reg.total_stake, ¶ms, &msgp, &sig_reg_list); + // // println!("list idx: {:?}", idx_by_mtidx); + // // println!("nb of idx: {:?}", idx_by_mtidx.len()); + // let mut vec_single_sig = + // BasicVerifier::valid_signatures_from_k_indices(¶ms, idx_by_mtidx, btm).unwrap(); + // vec_single_sig.sort_unstable(); let now = Instant::now(); - let (idx_by_mtidx, btm) = - BasicVerifier::get_k_indices(&closed_reg.total_stake, ¶ms, &msgp, &sig_reg_list); - // println!("list idx: {:?}", idx_by_mtidx); - // println!("nb of idx: {:?}", idx_by_mtidx.len()); - let mut vec_single_sig = - BasicVerifier::valid_signatures_from_k_indices(¶ms, idx_by_mtidx, btm).unwrap(); - vec_single_sig.sort_unstable(); + for _ in 0..100 { + let _ = BasicVerifier::reworked_select_valid_signatures_for_k_indices( + &closed_reg.total_stake, + ¶ms, + &msgp, + &sig_reg_list, + ) + .unwrap(); + } println!( "Time to get the indices NEW: {:?} ms.", - (Instant::now() - now).as_millis() + ((Instant::now() - now)/100).as_millis() ); + let mut new_unique_sigs = BasicVerifier::reworked_select_valid_signatures_for_k_indices( + &closed_reg.total_stake, + ¶ms, + &msgp, + &sig_reg_list, + ) + .unwrap(); + new_unique_sigs.sort_unstable(); let now = Instant::now(); + for _ in 0..100 { + let _ = BasicVerifier::reworked_select_valid_signatures_for_k_indices( + &closed_reg.total_stake, + ¶ms, + &msgp, + &sig_reg_list, + ) + .unwrap(); + } + println!( + "Time to get the indices OLD: {:?} ms.", + ((Instant::now() - now)/100).as_millis() + ); + let mut unique_sigs = BasicVerifier::select_valid_signatures_for_k_indices( &closed_reg.total_stake, ¶ms, @@ -108,14 +139,10 @@ pub fn select_sig_test() { ) .unwrap(); unique_sigs.sort_unstable(); - println!( - "Time to get the indices OLD: {:?} ms.", - (Instant::now() - now).as_millis() - ); // println!("bytes sigma: {:?}", vec_single_sig[0].sig.sigma.to_bytes().len()); - for sig_new in vec_single_sig.iter_mut() { + for sig_new in new_unique_sigs.iter_mut() { sig_new.sig.indexes.sort(); // println!("sig new: {:?}", sig_new.sig.indexes); } @@ -125,7 +152,7 @@ pub fn select_sig_test() { // println!("sig old: {:?}", sig_old.sig.indexes); } - let nb_bytes_new: usize = vec_single_sig.iter().map(|s| s.to_bytes().len()).sum(); + let nb_bytes_new: usize = new_unique_sigs.iter().map(|s| s.to_bytes().len()).sum(); let nb_bytes_old: usize = unique_sigs.iter().map(|s| s.to_bytes().len()).sum(); println!("nb_bytes_new: {:?}", nb_bytes_new); @@ -139,7 +166,7 @@ pub fn select_sig_test() { println!("batch_proof_old: {:?}", mt_index_list); let batch_proof_old = closed_reg.merkle_tree.get_batched_path(mt_index_list); - let mt_index_list = vec_single_sig + let mt_index_list = new_unique_sigs .iter() .map(|sig_reg| sig_reg.sig.signer_index as usize) .collect::>();