From 9f09a2aace8f3895872567cb5d031d1aa8073882 Mon Sep 17 00:00:00 2001
From: Daniel Lubarov <daniel@lubarov.com>
Date: Wed, 12 Jan 2022 16:25:12 -0800
Subject: [PATCH] Add Merkle tree benchmark (#429)

And one for a single Keccak hash
---
 plonky2/Cargo.toml             |  4 ++++
 plonky2/benches/hashing.rs     | 15 +++++++++++++++
 plonky2/benches/merkle.rs      | 35 ++++++++++++++++++++++++++++++++++
 plonky2/src/hash/hash_types.rs | 12 ++++++++++++
 4 files changed, 66 insertions(+)
 create mode 100644 plonky2/benches/merkle.rs

diff --git a/plonky2/Cargo.toml b/plonky2/Cargo.toml
index 54cf5c1f..bce9f400 100644
--- a/plonky2/Cargo.toml
+++ b/plonky2/Cargo.toml
@@ -48,6 +48,10 @@ harness = false
 name = "hashing"
 harness = false
 
+[[bench]]
+name = "merkle"
+harness = false
+
 [[bench]]
 name = "transpose"
 harness = false
diff --git a/plonky2/benches/hashing.rs b/plonky2/benches/hashing.rs
index b1193516..2d17dba1 100644
--- a/plonky2/benches/hashing.rs
+++ b/plonky2/benches/hashing.rs
@@ -1,10 +1,14 @@
+#![allow(incomplete_features)]
 #![feature(generic_const_exprs)]
 
 use criterion::{criterion_group, criterion_main, BatchSize, Criterion};
 use plonky2::field::goldilocks_field::GoldilocksField;
 use plonky2::hash::gmimc::GMiMC;
+use plonky2::hash::hash_types::{BytesHash, RichField};
 use plonky2::hash::hashing::SPONGE_WIDTH;
+use plonky2::hash::keccak::KeccakHash;
 use plonky2::hash::poseidon::Poseidon;
+use plonky2::plonk::config::Hasher;
 use tynm::type_name;
 
 pub(crate) fn bench_gmimc<F: GMiMC<WIDTH>, const WIDTH: usize>(c: &mut Criterion) {
@@ -17,6 +21,16 @@ pub(crate) fn bench_gmimc<F: GMiMC<WIDTH>, const WIDTH: usize>(c: &mut Criterion
     });
 }
 
+pub(crate) fn bench_keccak<F: RichField>(c: &mut Criterion) {
+    c.bench_function("keccak256", |b| {
+        b.iter_batched(
+            || (BytesHash::<32>::rand(), BytesHash::<32>::rand()),
+            |(left, right)| <KeccakHash<32> as Hasher<F>>::two_to_one(left, right),
+            BatchSize::SmallInput,
+        )
+    });
+}
+
 pub(crate) fn bench_poseidon<F: Poseidon>(c: &mut Criterion) {
     c.bench_function(
         &format!("poseidon<{}, {}>", type_name::<F>(), SPONGE_WIDTH),
@@ -33,6 +47,7 @@ pub(crate) fn bench_poseidon<F: Poseidon>(c: &mut Criterion) {
 fn criterion_benchmark(c: &mut Criterion) {
     bench_gmimc::<GoldilocksField, 12>(c);
     bench_poseidon::<GoldilocksField>(c);
+    bench_keccak::<GoldilocksField>(c);
 }
 
 criterion_group!(benches, criterion_benchmark);
diff --git a/plonky2/benches/merkle.rs b/plonky2/benches/merkle.rs
new file mode 100644
index 00000000..7445682b
--- /dev/null
+++ b/plonky2/benches/merkle.rs
@@ -0,0 +1,35 @@
+use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
+use plonky2::field::goldilocks_field::GoldilocksField;
+use plonky2::hash::hash_types::RichField;
+use plonky2::hash::keccak::KeccakHash;
+use plonky2::hash::merkle_tree::MerkleTree;
+use plonky2::hash::poseidon::PoseidonHash;
+use plonky2::plonk::config::Hasher;
+use tynm::type_name;
+
+const ELEMS_PER_LEAF: usize = 135;
+
+pub(crate) fn bench_merkle_tree<F: RichField, H: Hasher<F>>(c: &mut Criterion) {
+    let mut group = c.benchmark_group(&format!(
+        "merkle-tree<{}, {}>",
+        type_name::<F>(),
+        type_name::<H>()
+    ));
+    group.sample_size(10);
+
+    for size_log in [13, 14, 15] {
+        let size = 1 << size_log;
+        group.bench_with_input(BenchmarkId::from_parameter(size), &size, |b, _| {
+            let leaves = vec![F::rand_vec(ELEMS_PER_LEAF); size];
+            b.iter(|| MerkleTree::<F, H>::new(leaves.clone(), 0));
+        });
+    }
+}
+
+fn criterion_benchmark(c: &mut Criterion) {
+    bench_merkle_tree::<GoldilocksField, PoseidonHash>(c);
+    bench_merkle_tree::<GoldilocksField, KeccakHash<25>>(c);
+}
+
+criterion_group!(benches, criterion_benchmark);
+criterion_main!(benches);
diff --git a/plonky2/src/hash/hash_types.rs b/plonky2/src/hash/hash_types.rs
index 1cfa7de8..82b8ceea 100644
--- a/plonky2/src/hash/hash_types.rs
+++ b/plonky2/src/hash/hash_types.rs
@@ -123,6 +123,18 @@ pub struct MerkleCapTarget(pub Vec<HashOutTarget>);
 #[derive(Eq, PartialEq, Copy, Clone, Debug)]
 pub struct BytesHash<const N: usize>(pub [u8; N]);
 
+impl<const N: usize> BytesHash<N> {
+    pub fn rand_from_rng<R: Rng>(rng: &mut R) -> Self {
+        let mut buf = [0; N];
+        rng.fill_bytes(&mut buf);
+        Self(buf)
+    }
+
+    pub fn rand() -> Self {
+        Self::rand_from_rng(&mut rand::thread_rng())
+    }
+}
+
 impl<F: RichField, const N: usize> GenericHashOut<F> for BytesHash<N> {
     fn to_bytes(&self) -> Vec<u8> {
         self.0.to_vec()