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