From 33984c6c236f3848cb3c33992c280499365cd051 Mon Sep 17 00:00:00 2001
From: "dustin.ray" <dustin.vaughan@spaceandtime.io>
Date: Thu, 22 Aug 2024 21:50:22 -0700
Subject: [PATCH] add benches

---
 Cargo.toml                    |  6 +++
 benches/hash_chain_benches.rs | 87 +++++++++++++++++++++++++++++++++++
 src/lib.rs                    |  6 +--
 3 files changed, 96 insertions(+), 3 deletions(-)
 create mode 100644 benches/hash_chain_benches.rs

diff --git a/Cargo.toml b/Cargo.toml
index f73353a..9030197 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -20,6 +20,7 @@ sha3 = "0.10.6"
 structopt = { version = "0.3.26", default-features = false }
 env_logger = "0.11.5"
 thiserror = "1.0.63"
+criterion = "0.3"
 
 [dev-dependencies]
 debug_print = { version = "1.0.0" }
@@ -34,3 +35,8 @@ plonky2_crypto = { git = "https://github.com/Lagrange-Labs/plonky2-crypto", bran
 
 [profile.dev]
 opt-level = 3
+
+[[bench]]
+name = "hash_chain_bench"
+path = "benches/hash_chain_benches.rs"
+harness = false
\ No newline at end of file
diff --git a/benches/hash_chain_benches.rs b/benches/hash_chain_benches.rs
new file mode 100644
index 0000000..55b209a
--- /dev/null
+++ b/benches/hash_chain_benches.rs
@@ -0,0 +1,87 @@
+use criterion::{black_box, criterion_group, criterion_main, Criterion};
+use hash_chain::HashChain;
+use plonky2::{
+    field::goldilocks_field::GoldilocksField,
+    plonk::{
+        circuit_builder::CircuitBuilder,
+        circuit_data::CircuitConfig,
+        config::{GenericConfig, PoseidonGoldilocksConfig},
+    },
+};
+
+fn hash_chain_proving_benchmark(c: &mut Criterion) {
+    const D: usize = 2;
+    type C = PoseidonGoldilocksConfig;
+    type F = <C as GenericConfig<D>>::F;
+
+    let config = CircuitConfig::standard_recursion_config();
+
+    // Power of two recursive step sizes
+    let step_sizes = [2 /*, 4, 8, 16, 32, 64 */];
+
+    let mut group = c.benchmark_group("HashChain Prover");
+
+    // Configure the group
+    group.sample_size(10);
+    for &steps in &step_sizes {
+        group.bench_function(&format!("hash_chain_{}_steps", steps), |b| {
+            b.iter(|| {
+                let mut circuit = black_box(CircuitBuilder::<F, D>::new(config.clone()));
+                let (_, _) = <CircuitBuilder<GoldilocksField, D> as HashChain<
+                    GoldilocksField,
+                    D,
+                    C,
+                >>::build_hash_chain_circuit(&mut circuit, steps)
+                .unwrap();
+            });
+        });
+    }
+}
+
+fn hash_chain_verification_benchmark(c: &mut Criterion) {
+    const D: usize = 2;
+    type C = PoseidonGoldilocksConfig;
+    type F = <C as GenericConfig<D>>::F;
+
+    let config = CircuitConfig::standard_recursion_config();
+
+    let mut group = c.benchmark_group("HashChain Verifier");
+
+    // Configure the group
+    group.sample_size(10);
+
+    // Power of two recursive step sizes
+    let step_sizes = [2 /*, 4, 8, 16, 32, 64 */];
+
+    for &steps in &step_sizes {
+        group.bench_function(&format!("hash_chain_verify_{}_steps", steps), |b| {
+            // Move the circuit and proof generation out of the iterated benchmark block
+            let mut circuit = CircuitBuilder::<F, D>::new(config.clone());
+            let (proof, circuit_map) = <CircuitBuilder<GoldilocksField, D> as HashChain<
+                GoldilocksField,
+                D,
+                C,
+            >>::build_hash_chain_circuit(&mut circuit, steps)
+            .unwrap();
+
+            b.iter(|| {
+                // Only verification is timed
+                let verification_result =
+                    black_box(<CircuitBuilder<GoldilocksField, D> as HashChain<
+                        GoldilocksField,
+                        D,
+                        C,
+                    >>::verify(proof.clone(), &circuit_map));
+
+                black_box(verification_result.unwrap());
+            });
+        });
+    }
+}
+
+criterion_group!(
+    benches,
+    hash_chain_proving_benchmark,
+    hash_chain_verification_benchmark
+);
+criterion_main!(benches);
diff --git a/src/lib.rs b/src/lib.rs
index 095ea42..9bdfe14 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -72,7 +72,7 @@ pub trait HashChain<F: RichField + Extendable<D>, const D: usize, C: GenericConf
 
     fn verify(
         proof: Proof<F, C, D>,
-        cyclic_circuit_data: CircuitMap<F, C, D>,
+        cyclic_circuit_data: &CircuitMap<F, C, D>,
     ) -> Result<(), HashChainError>;
 
     fn check_cyclic_proof_layer(
@@ -397,7 +397,7 @@ where
     // for zero-knowledge and should not be considered private.
     fn verify(
         proof: ProofWithPublicInputs<F, C, D>,
-        cyclic_circuit_data: CircuitData<F, C, D>,
+        cyclic_circuit_data: &CircuitData<F, C, D>,
     ) -> Result<(), HashChainError> {
         // Use the given hash permutation from plonky2 to verify
         // that the repeated hash is computed correctly.
@@ -466,7 +466,7 @@ mod tests {
         let result =
             <CircuitBuilder<GoldilocksField, D> as HashChain<GoldilocksField, D, C>>::verify(
                 proof,
-                circuit_map,
+                &circuit_map,
             );
         assert!(result.is_ok())
     }