1//! This module implements useful functions for the permutation check protocol.
  2
  3use crate::PolyIOPErrors;
  4use ark_ff::PrimeField;
  5use ark_poly::DenseMultilinearExtension;
  6use ark_std::{end_timer, rand::RngCore, start_timer};
  7
  8/// Returns the evaluations of three MLEs:
  9/// - prod(0,x)
 10/// - numerator
 11/// - denominator
 12///
 13/// where
 14/// - `prod(0,x) := prod(0, x1, …, xn)` which is the MLE over the
 15/// evaluations of the following polynomial on the boolean hypercube {0,1}^n:
 16///
 17///  (f(x) + \beta s_id(x) + \gamma)/(g(x) + \beta s_perm(x) + \gamma)
 18///
 19///  where
 20///  - beta and gamma are challenges
 21///  - f(x), g(x), s_id(x), s_perm(x) are mle-s
 22///
 23/// - numerator is the MLE for `f(x) + \beta s_id(x) + \gamma`
 24/// - denominator is the MLE for `g(x) + \beta s_perm(x) + \gamma`
 25///
 26/// The caller needs to check num_vars matches in f/g/s_id/s_perm
 27/// Cost: linear in N.
 28///
 29/// TODO: replace `s_perm` with the merged poly `s`.
 30#[allow(clippy::type_complexity)]
 31pub(super) fn compute_prod_0<F: PrimeField>(
 32    beta: &F,
 33    gamma: &F,
 34    fx: &DenseMultilinearExtension<F>,
 35    gx: &DenseMultilinearExtension<F>,
 36    s_perm: &DenseMultilinearExtension<F>,
 37) -> Result<(Vec<F>, Vec<F>, Vec<F>), PolyIOPErrors> {
 38    let start = start_timer!(|| "compute prod(1,x)");
 39
 40    let num_vars = fx.num_vars;
 41    let mut prod_0x_evals = vec![];
 42    let mut numerator_evals = vec![];
 43    let mut denominator_evals = vec![];
 44
 45    // TODO: remove this line after replacing `s_perm` with `s`.
 46    let s_id = identity_permutation_mle::<F>(num_vars);
 47
 48    for (&fi, (&gi, (&s_id_i, &s_perm_i))) in
 49        fx.iter().zip(gx.iter().zip(s_id.iter().zip(s_perm.iter())))
 50    {
 51        let numerator = fi + *beta * s_id_i + gamma;
 52        let denominator = gi + *beta * s_perm_i + gamma;
 53
 54        prod_0x_evals.push(numerator / denominator);
 55        numerator_evals.push(numerator);
 56        denominator_evals.push(denominator);
 57    }
 58
 59    end_timer!(start);
 60    Ok((prod_0x_evals, numerator_evals, denominator_evals))
 61}
 62
 63/// An MLE that represent an identity permutation: `f(index) \mapto index`
 64pub fn identity_permutation_mle<F: PrimeField>(num_vars: usize) -> DenseMultilinearExtension<F> {
 65    let s_id_vec = (0..1u64 << num_vars).map(F::from).collect();
 66    DenseMultilinearExtension::from_evaluations_vec(num_vars, s_id_vec)
 67}
 68
 69/// An MLE that represent a random permutation
 70pub fn random_permutation_mle<F: PrimeField, R: RngCore>(
 71    num_vars: usize,
 72    rng: &mut R,
 73) -> DenseMultilinearExtension<F> {
 74    let len = 1u64 << num_vars;
 75    let mut s_id_vec: Vec<F> = (0..len).map(F::from).collect();
 76    let mut s_perm_vec = vec![];
 77    for _ in 0..len {
 78        let index = rng.next_u64() as usize % s_id_vec.len();
 79        s_perm_vec.push(s_id_vec.remove(index));
 80    }
 81    DenseMultilinearExtension::from_evaluations_vec(num_vars, s_perm_vec)
 82}
 83
 84/// Helper function of the IOP.
 85///
 86/// Input:
 87/// - prod(x)
 88///
 89/// Output: the following 4 polynomials
 90/// - prod(0, x)
 91/// - prod(1, x)
 92/// - prod(x, 0)
 93/// - prod(x, 1)
 94pub(super) fn build_prod_partial_eval<F: PrimeField>(
 95    prod_x: &DenseMultilinearExtension<F>,
 96) -> Result<[DenseMultilinearExtension<F>; 4], PolyIOPErrors> {
 97    let start = start_timer!(|| "build prod polynomial");
 98
 99    let prod_x_eval = &prod_x.evaluations;
100    let num_vars = prod_x.num_vars - 1;
101
102    // prod(0, x)
103    let prod_0_x =
104        DenseMultilinearExtension::from_evaluations_slice(num_vars, &prod_x_eval[0..1 << num_vars]);
105    // prod(1, x)
106    let prod_1_x = DenseMultilinearExtension::from_evaluations_slice(
107        num_vars,
108        &prod_x_eval[1 << num_vars..1 << (num_vars + 1)],
109    );
110
111    // ===================================
112    // prod(x, 0) and prod(x, 1)
113    // ===================================
114    //
115    // now we compute eval_x0 and eval_x1
116    // eval_0x will be the odd coefficients of eval
117    // and eval_1x will be the even coefficients of eval
118    let mut eval_x0 = vec![];
119    let mut eval_x1 = vec![];
120    for (x, &prod_x) in prod_x_eval.iter().enumerate() {
121        if x & 1 == 0 {
122            eval_x0.push(prod_x);
123        } else {
124            eval_x1.push(prod_x);
125        }
126    }
127    let prod_x_0 = DenseMultilinearExtension::from_evaluations_vec(num_vars, eval_x0);
128    let prod_x_1 = DenseMultilinearExtension::from_evaluations_vec(num_vars, eval_x1);
129
130    end_timer!(start);
131
132    Ok([prod_0_x, prod_1_x, prod_x_0, prod_x_1])
133}
134
135#[cfg(test)]
136mod test {
137    use super::*;
138    use crate::utils::bit_decompose;
139    use ark_bls12_381::Fr;
140    use ark_ff::UniformRand;
141    use ark_poly::MultilinearExtension;
142    use ark_std::test_rng;
143
144    #[test]
145    fn test_compute_prod_0() -> Result<(), PolyIOPErrors> {
146        let mut rng = test_rng();
147
148        for num_vars in 2..6 {
149            let f = DenseMultilinearExtension::rand(num_vars, &mut rng);
150            let g = DenseMultilinearExtension::rand(num_vars, &mut rng);
151
152            let s_id = identity_permutation_mle::<Fr>(num_vars);
153            let s_perm = random_permutation_mle(num_vars, &mut rng);
154
155            let beta = Fr::rand(&mut rng);
156            let gamma = Fr::rand(&mut rng);
157
158            let (prod_0, numerator, denominator) = compute_prod_0(&beta, &gamma, &f, &g, &s_perm)?;
159            let prod_0 = DenseMultilinearExtension::from_evaluations_vec(num_vars, prod_0);
160            let numerator = DenseMultilinearExtension::from_evaluations_vec(num_vars, numerator);
161            let denominator =
162                DenseMultilinearExtension::from_evaluations_vec(num_vars, denominator);
163
164            for i in 0..1 << num_vars {
165                let r: Vec<Fr> = bit_decompose(i, num_vars)
166                    .iter()
167                    .map(|&x| Fr::from(x))
168                    .collect();
169
170                let prod_0_eval = prod_0.evaluate(&r).unwrap();
171                let numerator_eval = numerator.evaluate(&r).unwrap();
172                let denominator_eval = denominator.evaluate(&r).unwrap();
173
174                let f_eval = f.evaluate(&r).unwrap();
175                let g_eval = g.evaluate(&r).unwrap();
176                let s_id_eval = s_id.evaluate(&r).unwrap();
177                let s_perm_eval = s_perm.evaluate(&r).unwrap();
178
179                let numerator_eval_rec = f_eval + beta * s_id_eval + gamma;
180                let denominator_eval_rec = g_eval + beta * s_perm_eval + gamma;
181                let prod_0_eval_rec = numerator_eval_rec / denominator_eval_rec;
182
183                assert_eq!(numerator_eval, numerator_eval_rec);
184                assert_eq!(denominator_eval, denominator_eval_rec);
185                assert_eq!(prod_0_eval, prod_0_eval_rec);
186            }
187        }
188        Ok(())
189    }
190}