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