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 sha256 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::crh::sha256::{constraints::Sha256Gadget, digest::Digest, Sha256};
 16    use ark_r1cs_std::fields::fp::FpVar;
 17    use ark_r1cs_std::{alloc::AllocVar, eq::EqGadget};
 18    use ark_r1cs_std::{
 19        boolean::Boolean,
 20        convert::{ToBitsGadget, ToBytesGadget},
 21        uint8::UInt8,
 22    };
 23    use ark_relations::r1cs::{
 24        ConstraintSynthesizer, ConstraintSystem, ConstraintSystemRef, SynthesisError,
 25    };
 26
 27    use crate::utils::tests::*;
 28
 29    /// Test circuit to be folded
 30    #[derive(Clone, Debug)]
 31    pub struct SHA256ChainCircuit<F: PrimeField, const N: usize, const HASHES_PER_STEP: usize> {
 32        z_0: Option<Vec<F>>,
 33        z_n: Option<Vec<F>>,
 34    }
 35    impl<F: PrimeField, const N: usize, const HASHES_PER_STEP: usize> ConstraintSynthesizer<F>
 36        for SHA256ChainCircuit<F, N, HASHES_PER_STEP>
 37    {
 38        fn generate_constraints(self, cs: ConstraintSystemRef<F>) -> Result<(), SynthesisError> {
 39            let z_0 = Vec::<FpVar<F>>::new_witness(cs.clone(), || {
 40                Ok(self.z_0.unwrap_or(vec![F::zero()]))
 41            })?;
 42            let z_n =
 43                Vec::<FpVar<F>>::new_input(cs.clone(), || Ok(self.z_n.unwrap_or(vec![F::zero()])))?;
 44            let mut z_i: Vec<FpVar<F>> = z_0.clone();
 45            for _ in 0..N {
 46                let mut b: Vec<UInt8<F>> = z_i
 47                    .iter()
 48                    .map(|f| UInt8::<F>::from_bits_le(&f.to_bits_le().unwrap()[..8]))
 49                    .collect::<Vec<_>>();
 50
 51                for _ in 0..HASHES_PER_STEP {
 52                    let mut sha256_var = Sha256Gadget::default();
 53                    sha256_var.update(&b).unwrap();
 54                    b = sha256_var.finalize()?.to_bytes_le()?;
 55                }
 56
 57                // update z_i = z_{i+1}
 58                z_i = b
 59                    .iter()
 60                    .map(|e| {
 61                        let bits = e.to_bits_le().unwrap();
 62                        Boolean::<F>::le_bits_to_fp(&bits).unwrap()
 63                    })
 64                    .collect();
 65            }
 66
 67            z_i.enforce_equal(&z_n)?;
 68            Ok(())
 69        }
 70    }
 71    // compute natively in rust the expected result
 72    fn rust_native_result(z_0: Vec<Fr>, n_steps: usize, hashes_per_step: usize) -> Vec<Fr> {
 73        let mut z_i: Vec<Fr> = z_0.clone();
 74        for _ in 0..n_steps {
 75            let mut b = f_vec_to_bytes(z_i.to_vec());
 76
 77            for _ in 0..hashes_per_step {
 78                let mut sha256 = Sha256::default();
 79                sha256.update(b);
 80                b = sha256.finalize().to_vec();
 81            }
 82
 83            z_i = bytes_to_f_vec(b.to_vec()).unwrap();
 84        }
 85        z_i.clone()
 86    }
 87
 88    #[test]
 89    fn full_flow() {
 90        // set how many iterations of the SHA256ChainCircuit circuit internal loop we want to
 91        // compute
 92        const N_STEPS: usize = 100;
 93        const HASHES_PER_STEP: usize = 10;
 94        println!("running the 'naive' SHA256ChainCircuit, with N_STEPS={}, HASHES_PER_STEP={}. Total hashes = {}", N_STEPS, HASHES_PER_STEP, N_STEPS* HASHES_PER_STEP);
 95
 96        // set the initial state
 97        // let z_0_aux: Vec<u32> = vec![0_u32; 32 * 8];
 98        let z_0_aux: Vec<u8> = vec![0_u8; 32];
 99        let z_0: Vec<Fr> = z_0_aux.iter().map(|v| Fr::from(*v)).collect::<Vec<Fr>>();
100
101        // run the N iterations 'natively' in rust to compute the expected `z_n`
102        let z_n = rust_native_result(z_0.clone(), N_STEPS, HASHES_PER_STEP);
103
104        let circuit = SHA256ChainCircuit::<Fr, N_STEPS, HASHES_PER_STEP> {
105            z_0: Some(z_0),
106            z_n: Some(z_n.clone()),
107        };
108
109        let cs = ConstraintSystem::<Fr>::new_ref();
110        circuit.clone().generate_constraints(cs.clone()).unwrap();
111        println!(
112            "number of constraints of the (naive) SHA256ChainCircuit with N_STEPS*HASHES_PER_STEP={} sha256 hashes in total: {}",
113            N_STEPS * HASHES_PER_STEP,
114            cs.num_constraints()
115        );
116
117        // now let's generate an actual Groth16 proof
118        let mut rng = rand::rngs::OsRng;
119        let (g16_pk, g16_vk) =
120            Groth16::<Bn254>::circuit_specific_setup(circuit.clone(), &mut rng).unwrap();
121
122        let start = Instant::now();
123        let proof = Groth16::<Bn254>::prove(&g16_pk, circuit.clone(), &mut rng).unwrap();
124        println!(
125            "Groth16 proof generation (for the naive SHA256ChainCircuit): {:?}",
126            start.elapsed()
127        );
128
129        let public_inputs = z_n;
130        let valid_proof = Groth16::<Bn254>::verify(&g16_vk, &public_inputs, &proof).unwrap();
131
132        assert!(valid_proof);
133
134        println!("finished running the 'naive' SHA256ChainCircuit, with N_STEPS={}, HASHES_PER_STEP={}. Total hashes = {}", N_STEPS, HASHES_PER_STEP, N_STEPS* HASHES_PER_STEP);
135    }
136}