1//! This module implements useful functions for the permutation check protocol.
  2
  3use crate::errors::PolyIOPErrors;
  4use ark_ff::PrimeField;
  5use ark_poly::DenseMultilinearExtension;
  6use ark_std::{end_timer, rand::RngCore, start_timer};
  7use std::rc::Rc;
  8
  9/// Returns the evaluations of three MLEs:
 10/// - prod(0,x)
 11/// - numerator
 12/// - denominator
 13///
 14/// where
 15/// - `prod(0,x) := prod(0, x1, …, xn)` which is the MLE over the
 16/// evaluations of the following polynomial on the boolean hypercube {0,1}^n:
 17///
 18///  (f(x) + \beta s_id(x) + \gamma)/(g(x) + \beta s_perm(x) + \gamma)
 19///
 20///  where
 21///  - beta and gamma are challenges
 22///  - f(x), g(x), s_id(x), s_perm(x) are mle-s
 23///
 24/// - numerator is the MLE for `f(x) + \beta s_id(x) + \gamma`
 25/// - denominator is the MLE for `g(x) + \beta s_perm(x) + \gamma`
 26///
 27/// The caller needs to check num_vars matches in f/g/s_id/s_perm
 28/// Cost: linear in N.
 29#[cfg(test)]
 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/// Returns the evaluations of two MLEs:
 64/// - numerator
 65/// - denominator
 66///
 67///  where
 68///  - beta and gamma are challenges
 69///  - f(x), g(x), s_id(x), s_perm(x) are mle-s
 70///
 71/// - numerator is the MLE for `f(x) + \beta s_id(x) + \gamma`
 72/// - denominator is the MLE for `g(x) + \beta s_perm(x) + \gamma`
 73#[allow(clippy::type_complexity)]
 74pub(super) fn computer_num_and_denom<F: PrimeField>(
 75    beta: &F,
 76    gamma: &F,
 77    fx: &DenseMultilinearExtension<F>,
 78    gx: &DenseMultilinearExtension<F>,
 79    s_perm: &DenseMultilinearExtension<F>,
 80) -> Result<
 81    (
 82        Rc<DenseMultilinearExtension<F>>,
 83        Rc<DenseMultilinearExtension<F>>,
 84    ),
 85    PolyIOPErrors,
 86> {
 87    let start = start_timer!(|| "compute numerator and denominator");
 88
 89    let num_vars = fx.num_vars;
 90    let mut numerator_evals = vec![];
 91    let mut denominator_evals = vec![];
 92    let s_id = identity_permutation_mle::<F>(num_vars);
 93
 94    for (&fi, (&gi, (&s_id_i, &s_perm_i))) in
 95        fx.iter().zip(gx.iter().zip(s_id.iter().zip(s_perm.iter())))
 96    {
 97        let numerator = fi + *beta * s_id_i + gamma;
 98        let denominator = gi + *beta * s_perm_i + gamma;
 99
100        numerator_evals.push(numerator);
101        denominator_evals.push(denominator);
102    }
103    let numerator = Rc::new(DenseMultilinearExtension::from_evaluations_vec(
104        num_vars,
105        numerator_evals,
106    ));
107    let denominator = Rc::new(DenseMultilinearExtension::from_evaluations_vec(
108        num_vars,
109        denominator_evals,
110    ));
111
112    end_timer!(start);
113    Ok((numerator, denominator))
114}
115
116/// An MLE that represent an identity permutation: `f(index) \mapto index`
117pub fn identity_permutation_mle<F: PrimeField>(
118    num_vars: usize,
119) -> Rc<DenseMultilinearExtension<F>> {
120    let s_id_vec = (0..1u64 << num_vars).map(F::from).collect();
121    Rc::new(DenseMultilinearExtension::from_evaluations_vec(
122        num_vars, s_id_vec,
123    ))
124}
125
126/// An MLE that represent a random permutation
127pub fn random_permutation_mle<F: PrimeField, R: RngCore>(
128    num_vars: usize,
129    rng: &mut R,
130) -> Rc<DenseMultilinearExtension<F>> {
131    let len = 1u64 << num_vars;
132    let mut s_id_vec: Vec<F> = (0..len).map(F::from).collect();
133    let mut s_perm_vec = vec![];
134    for _ in 0..len {
135        let index = rng.next_u64() as usize % s_id_vec.len();
136        s_perm_vec.push(s_id_vec.remove(index));
137    }
138    Rc::new(DenseMultilinearExtension::from_evaluations_vec(
139        num_vars, s_perm_vec,
140    ))
141}
142
143/// Helper function of the IOP.
144///
145/// Input:
146/// - prod(x)
147///
148/// Output: the following 4 polynomials
149/// - prod(0, x)
150/// - prod(1, x)
151/// - prod(x, 0)
152/// - prod(x, 1)
153#[cfg(test)]
154pub(super) fn build_prod_partial_eval<F: PrimeField>(
155    prod_x: &Rc<DenseMultilinearExtension<F>>,
156) -> Result<[Rc<DenseMultilinearExtension<F>>; 4], PolyIOPErrors> {
157    let start = start_timer!(|| "build prod polynomial");
158
159    let prod_x_eval = &prod_x.evaluations;
160    let num_vars = prod_x.num_vars - 1;
161
162    // prod(0, x)
163    let prod_0_x = Rc::new(DenseMultilinearExtension::from_evaluations_slice(
164        num_vars,
165        &prod_x_eval[0..1 << num_vars],
166    ));
167    // prod(1, x)
168    let prod_1_x = Rc::new(DenseMultilinearExtension::from_evaluations_slice(
169        num_vars,
170        &prod_x_eval[1 << num_vars..1 << (num_vars + 1)],
171    ));
172
173    // ===================================
174    // prod(x, 0) and prod(x, 1)
175    // ===================================
176    //
177    // now we compute eval_x0 and eval_x1
178    // eval_0x will be the odd coefficients of eval
179    // and eval_1x will be the even coefficients of eval
180    let mut eval_x0 = vec![];
181    let mut eval_x1 = vec![];
182    for (x, &prod_x) in prod_x_eval.iter().enumerate() {
183        if x & 1 == 0 {
184            eval_x0.push(prod_x);
185        } else {
186            eval_x1.push(prod_x);
187        }
188    }
189    let prod_x_0 = Rc::new(DenseMultilinearExtension::from_evaluations_vec(
190        num_vars, eval_x0,
191    ));
192    let prod_x_1 = Rc::new(DenseMultilinearExtension::from_evaluations_vec(
193        num_vars, eval_x1,
194    ));
195
196    end_timer!(start);
197
198    Ok([prod_0_x, prod_1_x, prod_x_0, prod_x_1])
199}
200
201#[cfg(test)]
202mod test {
203    use super::*;
204    use crate::utils::bit_decompose;
205    use ark_bls12_381::Fr;
206    use ark_ff::UniformRand;
207    use ark_poly::MultilinearExtension;
208    use ark_std::test_rng;
209
210    #[test]
211    fn test_compute_prod_0() -> Result<(), PolyIOPErrors> {
212        let mut rng = test_rng();
213
214        for num_vars in 2..6 {
215            let f = DenseMultilinearExtension::rand(num_vars, &mut rng);
216            let g = DenseMultilinearExtension::rand(num_vars, &mut rng);
217
218            let s_id = identity_permutation_mle::<Fr>(num_vars);
219            let s_perm = random_permutation_mle(num_vars, &mut rng);
220
221            let beta = Fr::rand(&mut rng);
222            let gamma = Fr::rand(&mut rng);
223
224            let (prod_0, numerator, denominator) = compute_prod_0(&beta, &gamma, &f, &g, &s_perm)?;
225            let prod_0 = DenseMultilinearExtension::from_evaluations_vec(num_vars, prod_0);
226            let numerator = DenseMultilinearExtension::from_evaluations_vec(num_vars, numerator);
227            let denominator =
228                DenseMultilinearExtension::from_evaluations_vec(num_vars, denominator);
229
230            for i in 0..1 << num_vars {
231                let r: Vec<Fr> = bit_decompose(i, num_vars)
232                    .iter()
233                    .map(|&x| Fr::from(x))
234                    .collect();
235
236                let prod_0_eval = prod_0.evaluate(&r).unwrap();
237                let numerator_eval = numerator.evaluate(&r).unwrap();
238                let denominator_eval = denominator.evaluate(&r).unwrap();
239
240                let f_eval = f.evaluate(&r).unwrap();
241                let g_eval = g.evaluate(&r).unwrap();
242                let s_id_eval = s_id.evaluate(&r).unwrap();
243                let s_perm_eval = s_perm.evaluate(&r).unwrap();
244
245                let numerator_eval_rec = f_eval + beta * s_id_eval + gamma;
246                let denominator_eval_rec = g_eval + beta * s_perm_eval + gamma;
247                let prod_0_eval_rec = numerator_eval_rec / denominator_eval_rec;
248
249                assert_eq!(numerator_eval, numerator_eval_rec);
250                assert_eq!(denominator_eval, denominator_eval_rec);
251                assert_eq!(prod_0_eval, prod_0_eval_rec);
252            }
253        }
254        Ok(())
255    }
256}