Added new functions to try replace the select k indices one.

mithril-stm/src/aggregate_signature/basic_verifier.rs

@@ -172,6 +172,105 @@ impl BasicVerifier {
         Err(AggregationError::NotEnoughSignatures(count, params.k))
     }
 
+    // Modification of the function select_valid_signatures_for_k_indices
+    // to try to improve readability
+    // Still not outputing the same results, more tests need to be done
+    pub fn modified_select_valid_signatures_for_k_indices(
+        total_stake: &Stake,
+        params: &Parameters,
+        msg: &[u8],
+        sigs: &[SingleSignatureWithRegisteredParty],
+    ) -> Result<Vec<SingleSignatureWithRegisteredParty>, 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)
+    }
+
+    // Get a set of k unique indices connected to indices in the MT
+    // Also creates a more explicit link between MT index and signature to reuse later
+    pub fn get_k_indices<'a>(
+        total_stake: &Stake,
+        params: &Parameters,
+        msg: &[u8],
+        sigs: &'a [SingleSignatureWithRegisteredParty],
+    ) -> (
+        HashMap<Index, Index>,
+        BTreeMap<Index, &'a SingleSignatureWithRegisteredParty>,
+    ) {
+        let mut sig_by_mt_index: BTreeMap<Index, &SingleSignatureWithRegisteredParty> =
+            BTreeMap::new();
+        let mut indices_by_mt_idx: HashMap<Index, Index> = 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 if (indices_by_mt_idx.len() as u64) < params.k {
+                    // 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)
+    }
+
+    // From a set of k unique indices together with Merkle Tree indices
+    // Creates a vector of signatures with correct indices
+    pub fn valid_signatures_from_k_indices(
+        params: &Parameters,
+        list_k_valid_indices: HashMap<Index, Index>,
+        sig_by_mt_index: BTreeMap<Index, &SingleSignatureWithRegisteredParty>,
+    ) -> Result<Vec<SingleSignatureWithRegisteredParty>, AggregationError> {
+        let mut valid_idx_for_mt_idx: HashMap<u64, Vec<u64>> = HashMap::new();
+
+        for (&valid_idx, &mt_id) in list_k_valid_indices.iter() {
+            if let Some(val) = valid_idx_for_mt_idx.get_mut(&mt_id) {
+                val.push(valid_idx);
+            } else {
+                valid_idx_for_mt_idx.insert(mt_id, vec![valid_idx]);
+            }
+        }
+        let mut uniques_sig: Vec<SingleSignatureWithRegisteredParty> = Vec::new();
+        let mut count_idx = 0;
+        for (mt_idx, indices) in valid_idx_for_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));
+            };
+            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).