poseidon hash bench

hash/plonky2/Cargo.toml

@@ -8,3 +8,5 @@ edition = "2021"
 [dependencies]
 anyhow = "1.0.79"
 plonky2 = "0.1.4"
+rand = "0.8.3"
+

hash/plonky2/src/bench/mod.rs

@@ -0,0 +1 @@
+pub mod poseidon;

hash/plonky2/src/bench/poseidon.rs

@@ -0,0 +1,64 @@
+use anyhow::Result;
+use plonky2::field::types::Field;
+// use plonky2::hash::hash_types::{HashOutTarget, RichField};
+use plonky2::field::goldilocks_field::GoldilocksField;
+// use plonky2::hash::keccak;
+// use plonky2::hash::keccak::KeccakHash;
+use plonky2::hash::poseidon::PoseidonHash;
+use plonky2::iop::witness::{PartialWitness, WitnessWrite};
+use plonky2::plonk::circuit_builder::CircuitBuilder;
+use plonky2::plonk::circuit_data::CircuitConfig;
+use plonky2::plonk::config::{/*AlgebraicHasher,*/ GenericConfig, PoseidonGoldilocksConfig};
+use rand::Rng;
+
+fn generate_data(size: usize) -> Vec<GoldilocksField> {
+    // let mut rng = rand::thread_rng();
+    // (0..size).map(|_| rng.gen()).collect()
+
+    let mut data: Vec<GoldilocksField> = Vec::new();
+    for _ in 0..(1<<size) {
+        let mut rng = rand::thread_rng();
+        let random_u64: u64 = rng.gen();
+        data.push(GoldilocksField::from_canonical_u64(random_u64));
+    }
+    // eprint!("data: {:?}", data);
+    data
+
+}
+
+
+pub fn poseidon_bench(depth: usize) -> Result<()> {
+
+    let data = generate_data(depth);
+
+    const D: usize = 2;
+    type C = PoseidonGoldilocksConfig;
+    type F = <C as GenericConfig<D>>::F;
+
+    let config = CircuitConfig::standard_recursion_config();
+    let mut builder = CircuitBuilder::<F, D>::new(config);
+
+    // The arithmetic circuit.
+    let initial = builder.add_virtual_targets(data.len());
+
+    let hash = builder.hash_or_noop::<PoseidonHash>(initial.clone());
+
+    // Public inputs are the initial value (provided below) and the result (which is generated).
+    builder.register_public_inputs(initial.clone().as_slice());
+    builder.register_public_input(hash.elements[0]);
+    builder.register_public_input(hash.elements[1]);
+    builder.register_public_input(hash.elements[2]);
+    builder.register_public_input(hash.elements[3]);
+
+    // Provide initial values.
+    let mut pw = PartialWitness::new();
+    pw.set_target_arr(initial.as_slice(), data.as_slice());
+
+
+    let data = builder.build::<C>();
+    let proof = data.prove(pw)?;
+
+
+    data.verify(proof)
+
+}

hash/plonky2/src/main.rs

@@ -1,47 +1,32 @@
-use anyhow::Result;
-use plonky2::field::types::Field;
-use plonky2::hash::hash_types::RichField;
-use plonky2::field::goldilocks_field::GoldilocksField;
-use plonky2::hash::keccak;
-use plonky2::hash::keccak::KeccakHash;
-use plonky2::hash::poseidon::PoseidonHash;
-use plonky2::iop::witness::{PartialWitness, WitnessWrite};
-use plonky2::plonk::circuit_builder::CircuitBuilder;
-use plonky2::plonk::circuit_data::CircuitConfig;
-use plonky2::plonk::config::{AlgebraicHasher, GenericConfig, PoseidonGoldilocksConfig};
+use std::process;
+mod bench;
+use bench::poseidon::poseidon_bench;
 
-fn main() -> Result<()> {
-    const D: usize = 2;
-    type C = PoseidonGoldilocksConfig;
-    type F = <C as GenericConfig<D>>::F;
+fn main() {
+    let args: Vec<String> = std::env::args().collect();
 
-    let config = CircuitConfig::standard_recursion_config();
-    let mut builder = CircuitBuilder::<F, D>::new(config);
+    if args.len() != 3 {
+        println!("Wrong number of arguments! The program expects two arguments: <hash_type> and <size>");
+        // Exit the program with a non-zero exit code
+        process::exit(1);
+    }
+    
+    let hash_type = &args[1];
+    let size = args[2].parse::<usize>().unwrap();
 
-    // The arithmetic circuit.
-    let initial = builder.add_virtual_target();
-    let hash = builder.hash_or_noop::<PoseidonHash>(vec![initial]);
+    match hash_type.as_str() {
 
-    // Public inputs are the initial value (provided below) and the result (which is generated).
-    builder.register_public_input(initial);
-    builder.register_public_input(hash.elements[0]);
-    builder.register_public_input(hash.elements[1]);
-    builder.register_public_input(hash.elements[2]);
-    builder.register_public_input(hash.elements[3]);
+        "poseidon" => {
+            println!("Running Poseidon: ");
+            eprintln!("Tree Depth: {:?}", size);
+            let _ = poseidon_bench(size);
+        }
 
-    // Provide initial values.
-    let mut pw = PartialWitness::new();
-    pw.set_target(initial, F::ONE);
+        _ => {
+            println!("Wrong Benchmark Name!");
+        }
+    }
 
-
-    let data = builder.build::<C>();
-    let proof = data.prove(pw)?;
-
-    println!(
-        "hash of {} is: {}",
-        proof.public_inputs[0], proof.public_inputs[1]
-    );
-
-    data.verify(proof)
-
-}
+    println!("All Done!");
+    
+}