Merge pull request #236 from zcash/bench-poseidon-2

Benchmark Poseidon gadget for rates {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,258 @@
+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, const L: usize>
+where
+    S: Spec<Fp, WIDTH, RATE> + Clone + Copy,
+{
+    message: Option<[Fp; L]>,
+    // 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, const L: usize> {
+    input: [Column<Advice>; L],
+    poseidon_config: Pow5Config<Fp, WIDTH, RATE>,
+}
+
+impl<S, const WIDTH: usize, const RATE: usize, const L: usize> Circuit<Fp>
+    for HashCircuit<S, WIDTH, RATE, L>
+where
+    S: Spec<Fp, WIDTH, RATE> + Copy + Clone,
+{
+    type Config = MyConfig<WIDTH, RATE, L>;
+    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::Synthesis),
+                    )?;
+                    Ok(CellValue::new(cell, value))
+                };
+
+                let message: Result<Vec<_>, Error> = (0..L).map(message_word).collect();
+                Ok(message?.try_into().unwrap())
+            },
+        )?;
+
+        let hasher = Hash::<_, _, S, _, WIDTH, RATE>::init(
+            chip,
+            layouter.namespace(|| "init"),
+            ConstantLength::<L>,
+        )?;
+        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::Synthesis),
+                )?;
+                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, const L: 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, L> {
+        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..L)
+        .map(|_| pallas::Base::random(rng))
+        .collect::<Vec<_>>()
+        .try_into()
+        .unwrap();
+    let output = poseidon::Hash::<_, S, _, WIDTH, RATE>::init(ConstantLength::<L>).hash(message);
+
+    let circuit = HashCircuit::<S, WIDTH, RATE, L> {
+        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, 2>("WIDTH = 3, RATE = 2", c);
+    bench_poseidon::<MySpec<9, 8>, 9, 8, 8>("WIDTH = 9, RATE = 8", c);
+    bench_poseidon::<MySpec<12, 11>, 12, 11, 11>("WIDTH = 12, RATE = 11", c);
+}
+
+criterion_group!(benches, criterion_benchmark);
+criterion_main!(benches);

benches/primitives.rs

@@ -21,7 +21,7 @@ fn bench_primitives(c: &mut Criterion) {
         let message = [pallas::Base::random(rng), pallas::Base::random(rng)];
 
         group.bench_function("2-to-1", |b| {
-            b.iter(|| poseidon::Hash::init(P128Pow5T3, ConstantLength).hash(message))
+            b.iter(|| poseidon::Hash::<_, P128Pow5T3, _, 3, 2>::init(ConstantLength).hash(message))
         });
     }
 

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;
@@ -239,9 +239,8 @@ impl plonk::Circuit<pallas::Base> for Circuit {
         let ecc_config = EccChip::configure(meta, advices, lagrange_coeffs, range_check.clone());
 
         // Configuration for the Poseidon hash.
-        let poseidon_config = PoseidonChip::configure(
+        let poseidon_config = PoseidonChip::configure::<poseidon::P128Pow5T3>(
             meta,
-            poseidon::P128Pow5T3,
             // We place the state columns after the partial_sbox column so that the
             // pad-and-add region can be layed out more efficiently.
             advices[6..9].try_into().unwrap(),

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

@@ -52,7 +52,6 @@ impl<F: FieldExt, const WIDTH: usize, const RATE: usize> Pow5Chip<F, WIDTH, RATE
     // necessary for the permutation.
     pub fn configure<S: Spec<F, WIDTH, RATE>>(
         meta: &mut ConstraintSystem<F>,
-        spec: S,
         state: [Column<Advice>; WIDTH],
         partial_sbox: Column<Advice>,
         rc_a: [Column<Fixed>; WIDTH],
@@ -65,7 +64,7 @@ impl<F: FieldExt, const WIDTH: usize, const RATE: usize> Pow5Chip<F, WIDTH, RATE
         assert!(S::partial_rounds() & 1 == 0);
         let half_full_rounds = S::full_rounds() / 2;
         let half_partial_rounds = S::partial_rounds() / 2;
-        let (round_constants, m_reg, m_inv) = spec.constants();
+        let (round_constants, m_reg, m_inv) = S::constants();
 
         // This allows state words to be initialized (by constraining them equal to fixed
         // values), and used in a permutation from an arbitrary region. rc_a is used in
@@ -200,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 }
     }
@@ -402,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)
@@ -426,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]);
 
@@ -616,40 +625,37 @@ mod tests {
         },
         primitives::poseidon::{self, ConstantLength, P128Pow5T3 as OrchardNullifier, Spec},
     };
+    use std::convert::TryInto;
+    use std::marker::PhantomData;
 
-    const WIDTH: usize = 3;
-    const RATE: usize = 2;
+    struct PermuteCircuit<S: Spec<Fp, WIDTH, RATE>, const WIDTH: usize, const RATE: usize>(
+        PhantomData<S>,
+    );
 
-    struct PermuteCircuit {}
-
-    impl Circuit<Fp> for PermuteCircuit {
+    impl<S: Spec<Fp, WIDTH, RATE>, const WIDTH: usize, const RATE: usize> Circuit<Fp>
+        for PermuteCircuit<S, WIDTH, RATE>
+    {
         type Config = Pow5Config<Fp, WIDTH, RATE>;
         type FloorPlanner = SimpleFloorPlanner;
 
         fn without_witnesses(&self) -> Self {
-            PermuteCircuit {}
+            PermuteCircuit::<S, WIDTH, RATE>(PhantomData)
         }
 
         fn configure(meta: &mut ConstraintSystem<Fp>) -> Pow5Config<Fp, WIDTH, RATE> {
-            let state = [
-                meta.advice_column(),
-                meta.advice_column(),
-                meta.advice_column(),
-            ];
+            let state = (0..WIDTH).map(|_| meta.advice_column()).collect::<Vec<_>>();
             let partial_sbox = meta.advice_column();
 
-            let rc_a = [
-                meta.fixed_column(),
-                meta.fixed_column(),
-                meta.fixed_column(),
-            ];
-            let rc_b = [
-                meta.fixed_column(),
-                meta.fixed_column(),
-                meta.fixed_column(),
-            ];
+            let rc_a = (0..WIDTH).map(|_| meta.fixed_column()).collect::<Vec<_>>();
+            let rc_b = (0..WIDTH).map(|_| meta.fixed_column()).collect::<Vec<_>>();
 
-            Pow5Chip::configure(meta, OrchardNullifier, state, partial_sbox, rc_a, rc_b)
+            Pow5Chip::configure::<S>(
+                meta,
+                state.try_into().unwrap(),
+                partial_sbox,
+                rc_a.try_into().unwrap(),
+                rc_b.try_into().unwrap(),
+            )
         }
 
         fn synthesize(
@@ -660,7 +666,7 @@ mod tests {
             let initial_state = layouter.assign_region(
                 || "prepare initial state",
                 |mut region| {
-                    let mut state_word = |i: usize| -> Result<_, Error> {
+                    let state_word = |i: usize| {
                         let value = Some(Fp::from(i as u64));
                         let var = region.assign_advice(
                             || format!("load state_{}", i),
@@ -671,22 +677,27 @@ mod tests {
                         Ok(StateWord { var, value })
                     };
 
-                    Ok([state_word(0)?, state_word(1)?, state_word(2)?])
+                    let state: Result<Vec<_>, Error> = (0..WIDTH).map(state_word).collect();
+                    Ok(state?.try_into().unwrap())
                 },
             )?;
 
             let chip = Pow5Chip::construct(config.clone());
             let final_state = <Pow5Chip<_, WIDTH, RATE> as PoseidonInstructions<
                 Fp,
-                OrchardNullifier,
+                S,
                 WIDTH,
-                2,
+                RATE,
             >>::permute(&chip, &mut layouter, &initial_state)?;
 
             // For the purpose of this test, compute the real final state inline.
-            let mut expected_final_state = [Fp::zero(), Fp::one(), Fp::from_u64(2)];
-            let (round_constants, mds, _) = OrchardNullifier.constants();
-            poseidon::permute::<_, OrchardNullifier, WIDTH, RATE>(
+            let mut expected_final_state = (0..WIDTH)
+                .map(|idx| Fp::from_u64(idx as u64))
+                .collect::<Vec<_>>()
+                .try_into()
+                .unwrap();
+            let (round_constants, mds, _) = S::constants();
+            poseidon::permute::<_, S, WIDTH, RATE>(
                 &mut expected_final_state,
                 &mds,
                 &round_constants,
@@ -705,9 +716,11 @@ mod tests {
                         region.constrain_equal(final_state[i].var, var)
                     };
 
-                    final_state_word(0)?;
-                    final_state_word(1)?;
-                    final_state_word(2)
+                    for i in 0..(WIDTH) {
+                        final_state_word(i)?;
+                    }
+
+                    Ok(())
                 },
             )
         }
@@ -716,49 +729,54 @@ mod tests {
     #[test]
     fn poseidon_permute() {
         let k = 6;
-        let circuit = PermuteCircuit {};
+        let circuit = PermuteCircuit::<OrchardNullifier, 3, 2>(PhantomData);
         let prover = MockProver::run(k, &circuit, vec![]).unwrap();
         assert_eq!(prover.verify(), Ok(()))
     }
 
-    #[derive(Default)]
-    struct HashCircuit {
-        message: Option<[Fp; 2]>,
+    struct HashCircuit<
+        S: Spec<Fp, WIDTH, RATE>,
+        const WIDTH: usize,
+        const RATE: usize,
+        const L: usize,
+    > {
+        message: Option<[Fp; L]>,
         // For the purpose of this test, witness the result.
         // TODO: Move this into an instance column.
         output: Option<Fp>,
+        _spec: PhantomData<S>,
     }
 
-    impl Circuit<Fp> for HashCircuit {
+    impl<S: Spec<Fp, WIDTH, RATE>, const WIDTH: usize, const RATE: usize, const L: usize>
+        Circuit<Fp> for HashCircuit<S, WIDTH, RATE, L>
+    {
         type Config = Pow5Config<Fp, WIDTH, RATE>;
         type FloorPlanner = SimpleFloorPlanner;
 
         fn without_witnesses(&self) -> Self {
-            Self::default()
+            Self {
+                message: None,
+                output: None,
+                _spec: PhantomData,
+            }
         }
 
         fn configure(meta: &mut ConstraintSystem<Fp>) -> Pow5Config<Fp, WIDTH, RATE> {
-            let state = [
-                meta.advice_column(),
-                meta.advice_column(),
-                meta.advice_column(),
-            ];
+            let state = (0..WIDTH).map(|_| meta.advice_column()).collect::<Vec<_>>();
             let partial_sbox = meta.advice_column();
 
-            let rc_a = [
-                meta.fixed_column(),
-                meta.fixed_column(),
-                meta.fixed_column(),
-            ];
-            let rc_b = [
-                meta.fixed_column(),
-                meta.fixed_column(),
-                meta.fixed_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]);
 
-            Pow5Chip::configure(meta, OrchardNullifier, state, partial_sbox, rc_a, rc_b)
+            Pow5Chip::configure::<S>(
+                meta,
+                state.try_into().unwrap(),
+                partial_sbox,
+                rc_a.try_into().unwrap(),
+                rc_b.try_into().unwrap(),
+            )
         }
 
         fn synthesize(
@@ -771,7 +789,7 @@ mod tests {
             let message = layouter.assign_region(
                 || "load message",
                 |mut region| {
-                    let mut message_word = |i: usize| -> Result<_, Error> {
+                    let message_word = |i: usize| {
                         let value = self.message.map(|message_vals| message_vals[i]);
                         let cell = region.assign_advice(
                             || format!("load message_{}", i),
@@ -782,14 +800,15 @@ mod tests {
                         Ok(CellValue::new(cell, value))
                     };
 
-                    Ok([message_word(0)?, message_word(1)?])
+                    let message: Result<Vec<_>, Error> = (0..L).map(message_word).collect();
+                    Ok(message?.try_into().unwrap())
                 },
             )?;
 
-            let hasher = Hash::<_, _, OrchardNullifier, _, WIDTH, 2>::init(
+            let hasher = Hash::<_, _, S, _, WIDTH, RATE>::init(
                 chip,
                 layouter.namespace(|| "init"),
-                ConstantLength::<2>,
+                ConstantLength::<L>,
             )?;
             let output = hasher.hash(layouter.namespace(|| "hash"), message)?;
 
@@ -811,12 +830,14 @@ mod tests {
     #[test]
     fn poseidon_hash() {
         let message = [Fp::rand(), Fp::rand()];
-        let output = poseidon::Hash::init(OrchardNullifier, ConstantLength::<2>).hash(message);
+        let output =
+            poseidon::Hash::<_, OrchardNullifier, _, 3, 2>::init(ConstantLength::<2>).hash(message);
 
         let k = 6;
-        let circuit = HashCircuit {
+        let circuit = HashCircuit::<OrchardNullifier, 3, 2, 2> {
             message: Some(message),
             output: Some(output),
+            _spec: PhantomData,
         };
         let prover = MockProver::run(k, &circuit, vec![]).unwrap();
         assert_eq!(prover.verify(), Ok(()))
@@ -829,12 +850,14 @@ mod tests {
                 pallas::Base::from_repr(tv.input[0]).unwrap(),
                 pallas::Base::from_repr(tv.input[1]).unwrap(),
             ];
-            let output = poseidon::Hash::init(OrchardNullifier, ConstantLength).hash(message);
+            let output =
+                poseidon::Hash::<_, OrchardNullifier, _, 3, 2>::init(ConstantLength).hash(message);
 
             let k = 6;
-            let circuit = HashCircuit {
+            let circuit = HashCircuit::<OrchardNullifier, 3, 2, 2> {
                 message: Some(message),
                 output: Some(output),
+                _spec: PhantomData,
             };
             let prover = MockProver::run(k, &circuit, vec![]).unwrap();
             assert_eq!(prover.verify(), Ok(()));
@@ -852,9 +875,10 @@ mod tests {
             .titled("Poseidon Chip Layout", ("sans-serif", 60))
             .unwrap();
 
-        let circuit = HashCircuit {
+        let circuit = HashCircuit::<OrchardNullifier, 3, 2, 2> {
             message: None,
             output: None,
+            _spec: PhantomData,
         };
         halo2::dev::CircuitLayout::default()
             .render(6, &circuit, &root)

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

@@ -30,7 +30,7 @@ pub(crate) type SpongeState<F, const RATE: usize> = [Option<F>; RATE];
 pub(crate) type Mds<F, const T: usize> = [[F; T]; T];
 
 /// A specification for a Poseidon permutation.
-pub trait Spec<F: FieldExt, const T: usize, const RATE: usize> {
+pub trait Spec<F: FieldExt, const T: usize, const RATE: usize>: fmt::Debug {
     /// The number of full rounds for this specification.
     ///
     /// This must be an even number.
@@ -47,10 +47,10 @@ pub trait Spec<F: FieldExt, const T: usize, const RATE: usize> {
     ///
     /// This is used by the default implementation of [`Spec::constants`]. If you are
     /// hard-coding the constants, you may leave this unimplemented.
-    fn secure_mds(&self) -> usize;
+    fn secure_mds() -> usize;
 
     /// Generates `(round_constants, mds, mds^-1)` corresponding to this specification.
-    fn constants(&self) -> (Vec<[F; T]>, Mds<F, T>, Mds<F, T>) {
+    fn constants() -> (Vec<[F; T]>, Mds<F, T>, Mds<F, T>) {
         let r_f = Self::full_rounds();
         let r_p = Self::partial_rounds();
 
@@ -69,7 +69,7 @@ pub trait Spec<F: FieldExt, const T: usize, const RATE: usize> {
             })
             .collect();
 
-        let (mds, mds_inv) = mds::generate_mds::<F, T>(&mut grain, self.secure_mds());
+        let (mds, mds_inv) = mds::generate_mds::<F, T>(&mut grain, Self::secure_mds());
 
         (round_constants, mds, mds_inv)
     }
@@ -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>),
@@ -167,11 +168,10 @@ pub(crate) struct Duplex<F: FieldExt, S: Spec<F, T, RATE>, const T: usize, const
 impl<F: FieldExt, S: Spec<F, T, RATE>, const T: usize, const RATE: usize> Duplex<F, S, T, RATE> {
     /// Constructs a new duplex sponge for the given Poseidon specification.
     pub(crate) fn new(
-        spec: S,
         initial_capacity_element: F,
         pad_and_add: Box<dyn Fn(&mut State<F, T>, &SpongeState<F, RATE>)>,
     ) -> Self {
-        let (round_constants, mds_matrix, _) = spec.constants();
+        let (round_constants, mds_matrix, _) = S::constants();
 
         let input = [None; RATE];
         let mut state = [F::zero(); T];
@@ -336,13 +336,9 @@ impl<
     > Hash<F, S, D, T, RATE>
 {
     /// Initializes a new hasher.
-    pub fn init(spec: S, domain: D) -> Self {
+    pub fn init(domain: D) -> Self {
         Hash {
-            duplex: Duplex::new(
-                spec,
-                domain.initial_capacity_element(),
-                domain.pad_and_add(),
-            ),
+            duplex: Duplex::new(domain.initial_capacity_element(), domain.pad_and_add()),
             domain,
         }
     }
@@ -371,9 +367,9 @@ mod tests {
     fn orchard_spec_equivalence() {
         let message = [pallas::Base::from_u64(6), pallas::Base::from_u64(42)];
 
-        let (round_constants, mds, _) = OrchardNullifier.constants();
+        let (round_constants, mds, _) = OrchardNullifier::constants();
 
-        let hasher = Hash::init(OrchardNullifier, ConstantLength);
+        let hasher = Hash::<_, OrchardNullifier, _, 3, 2>::init(ConstantLength);
         let result = hasher.hash(message);
 
         // The result should be equivalent to just directly applying the permutation and

src/primitives/poseidon/p128pow5t3.rs

@@ -25,11 +25,11 @@ impl Spec<Fp, 3, 2> for P128Pow5T3 {
         val.pow_vartime(&[5])
     }
 
-    fn secure_mds(&self) -> usize {
+    fn secure_mds() -> usize {
         unimplemented!()
     }
 
-    fn constants(&self) -> (Vec<[Fp; 3]>, Mds<Fp, 3>, Mds<Fp, 3>) {
+    fn constants() -> (Vec<[Fp; 3]>, Mds<Fp, 3>, Mds<Fp, 3>) {
         (
             super::fp::ROUND_CONSTANTS[..].to_vec(),
             super::fp::MDS,
@@ -51,11 +51,11 @@ impl Spec<Fq, 3, 2> for P128Pow5T3 {
         val.pow_vartime(&[5])
     }
 
-    fn secure_mds(&self) -> usize {
+    fn secure_mds() -> usize {
         unimplemented!()
     }
 
-    fn constants(&self) -> (Vec<[Fq; 3]>, Mds<Fq, 3>, Mds<Fq, 3>) {
+    fn constants() -> (Vec<[Fq; 3]>, Mds<Fq, 3>, Mds<Fq, 3>) {
         (
             super::fq::ROUND_CONSTANTS[..].to_vec(),
             super::fq::MDS,
@@ -80,21 +80,15 @@ mod tests {
     /// The same Poseidon specification as poseidon::P128Pow5T3, but constructed
     /// such that its constants will be generated at runtime.
     #[derive(Debug)]
-    pub struct P128Pow5T3Gen<F: FieldExt> {
-        secure_mds: usize,
-        _field: PhantomData<F>,
-    }
+    pub struct P128Pow5T3Gen<F: FieldExt, const SECURE_MDS: usize>(PhantomData<F>);
 
-    impl<F: FieldExt> P128Pow5T3Gen<F> {
-        pub fn new(secure_mds: usize) -> Self {
-            P128Pow5T3Gen {
-                secure_mds,
-                _field: PhantomData::default(),
-            }
+    impl<F: FieldExt, const SECURE_MDS: usize> P128Pow5T3Gen<F, SECURE_MDS> {
+        pub fn new() -> Self {
+            P128Pow5T3Gen(PhantomData::default())
         }
     }
 
-    impl<F: FieldExt> Spec<F, 3, 2> for P128Pow5T3Gen<F> {
+    impl<F: FieldExt, const SECURE_MDS: usize> Spec<F, 3, 2> for P128Pow5T3Gen<F, SECURE_MDS> {
         fn full_rounds() -> usize {
             8
         }
@@ -107,8 +101,8 @@ mod tests {
             val.pow_vartime(&[5])
         }
 
-        fn secure_mds(&self) -> usize {
-            self.secure_mds
+        fn secure_mds() -> usize {
+            SECURE_MDS
         }
     }
 
@@ -119,8 +113,7 @@ mod tests {
             expected_mds: [[F; 3]; 3],
             expected_mds_inv: [[F; 3]; 3],
         ) {
-            let poseidon = P128Pow5T3Gen::<F>::new(0);
-            let (round_constants, mds, mds_inv) = poseidon.constants();
+            let (round_constants, mds, mds_inv) = P128Pow5T3Gen::<F, 0>::constants();
 
             for (actual, expected) in round_constants
                 .iter()
@@ -196,7 +189,7 @@ mod tests {
                 ]),
             ];
 
-            permute::<Fp, P128Pow5T3Gen<Fp>, 3, 2>(&mut input, &fp::MDS, &fp::ROUND_CONSTANTS);
+            permute::<Fp, P128Pow5T3Gen<Fp, 0>, 3, 2>(&mut input, &fp::MDS, &fp::ROUND_CONSTANTS);
             assert_eq!(input, expected_output);
         }
 
@@ -247,7 +240,7 @@ mod tests {
                 ]),
             ];
 
-            permute::<Fq, P128Pow5T3Gen<Fq>, 3, 2>(&mut input, &fq::MDS, &fq::ROUND_CONSTANTS);
+            permute::<Fq, P128Pow5T3Gen<Fq, 0>, 3, 2>(&mut input, &fq::MDS, &fq::ROUND_CONSTANTS);
             assert_eq!(input, expected_output);
         }
     }
@@ -255,7 +248,7 @@ mod tests {
     #[test]
     fn permute_test_vectors() {
         {
-            let (round_constants, mds, _) = super::P128Pow5T3.constants();
+            let (round_constants, mds, _) = super::P128Pow5T3::constants();
 
             for tv in crate::primitives::poseidon::test_vectors::fp::permute() {
                 let mut state = [
@@ -273,7 +266,7 @@ mod tests {
         }
 
         {
-            let (round_constants, mds, _) = super::P128Pow5T3.constants();
+            let (round_constants, mds, _) = super::P128Pow5T3::constants();
 
             for tv in crate::primitives::poseidon::test_vectors::fq::permute() {
                 let mut state = [
@@ -299,7 +292,7 @@ mod tests {
                 Fp::from_repr(tv.input[1]).unwrap(),
             ];
 
-            let result = Hash::init(super::P128Pow5T3, ConstantLength).hash(message);
+            let result = Hash::<_, super::P128Pow5T3, _, 3, 2>::init(ConstantLength).hash(message);
 
             assert_eq!(result.to_repr(), tv.output);
         }
@@ -310,7 +303,7 @@ mod tests {
                 Fq::from_repr(tv.input[1]).unwrap(),
             ];
 
-            let result = Hash::init(super::P128Pow5T3, ConstantLength).hash(message);
+            let result = Hash::<_, super::P128Pow5T3, _, 3, 2>::init(ConstantLength).hash(message);
 
             assert_eq!(result.to_repr(), tv.output);
         }

src/spec.rs

@@ -212,7 +212,8 @@ pub(crate) fn diversify_hash(d: &[u8; 11]) -> NonIdentityPallasPoint {
 ///
 /// [concreteprfs]: https://zips.z.cash/protocol/nu5.pdf#concreteprfs
 pub(crate) fn prf_nf(nk: pallas::Base, rho: pallas::Base) -> pallas::Base {
-    poseidon::Hash::init(poseidon::P128Pow5T3, poseidon::ConstantLength).hash([nk, rho])
+    poseidon::Hash::<_, poseidon::P128Pow5T3, _, 3, 2>::init(poseidon::ConstantLength)
+        .hash([nk, rho])
 }
 
 /// Defined in [Zcash Protocol Spec § 5.4.5.5: Orchard Key Agreement][concreteorchardkeyagreement].