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 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<
 38    (
 39        DenseMultilinearExtension<F>,
 40        DenseMultilinearExtension<F>,
 41        DenseMultilinearExtension<F>,
 42    ),
 43    PolyIOPErrors,
 44> {
 45    let start = start_timer!(|| "compute prod(1,x)");
 46
 47    let num_vars = fx.num_vars;
 48    let mut prod_0x_evals = vec![];
 49    let mut numerator_evals = vec![];
 50    let mut denominator_evals = vec![];
 51
 52    // TODO: remove this line after replacing `s_perm` with `s`.
 53    let s_id = identity_permutation_mle::<F>(num_vars);
 54
 55    for (&fi, (&gi, (&s_id_i, &s_perm_i))) in
 56        fx.iter().zip(gx.iter().zip(s_id.iter().zip(s_perm.iter())))
 57    {
 58        let numerator = fi + *beta * s_id_i + gamma;
 59        let denominator = gi + *beta * s_perm_i + gamma;
 60
 61        prod_0x_evals.push(numerator / denominator);
 62        numerator_evals.push(numerator);
 63        denominator_evals.push(denominator);
 64    }
 65
 66    let prod_0x = DenseMultilinearExtension::from_evaluations_vec(num_vars, prod_0x_evals);
 67    let numerator = DenseMultilinearExtension::from_evaluations_vec(num_vars, numerator_evals);
 68    let denominator = DenseMultilinearExtension::from_evaluations_vec(num_vars, denominator_evals);
 69
 70    end_timer!(start);
 71    Ok((prod_0x, numerator, denominator))
 72}
 73
 74/// An MLE that represent an identity permutation: `f(index) \mapto index`
 75pub fn identity_permutation_mle<F: PrimeField>(num_vars: usize) -> DenseMultilinearExtension<F> {
 76    let s_id_vec = (0..1u64 << num_vars).map(F::from).collect();
 77    DenseMultilinearExtension::from_evaluations_vec(num_vars, s_id_vec)
 78}
 79
 80/// An MLE that represent a random permutation
 81pub fn random_permutation_mle<F: PrimeField, R: RngCore>(
 82    num_vars: usize,
 83    rng: &mut R,
 84) -> DenseMultilinearExtension<F> {
 85    let len = 1u64 << num_vars;
 86    let mut s_id_vec: Vec<F> = (0..len).map(F::from).collect();
 87    let mut s_perm_vec = vec![];
 88    for _ in 0..len {
 89        let index = rng.next_u64() as usize % s_id_vec.len();
 90        s_perm_vec.push(s_id_vec.remove(index));
 91    }
 92    DenseMultilinearExtension::from_evaluations_vec(num_vars, s_perm_vec)
 93}
 94
 95#[cfg(test)]
 96mod test {
 97    use super::*;
 98    use crate::utils::bit_decompose;
 99    use ark_bls12_381::Fr;
100    use ark_ff::UniformRand;
101    use ark_poly::MultilinearExtension;
102    use ark_std::test_rng;
103
104    #[test]
105    fn test_compute_prod_0() -> Result<(), PolyIOPErrors> {
106        let mut rng = test_rng();
107
108        for num_vars in 2..6 {
109            let f = DenseMultilinearExtension::rand(num_vars, &mut rng);
110            let g = DenseMultilinearExtension::rand(num_vars, &mut rng);
111
112            let s_id = identity_permutation_mle::<Fr>(num_vars);
113            let s_perm = random_permutation_mle(num_vars, &mut rng);
114
115            let beta = Fr::rand(&mut rng);
116            let gamma = Fr::rand(&mut rng);
117
118            let (prod_0, numerator, denominator) = compute_prod_0(&beta, &gamma, &f, &g, &s_perm)?;
119
120            for i in 0..1 << num_vars {
121                let r: Vec<Fr> = bit_decompose(i, num_vars)
122                    .iter()
123                    .map(|&x| Fr::from(x))
124                    .collect();
125
126                let prod_0_eval = prod_0.evaluate(&r).unwrap();
127                let numerator_eval = numerator.evaluate(&r).unwrap();
128                let denominator_eval = denominator.evaluate(&r).unwrap();
129
130                let f_eval = f.evaluate(&r).unwrap();
131                let g_eval = g.evaluate(&r).unwrap();
132                let s_id_eval = s_id.evaluate(&r).unwrap();
133                let s_perm_eval = s_perm.evaluate(&r).unwrap();
134
135                let numerator_eval_rec = f_eval + beta * s_id_eval + gamma;
136                let denominator_eval_rec = g_eval + beta * s_perm_eval + gamma;
137                let prod_0_eval_rec = numerator_eval_rec / denominator_eval_rec;
138
139                assert_eq!(numerator_eval, numerator_eval_rec);
140                assert_eq!(denominator_eval, denominator_eval_rec);
141                assert_eq!(prod_0_eval, prod_0_eval_rec);
142            }
143        }
144        Ok(())
145    }
146}