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