Poseidon benchmark that also runs in Wasm

.gitignore

@@ -18,6 +18,7 @@ idea
 *.sage.py
 params
 _build
+build
 .merlin
 *.swp
 *.swo

bench-wasm.sh

@@ -0,0 +1,60 @@
+#!/bin/bash
+
+# Default benchmark name
+BENCH_NAME=""
+
+# Parse the optional --bench=<NAME> argument
+for i in "$@"; do
+    case $i in
+        --bench=*)
+        BENCH_NAME="${i#*=}"
+        shift # Remove --bench=<NAME> from processing
+        ;;
+        *)
+        # unknown option
+        ;;
+    esac
+done
+
+# Throw an error if --bench was not provided
+if [ -z "$BENCH_NAME" ]; then
+    echo "Error: You must specify a benchmark using --bench=<NAME>"
+    exit 1
+fi
+
+# Function to delete any old .wasm files related to the current benchmark
+cleanup_old_wasm_files() {
+    WASM_FILES=$(find target/wasm32-wasi/release/deps/ -type f -name "*$BENCH_NAME*.wasm")
+
+    if [ -n "$WASM_FILES" ]; then
+        echo "Cleaning up old WASM files for benchmark '$BENCH_NAME'..."
+        rm -f $WASM_FILES
+    else
+        echo "No old WASM files found for benchmark '$BENCH_NAME'."
+    fi
+}
+
+# Call the cleanup function
+cleanup_old_wasm_files
+
+# Build the WASM benchmark with cargo-wasi
+echo "Building benchmark '$BENCH_NAME' with cargo-wasi..."
+cargo wasi build --bench="$BENCH_NAME" --release
+
+# Function to find the correct WASM file
+find_wasm_file() {
+    # Search for the WASM file corresponding to the benchmark name
+    WASM_FILE=$(find target/wasm32-wasi/release/deps/ -type f -name "*$BENCH_NAME*.wasm" | head -n 1)
+
+    if [ -z "$WASM_FILE" ]; then
+        echo "Error: No WASM file found for benchmark '$BENCH_NAME'."
+        exit 1
+    fi
+}
+
+# Call the function to find the correct WASM file
+find_wasm_file
+
+echo "Running benchmark at $WASM_FILE with wasmer-js..."
+# Run the WASM file with wasmer-js
+wasmer-js run --dir=. "$WASM_FILE" -- --bench

poseidon/Cargo.toml

@@ -33,7 +33,15 @@ ocaml-gen = { workspace = true, optional = true }
 [dev-dependencies]
 serde_json.workspace = true
 hex.workspace = true
+criterion = { version = "0.3", default-features = false, features = [
+  "cargo_bench_support",
+  "html_reports",
+] }
 
 [features]
 default = []
-ocaml_types = [ "ocaml", "ocaml-gen", ]
+ocaml_types = ["ocaml", "ocaml-gen"]
+
+[[bench]]
+name = "poseidon_bench"
+harness = false

poseidon/README.md

@@ -1,4 +1,4 @@
-# Oracle
+# Poseidon
 
 ## Test vectors
 
@@ -17,3 +17,33 @@ cargo run -p export_test_vectors -- B10 legacy -
 cargo run -p export_test_vectors -- b10 legacy legacy.json
 cargo run -p export_test_vectors -- hex kimchi kimchi.json
 ```
+
+## Benchmark
+
+This folder contains a Poseidon benchmark `poseidon_bench`.
+
+To run the benchmark natively, do:
+
+```sh
+cargo bench --bench=poseidon_bench
+```
+
+It can also be run in WebAssembly (executed by Node.js), with the following prerequisites:
+
+- Add the `wasm32-wasi` target
+- Install `cargo-wasi`
+- Install the wasmer JS CLI
+
+```sh
+rustup target add wasm32-wasi
+cargo install cargo-wasi
+npm install -g @wasmer/cli
+```
+
+Now, you can run this and other benchmarks in Wasm, using the following from the repository root:
+
+```sh
+./bench-wasm.sh --bench=poseidon_bench
+```
+
+For this to work, the filename of your benchmark has to be the same as the benchmark name!

poseidon/benches/poseidon_bench.rs

@@ -0,0 +1,30 @@
+use criterion::{criterion_group, criterion_main, Criterion};
+use mina_curves::pasta::Fp;
+use mina_poseidon::{
+    constants::PlonkSpongeConstantsKimchi,
+    pasta::fp_kimchi as SpongeParametersKimchi,
+    poseidon::{ArithmeticSponge as Poseidon, Sponge},
+};
+
+pub fn bench_poseidon_kimchi(c: &mut Criterion) {
+    let mut group = c.benchmark_group("Poseidon");
+    group.sample_size(100);
+
+    // Chain of hashes, starting from a random value
+    group.bench_function("poseidon_hash_kimchi", |b| {
+        let mut hash: Fp = rand::random();
+        let mut poseidon = Poseidon::<Fp, PlonkSpongeConstantsKimchi>::new(
+            SpongeParametersKimchi::static_params(),
+        );
+
+        b.iter(|| {
+            poseidon.absorb(&[hash]);
+            hash = poseidon.squeeze();
+        })
+    });
+
+    group.finish();
+}
+
+criterion_group!(benches, bench_poseidon_kimchi);
+criterion_main!(benches);