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 two MLEs:
10/// - numerator
11/// - denominator
12///
13///  where
14///  - beta and gamma are challenges
15///  - f(x), g(x), s_id(x), s_perm(x) are mle-s
16///
17/// - numerator is the MLE for `f(x) + \beta s_id(x) + \gamma`
18/// - denominator is the MLE for `g(x) + \beta s_perm(x) + \gamma`
19#[allow(clippy::type_complexity)]
20pub(super) fn computer_num_and_denom<F: PrimeField>(
21    beta: &F,
22    gamma: &F,
23    fx: &DenseMultilinearExtension<F>,
24    gx: &DenseMultilinearExtension<F>,
25    s_perm: &DenseMultilinearExtension<F>,
26) -> Result<
27    (
28        Rc<DenseMultilinearExtension<F>>,
29        Rc<DenseMultilinearExtension<F>>,
30    ),
31    PolyIOPErrors,
32> {
33    let start = start_timer!(|| "compute numerator and denominator");
34
35    let num_vars = fx.num_vars;
36    let mut numerator_evals = vec![];
37    let mut denominator_evals = vec![];
38    let s_id = identity_permutation_mle::<F>(num_vars);
39
40    for (&fi, (&gi, (&s_id_i, &s_perm_i))) in
41        fx.iter().zip(gx.iter().zip(s_id.iter().zip(s_perm.iter())))
42    {
43        let numerator = fi + *beta * s_id_i + gamma;
44        let denominator = gi + *beta * s_perm_i + gamma;
45
46        numerator_evals.push(numerator);
47        denominator_evals.push(denominator);
48    }
49    let numerator = Rc::new(DenseMultilinearExtension::from_evaluations_vec(
50        num_vars,
51        numerator_evals,
52    ));
53    let denominator = Rc::new(DenseMultilinearExtension::from_evaluations_vec(
54        num_vars,
55        denominator_evals,
56    ));
57
58    end_timer!(start);
59    Ok((numerator, denominator))
60}
61
62/// An MLE that represent an identity permutation: `f(index) \mapto index`
63pub fn identity_permutation_mle<F: PrimeField>(
64    num_vars: usize,
65) -> Rc<DenseMultilinearExtension<F>> {
66    let s_id_vec = (0..1u64 << num_vars).map(F::from).collect();
67    Rc::new(DenseMultilinearExtension::from_evaluations_vec(
68        num_vars, s_id_vec,
69    ))
70}
71
72/// An MLE that represent a random permutation
73pub fn random_permutation_mle<F: PrimeField, R: RngCore>(
74    num_vars: usize,
75    rng: &mut R,
76) -> Rc<DenseMultilinearExtension<F>> {
77    let len = 1u64 << num_vars;
78    let mut s_id_vec: Vec<F> = (0..len).map(F::from).collect();
79    let mut s_perm_vec = vec![];
80    for _ in 0..len {
81        let index = rng.next_u64() as usize % s_id_vec.len();
82        s_perm_vec.push(s_id_vec.remove(index));
83    }
84    Rc::new(DenseMultilinearExtension::from_evaluations_vec(
85        num_vars, s_perm_vec,
86    ))
87}