Benchmark proof creation and verification for RATE = 2, 8, 11.

Cargo.toml

@@ -77,6 +77,10 @@ harness = false
 name = "circuit"
 harness = false
 
+[[bench]]
+name = "poseidon"
+harness = false
+
 [profile.release]
 debug = true
 

benches/poseidon.rs

@@ -0,0 +1,255 @@
+use ff::Field;
+use halo2::{
+    circuit::{Layouter, SimpleFloorPlanner},
+    pasta::Fp,
+    plonk::{
+        create_proof, keygen_pk, keygen_vk, verify_proof, Advice, Circuit, Column,
+        ConstraintSystem, Error,
+    },
+    poly::commitment::Params,
+    transcript::{Blake2bRead, Blake2bWrite, Challenge255},
+};
+use pasta_curves::{pallas, vesta};
+
+use orchard::{
+    circuit::gadget::{
+        poseidon::{Hash, Pow5Chip, Pow5Config},
+        utilities::{CellValue, Var},
+    },
+    primitives::poseidon::{self, ConstantLength, Spec},
+};
+use std::convert::TryInto;
+use std::marker::PhantomData;
+
+use criterion::{criterion_group, criterion_main, Criterion};
+use rand::rngs::OsRng;
+
+#[derive(Clone, Copy)]
+struct HashCircuit<S, const WIDTH: usize, const RATE: usize>
+where
+    S: Spec<Fp, WIDTH, RATE> + Clone + Copy,
+{
+    message: Option<[Fp; RATE]>,
+    // For the purpose of this test, witness the result.
+    // TODO: Move this into an instance column.
+    output: Option<Fp>,
+    _spec: PhantomData<S>,
+}
+
+#[derive(Debug, Clone)]
+struct MyConfig<const WIDTH: usize, const RATE: usize> {
+    input: [Column<Advice>; RATE],
+    poseidon_config: Pow5Config<Fp, WIDTH, RATE>,
+}
+
+impl<S, const WIDTH: usize, const RATE: usize> Circuit<Fp> for HashCircuit<S, WIDTH, RATE>
+where
+    S: Spec<Fp, WIDTH, RATE> + Copy + Clone,
+{
+    type Config = MyConfig<WIDTH, RATE>;
+    type FloorPlanner = SimpleFloorPlanner;
+
+    fn without_witnesses(&self) -> Self {
+        Self {
+            message: None,
+            output: None,
+            _spec: PhantomData,
+        }
+    }
+
+    fn configure(meta: &mut ConstraintSystem<Fp>) -> Self::Config {
+        let state = (0..WIDTH).map(|_| meta.advice_column()).collect::<Vec<_>>();
+        let partial_sbox = meta.advice_column();
+
+        let rc_a = (0..WIDTH).map(|_| meta.fixed_column()).collect::<Vec<_>>();
+        let rc_b = (0..WIDTH).map(|_| meta.fixed_column()).collect::<Vec<_>>();
+
+        meta.enable_constant(rc_b[0]);
+
+        Self::Config {
+            input: state[..RATE].try_into().unwrap(),
+            poseidon_config: Pow5Chip::configure::<S>(
+                meta,
+                state.try_into().unwrap(),
+                partial_sbox,
+                rc_a.try_into().unwrap(),
+                rc_b.try_into().unwrap(),
+            ),
+        }
+    }
+
+    fn synthesize(
+        &self,
+        config: Self::Config,
+        mut layouter: impl Layouter<Fp>,
+    ) -> Result<(), Error> {
+        let chip = Pow5Chip::construct(config.poseidon_config.clone());
+
+        let message = layouter.assign_region(
+            || "load message",
+            |mut region| {
+                let message_word = |i: usize| {
+                    let value = self.message.map(|message_vals| message_vals[i]);
+                    let cell = region.assign_advice(
+                        || format!("load message_{}", i),
+                        config.input[i],
+                        0,
+                        || value.ok_or(Error::SynthesisError),
+                    )?;
+                    Ok(CellValue::new(cell, value))
+                };
+
+                let message: Result<Vec<_>, Error> = (0..RATE).map(message_word).collect();
+                Ok(message?.try_into().unwrap())
+            },
+        )?;
+
+        let hasher = Hash::<_, _, S, _, WIDTH, RATE>::init(
+            chip,
+            layouter.namespace(|| "init"),
+            ConstantLength::<RATE>,
+        )?;
+        let output = hasher.hash(layouter.namespace(|| "hash"), message)?;
+
+        layouter.assign_region(
+            || "constrain output",
+            |mut region| {
+                let expected_var = region.assign_advice(
+                    || "load output",
+                    config.input[0],
+                    0,
+                    || self.output.ok_or(Error::SynthesisError),
+                )?;
+                region.constrain_equal(output.cell(), expected_var)
+            },
+        )
+    }
+}
+
+#[derive(Debug, Clone, Copy)]
+struct MySpec<const WIDTH: usize, const RATE: usize>;
+
+impl Spec<Fp, 3, 2> for MySpec<3, 2> {
+    fn full_rounds() -> usize {
+        8
+    }
+
+    fn partial_rounds() -> usize {
+        56
+    }
+
+    fn sbox(val: Fp) -> Fp {
+        val.pow_vartime(&[5])
+    }
+
+    fn secure_mds() -> usize {
+        0
+    }
+}
+
+impl Spec<Fp, 9, 8> for MySpec<9, 8> {
+    fn full_rounds() -> usize {
+        8
+    }
+
+    fn partial_rounds() -> usize {
+        56
+    }
+
+    fn sbox(val: Fp) -> Fp {
+        val.pow_vartime(&[5])
+    }
+
+    fn secure_mds() -> usize {
+        0
+    }
+}
+
+impl Spec<Fp, 12, 11> for MySpec<12, 11> {
+    fn full_rounds() -> usize {
+        8
+    }
+
+    fn partial_rounds() -> usize {
+        56
+    }
+
+    fn sbox(val: Fp) -> Fp {
+        val.pow_vartime(&[5])
+    }
+
+    fn secure_mds() -> usize {
+        0
+    }
+}
+
+const K: u32 = 6;
+
+fn bench_poseidon<S, const WIDTH: usize, const RATE: usize>(name: &str, c: &mut Criterion)
+where
+    S: Spec<Fp, WIDTH, RATE> + Copy + Clone,
+{
+    // Initialize the polynomial commitment parameters
+    let params: Params<vesta::Affine> = Params::new(K);
+
+    let empty_circuit = HashCircuit::<S, WIDTH, RATE> {
+        message: None,
+        output: None,
+        _spec: PhantomData,
+    };
+
+    // Initialize the proving key
+    let vk = keygen_vk(&params, &empty_circuit).expect("keygen_vk should not fail");
+    let pk = keygen_pk(&params, vk, &empty_circuit).expect("keygen_pk should not fail");
+
+    let prover_name = name.to_string() + "-prover";
+    let verifier_name = name.to_string() + "-verifier";
+
+    let rng = OsRng;
+    let message = (0..RATE)
+        .map(|_| pallas::Base::random(rng))
+        .collect::<Vec<_>>()
+        .try_into()
+        .unwrap();
+    let output = poseidon::Hash::<_, S, _, WIDTH, RATE>::init(ConstantLength::<RATE>).hash(message);
+
+    let circuit = HashCircuit::<S, WIDTH, RATE> {
+        message: Some(message),
+        output: Some(output),
+        _spec: PhantomData,
+    };
+
+    c.bench_function(&prover_name, |b| {
+        b.iter(|| {
+            // Create a proof
+            let mut transcript = Blake2bWrite::<_, _, Challenge255<_>>::init(vec![]);
+            create_proof(&params, &pk, &[circuit], &[&[]], &mut transcript)
+                .expect("proof generation should not fail")
+        })
+    });
+
+    // Create a proof
+    let mut transcript = Blake2bWrite::<_, _, Challenge255<_>>::init(vec![]);
+    create_proof(&params, &pk, &[circuit], &[&[]], &mut transcript)
+        .expect("proof generation should not fail");
+    let proof = transcript.finalize();
+
+    c.bench_function(&verifier_name, |b| {
+        b.iter(|| {
+            let msm = params.empty_msm();
+            let mut transcript = Blake2bRead::<_, _, Challenge255<_>>::init(&proof[..]);
+            let guard = verify_proof(&params, pk.get_vk(), msm, &[&[]], &mut transcript).unwrap();
+            let msm = guard.clone().use_challenges();
+            assert!(msm.eval());
+        });
+    });
+}
+
+fn criterion_benchmark(c: &mut Criterion) {
+    bench_poseidon::<MySpec<3, 2>, 3, 2>("WIDTH = 3, RATE = 2", c);
+    bench_poseidon::<MySpec<9, 8>, 9, 8>("WIDTH = 9, RATE = 8", c);
+    bench_poseidon::<MySpec<12, 11>, 12, 11>("WIDTH = 12, RATE = 11", c);
+}
+
+criterion_group!(benches, criterion_benchmark);
+criterion_main!(benches);

src/circuit.rs

@@ -57,7 +57,7 @@ use std::convert::TryInto;
 
 use self::gadget::utilities::lookup_range_check::LookupRangeCheckConfig;
 
-pub(crate) mod gadget;
+pub mod gadget;
 
 /// Size of the Orchard circuit.
 const K: u32 = 11;

src/circuit/gadget.rs

@@ -1,3 +1,5 @@
+//! Gadgets used in the Orchard circuit.
+
 use pasta_curves::pallas;
 
 use ecc::chip::EccChip;
@@ -5,9 +7,9 @@ use poseidon::Pow5Chip as PoseidonChip;
 use sinsemilla::{chip::SinsemillaChip, merkle::chip::MerkleChip};
 
 pub(crate) mod ecc;
-pub(crate) mod poseidon;
+pub mod poseidon;
 pub(crate) mod sinsemilla;
-pub(crate) mod utilities;
+pub mod utilities;
 
 impl super::Config {
     pub(super) fn ecc_chip(&self) -> EccChip {

src/circuit/gadget/poseidon.rs

@@ -80,11 +80,13 @@ impl<
         const RATE: usize,
     > Word<F, PoseidonChip, S, T, RATE>
 {
-    pub(crate) fn inner(&self) -> PoseidonChip::Word {
+    /// The word contained in this gadget.
+    pub fn inner(&self) -> PoseidonChip::Word {
         self.inner
     }
 
-    pub(crate) fn from_inner(inner: PoseidonChip::Word) -> Self {
+    /// Construct a [`Word`] gadget from the inner word.
+    pub fn from_inner(inner: PoseidonChip::Word) -> Self {
         Self { inner }
     }
 }
@@ -109,6 +111,7 @@ fn poseidon_duplex<
 }
 
 /// A Poseidon duplex sponge.
+#[derive(Debug)]
 pub struct Duplex<
     F: FieldExt,
     PoseidonChip: PoseidonDuplexInstructions<F, S, T, RATE>,
@@ -210,6 +213,7 @@ impl<
 }
 
 /// A Poseidon hash function, built around a duplex sponge.
+#[derive(Debug)]
 pub struct Hash<
     F: FieldExt,
     PoseidonChip: PoseidonDuplexInstructions<F, S, T, RATE>,

src/circuit/gadget/poseidon/pow5.rs

@@ -199,6 +199,7 @@ impl<F: FieldExt, const WIDTH: usize, const RATE: usize> Pow5Chip<F, WIDTH, RATE
         }
     }
 
+    /// Construct a [`Pow5Chip`].
     pub fn construct(config: Pow5Config<F, WIDTH, RATE>) -> Self {
         Pow5Chip { config }
     }
@@ -401,18 +402,13 @@ impl<F: FieldExt, S: Spec<F, WIDTH, RATE>, const WIDTH: usize, const RATE: usize
     }
 }
 
+/// A word in the Poseidon state.
 #[derive(Clone, Copy, Debug)]
 pub struct StateWord<F: FieldExt> {
     var: Cell,
     value: Option<F>,
 }
 
-impl<F: FieldExt> StateWord<F> {
-    pub fn new(var: Cell, value: Option<F>) -> Self {
-        Self { var, value }
-    }
-}
-
 impl<F: FieldExt> From<StateWord<F>> for CellValue<F> {
     fn from(state_word: StateWord<F>) -> CellValue<F> {
         CellValue::new(state_word.var, state_word.value)
@@ -425,6 +421,20 @@ impl<F: FieldExt> From<CellValue<F>> for StateWord<F> {
     }
 }
 
+impl<F: FieldExt> Var<F> for StateWord<F> {
+    fn new(var: Cell, value: Option<F>) -> Self {
+        Self { var, value }
+    }
+
+    fn cell(&self) -> Cell {
+        self.var
+    }
+
+    fn value(&self) -> Option<F> {
+        self.value
+    }
+}
+
 #[derive(Debug)]
 struct Pow5State<F: FieldExt, const WIDTH: usize>([StateWord<F>; WIDTH]);
 

src/circuit/gadget/utilities.rs

@@ -1,3 +1,5 @@
+//! Utility gadgets.
+
 use ff::PrimeFieldBits;
 use halo2::{
     circuit::{Cell, Layouter, Region},
@@ -17,9 +19,15 @@ pub struct CellValue<F: FieldExt> {
     value: Option<F>,
 }
 
+/// Trait for a variable in the circuit.
 pub trait Var<F: FieldExt>: Copy + Clone + std::fmt::Debug {
+    /// Construct a new variable.
     fn new(cell: Cell, value: Option<F>) -> Self;
+
+    /// The cell at which this variable was allocated.
     fn cell(&self) -> Cell;
+
+    /// The value allocated to this variable.
     fn value(&self) -> Option<F>;
 }
 
@@ -37,9 +45,12 @@ impl<F: FieldExt> Var<F> for CellValue<F> {
     }
 }
 
+/// Trait for utilities used across circuits.
 pub trait UtilitiesInstructions<F: FieldExt> {
+    /// Variable in the circuit.
     type Var: Var<F>;
 
+    /// Load a variable.
     fn load_private(
         &self,
         mut layouter: impl Layouter<F>,
@@ -88,7 +99,7 @@ where
     Ok(CellValue::new(cell, copy.value))
 }
 
-pub fn transpose_option_array<T: Copy + std::fmt::Debug, const LEN: usize>(
+pub(crate) fn transpose_option_array<T: Copy + std::fmt::Debug, const LEN: usize>(
     option_array: Option<[T; LEN]>,
 ) -> [Option<T>; LEN] {
     let mut ret = [None; LEN];

src/primitives/poseidon.rs

@@ -141,6 +141,7 @@ fn poseidon_duplex<F: FieldExt, S: Spec<F, T, RATE>, const T: usize, const RATE:
     output
 }
 
+#[derive(Debug)]
 pub(crate) enum Sponge<F, const RATE: usize> {
     Absorbing(SpongeState<F, RATE>),
     Squeezing(SpongeState<F, RATE>),