1// Copyright (c) 2022 Espresso Systems (espressosys.com)
  2// This file is part of the Jellyfish library.
  3
  4// You should have received a copy of the MIT License
  5// along with the Jellyfish library. If not, see <https://mit-license.org/>.
  6
  7//! Implementing Structured Reference Strings for multilinear polynomial KZG
  8use crate::pcs::{
  9    multilinear_kzg::util::{eq_eval, eq_extension},
 10    prelude::PCSError,
 11    StructuredReferenceString,
 12};
 13use ark_ec::{msm::FixedBaseMSM, AffineCurve, PairingEngine, ProjectiveCurve};
 14use ark_ff::{Field, PrimeField, Zero};
 15use ark_poly::DenseMultilinearExtension;
 16use ark_serialize::{CanonicalDeserialize, CanonicalSerialize, Read, SerializationError, Write};
 17use ark_std::{
 18    collections::LinkedList,
 19    end_timer, format,
 20    rand::{CryptoRng, RngCore},
 21    start_timer,
 22    string::ToString,
 23    vec::Vec,
 24    UniformRand,
 25};
 26use core::iter::FromIterator;
 27
 28/// Evaluations over {0,1}^n for G1 or G2
 29#[derive(CanonicalSerialize, CanonicalDeserialize, Clone, Debug)]
 30pub struct Evaluations<C: AffineCurve> {
 31    /// The evaluations.
 32    pub evals: Vec<C>,
 33}
 34
 35/// Universal Parameter
 36#[derive(CanonicalSerialize, CanonicalDeserialize, Clone, Debug)]
 37pub struct MultilinearUniversalParams<E: PairingEngine> {
 38    /// prover parameters
 39    pub prover_param: MultilinearProverParam<E>,
 40    /// h^randomness: h^t1, h^t2, ..., **h^{t_nv}**
 41    pub h_mask: Vec<E::G2Affine>,
 42}
 43
 44/// Prover Parameters
 45#[derive(CanonicalSerialize, CanonicalDeserialize, Clone, Debug)]
 46pub struct MultilinearProverParam<E: PairingEngine> {
 47    /// number of variables
 48    pub num_vars: usize,
 49    /// `pp_{0}`, `pp_{1}`, ...,pp_{nu_vars} defined
 50    /// by XZZPD19 where pp_{nv-0}=g and
 51    /// pp_{nv-i}=g^{eq((t_1,..t_i),(X_1,..X_i))}
 52    pub powers_of_g: Vec<Evaluations<E::G1Affine>>,
 53    /// generator for G1
 54    pub g: E::G1Affine,
 55    /// generator for G2
 56    pub h: E::G2Affine,
 57}
 58
 59/// Verifier Parameters
 60#[derive(CanonicalSerialize, CanonicalDeserialize, Clone, Debug)]
 61pub struct MultilinearVerifierParam<E: PairingEngine> {
 62    /// number of variables
 63    pub num_vars: usize,
 64    /// generator of G1
 65    pub g: E::G1Affine,
 66    /// generator of G2
 67    pub h: E::G2Affine,
 68    /// h^randomness: h^t1, h^t2, ..., **h^{t_nv}**
 69    pub h_mask: Vec<E::G2Affine>,
 70}
 71
 72impl<E: PairingEngine> StructuredReferenceString<E> for MultilinearUniversalParams<E> {
 73    type ProverParam = MultilinearProverParam<E>;
 74    type VerifierParam = MultilinearVerifierParam<E>;
 75
 76    /// Extract the prover parameters from the public parameters.
 77    fn extract_prover_param(&self, supported_num_vars: usize) -> Self::ProverParam {
 78        let to_reduce = self.prover_param.num_vars - supported_num_vars;
 79
 80        Self::ProverParam {
 81            powers_of_g: self.prover_param.powers_of_g[to_reduce..].to_vec(),
 82            g: self.prover_param.g,
 83            h: self.prover_param.h,
 84            num_vars: supported_num_vars,
 85        }
 86    }
 87
 88    /// Extract the verifier parameters from the public parameters.
 89    fn extract_verifier_param(&self, supported_num_vars: usize) -> Self::VerifierParam {
 90        let to_reduce = self.prover_param.num_vars - supported_num_vars;
 91        Self::VerifierParam {
 92            num_vars: supported_num_vars,
 93            g: self.prover_param.g,
 94            h: self.prover_param.h,
 95            h_mask: self.h_mask[to_reduce..].to_vec(),
 96        }
 97    }
 98
 99    /// Trim the universal parameters to specialize the public parameters
100    /// for multilinear polynomials to the given `supported_num_vars`, and
101    /// returns committer key and verifier key. `supported_num_vars` should
102    /// be in range `1..=params.num_vars`
103    fn trim(
104        &self,
105        supported_num_vars: usize,
106    ) -> Result<(Self::ProverParam, Self::VerifierParam), PCSError> {
107        if supported_num_vars > self.prover_param.num_vars {
108            return Err(PCSError::InvalidParameters(format!(
109                "SRS does not support target number of vars {}",
110                supported_num_vars
111            )));
112        }
113
114        let to_reduce = self.prover_param.num_vars - supported_num_vars;
115        let ck = Self::ProverParam {
116            powers_of_g: self.prover_param.powers_of_g[to_reduce..].to_vec(),
117            g: self.prover_param.g,
118            h: self.prover_param.h,
119            num_vars: supported_num_vars,
120        };
121        let vk = Self::VerifierParam {
122            num_vars: supported_num_vars,
123            g: self.prover_param.g,
124            h: self.prover_param.h,
125            h_mask: self.h_mask[to_reduce..].to_vec(),
126        };
127        Ok((ck, vk))
128    }
129
130    /// Build SRS for testing.
131    /// WARNING: THIS FUNCTION IS FOR TESTING PURPOSE ONLY.
132    /// THE OUTPUT SRS SHOULD NOT BE USED IN PRODUCTION.
133    fn gen_srs_for_testing<R: RngCore + CryptoRng>(
134        rng: &mut R,
135        num_vars: usize,
136    ) -> Result<Self, PCSError> {
137        if num_vars == 0 {
138            return Err(PCSError::InvalidParameters(
139                "constant polynomial not supported".to_string(),
140            ));
141        }
142
143        let total_timer = start_timer!(|| "SRS generation");
144
145        let pp_generation_timer = start_timer!(|| "Prover Param generation");
146
147        let g = E::G1Projective::rand(rng);
148        let h = E::G2Projective::rand(rng);
149
150        let mut powers_of_g = Vec::new();
151
152        let t: Vec<_> = (0..num_vars).map(|_| E::Fr::rand(rng)).collect();
153        let scalar_bits = E::Fr::size_in_bits();
154
155        let mut eq: LinkedList<DenseMultilinearExtension<E::Fr>> =
156            LinkedList::from_iter(eq_extension(&t).into_iter());
157        let mut eq_arr = LinkedList::new();
158        let mut base = eq.pop_back().unwrap().evaluations;
159
160        for i in (0..num_vars).rev() {
161            eq_arr.push_front(remove_dummy_variable(&base, i)?);
162            if i != 0 {
163                let mul = eq.pop_back().unwrap().evaluations;
164                base = base
165                    .into_iter()
166                    .zip(mul.into_iter())
167                    .map(|(a, b)| a * b)
168                    .collect();
169            }
170        }
171
172        let mut pp_powers = Vec::new();
173        let mut total_scalars = 0;
174        for i in 0..num_vars {
175            let eq = eq_arr.pop_front().unwrap();
176            let pp_k_powers = (0..(1 << (num_vars - i))).map(|x| eq[x]);
177            pp_powers.extend(pp_k_powers);
178            total_scalars += 1 << (num_vars - i);
179        }
180        let window_size = FixedBaseMSM::get_mul_window_size(total_scalars);
181        let g_table = FixedBaseMSM::get_window_table(scalar_bits, window_size, g);
182
183        let pp_g = E::G1Projective::batch_normalization_into_affine(
184            &FixedBaseMSM::multi_scalar_mul(scalar_bits, window_size, &g_table, &pp_powers),
185        );
186
187        let mut start = 0;
188        for i in 0..num_vars {
189            let size = 1 << (num_vars - i);
190            let pp_k_g = Evaluations {
191                evals: pp_g[start..(start + size)].to_vec(),
192            };
193            // check correctness of pp_k_g
194            let t_eval_0 = eq_eval(&vec![E::Fr::zero(); num_vars - i], &t[i..num_vars])?;
195            assert_eq!(g.mul(t_eval_0.into_repr()).into_affine(), pp_k_g.evals[0]);
196            powers_of_g.push(pp_k_g);
197            start += size;
198        }
199        let gg = Evaluations {
200            evals: [g.into_affine()].to_vec(),
201        };
202        powers_of_g.push(gg);
203
204        let pp = Self::ProverParam {
205            num_vars,
206            g: g.into_affine(),
207            h: h.into_affine(),
208            powers_of_g,
209        };
210
211        end_timer!(pp_generation_timer);
212
213        let vp_generation_timer = start_timer!(|| "VP generation");
214        let h_mask = {
215            let window_size = FixedBaseMSM::get_mul_window_size(num_vars);
216            let h_table = FixedBaseMSM::get_window_table(scalar_bits, window_size, h);
217            E::G2Projective::batch_normalization_into_affine(&FixedBaseMSM::multi_scalar_mul(
218                scalar_bits,
219                window_size,
220                &h_table,
221                &t,
222            ))
223        };
224        end_timer!(vp_generation_timer);
225        end_timer!(total_timer);
226        Ok(Self {
227            prover_param: pp,
228            h_mask,
229        })
230    }
231}
232
233/// fix first `pad` variables of `poly` represented in evaluation form to zero
234fn remove_dummy_variable<F: Field>(poly: &[F], pad: usize) -> Result<Vec<F>, PCSError> {
235    if pad == 0 {
236        return Ok(poly.to_vec());
237    }
238    if !poly.len().is_power_of_two() {
239        return Err(PCSError::InvalidParameters(
240            "Size of polynomial should be power of two.".to_string(),
241        ));
242    }
243    let nv = ark_std::log2(poly.len()) as usize - pad;
244    Ok((0..(1 << nv)).map(|x| poly[x << pad]).collect())
245}
246
247#[cfg(test)]
248mod tests {
249    use super::*;
250    use ark_bls12_381::Bls12_381;
251    use ark_std::test_rng;
252    type E = Bls12_381;
253
254    #[test]
255    fn test_srs_gen() -> Result<(), PCSError> {
256        let mut rng = test_rng();
257        for nv in 4..10 {
258            let _ = MultilinearUniversalParams::<E>::gen_srs_for_testing(&mut rng, nv)?;
259        }
260
261        Ok(())
262    }
263}