feat: benchmark MKMapProof generation for multiple leaves

Author: jpraynaud

mithril-common/benches/merkle_map.rs

@@ -2,7 +2,7 @@ use std::ops::Range;
 
 use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
 use mithril_common::{
-    crypto_helper::{MKMap, MKMapNode, MKMapProof, MKMapValue, MKTree},
+    crypto_helper::{MKMap, MKMapNode, MKMapProof, MKMapValue, MKTree, MKTreeNode},
     entities::BlockRange,
 };
 
@@ -23,6 +23,7 @@ const TOTAL_TRANSACTIONS_BENCHES: &[u64] = &[
 ];
 const BLOCK_RANGE_LENGTH_BENCH: u64 = 15;
 const TOTAL_TRANSACTIONS_PER_BLOCK: u64 = 50;
+const MAX_TRANSACTIONS_PER_PROOF_BENCH: u64 = 100;
 
 fn generate_block_ranges_nodes_iterator(
     total_transactions: u64,
@@ -32,6 +33,10 @@ fn generate_block_ranges_nodes_iterator(
 ) -> impl Iterator<Item = (BlockRange, MKMapNode<BlockRange>)> {
     let total_block_ranges =
         total_transactions / (total_transactions_per_block * block_range_length);
+    assert!(
+        total_block_ranges > 0,
+        "total_block_ranges should be strictly greater than 0"
+    );
     (0..total_block_ranges).map(move |block_range_index| {
         let block_range = BlockRange::new(
             block_range_index * total_transactions_per_block * block_range_length,
@@ -55,15 +60,15 @@ fn generate_merkle_map_compressed(
     block_ranges_nodes_iterator: impl Iterator<Item = (BlockRange, MKMapNode<BlockRange>)>,
 ) -> MKMap<BlockRange, MKMapNode<BlockRange>> {
     let mut mk_map = MKMap::new(&[]).unwrap();
-    for (block_range, mk_tree) in block_ranges_nodes_iterator {
+    for (block_range, mk_map_node) in block_ranges_nodes_iterator {
         mk_map
-            .insert(block_range, mk_tree.compute_root().unwrap().into())
+            .insert(block_range, mk_map_node.compute_root().unwrap().into())
             .unwrap();
     }
     mk_map
 }
 
-fn generate_merkle_map_proof(
+fn generate_merkle_map_root(
     block_ranges_nodes_iterator: impl Iterator<Item = (BlockRange, MKMapNode<BlockRange>)>,
     mk_map_compressed: &MKMap<BlockRange, MKMapNode<BlockRange>>,
 ) -> MKMapProof<BlockRange> {
@@ -72,7 +77,7 @@ fn generate_merkle_map_proof(
     let mktree_to_prove = if let MKMapNode::Tree(mktree_to_prove) = mk_map_node_to_prove {
         mktree_to_prove
     } else {
-        panic!("Expected MKMapNode::TreeNode");
+        panic!("Expected MKMapNode::Tree");
     };
     let leaves_to_prove = mktree_to_prove
         .leaves()
@@ -89,9 +94,43 @@ fn generate_merkle_map_proof(
     mk_map_compressed.compute_proof(&leaves_to_prove).unwrap()
 }
 
-fn create_merkle_map_compressed_benches(c: &mut Criterion) {
+fn generate_merkle_map_proof(
+    block_ranges_nodes_iterator: impl Iterator<Item = (BlockRange, MKMapNode<BlockRange>)>,
+    mk_map_compressed: &MKMap<BlockRange, MKMapNode<BlockRange>>,
+    total_proofs: u64,
+) -> MKMapProof<BlockRange> {
+    let mut leaves_to_prove_all: Vec<MKTreeNode> = vec![];
+    let mut mk_map_compressed = mk_map_compressed.clone();
+    for (mk_map_key_to_prove, mk_map_node_to_prove) in &block_ranges_nodes_iterator
+        .take(total_proofs as usize)
+        .collect::<Vec<_>>()
+    {
+        let mktree_to_prove = if let MKMapNode::Tree(mktree_to_prove) = mk_map_node_to_prove {
+            mktree_to_prove
+        } else {
+            panic!("Expected MKMapNode::Tree");
+        };
+        let leaves_to_prove = mktree_to_prove
+            .leaves()
+            .into_iter()
+            .take(1)
+            .collect::<Vec<_>>();
+        leaves_to_prove_all.extend(leaves_to_prove);
+        mk_map_compressed
+            .insert(
+                mk_map_key_to_prove.to_owned(),
+                mk_map_node_to_prove.to_owned(),
+            )
+            .unwrap();
+    }
+    mk_map_compressed
+        .compute_proof(&leaves_to_prove_all)
+        .unwrap()
+}
+
+fn create_merkle_map_root_benches(c: &mut Criterion) {
     let mut group = c.benchmark_group(format!(
-        "create_merkle_map_compressed(blocks_ranges_length={BLOCK_RANGE_LENGTH_BENCH},txs_per_block={TOTAL_TRANSACTIONS_PER_BLOCK})"
+        "create_merkle_map_root(blocks_ranges_length={BLOCK_RANGE_LENGTH_BENCH},txs_per_block={TOTAL_TRANSACTIONS_PER_BLOCK})"
     ));
     for total_leaves in TOTAL_TRANSACTIONS_BENCHES.iter() {
         let mk_trees_by_block_range_iterator = generate_block_ranges_nodes_iterator(
@@ -106,8 +145,13 @@ fn create_merkle_map_compressed_benches(c: &mut Criterion) {
             total_leaves,
             |b, &_total_leaves| {
                 b.iter(|| {
-                    let mk_map_compressed_clone = mk_map_compressed.clone();
-                    mk_map_compressed_clone.compute_root().unwrap();
+                    let mk_trees_by_block_range_iterator = generate_block_ranges_nodes_iterator(
+                        *total_leaves,
+                        TOTAL_TRANSACTIONS_PER_BLOCK,
+                        BLOCK_RANGE_LENGTH_BENCH,
+                        1,
+                    );
+                    generate_merkle_map_root(mk_trees_by_block_range_iterator, &mk_map_compressed);
                 });
             },
         );
@@ -117,14 +161,14 @@ fn create_merkle_map_compressed_benches(c: &mut Criterion) {
 
 fn create_merkle_map_proof_benches(c: &mut Criterion) {
     let mut group = c.benchmark_group(format!(
-        "create_merkle_map_proof_(blocks_ranges_length={BLOCK_RANGE_LENGTH_BENCH},txs_per_block={TOTAL_TRANSACTIONS_PER_BLOCK})"
+        "create_merkle_map_proof_(blocks_ranges_length={BLOCK_RANGE_LENGTH_BENCH},txs_per_block={TOTAL_TRANSACTIONS_PER_BLOCK},txs_per_proof={MAX_TRANSACTIONS_PER_PROOF_BENCH})"
     ));
     for total_leaves in TOTAL_TRANSACTIONS_BENCHES.iter() {
         let mk_trees_by_block_range_iterator = generate_block_ranges_nodes_iterator(
             *total_leaves,
             TOTAL_TRANSACTIONS_PER_BLOCK,
             BLOCK_RANGE_LENGTH_BENCH,
-            1,
+            MAX_TRANSACTIONS_PER_PROOF_BENCH,
         );
         let mk_map_compressed = generate_merkle_map_compressed(mk_trees_by_block_range_iterator);
 
@@ -137,9 +181,13 @@ fn create_merkle_map_proof_benches(c: &mut Criterion) {
                         *total_leaves,
                         TOTAL_TRANSACTIONS_PER_BLOCK,
                         BLOCK_RANGE_LENGTH_BENCH,
-                        1,
+                        MAX_TRANSACTIONS_PER_PROOF_BENCH,
+                    );
+                    generate_merkle_map_proof(
+                        mk_trees_by_block_range_iterator,
+                        &mk_map_compressed,
+                        MAX_TRANSACTIONS_PER_PROOF_BENCH,
                     );
-                    generate_merkle_map_proof(mk_trees_by_block_range_iterator, &mk_map_compressed);
                 });
             },
         );
@@ -149,7 +197,7 @@ fn create_merkle_map_proof_benches(c: &mut Criterion) {
 
 fn verify_merkle_map_proof_benches(c: &mut Criterion) {
     let mut group = c.benchmark_group(format!(
-        "verify_merkle_map_proof_(blocks_ranges_length={BLOCK_RANGE_LENGTH_BENCH},txs_per_block={TOTAL_TRANSACTIONS_PER_BLOCK})"
+        "verify_merkle_map_proof_(blocks_ranges_length={BLOCK_RANGE_LENGTH_BENCH},txs_per_block={TOTAL_TRANSACTIONS_PER_BLOCK},txs_per_proof={MAX_TRANSACTIONS_PER_PROOF_BENCH})"
     ));
     for total_leaves in TOTAL_TRANSACTIONS_BENCHES.iter() {
         let mk_trees_by_block_range_iterator = generate_block_ranges_nodes_iterator(
@@ -163,10 +211,13 @@ fn verify_merkle_map_proof_benches(c: &mut Criterion) {
             *total_leaves,
             TOTAL_TRANSACTIONS_PER_BLOCK,
             BLOCK_RANGE_LENGTH_BENCH,
-            1,
+            MAX_TRANSACTIONS_PER_PROOF_BENCH,
+        );
+        let mk_map_proof = generate_merkle_map_proof(
+            mk_trees_by_block_range_iterator,
+            &mk_map_compressed,
+            MAX_TRANSACTIONS_PER_PROOF_BENCH,
         );
-        let mk_map_proof =
-            generate_merkle_map_proof(mk_trees_by_block_range_iterator, &mk_map_compressed);
 
         group.bench_with_input(
             BenchmarkId::from_parameter(total_leaves),
@@ -183,7 +234,7 @@ criterion_group!(
     name = benches;
     config = Criterion::default().sample_size(10);
     targets =
-        create_merkle_map_compressed_benches,
+        create_merkle_map_root_benches,
         create_merkle_map_proof_benches,
         verify_merkle_map_proof_benches
 );

mithril-common/src/crypto_helper/merkle_map.rs

@@ -122,6 +122,9 @@ impl<K: MKMapKey, V: MKMapValue<K>> MKMap<K, V> {
         &self,
         leaves: &[T],
     ) -> StdResult<MKMapProof<K>> {
+        if leaves.is_empty() {
+            return Err(anyhow!("MKMap could not compute proof for empty leaves"));
+        }
         let leaves_by_keys: HashMap<K, Vec<MKTreeNode>> = leaves
             .iter()
             .filter_map(|leaf| match self.contains(&leaf.to_owned().into()) {
@@ -537,6 +540,20 @@ mod tests {
         );
     }
 
+    #[test]
+    fn test_mk_map_should_not_compute_proof_for_no_leaves() {
+        let entries = generate_merkle_trees(100000, 100);
+        let mktree_nodes_to_certify: &[MKTreeNode] = &[];
+        let merkle_tree_node_entries = &entries
+            .into_iter()
+            .map(|(range, mktree)| (range.to_owned(), mktree.into()))
+            .collect::<Vec<(_, MKMapNode<_>)>>();
+        let mk_map_full = MKMap::new(merkle_tree_node_entries.as_slice()).unwrap();
+        mk_map_full
+            .compute_proof(mktree_nodes_to_certify)
+            .expect_err("MKMap should not compute proof for no leaves");
+    }
+
     #[test]
     fn test_mk_map_should_compute_and_verify_valid_proof() {
         let entries = generate_merkle_trees(100000, 100);