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