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::{bit::bit_decompose, 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> + ProtoGalaxyTranscript<C>),
 39        r1cs: &R1CS<C::ScalarField>,
 40        // running instance
 41        instance: &CommittedInstance<C>,
 42        w: &Witness<C::ScalarField>,
 43        // incomming 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        ),
 53        Error,
 54    > {
 55        if vec_instances.len() != vec_w.len() {
 56            return Err(Error::NotSameLength(
 57                "vec_instances.len()".to_string(),
 58                vec_instances.len(),
 59                "vec_w.len()".to_string(),
 60                vec_w.len(),
 61            ));
 62        }
 63        let d = 2; // for the moment hardcoded to 2 since it only supports R1CS
 64        let k = vec_instances.len();
 65        let t = instance.betas.len();
 66        let n = r1cs.A.n_cols;
 67        if w.w.len() != n {
 68            return Err(Error::NotSameLength(
 69                "w.w.len()".to_string(),
 70                w.w.len(),
 71                "n".to_string(),
 72                n,
 73            ));
 74        }
 75        if log2(n) as usize != t {
 76            return Err(Error::NotEqual);
 77        }
 78        if !(k + 1).is_power_of_two() {
 79            return Err(Error::ProtoGalaxy(ProtoGalaxyError::WrongNumInstances(k)));
 80        }
 81
 82        // absorb the committed instances
 83        transcript.absorb_committed_instance(instance)?;
 84        for ci in vec_instances.iter() {
 85            transcript.absorb_committed_instance(ci)?;
 86        }
 87
 88        let delta = transcript.get_challenge();
 89        let deltas = exponential_powers(delta, t);
 90
 91        let f_w = eval_f(r1cs, &w.w)?;
 92
 93        // F(X)
 94        let F_X: SparsePolynomial<C::ScalarField> =
 95            calc_f_from_btree(&f_w, &instance.betas, &deltas).expect("Error calculating F[x]");
 96        let F_X_dense = DensePolynomial::from(F_X.clone());
 97        transcript.absorb_vec(&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            },
117            w,
118        )?;
119
120        let ws: Vec<Vec<C::ScalarField>> = std::iter::once(w.w.clone())
121            .chain(
122                vec_w
123                    .iter()
124                    .map(|wj| {
125                        if wj.w.len() != n {
126                            return Err(Error::NotSameLength(
127                                "wj.w.len()".to_string(),
128                                wj.w.len(),
129                                "n".to_string(),
130                                n,
131                            ));
132                        }
133                        Ok(wj.w.clone())
134                    })
135                    .collect::<Result<Vec<Vec<C::ScalarField>>, Error>>()?,
136            )
137            .collect::<Vec<Vec<C::ScalarField>>>();
138
139        let H =
140            GeneralEvaluationDomain::<C::ScalarField>::new(k + 1).ok_or(Error::NewDomainFail)?;
141        let G_domain = GeneralEvaluationDomain::<C::ScalarField>::new((d * k) + 1)
142            .ok_or(Error::NewDomainFail)?;
143        let L_X: Vec<DensePolynomial<C::ScalarField>> = lagrange_polys(H);
144
145        // K(X) computation in a naive way, next iterations will compute K(X) as described in Claim
146        // 4.5 of the paper.
147        let mut G_evals: Vec<C::ScalarField> = vec![C::ScalarField::zero(); G_domain.size()];
148        for (hi, h) in G_domain.elements().enumerate() {
149            // each iteration evaluates G(h)
150            // inner = L_0(x) * w + \sum_k L_i(x) * w_j
151            let mut inner: Vec<C::ScalarField> = vec![C::ScalarField::zero(); ws[0].len()];
152            for (i, w) in ws.iter().enumerate() {
153                // Li_w_h = (Li(X)*wj)(h) = Li(h) * wj
154                let mut Liw_h: Vec<C::ScalarField> = vec![C::ScalarField::zero(); w.len()];
155                for (j, wj) in w.iter().enumerate() {
156                    Liw_h[j] = (&L_X[i] * *wj).evaluate(&h);
157                }
158
159                for j in 0..inner.len() {
160                    inner[j] += Liw_h[j];
161                }
162            }
163            let f_ev = eval_f(r1cs, &inner)?;
164
165            let mut Gsum = C::ScalarField::zero();
166            for (i, f_ev_i) in f_ev.iter().enumerate() {
167                let pow_i_betas = pow_i(i, &betas_star);
168                let curr = pow_i_betas * f_ev_i;
169                Gsum += curr;
170            }
171            G_evals[hi] = Gsum;
172        }
173        let G_X: DensePolynomial<C::ScalarField> =
174            Evaluations::<C::ScalarField>::from_vec_and_domain(G_evals, G_domain).interpolate();
175        let Z_X: DensePolynomial<C::ScalarField> = H.vanishing_polynomial().into();
176        // K(X) = (G(X) - F(alpha)*L_0(X)) / Z(X)
177        // Notice that L0(X)*F(a) will be 0 in the native case (the instance of the first folding
178        // iteration case).
179        let L0_e = &L_X[0] * F_alpha;
180        let G_L0e = &G_X - &L0_e;
181        // Pending optimization: move division by Z_X to the prev loop
182        let (K_X, remainder) = G_L0e.divide_by_vanishing_poly(H).ok_or(Error::ProtoGalaxy(
183            ProtoGalaxyError::CouldNotDivideByVanishing,
184        ))?;
185        if !remainder.is_zero() {
186            return Err(Error::ProtoGalaxy(ProtoGalaxyError::RemainderNotZero));
187        }
188
189        transcript.absorb_vec(&K_X.coeffs);
190
191        let gamma = transcript.get_challenge();
192
193        let e_star =
194            F_alpha * L_X[0].evaluate(&gamma) + Z_X.evaluate(&gamma) * K_X.evaluate(&gamma);
195
196        let mut phi_star: C = instance.phi * L_X[0].evaluate(&gamma);
197        for i in 0..k {
198            phi_star += vec_instances[i].phi * L_X[i + 1].evaluate(&gamma);
199        }
200        let mut w_star: Vec<C::ScalarField> = vec_scalar_mul(&w.w, &L_X[0].evaluate(&gamma));
201        let mut r_w_star: C::ScalarField = w.r_w * L_X[0].evaluate(&gamma);
202        for i in 0..k {
203            let L_X_at_i1 = L_X[i + 1].evaluate(&gamma);
204            w_star = vec_add(&w_star, &vec_scalar_mul(&vec_w[i].w, &L_X_at_i1))?;
205            r_w_star += vec_w[i].r_w * L_X_at_i1;
206        }
207
208        Ok((
209            CommittedInstance {
210                betas: betas_star,
211                phi: phi_star,
212                e: e_star,
213            },
214            Witness {
215                w: w_star,
216                r_w: r_w_star,
217            },
218            F_X_dense.coeffs,
219            K_X.coeffs,
220        ))
221    }
222
223    /// implements the non-interactive Verifier from the folding scheme described in section 4
224    pub fn verify(
225        transcript: &mut (impl Transcript<C> + ProtoGalaxyTranscript<C>),
226        r1cs: &R1CS<C::ScalarField>,
227        // running instance
228        instance: &CommittedInstance<C>,
229        // incomming instances
230        vec_instances: &[CommittedInstance<C>],
231        // polys from P
232        F_coeffs: Vec<C::ScalarField>,
233        K_coeffs: Vec<C::ScalarField>,
234    ) -> Result<CommittedInstance<C>, Error> {
235        let t = instance.betas.len();
236        let n = r1cs.A.n_cols;
237
238        // absorb the committed instances
239        transcript.absorb_committed_instance(instance)?;
240        for ci in vec_instances.iter() {
241            transcript.absorb_committed_instance(ci)?;
242        }
243
244        let delta = transcript.get_challenge();
245        let deltas = exponential_powers(delta, t);
246
247        transcript.absorb_vec(&F_coeffs);
248
249        let alpha = transcript.get_challenge();
250        let alphas = all_powers(alpha, n);
251
252        // F(alpha) = e + \sum_t F_i * alpha^i
253        let mut F_alpha = instance.e;
254        for (i, F_i) in F_coeffs.iter().skip(1).enumerate() {
255            F_alpha += *F_i * alphas[i + 1];
256        }
257
258        let betas_star = betas_star(&instance.betas, &deltas, alpha);
259
260        let k = vec_instances.len();
261        let H =
262            GeneralEvaluationDomain::<C::ScalarField>::new(k + 1).ok_or(Error::NewDomainFail)?;
263        let L_X: Vec<DensePolynomial<C::ScalarField>> = lagrange_polys(H);
264        let Z_X: DensePolynomial<C::ScalarField> = H.vanishing_polynomial().into();
265        let K_X: DensePolynomial<C::ScalarField> =
266            DensePolynomial::<C::ScalarField>::from_coefficients_vec(K_coeffs);
267
268        transcript.absorb_vec(&K_X.coeffs);
269
270        let gamma = transcript.get_challenge();
271
272        let e_star =
273            F_alpha * L_X[0].evaluate(&gamma) + Z_X.evaluate(&gamma) * K_X.evaluate(&gamma);
274
275        let mut phi_star: C = instance.phi * L_X[0].evaluate(&gamma);
276        for i in 0..k {
277            phi_star += vec_instances[i].phi * L_X[i + 1].evaluate(&gamma);
278        }
279
280        // return the folded instance
281        Ok(CommittedInstance {
282            betas: betas_star,
283            phi: phi_star,
284            e: e_star,
285        })
286    }
287}
288
289// naive impl of pow_i for betas, assuming that betas=(b, b^2, b^4, ..., b^{2^{t-1}})
290fn pow_i<F: PrimeField>(i: usize, betas: &Vec<F>) -> F {
291    // WIP check if makes more sense to do it with ifs instead of arithmetic
292
293    let n = 2_u64.pow(betas.len() as u32);
294    let b = bit_decompose(i as u64, n as usize);
295
296    let mut r: F = F::one();
297    for (j, beta_j) in betas.iter().enumerate() {
298        let mut b_j = F::zero();
299        if b[j] {
300            b_j = F::one();
301        }
302        r *= (F::one() - b_j) + b_j * beta_j;
303    }
304    r
305}
306
307/// calculates F[x] using the optimized binary-tree technique
308/// described in Claim 4.4
309/// of [Protogalaxy](https://eprint.iacr.org/2023/1106.pdf)
310fn calc_f_from_btree<F: PrimeField>(
311    fw: &[F],
312    betas: &[F],
313    deltas: &[F],
314) -> Result<SparsePolynomial<F>, Error> {
315    let fw_len = fw.len();
316    let betas_len = betas.len();
317    let deltas_len = deltas.len();
318
319    // ensure our binary tree is full
320    if !fw_len.is_power_of_two() {
321        return Err(Error::ProtoGalaxy(ProtoGalaxyError::BTreeNotFull(fw_len)));
322    }
323
324    if betas_len != deltas_len {
325        return Err(Error::ProtoGalaxy(ProtoGalaxyError::WrongLenBetas(
326            betas_len, deltas_len,
327        )));
328    }
329
330    let mut layers: Vec<Vec<SparsePolynomial<F>>> = Vec::new();
331    let leaves: Vec<SparsePolynomial<F>> = fw
332        .iter()
333        .copied()
334        .map(|e| SparsePolynomial::<F>::from_coefficients_slice(&[(0, e)]))
335        .collect();
336    layers.push(leaves.to_vec());
337    let mut currentNodes = leaves.clone();
338    while currentNodes.len() > 1 {
339        let index = layers.len();
340        layers.push(vec![]);
341        for (i, ni) in currentNodes.iter().enumerate().step_by(2) {
342            let left = ni.clone();
343            let right = SparsePolynomial::<F>::from_coefficients_vec(vec![
344                (0, betas[layers.len() - 2]),
345                (1, deltas[layers.len() - 2]),
346            ])
347            .mul(&currentNodes[i + 1]);
348
349            layers[index].push(left + right);
350        }
351        currentNodes = layers[index].clone();
352    }
353    let root_index = layers.len() - 1;
354    Ok(layers[root_index][0].clone())
355}
356
357// lagrange_polys method from caulk: https://github.com/caulk-crypto/caulk/tree/8210b51fb8a9eef4335505d1695c44ddc7bf8170/src/multi/setup.rs#L300
358fn lagrange_polys<F: PrimeField>(domain_n: GeneralEvaluationDomain<F>) -> Vec<DensePolynomial<F>> {
359    let mut lagrange_polynomials: Vec<DensePolynomial<F>> = Vec::new();
360    for i in 0..domain_n.size() {
361        let evals: Vec<F> = cfg_into_iter!(0..domain_n.size())
362            .map(|k| if k == i { F::one() } else { F::zero() })
363            .collect();
364        lagrange_polynomials.push(Evaluations::from_vec_and_domain(evals, domain_n).interpolate());
365    }
366    lagrange_polynomials
367}
368
369// f(w) in R1CS context. For the moment we use R1CS, in the future we will abstract this with a
370// trait
371fn eval_f<F: PrimeField>(r1cs: &R1CS<F>, w: &[F]) -> Result<Vec<F>, Error> {
372    let Az = mat_vec_mul_sparse(&r1cs.A, w)?;
373    let Bz = mat_vec_mul_sparse(&r1cs.B, w)?;
374    let Cz = mat_vec_mul_sparse(&r1cs.C, w)?;
375    let AzBz = hadamard(&Az, &Bz)?;
376    vec_sub(&AzBz, &Cz)
377}
378
379#[cfg(test)]
380mod tests {
381    use super::*;
382    use ark_pallas::{Fr, Projective};
383    use ark_std::UniformRand;
384
385    use crate::ccs::r1cs::tests::{get_test_r1cs, get_test_z};
386    use crate::commitment::{pedersen::Pedersen, CommitmentProver};
387    use crate::transcript::poseidon::{poseidon_test_config, PoseidonTranscript};
388
389    pub(crate) fn check_instance<C: CurveGroup>(
390        r1cs: &R1CS<C::ScalarField>,
391        instance: &CommittedInstance<C>,
392        w: &Witness<C::ScalarField>,
393    ) -> Result<(), Error> {
394        if instance.betas.len() != log2(w.w.len()) as usize {
395            return Err(Error::NotSameLength(
396                "instance.betas.len()".to_string(),
397                instance.betas.len(),
398                "log2(w.w.len())".to_string(),
399                log2(w.w.len()) as usize,
400            ));
401        }
402
403        let f_w = eval_f(r1cs, &w.w)?; // f(w)
404
405        let mut r = C::ScalarField::zero();
406        for (i, f_w_i) in f_w.iter().enumerate() {
407            r += pow_i(i, &instance.betas) * f_w_i;
408        }
409        if instance.e == r {
410            return Ok(());
411        }
412        Err(Error::NotSatisfied)
413    }
414
415    #[test]
416    fn test_pow_i() {
417        let mut rng = ark_std::test_rng();
418        let t = 4;
419        let n = 16;
420        let beta = Fr::rand(&mut rng);
421        let betas = exponential_powers(beta, t);
422        let not_betas = all_powers(beta, n);
423
424        #[allow(clippy::needless_range_loop)]
425        for i in 0..n {
426            assert_eq!(pow_i(i, &betas), not_betas[i]);
427        }
428    }
429
430    #[test]
431    fn test_eval_f() {
432        let r1cs = get_test_r1cs::<Fr>();
433        let mut z = get_test_z::<Fr>(3);
434
435        let f_w = eval_f(&r1cs, &z).unwrap();
436        assert!(is_zero_vec(&f_w));
437
438        z[1] = Fr::from(111);
439        let f_w = eval_f(&r1cs, &z).unwrap();
440        assert!(!is_zero_vec(&f_w));
441    }
442
443    // k represents the number of instances to be fold, appart from the running instance
444    #[allow(clippy::type_complexity)]
445    fn prepare_inputs(
446        k: usize,
447    ) -> (
448        Witness<Fr>,
449        CommittedInstance<Projective>,
450        Vec<Witness<Fr>>,
451        Vec<CommittedInstance<Projective>>,
452    ) {
453        let mut rng = ark_std::test_rng();
454        let pedersen_params = Pedersen::<Projective>::new_params(&mut rng, 100); // 100 is wip, will get it from actual vec
455
456        let z = get_test_z::<Fr>(3);
457        let mut zs: Vec<Vec<Fr>> = Vec::new();
458        for i in 0..k {
459            let z_i = get_test_z::<Fr>(i + 4);
460            zs.push(z_i);
461        }
462
463        let n = z.len();
464        let t = log2(n) as usize;
465
466        let beta = Fr::rand(&mut rng);
467        let betas = exponential_powers(beta, t);
468
469        let witness = Witness::<Fr> {
470            w: z.clone(),
471            r_w: Fr::rand(&mut rng),
472        };
473        let phi =
474            Pedersen::<Projective>::commit(&pedersen_params, &witness.w, &witness.r_w).unwrap();
475        let instance = CommittedInstance::<Projective> {
476            phi,
477            betas: betas.clone(),
478            e: Fr::zero(),
479        };
480        // same for the other instances
481        let mut witnesses: Vec<Witness<Fr>> = Vec::new();
482        let mut instances: Vec<CommittedInstance<Projective>> = Vec::new();
483        #[allow(clippy::needless_range_loop)]
484        for i in 0..k {
485            let witness_i = Witness::<Fr> {
486                w: zs[i].clone(),
487                r_w: Fr::rand(&mut rng),
488            };
489            let phi_i =
490                Pedersen::<Projective>::commit(&pedersen_params, &witness_i.w, &witness_i.r_w)
491                    .unwrap();
492            let instance_i = CommittedInstance::<Projective> {
493                phi: phi_i,
494                betas: betas.clone(),
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_test_config::<Fr>();
512        let mut transcript_p = PoseidonTranscript::<Projective>::new(&poseidon_config);
513        let mut transcript_v = PoseidonTranscript::<Projective>::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        // veriier
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_test_config::<Fr>();
552        let mut transcript_p = PoseidonTranscript::<Projective>::new(&poseidon_config);
553        let mut transcript_v = PoseidonTranscript::<Projective>::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            // veriier
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}