Support optimized witness generation for circuit2.

Cargo.toml

@@ -10,7 +10,7 @@ resolver = "2"
 rust-version = "1.62.1"
 
 [dependencies]
-bellperson = { version = "0.25", default-features = false }
+bellperson = { git = "https://github.com/filecoin-project/bellperson", rev = "bd31feeb88565e66aaf0f2a28c15f4f0a7284710" }
 blake2s_simd = "0.5"
 blstrs = { version = "0.7.0", optional = true }
 byteorder = "1"

benches/synthesis.rs

@@ -1,52 +1,65 @@
 use crate::poseidon::{Arity, PoseidonConstants};
 use bellperson::gadgets::num::AllocatedNum;
 use bellperson::util_cs::bench_cs::BenchCS;
-use bellperson::{Circuit, ConstraintSystem, SynthesisError};
+use bellperson::util_cs::witness_cs::WitnessCS;
+use bellperson::{ConstraintSystem, SynthesisError};
 use blstrs::Scalar as Fr;
 use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
 use ff::Field;
 use generic_array::typenum;
 use neptune::circuit::{poseidon_hash_circuit, CircuitType};
+use neptune::circuit2::poseidon_hash_allocated_mut;
 use neptune::*;
 use rand::thread_rng;
 use std::marker::PhantomData;
 
 struct BenchCircuit<'a, A: Arity<Fr>> {
     n: usize,
-    circuit_type: &'a CircuitType,
+    circuit_type: Option<&'a CircuitType>,
     _a: PhantomData<A>,
 }
 
-impl<A: Arity<Fr>> Circuit<Fr> for BenchCircuit<'_, A> {
-    fn synthesize<CS: ConstraintSystem<Fr>>(self, mut cs: &mut CS) -> Result<(), SynthesisError> {
-        let mut rng = thread_rng();
-        let arity = A::to_usize();
-        let constants = PoseidonConstants::<Fr, A>::new();
-
+impl<A: Arity<Fr>> BenchCircuit<'_, A> {
+    fn synthesize<CS: ConstraintSystem<Fr>>(
+        self,
+        cs: &mut CS,
+        data: Vec<AllocatedNum<Fr>>,
+        constants: &PoseidonConstants<Fr, A>,
+    ) -> Result<(), SynthesisError> {
         for _ in 0..self.n {
-            let mut i = 0;
-            let mut fr_data = vec![Fr::random(&mut rng); arity];
-            let data: Vec<AllocatedNum<Fr>> = (0..arity)
-                .enumerate()
-                .map(|_| {
-                    let fr = Fr::random(&mut rng);
-                    fr_data[i] = fr;
-                    i += 1;
-                    AllocatedNum::alloc(cs.namespace(|| format!("data {}", i)), || Ok(fr)).unwrap()
-                })
-                .collect::<Vec<_>>();
-            let _ = poseidon_hash_circuit(&mut cs, *self.circuit_type, data, &constants)
-                .expect("poseidon hashing failed");
+            let _ = if self.circuit_type.is_some() {
+                poseidon_hash_circuit(cs, *self.circuit_type.unwrap(), data.clone(), constants)
+                    .expect("poseidon hashing failed");
+            } else {
+                poseidon_hash_allocated_mut(cs, data.clone(), constants)
+                    .expect("poseidon hashing failed");
+            };
         }
         Ok(())
     }
+
+    fn data<CS: ConstraintSystem<Fr>>(cs: &mut CS) -> Vec<AllocatedNum<Fr>> {
+        let mut rng = thread_rng();
+        let arity = A::to_usize();
+
+        let mut fr_data = vec![Fr::random(&mut rng); arity];
+        let data: Vec<AllocatedNum<Fr>> = (0..arity)
+            .map(|i| {
+                let fr = Fr::random(&mut rng);
+                fr_data[i] = fr;
+                AllocatedNum::alloc(cs.namespace(|| format!("data {}", i)), || Ok(fr)).unwrap()
+            })
+            .collect::<Vec<_>>();
+        data
+    }
 }
 
 fn bench_synthesis<A>(c: &mut Criterion)
 where
     A: Arity<Fr>,
 {
     let mut group = c.benchmark_group(format!("synthesis-{}", A::to_usize()));
+    let constants = PoseidonConstants::<Fr, A>::new();
     for i in 0..4 {
         let num_hashes = 10usize.pow(i);
         for circuit_type in &[CircuitType::Legacy, CircuitType::OptimalAllocated] {
@@ -57,19 +70,39 @@ where
                 ),
                 &num_hashes,
                 |b, n| {
+                    let mut cs = BenchCS::<Fr>::new();
+                    let data = BenchCircuit::<A>::data(&mut cs);
                     b.iter(|| {
-                        let mut cs = BenchCS::<Fr>::new();
                         let circuit = BenchCircuit::<A> {
                             n: *n,
-                            circuit_type,
+                            circuit_type: Some(circuit_type),
                             _a: PhantomData::<A>,
                         };
-                        circuit.synthesize(&mut cs)
+                        circuit.synthesize(&mut cs, data.clone(), &constants)
                     })
                 },
             );
         }
-        // num_hashes *= 10;
+
+        group.bench_with_input(
+            BenchmarkId::new(
+                "hash_allocated_witness",
+                format!("arity: {}, count: {}", A::to_usize(), num_hashes),
+            ),
+            &num_hashes,
+            |b, n| {
+                let mut cs = WitnessCS::<Fr>::new();
+                let data = BenchCircuit::<A>::data(&mut cs);
+                b.iter(|| {
+                    let circuit = BenchCircuit::<A> {
+                        n: *n,
+                        circuit_type: None,
+                        _a: PhantomData::<A>,
+                    };
+                    circuit.synthesize(&mut cs, data.clone(), &constants)
+                })
+            },
+        );
     }
 }
 

src/circuit.rs

@@ -2,6 +2,7 @@
 use std::ops::{AddAssign, MulAssign};
 
 use crate::circuit2::poseidon_hash_allocated;
+use crate::circuit2_witness::poseidon_hash_allocated_witness;
 use crate::hash_type::HashType;
 use crate::matrix::Matrix;
 use crate::mds::SparseMatrix;
@@ -30,7 +31,7 @@ impl CircuitType {
 
 /// Convenience function to potentially ease upgrade transition from legacy to optimal circuits.
 pub fn poseidon_hash_circuit<CS, Scalar, A>(
-    cs: CS,
+    cs: &mut CS,
     circuit_type: CircuitType,
     preimage: Vec<AllocatedNum<Scalar>>,
     constants: &PoseidonConstants<Scalar, A>,

src/circuit2.rs

@@ -1,6 +1,7 @@
 /// The `circuit2` module implements the optimal Poseidon hash circuit.
 use std::ops::{AddAssign, MulAssign};
 
+use crate::circuit2_witness::poseidon_hash_allocated_witness;
 use crate::hash_type::HashType;
 use crate::matrix::Matrix;
 use crate::mds::SparseMatrix;
@@ -434,9 +435,9 @@ where
     }
 }
 
-/// Create circuit for Poseidon hash, returning an unallocated `Num` at the cost of one constraint.
+/// Create circuit for Poseidon hash, returning an allocated `Num` at the cost of one constraint.
 pub fn poseidon_hash_allocated<CS, Scalar, A>(
-    cs: CS,
+    mut cs: CS,
     preimage: Vec<AllocatedNum<Scalar>>,
     constants: &PoseidonConstants<Scalar, A>,
 ) -> Result<AllocatedNum<Scalar>, SynthesisError>
@@ -445,24 +446,41 @@ where
     Scalar: PrimeField,
     A: Arity<Scalar>,
 {
-    let arity = A::to_usize();
-    let tag_element = Elt::num_from_fr::<CS>(constants.domain_tag);
-    let mut elements = Vec::with_capacity(arity + 1);
-    elements.push(tag_element);
-    elements.extend(preimage.into_iter().map(Elt::Allocated));
+    poseidon_hash_allocated_mut(&mut cs, preimage, constants)
+}
 
-    if let HashType::ConstantLength(length) = constants.hash_type {
-        assert!(length <= arity, "illegal length: constants are malformed");
-        // Add zero-padding.
-        for _ in 0..(arity - length) {
-            let elt = Elt::Num(num::Num::zero());
-            elements.push(elt);
+/// Create circuit for Poseidon hash, returning an allocated `Num` at the cost of one constraint.
+pub fn poseidon_hash_allocated_mut<CS, Scalar, A>(
+    cs: &mut CS,
+    preimage: Vec<AllocatedNum<Scalar>>,
+    constants: &PoseidonConstants<Scalar, A>,
+) -> Result<AllocatedNum<Scalar>, SynthesisError>
+where
+    CS: ConstraintSystem<Scalar>,
+    Scalar: PrimeField,
+    A: Arity<Scalar>,
+{
+    if cs.is_witness_generator() {
+        poseidon_hash_allocated_witness(cs, &preimage, constants)
+    } else {
+        let arity = A::to_usize();
+        let tag_element = Elt::num_from_fr::<CS>(constants.domain_tag);
+        let mut elements = Vec::with_capacity(arity + 1);
+        elements.push(tag_element);
+        elements.extend(preimage.into_iter().map(Elt::Allocated));
+
+        if let HashType::ConstantLength(length) = constants.hash_type {
+            assert!(length <= arity, "illegal length: constants are malformed");
+            // Add zero-padding.
+            for _ in 0..(arity - length) {
+                let elt = Elt::Num(num::Num::zero());
+                elements.push(elt);
+            }
         }
-    }
+        let mut p = PoseidonCircuit2::new(elements, constants);
 
-    let mut p = PoseidonCircuit2::new(elements, constants);
-
-    p.hash_to_allocated(cs)
+        p.hash_to_allocated(cs)
+    }
 }
 
 /// Create circuit for Poseidon hash, minimizing constraints by returning an unallocated `Num`.
@@ -557,7 +575,6 @@ where
         let mut tmp = elt.val().ok_or(SynthesisError::AssignmentMissing)?;
         tmp.add_assign(&to_add);
         tmp = tmp.square();
-
         Ok(tmp)
     })?;