1/// This example does the hash chain but in the naive approach: instead of using folding, it does a
  2/// big circuit containing n instantiations of the Poseidon constraints.
  3
  4#[cfg(test)]
  5mod tests {
  6    use ark_bn254::{Bn254, Fr};
  7
  8    use ark_groth16::Groth16;
  9    use ark_snark::SNARK;
 10
 11    use ark_ff::PrimeField;
 12
 13    use std::time::Instant;
 14
 15    use ark_crypto_primitives::sponge::{
 16        constraints::CryptographicSpongeVar,
 17        poseidon::{constraints::PoseidonSpongeVar, PoseidonConfig, PoseidonSponge},
 18        Absorb, CryptographicSponge,
 19    };
 20    use ark_r1cs_std::fields::fp::FpVar;
 21    use ark_r1cs_std::{alloc::AllocVar, eq::EqGadget};
 22    use ark_r1cs_std::{bits::uint8::UInt8, boolean::Boolean, ToBitsGadget, ToBytesGadget};
 23    use ark_relations::r1cs::{
 24        ConstraintSynthesizer, ConstraintSystem, ConstraintSystemRef, SynthesisError,
 25    };
 26
 27    use folding_schemes::transcript::poseidon::poseidon_canonical_config;
 28
 29    use crate::utils::tests::*;
 30
 31    /// Test circuit to be folded
 32    #[derive(Clone, Debug)]
 33    pub struct PoseidonChainCircuit<F: PrimeField, const N: usize, const HASHES_PER_STEP: usize> {
 34        z_0: Option<Vec<F>>,
 35        z_n: Option<Vec<F>>,
 36        config: PoseidonConfig<F>,
 37    }
 38    impl<F: PrimeField, const N: usize, const HASHES_PER_STEP: usize> ConstraintSynthesizer<F>
 39        for PoseidonChainCircuit<F, N, HASHES_PER_STEP>
 40    {
 41        fn generate_constraints(self, cs: ConstraintSystemRef<F>) -> Result<(), SynthesisError> {
 42            let z_0 = Vec::<FpVar<F>>::new_witness(cs.clone(), || {
 43                Ok(self.z_0.unwrap_or(vec![F::zero()]))
 44            })?;
 45            let z_n =
 46                Vec::<FpVar<F>>::new_input(cs.clone(), || Ok(self.z_n.unwrap_or(vec![F::zero()])))?;
 47
 48            let mut sponge = PoseidonSpongeVar::<F>::new(cs.clone(), &self.config);
 49            let mut z_i: Vec<FpVar<F>> = z_0.clone();
 50            for _ in 0..N {
 51                for _ in 0..HASHES_PER_STEP {
 52                    sponge.absorb(&z_i)?;
 53                    z_i = sponge.squeeze_field_elements(1)?;
 54                }
 55            }
 56
 57            z_i.enforce_equal(&z_n)?;
 58            Ok(())
 59        }
 60    }
 61    // compute natively in rust the expected result
 62    fn rust_native_result(
 63        poseidon_config: &PoseidonConfig<Fr>,
 64        z_0: Vec<Fr>,
 65        n_steps: usize,
 66        hashes_per_step: usize,
 67    ) -> Vec<Fr> {
 68        let mut z_i: Vec<Fr> = z_0.clone();
 69        for _ in 0..n_steps {
 70            let mut sponge = PoseidonSponge::<Fr>::new(&poseidon_config);
 71
 72            for _ in 0..hashes_per_step {
 73                sponge.absorb(&z_i);
 74                z_i = sponge.squeeze_field_elements(1);
 75            }
 76        }
 77        z_i.clone()
 78    }
 79
 80    #[test]
 81    fn full_flow() {
 82        // set how many iterations of the PoseidonChainCircuit circuit internal loop we want to
 83        // compute
 84        const N_STEPS: usize = 10;
 85        const HASHES_PER_STEP: usize = 400;
 86        println!("running the 'naive' PoseidonChainCircuit, with N_STEPS={}, HASHES_PER_STEP={}. Total hashes = {}", N_STEPS, HASHES_PER_STEP, N_STEPS* HASHES_PER_STEP);
 87
 88        let poseidon_config = poseidon_canonical_config::<Fr>();
 89
 90        // set the initial state
 91        // let z_0_aux: Vec<u32> = vec![0_u32; 32 * 8];
 92        let z_0_aux: Vec<u8> = vec![0_u8; 32];
 93        let z_0: Vec<Fr> = z_0_aux.iter().map(|v| Fr::from(*v)).collect::<Vec<Fr>>();
 94
 95        // run the N iterations 'natively' in rust to compute the expected `z_n`
 96        let z_n = rust_native_result(&poseidon_config, z_0.clone(), N_STEPS, HASHES_PER_STEP);
 97
 98        let circuit = PoseidonChainCircuit::<Fr, N_STEPS, HASHES_PER_STEP> {
 99            z_0: Some(z_0),
100            z_n: Some(z_n.clone()),
101            config: poseidon_config,
102        };
103
104        let cs = ConstraintSystem::<Fr>::new_ref();
105        circuit.clone().generate_constraints(cs.clone()).unwrap();
106        println!(
107            "number of constraints of the (naive) PoseidonChainCircuit with N_STEPS*HASHES_PER_STEP={} poseidon hashes in total: {} (num constraints)",
108            N_STEPS * HASHES_PER_STEP,
109            cs.num_constraints()
110        );
111
112        // now let's generate an actual Groth16 proof
113        let mut rng = rand::rngs::OsRng;
114        let (g16_pk, g16_vk) =
115            Groth16::<Bn254>::circuit_specific_setup(circuit.clone(), &mut rng).unwrap();
116
117        let start = Instant::now();
118        let proof = Groth16::<Bn254>::prove(&g16_pk, circuit.clone(), &mut rng).unwrap();
119        println!(
120            "Groth16 proof generation (for the naive PoseidonChainCircuit): {:?}",
121            start.elapsed()
122        );
123
124        let public_inputs = z_n;
125        let valid_proof = Groth16::<Bn254>::verify(&g16_vk, &public_inputs, &proof).unwrap();
126
127        assert!(valid_proof);
128
129        println!("finished running the 'naive' PoseidonChainCircuit, with N_STEPS={}, HASHES_PER_STEP={}. Total hashes = {}", N_STEPS, HASHES_PER_STEP, N_STEPS* HASHES_PER_STEP);
130    }
131}