1/// Implements the scheme described in [ProtoGalaxy](https://eprint.iacr.org/2023/1106.pdf)
  2use ark_crypto_primitives::sponge::Absorb;
  3use ark_ec::{CurveGroup, Group};
  4use ark_ff::PrimeField;
  5use ark_poly::{
  6    univariate::{DensePolynomial, SparsePolynomial},
  7    DenseUVPolynomial, EvaluationDomain, Evaluations, GeneralEvaluationDomain, Polynomial,
  8};
  9use ark_std::{cfg_into_iter, log2, Zero};
 10use rayon::iter::{IntoParallelIterator, ParallelIterator};
 11use std::marker::PhantomData;
 12
 13use super::utils::{all_powers, betas_star, exponential_powers};
 14use super::ProtoGalaxyError;
 15use super::{CommittedInstance, Witness};
 16
 17use crate::arith::r1cs::R1CS;
 18use crate::transcript::Transcript;
 19use crate::utils::vec::*;
 20use crate::utils::virtual_polynomial::bit_decompose;
 21use crate::Error;
 22
 23#[derive(Clone, Debug)]
 24/// Implements the protocol described in section 4 of
 25/// [ProtoGalaxy](https://eprint.iacr.org/2023/1106.pdf)
 26pub struct Folding<C: CurveGroup> {
 27    _phantom: PhantomData<C>,
 28}
 29impl<C: CurveGroup> Folding<C>
 30where
 31    <C as Group>::ScalarField: Absorb,
 32    <C as CurveGroup>::BaseField: Absorb,
 33{
 34    #![allow(clippy::type_complexity)]
 35    /// implements the non-interactive Prover from the folding scheme described in section 4
 36    pub fn prove(
 37        transcript: &mut impl Transcript<C::ScalarField>,
 38        r1cs: &R1CS<C::ScalarField>,
 39        // running instance
 40        instance: &CommittedInstance<C>,
 41        w: &Witness<C::ScalarField>,
 42        // incoming instances
 43        vec_instances: &[CommittedInstance<C>],
 44        vec_w: &[Witness<C::ScalarField>],
 45    ) -> Result<
 46        (
 47            CommittedInstance<C>,
 48            Witness<C::ScalarField>,
 49            Vec<C::ScalarField>, // F_X coeffs
 50            Vec<C::ScalarField>, // K_X coeffs
 51        ),
 52        Error,
 53    > {
 54        if vec_instances.len() != vec_w.len() {
 55            return Err(Error::NotSameLength(
 56                "vec_instances.len()".to_string(),
 57                vec_instances.len(),
 58                "vec_w.len()".to_string(),
 59                vec_w.len(),
 60            ));
 61        }
 62        let d = 2; // for the moment hardcoded to 2 since it only supports R1CS
 63        let k = vec_instances.len();
 64        let t = instance.betas.len();
 65        let n = r1cs.A.n_cols;
 66
 67        let z = [vec![instance.u], instance.x.clone(), w.w.clone()].concat();
 68
 69        if z.len() != n {
 70            return Err(Error::NotSameLength(
 71                "z.len()".to_string(),
 72                z.len(),
 73                "n".to_string(),
 74                n,
 75            ));
 76        }
 77        if log2(n) as usize != t {
 78            return Err(Error::NotEqual);
 79        }
 80        if !(k + 1).is_power_of_two() {
 81            return Err(Error::ProtoGalaxy(ProtoGalaxyError::WrongNumInstances(k)));
 82        }
 83
 84        // absorb the committed instances
 85        transcript.absorb(instance);
 86        transcript.absorb(&vec_instances);
 87
 88        let delta = transcript.get_challenge();
 89        let deltas = exponential_powers(delta, t);
 90
 91        let f_z = eval_f(r1cs, &z)?;
 92
 93        // F(X)
 94        let F_X: SparsePolynomial<C::ScalarField> =
 95            calc_f_from_btree(&f_z, &instance.betas, &deltas).expect("Error calculating F[x]");
 96        let F_X_dense = DensePolynomial::from(F_X.clone());
 97        transcript.absorb(&F_X_dense.coeffs);
 98
 99        let alpha = transcript.get_challenge();
100
101        // eval F(alpha)
102        let F_alpha = F_X.evaluate(&alpha);
103
104        // betas*
105        let betas_star = betas_star(&instance.betas, &deltas, alpha);
106
107        // sanity check: check that the new randomized instance (the original instance but with
108        // 'refreshed' randomness) satisfies the relation.
109        #[cfg(test)]
110        tests::check_instance(
111            r1cs,
112            &CommittedInstance {
113                phi: instance.phi,
114                betas: betas_star.clone(),
115                e: F_alpha,
116                u: instance.u,
117                x: instance.x.clone(),
118            },
119            w,
120        )?;
121
122        let zs: Vec<Vec<C::ScalarField>> = std::iter::once(z.clone())
123            .chain(
124                vec_w
125                    .iter()
126                    .zip(vec_instances)
127                    .map(|(wj, uj)| {
128                        let zj = [vec![uj.u], uj.x.clone(), wj.w.clone()].concat();
129                        if zj.len() != n {
130                            return Err(Error::NotSameLength(
131                                "zj.len()".to_string(),
132                                zj.len(),
133                                "n".to_string(),
134                                n,
135                            ));
136                        }
137                        Ok(zj)
138                    })
139                    .collect::<Result<Vec<Vec<C::ScalarField>>, Error>>()?,
140            )
141            .collect::<Vec<Vec<C::ScalarField>>>();
142
143        let H =
144            GeneralEvaluationDomain::<C::ScalarField>::new(k + 1).ok_or(Error::NewDomainFail)?;
145        let G_domain = GeneralEvaluationDomain::<C::ScalarField>::new((d * k) + 1)
146            .ok_or(Error::NewDomainFail)?;
147        let L_X: Vec<DensePolynomial<C::ScalarField>> = lagrange_polys(H);
148
149        // K(X) computation in a naive way, next iterations will compute K(X) as described in Claim
150        // 4.5 of the paper.
151        let mut G_evals: Vec<C::ScalarField> = vec![C::ScalarField::zero(); G_domain.size()];
152        for (hi, h) in G_domain.elements().enumerate() {
153            // each iteration evaluates G(h)
154            // inner = L_0(x) * z + \sum_k L_i(x) * z_j
155            let mut inner: Vec<C::ScalarField> = vec![C::ScalarField::zero(); zs[0].len()];
156            for (i, z) in zs.iter().enumerate() {
157                // Li_z_h = (Li(X)*zj)(h) = Li(h) * zj
158                let mut Liz_h: Vec<C::ScalarField> = vec![C::ScalarField::zero(); z.len()];
159                for (j, zj) in z.iter().enumerate() {
160                    Liz_h[j] = (&L_X[i] * *zj).evaluate(&h);
161                }
162
163                for j in 0..inner.len() {
164                    inner[j] += Liz_h[j];
165                }
166            }
167            let f_ev = eval_f(r1cs, &inner)?;
168
169            let mut Gsum = C::ScalarField::zero();
170            for (i, f_ev_i) in f_ev.iter().enumerate() {
171                let pow_i_betas = pow_i(i, &betas_star);
172                let curr = pow_i_betas * f_ev_i;
173                Gsum += curr;
174            }
175            G_evals[hi] = Gsum;
176        }
177        let G_X: DensePolynomial<C::ScalarField> =
178            Evaluations::<C::ScalarField>::from_vec_and_domain(G_evals, G_domain).interpolate();
179        let Z_X: DensePolynomial<C::ScalarField> = H.vanishing_polynomial().into();
180        // K(X) = (G(X) - F(alpha)*L_0(X)) / Z(X)
181        // Notice that L0(X)*F(a) will be 0 in the native case (the instance of the first folding
182        // iteration case).
183        let L0_e = &L_X[0] * F_alpha;
184        let G_L0e = &G_X - &L0_e;
185        // Pending optimization: move division by Z_X to the prev loop
186        let (K_X, remainder) = G_L0e.divide_by_vanishing_poly(H).ok_or(Error::ProtoGalaxy(
187            ProtoGalaxyError::CouldNotDivideByVanishing,
188        ))?;
189        if !remainder.is_zero() {
190            return Err(Error::ProtoGalaxy(ProtoGalaxyError::RemainderNotZero));
191        }
192
193        transcript.absorb(&K_X.coeffs);
194
195        let gamma = transcript.get_challenge();
196
197        let L_X_evals = L_X
198            .iter()
199            .take(k + 1)
200            .map(|L| L.evaluate(&gamma))
201            .collect::<Vec<_>>();
202
203        let e_star = F_alpha * L_X_evals[0] + Z_X.evaluate(&gamma) * K_X.evaluate(&gamma);
204        let mut w_star = vec_scalar_mul(&w.w, &L_X_evals[0]);
205        let mut r_w_star = w.r_w * L_X_evals[0];
206        let mut phi_star = instance.phi * L_X_evals[0];
207        let mut u_star = instance.u * L_X_evals[0];
208        let mut x_star = vec_scalar_mul(&instance.x, &L_X_evals[0]);
209        for i in 0..k {
210            w_star = vec_add(&w_star, &vec_scalar_mul(&vec_w[i].w, &L_X_evals[i + 1]))?;
211            r_w_star += vec_w[i].r_w * L_X_evals[i + 1];
212            phi_star += vec_instances[i].phi * L_X_evals[i + 1];
213            u_star += vec_instances[i].u * L_X_evals[i + 1];
214            x_star = vec_add(
215                &x_star,
216                &vec_scalar_mul(&vec_instances[i].x, &L_X_evals[i + 1]),
217            )?;
218        }
219
220        Ok((
221            CommittedInstance {
222                betas: betas_star,
223                phi: phi_star,
224                e: e_star,
225                u: u_star,
226                x: x_star,
227            },
228            Witness {
229                w: w_star,
230                r_w: r_w_star,
231            },
232            F_X_dense.coeffs,
233            K_X.coeffs,
234        ))
235    }
236
237    /// implements the non-interactive Verifier from the folding scheme described in section 4
238    pub fn verify(
239        transcript: &mut impl Transcript<C::ScalarField>,
240        r1cs: &R1CS<C::ScalarField>,
241        // running instance
242        instance: &CommittedInstance<C>,
243        // incoming instances
244        vec_instances: &[CommittedInstance<C>],
245        // polys from P
246        F_coeffs: Vec<C::ScalarField>,
247        K_coeffs: Vec<C::ScalarField>,
248    ) -> Result<CommittedInstance<C>, Error> {
249        let t = instance.betas.len();
250        let n = r1cs.A.n_cols;
251
252        // absorb the committed instances
253        transcript.absorb(instance);
254        transcript.absorb(&vec_instances);
255
256        let delta = transcript.get_challenge();
257        let deltas = exponential_powers(delta, t);
258
259        transcript.absorb(&F_coeffs);
260
261        let alpha = transcript.get_challenge();
262        let alphas = all_powers(alpha, n);
263
264        // F(alpha) = e + \sum_t F_i * alpha^i
265        let mut F_alpha = instance.e;
266        for (i, F_i) in F_coeffs.iter().skip(1).enumerate() {
267            F_alpha += *F_i * alphas[i + 1];
268        }
269
270        let betas_star = betas_star(&instance.betas, &deltas, alpha);
271
272        let k = vec_instances.len();
273        let H =
274            GeneralEvaluationDomain::<C::ScalarField>::new(k + 1).ok_or(Error::NewDomainFail)?;
275        let L_X: Vec<DensePolynomial<C::ScalarField>> = lagrange_polys(H);
276        let Z_X: DensePolynomial<C::ScalarField> = H.vanishing_polynomial().into();
277        let K_X: DensePolynomial<C::ScalarField> =
278            DensePolynomial::<C::ScalarField>::from_coefficients_vec(K_coeffs);
279
280        transcript.absorb(&K_X.coeffs);
281
282        let gamma = transcript.get_challenge();
283
284        let L_X_evals = L_X
285            .iter()
286            .take(k + 1)
287            .map(|L| L.evaluate(&gamma))
288            .collect::<Vec<_>>();
289
290        let e_star = F_alpha * L_X_evals[0] + Z_X.evaluate(&gamma) * K_X.evaluate(&gamma);
291
292        let mut phi_star = instance.phi * L_X_evals[0];
293        let mut u_star = instance.u * L_X_evals[0];
294        let mut x_star = vec_scalar_mul(&instance.x, &L_X_evals[0]);
295        for i in 0..k {
296            phi_star += vec_instances[i].phi * L_X_evals[i + 1];
297            u_star += vec_instances[i].u * L_X_evals[i + 1];
298            x_star = vec_add(
299                &x_star,
300                &vec_scalar_mul(&vec_instances[i].x, &L_X_evals[i + 1]),
301            )?;
302        }
303
304        // return the folded instance
305        Ok(CommittedInstance {
306            betas: betas_star,
307            phi: phi_star,
308            e: e_star,
309            u: u_star,
310            x: x_star,
311        })
312    }
313}
314
315// naive impl of pow_i for betas, assuming that betas=(b, b^2, b^4, ..., b^{2^{t-1}})
316fn pow_i<F: PrimeField>(i: usize, betas: &[F]) -> F {
317    // WIP check if makes more sense to do it with ifs instead of arithmetic
318
319    let n = 2_u64.pow(betas.len() as u32);
320    let b = bit_decompose(i as u64, n as usize);
321
322    let mut r: F = F::one();
323    for (j, beta_j) in betas.iter().enumerate() {
324        let mut b_j = F::zero();
325        if b[j] {
326            b_j = F::one();
327        }
328        r *= (F::one() - b_j) + b_j * beta_j;
329    }
330    r
331}
332
333/// calculates F[x] using the optimized binary-tree technique
334/// described in Claim 4.4
335/// of [ProtoGalaxy](https://eprint.iacr.org/2023/1106.pdf)
336fn calc_f_from_btree<F: PrimeField>(
337    fw: &[F],
338    betas: &[F],
339    deltas: &[F],
340) -> Result<SparsePolynomial<F>, Error> {
341    let fw_len = fw.len();
342    let betas_len = betas.len();
343    let deltas_len = deltas.len();
344
345    // ensure our binary tree is full
346    if !fw_len.is_power_of_two() {
347        return Err(Error::ProtoGalaxy(ProtoGalaxyError::BTreeNotFull(fw_len)));
348    }
349
350    if betas_len != deltas_len {
351        return Err(Error::ProtoGalaxy(ProtoGalaxyError::WrongLenBetas(
352            betas_len, deltas_len,
353        )));
354    }
355
356    let mut layers: Vec<Vec<SparsePolynomial<F>>> = Vec::new();
357    let leaves: Vec<SparsePolynomial<F>> = fw
358        .iter()
359        .copied()
360        .map(|e| SparsePolynomial::<F>::from_coefficients_slice(&[(0, e)]))
361        .collect();
362    layers.push(leaves.to_vec());
363    let mut currentNodes = leaves.clone();
364    while currentNodes.len() > 1 {
365        let index = layers.len();
366        layers.push(vec![]);
367        for (i, ni) in currentNodes.iter().enumerate().step_by(2) {
368            let left = ni.clone();
369            let right = SparsePolynomial::<F>::from_coefficients_vec(vec![
370                (0, betas[layers.len() - 2]),
371                (1, deltas[layers.len() - 2]),
372            ])
373            .mul(&currentNodes[i + 1]);
374
375            layers[index].push(left + right);
376        }
377        currentNodes = layers[index].clone();
378    }
379    let root_index = layers.len() - 1;
380    Ok(layers[root_index][0].clone())
381}
382
383// lagrange_polys method from caulk: https://github.com/caulk-crypto/caulk/tree/8210b51fb8a9eef4335505d1695c44ddc7bf8170/src/multi/setup.rs#L300
384pub fn lagrange_polys<F: PrimeField>(
385    domain_n: GeneralEvaluationDomain<F>,
386) -> Vec<DensePolynomial<F>> {
387    let mut lagrange_polynomials: Vec<DensePolynomial<F>> = Vec::new();
388    for i in 0..domain_n.size() {
389        let evals: Vec<F> = cfg_into_iter!(0..domain_n.size())
390            .map(|k| if k == i { F::one() } else { F::zero() })
391            .collect();
392        lagrange_polynomials.push(Evaluations::from_vec_and_domain(evals, domain_n).interpolate());
393    }
394    lagrange_polynomials
395}
396
397// f(w) in R1CS context. For the moment we use R1CS, in the future we will abstract this with a
398// trait
399fn eval_f<F: PrimeField>(r1cs: &R1CS<F>, z: &[F]) -> Result<Vec<F>, Error> {
400    let Az = mat_vec_mul(&r1cs.A, z)?;
401    let Bz = mat_vec_mul(&r1cs.B, z)?;
402    let Cz = mat_vec_mul(&r1cs.C, z)?;
403    let AzBz = hadamard(&Az, &Bz)?;
404    vec_sub(&AzBz, &Cz)
405}
406
407#[cfg(test)]
408pub mod tests {
409    use super::*;
410    use ark_crypto_primitives::sponge::poseidon::PoseidonSponge;
411    use ark_crypto_primitives::sponge::CryptographicSponge;
412    use ark_pallas::{Fr, Projective};
413    use ark_std::{rand::Rng, UniformRand};
414
415    use crate::arith::r1cs::tests::{get_test_r1cs, get_test_z, get_test_z_split};
416    use crate::commitment::{pedersen::Pedersen, CommitmentScheme};
417    use crate::transcript::poseidon::poseidon_canonical_config;
418
419    pub(crate) fn check_instance<C: CurveGroup>(
420        r1cs: &R1CS<C::ScalarField>,
421        instance: &CommittedInstance<C>,
422        w: &Witness<C::ScalarField>,
423    ) -> Result<(), Error> {
424        let z = [vec![instance.u], instance.x.clone(), w.w.clone()].concat();
425
426        if instance.betas.len() != log2(z.len()) as usize {
427            return Err(Error::NotSameLength(
428                "instance.betas.len()".to_string(),
429                instance.betas.len(),
430                "log2(z.len())".to_string(),
431                log2(z.len()) as usize,
432            ));
433        }
434
435        let f_z = eval_f(r1cs, &z)?; // f(z)
436
437        let mut r = C::ScalarField::zero();
438        for (i, f_z_i) in f_z.iter().enumerate() {
439            r += pow_i(i, &instance.betas) * f_z_i;
440        }
441        if instance.e == r {
442            return Ok(());
443        }
444        Err(Error::NotSatisfied)
445    }
446
447    #[test]
448    fn test_pow_i() {
449        let mut rng = ark_std::test_rng();
450        let t = 4;
451        let n = 16;
452        let beta = Fr::rand(&mut rng);
453        let betas = exponential_powers(beta, t);
454        let not_betas = all_powers(beta, n);
455
456        #[allow(clippy::needless_range_loop)]
457        for i in 0..n {
458            assert_eq!(pow_i(i, &betas), not_betas[i]);
459        }
460    }
461
462    #[test]
463    fn test_eval_f() {
464        let mut rng = ark_std::test_rng();
465        let r1cs = get_test_r1cs::<Fr>();
466        let mut z = get_test_z::<Fr>(rng.gen::<u16>() as usize);
467
468        let f_w = eval_f(&r1cs, &z).unwrap();
469        assert!(is_zero_vec(&f_w));
470
471        z[1] = Fr::from(111);
472        let f_w = eval_f(&r1cs, &z).unwrap();
473        assert!(!is_zero_vec(&f_w));
474    }
475
476    // k represents the number of instances to be fold, apart from the running instance
477    #[allow(clippy::type_complexity)]
478    pub fn prepare_inputs(
479        k: usize,
480    ) -> (
481        Witness<Fr>,
482        CommittedInstance<Projective>,
483        Vec<Witness<Fr>>,
484        Vec<CommittedInstance<Projective>>,
485    ) {
486        let mut rng = ark_std::test_rng();
487
488        let (u, x, w) = get_test_z_split::<Fr>(rng.gen::<u16>() as usize);
489
490        let (pedersen_params, _) = Pedersen::<Projective>::setup(&mut rng, w.len()).unwrap();
491
492        let n = 1 + x.len() + w.len();
493        let t = log2(n) as usize;
494
495        let beta = Fr::rand(&mut rng);
496        let betas = exponential_powers(beta, t);
497
498        let witness = Witness::<Fr> {
499            w,
500            r_w: Fr::rand(&mut rng),
501        };
502        let phi = Pedersen::<Projective, true>::commit(&pedersen_params, &witness.w, &witness.r_w)
503            .unwrap();
504        let instance = CommittedInstance::<Projective> {
505            phi,
506            betas: betas.clone(),
507            e: Fr::zero(),
508            u,
509            x,
510        };
511        // same for the other instances
512        let mut witnesses: Vec<Witness<Fr>> = Vec::new();
513        let mut instances: Vec<CommittedInstance<Projective>> = Vec::new();
514        #[allow(clippy::needless_range_loop)]
515        for _ in 0..k {
516            let (u_i, x_i, w_i) = get_test_z_split::<Fr>(rng.gen::<u16>() as usize);
517            let witness_i = Witness::<Fr> {
518                w: w_i,
519                r_w: Fr::rand(&mut rng),
520            };
521            let phi_i = Pedersen::<Projective, true>::commit(
522                &pedersen_params,
523                &witness_i.w,
524                &witness_i.r_w,
525            )
526            .unwrap();
527            let instance_i = CommittedInstance::<Projective> {
528                phi: phi_i,
529                betas: vec![],
530                e: Fr::zero(),
531                u: u_i,
532                x: x_i,
533            };
534            witnesses.push(witness_i);
535            instances.push(instance_i);
536        }
537
538        (witness, instance, witnesses, instances)
539    }
540
541    #[test]
542    fn test_fold() {
543        let k = 7;
544        let (witness, instance, witnesses, instances) = prepare_inputs(k);
545        let r1cs = get_test_r1cs::<Fr>();
546
547        // init Prover & Verifier's transcript
548        let poseidon_config = poseidon_canonical_config::<Fr>();
549        let mut transcript_p = PoseidonSponge::<Fr>::new(&poseidon_config);
550        let mut transcript_v = PoseidonSponge::<Fr>::new(&poseidon_config);
551
552        let (folded_instance, folded_witness, F_coeffs, K_coeffs) = Folding::<Projective>::prove(
553            &mut transcript_p,
554            &r1cs,
555            &instance,
556            &witness,
557            &instances,
558            &witnesses,
559        )
560        .unwrap();
561
562        // verifier
563        let folded_instance_v = Folding::<Projective>::verify(
564            &mut transcript_v,
565            &r1cs,
566            &instance,
567            &instances,
568            F_coeffs,
569            K_coeffs,
570        )
571        .unwrap();
572
573        // check that prover & verifier folded instances are the same values
574        assert_eq!(folded_instance.phi, folded_instance_v.phi);
575        assert_eq!(folded_instance.betas, folded_instance_v.betas);
576        assert_eq!(folded_instance.e, folded_instance_v.e);
577        assert!(!folded_instance.e.is_zero());
578
579        // check that the folded instance satisfies the relation
580        check_instance(&r1cs, &folded_instance, &folded_witness).unwrap();
581    }
582
583    #[test]
584    fn test_fold_various_iterations() {
585        let r1cs = get_test_r1cs::<Fr>();
586
587        // init Prover & Verifier's transcript
588        let poseidon_config = poseidon_canonical_config::<Fr>();
589        let mut transcript_p = PoseidonSponge::<Fr>::new(&poseidon_config);
590        let mut transcript_v = PoseidonSponge::<Fr>::new(&poseidon_config);
591
592        let (mut running_witness, mut running_instance, _, _) = prepare_inputs(0);
593
594        // fold k instances on each of num_iters iterations
595        let k = 7;
596        let num_iters = 10;
597        for _ in 0..num_iters {
598            // generate the instances to be fold
599            let (_, _, witnesses, instances) = prepare_inputs(k);
600
601            let (folded_instance, folded_witness, F_coeffs, K_coeffs) =
602                Folding::<Projective>::prove(
603                    &mut transcript_p,
604                    &r1cs,
605                    &running_instance,
606                    &running_witness,
607                    &instances,
608                    &witnesses,
609                )
610                .unwrap();
611
612            // verifier
613            let folded_instance_v = Folding::<Projective>::verify(
614                &mut transcript_v,
615                &r1cs,
616                &running_instance,
617                &instances,
618                F_coeffs,
619                K_coeffs,
620            )
621            .unwrap();
622
623            // check that prover & verifier folded instances are the same values
624            assert_eq!(folded_instance, folded_instance_v);
625
626            assert!(!folded_instance.e.is_zero());
627
628            // check that the folded instance satisfies the relation
629            check_instance(&r1cs, &folded_instance, &folded_witness).unwrap();
630
631            running_witness = folded_witness;
632            running_instance = folded_instance;
633        }
634    }
635}