From ddf38f41f27dd510c78f10edec1e63ffe7e4f2e8 Mon Sep 17 00:00:00 2001
From: Roman Krasiuk <rokrassyuk@gmail.com>
Date: Tue, 2 May 2023 17:28:26 +0300
Subject: [PATCH] bench(proofs): root calculations (#2513)

---
 crates/primitives/Cargo.toml           |  4 ++
 crates/primitives/benches/trie_root.rs | 92 ++++++++++++++++++++++++++
 crates/primitives/src/proofs.rs        | 11 +++
 crates/primitives/src/receipt.rs       |  2 +-
 4 files changed, 108 insertions(+), 1 deletion(-)
 create mode 100644 crates/primitives/benches/trie_root.rs

diff --git a/crates/primitives/Cargo.toml b/crates/primitives/Cargo.toml
index 386d22a1c..8ecd1434f 100644
--- a/crates/primitives/Cargo.toml
+++ b/crates/primitives/Cargo.toml
@@ -103,3 +103,7 @@ test-utils = []
 [[bench]]
 name = "recover_ecdsa_crit"
 harness = false
+
+[[bench]]
+name = "trie_root"
+harness = false
diff --git a/crates/primitives/benches/trie_root.rs b/crates/primitives/benches/trie_root.rs
new file mode 100644
index 000000000..fe0bd2f84
--- /dev/null
+++ b/crates/primitives/benches/trie_root.rs
@@ -0,0 +1,92 @@
+use criterion::{black_box, criterion_group, criterion_main, Criterion};
+use proptest::{
+    prelude::*,
+    strategy::{Strategy, ValueTree},
+    test_runner::TestRunner,
+};
+use reth_primitives::{proofs::KeccakHasher, ReceiptWithBloom, H256};
+
+/// Benchmarks different implementations of the root calculation.
+pub fn trie_root_benchmark(c: &mut Criterion) {
+    let mut group = c.benchmark_group("Receipts root calculation");
+
+    for size in [10, 100, 1_000] {
+        let group_name =
+            |description: &str| format!("receipts root | size: {size} | {description}");
+
+        let (test_data, expected) = generate_test_data(size);
+        use implementations::*;
+
+        group.bench_function(group_name("triehash::ordered_trie_root"), |b| {
+            b.iter(|| {
+                let receipts = test_data.clone();
+                let result = black_box(trie_hash_ordered_trie_root(receipts.into_iter()));
+                assert_eq!(result, expected);
+            });
+        });
+
+        group.bench_function(group_name("HashBuilder"), |b| {
+            b.iter(|| {
+                let receipts = test_data.clone();
+                let result = black_box(hash_builder_root(receipts));
+                assert_eq!(result, expected);
+            });
+        });
+    }
+}
+
+fn generate_test_data(size: usize) -> (Vec<ReceiptWithBloom>, H256) {
+    let receipts = prop::collection::vec(any::<ReceiptWithBloom>(), size)
+        .new_tree(&mut TestRunner::new(ProptestConfig::default()))
+        .unwrap()
+        .current();
+    let root = implementations::hash_builder_root(receipts.clone());
+    (receipts, root)
+}
+
+criterion_group! {
+    name = benches;
+    config = Criterion::default();
+    targets = trie_root_benchmark
+}
+criterion_main!(benches);
+
+mod implementations {
+    use super::*;
+    use bytes::BytesMut;
+    use reth_primitives::{
+        proofs::adjust_index_for_rlp,
+        trie::{HashBuilder, Nibbles},
+    };
+    use reth_rlp::Encodable;
+    use std::vec::IntoIter;
+
+    pub fn trie_hash_ordered_trie_root(receipts: IntoIter<ReceiptWithBloom>) -> H256 {
+        triehash::ordered_trie_root::<KeccakHasher, _>(receipts.map(|receipt| {
+            let mut receipt_rlp = Vec::new();
+            receipt.encode_inner(&mut receipt_rlp, false);
+            receipt_rlp
+        }))
+    }
+
+    pub fn hash_builder_root(receipts: Vec<ReceiptWithBloom>) -> H256 {
+        let mut index_buffer = BytesMut::new();
+        let mut value_buffer = BytesMut::new();
+
+        let mut hb = HashBuilder::default();
+        let receipts_len = receipts.len();
+        for i in 0..receipts_len {
+            let index = adjust_index_for_rlp(i, receipts_len);
+
+            index_buffer.clear();
+            index.encode(&mut index_buffer);
+
+            value_buffer.clear();
+            receipts[index].encode_inner(&mut value_buffer, false);
+
+            hb.add_leaf(Nibbles::unpack(&index_buffer), &value_buffer);
+        }
+
+        hb.root()
+    }
+}
diff --git a/crates/primitives/src/proofs.rs b/crates/primitives/src/proofs.rs
index 22419f8fb..48196cab8 100644
--- a/crates/primitives/src/proofs.rs
+++ b/crates/primitives/src/proofs.rs
@@ -33,6 +33,17 @@ impl Hasher for KeccakHasher {
     }
 }
 
+/// Adjust the index of an item for rlp encoding.
+pub const fn adjust_index_for_rlp(i: usize, len: usize) -> usize {
+    if i > 0x7f {
+        i
+    } else if i == 0x7f || i + 1 == len {
+        0
+    } else {
+        i + 1
+    }
+}
+
 /// Calculate a transaction root.
 ///
 /// `(rlp(index), encoded(tx))` pairs.
diff --git a/crates/primitives/src/receipt.rs b/crates/primitives/src/receipt.rs
index 62920da57..66b4a82ac 100644
--- a/crates/primitives/src/receipt.rs
+++ b/crates/primitives/src/receipt.rs
@@ -42,8 +42,8 @@ impl From<Receipt> for ReceiptWithBloom {
 }
 
 /// [`Receipt`] with calculated bloom filter.
-#[derive(Clone, Debug, PartialEq, Eq, Default)]
 #[main_codec]
+#[derive(Clone, Debug, PartialEq, Eq, Default)]
 pub struct ReceiptWithBloom {
     /// Bloom filter build from logs.
     pub bloom: Bloom,