1//! This module implements useful functions for the product check protocol.
  2
  3use crate::{errors::PolyIOPErrors, structs::IOPProof, zero_check::ZeroCheck, PolyIOP};
  4use arithmetic::{get_index, VirtualPolynomial};
  5use ark_ff::PrimeField;
  6use ark_poly::DenseMultilinearExtension;
  7use ark_std::{end_timer, start_timer};
  8use rayon::prelude::{IntoParallelRefIterator, ParallelIterator};
  9use std::rc::Rc;
 10use transcript::IOPTranscript;
 11
 12/// Compute the product polynomial `prod(x)` where
 13///
 14///  - `prod(0,x) := prod(0, x1, …, xn)` is the MLE over the
 15/// evaluations of `f(x)/g(x)` on the boolean hypercube {0,1}^n
 16///
 17/// - `prod(1,x)` is a MLE over the evaluations of `prod(x, 0) * prod(x, 1)`
 18/// on the boolean hypercube {0,1}^n
 19///
 20/// The caller needs to check num_vars matches in f and g
 21/// Cost: linear in N.
 22pub(super) fn compute_product_poly<F: PrimeField>(
 23    fx: &Rc<DenseMultilinearExtension<F>>,
 24    gx: &Rc<DenseMultilinearExtension<F>>,
 25) -> Result<Rc<DenseMultilinearExtension<F>>, PolyIOPErrors> {
 26    let start = start_timer!(|| "compute evaluations of prod polynomial");
 27    let num_vars = fx.num_vars;
 28
 29    // ===================================
 30    // prod(0, x)
 31    // ===================================
 32    let prod_0x_eval = compute_prod_0(fx, gx)?;
 33
 34    // ===================================
 35    // prod(1, x)
 36    // ===================================
 37    //
 38    // `prod(1, x)` can be computed via recursing the following formula for 2^n-1
 39    // times
 40    //
 41    // `prod(1, x_1, ..., x_n) :=
 42    //      prod(x_1, x_2, ..., x_n, 0) * prod(x_1, x_2, ..., x_n, 1)`
 43    //
 44    // At any given step, the right hand side of the equation
 45    // is available via either eval_0x or the current view of eval_1x
 46    let mut prod_1x_eval = vec![];
 47    for x in 0..(1 << num_vars) - 1 {
 48        // sign will decide if the evaluation should be looked up from eval_0x or
 49        // eval_1x; x_zero_index is the index for the evaluation (x_2, ..., x_n,
 50        // 0); x_one_index is the index for the evaluation (x_2, ..., x_n, 1);
 51        let (x_zero_index, x_one_index, sign) = get_index(x, num_vars);
 52        if !sign {
 53            prod_1x_eval.push(prod_0x_eval[x_zero_index] * prod_0x_eval[x_one_index]);
 54        } else {
 55            // sanity check: if we are trying to look up from the eval_1x table,
 56            // then the target index must already exist
 57            if x_zero_index >= prod_1x_eval.len() || x_one_index >= prod_1x_eval.len() {
 58                return Err(PolyIOPErrors::ShouldNotArrive);
 59            }
 60            prod_1x_eval.push(prod_1x_eval[x_zero_index] * prod_1x_eval[x_one_index]);
 61        }
 62    }
 63
 64    // prod(1, 1, ..., 1) := 0
 65    prod_1x_eval.push(F::zero());
 66
 67    // ===================================
 68    // prod(x)
 69    // ===================================
 70    // prod(x)'s evaluation is indeed `e := [eval_0x[..], eval_1x[..]].concat()`
 71    let eval = [prod_0x_eval.as_slice(), prod_1x_eval.as_slice()].concat();
 72
 73    let prod_x = Rc::new(DenseMultilinearExtension::from_evaluations_vec(
 74        num_vars + 1,
 75        eval,
 76    ));
 77
 78    end_timer!(start);
 79    Ok(prod_x)
 80}
 81
 82/// generate the zerocheck proof for the virtual polynomial
 83///    prod(1, x) - prod(x, 0) * prod(x, 1) + alpha * (f(x) - prod(0, x) * g(x))
 84///
 85/// Returns proof and Q(x) for testing purpose.
 86///
 87/// Cost: O(N)
 88pub(super) fn prove_zero_check<F: PrimeField>(
 89    fx: &Rc<DenseMultilinearExtension<F>>,
 90    gx: &Rc<DenseMultilinearExtension<F>>,
 91    prod_x: &Rc<DenseMultilinearExtension<F>>,
 92    alpha: &F,
 93    transcript: &mut IOPTranscript<F>,
 94) -> Result<(IOPProof<F>, VirtualPolynomial<F>), PolyIOPErrors> {
 95    let start = start_timer!(|| "zerocheck in product check");
 96
 97    let prod_partial_evals = build_prod_partial_eval(prod_x)?;
 98    let prod_0x = prod_partial_evals[0].clone();
 99    let prod_1x = prod_partial_evals[1].clone();
100    let prod_x0 = prod_partial_evals[2].clone();
101    let prod_x1 = prod_partial_evals[3].clone();
102
103    // compute g(x) * prod(0, x) * alpha
104    let mut q_x = VirtualPolynomial::new_from_mle(gx, F::one());
105    q_x.mul_by_mle(prod_0x, *alpha)?;
106
107    //   g(x) * prod(0, x) * alpha
108    // - f(x) * alpha
109    q_x.add_mle_list([fx.clone()], -*alpha)?;
110
111    // Q(x) := prod(1,x) - prod(x, 0) * prod(x, 1)
112    //       + alpha * (
113    //             g(x) * prod(0, x)
114    //           - f(x))
115    q_x.add_mle_list([prod_x0, prod_x1], -F::one())?;
116    q_x.add_mle_list([prod_1x], F::one())?;
117
118    let iop_proof = <PolyIOP<F> as ZeroCheck<F>>::prove(&q_x, transcript)?;
119
120    end_timer!(start);
121    Ok((iop_proof, q_x))
122}
123
124/// Helper function of the IOP.
125///
126/// Input:
127/// - prod(x)
128///
129/// Output: the following 4 polynomials
130/// - prod(0, x)
131/// - prod(1, x)
132/// - prod(x, 0)
133/// - prod(x, 1)
134fn build_prod_partial_eval<F: PrimeField>(
135    prod_x: &Rc<DenseMultilinearExtension<F>>,
136) -> Result<[Rc<DenseMultilinearExtension<F>>; 4], PolyIOPErrors> {
137    let start = start_timer!(|| "build partial prod polynomial");
138
139    let prod_x_eval = &prod_x.evaluations;
140    let num_vars = prod_x.num_vars - 1;
141
142    // prod(0, x)
143    let prod_0_x = Rc::new(DenseMultilinearExtension::from_evaluations_slice(
144        num_vars,
145        &prod_x_eval[0..1 << num_vars],
146    ));
147    // prod(1, x)
148    let prod_1_x = Rc::new(DenseMultilinearExtension::from_evaluations_slice(
149        num_vars,
150        &prod_x_eval[1 << num_vars..1 << (num_vars + 1)],
151    ));
152
153    // ===================================
154    // prod(x, 0) and prod(x, 1)
155    // ===================================
156    //
157    // now we compute eval_x0 and eval_x1
158    // eval_0x will be the odd coefficients of eval
159    // and eval_1x will be the even coefficients of eval
160    let mut eval_x0 = vec![];
161    let mut eval_x1 = vec![];
162    for (x, &prod_x) in prod_x_eval.iter().enumerate() {
163        if x & 1 == 0 {
164            eval_x0.push(prod_x);
165        } else {
166            eval_x1.push(prod_x);
167        }
168    }
169    let prod_x_0 = Rc::new(DenseMultilinearExtension::from_evaluations_vec(
170        num_vars, eval_x0,
171    ));
172    let prod_x_1 = Rc::new(DenseMultilinearExtension::from_evaluations_vec(
173        num_vars, eval_x1,
174    ));
175
176    end_timer!(start);
177
178    Ok([prod_0_x, prod_1_x, prod_x_0, prod_x_1])
179}
180
181/// Returns the evaluations of
182/// - `prod(0,x) := prod(0, x1, …, xn)` which is the MLE over the
183/// evaluations of f(x)/g(x) on the boolean hypercube {0,1}^n:
184///
185/// The caller needs to check num_vars matches in f/g
186/// Cost: linear in N.
187fn compute_prod_0<F: PrimeField>(
188    fx: &DenseMultilinearExtension<F>,
189    gx: &DenseMultilinearExtension<F>,
190) -> Result<Vec<F>, PolyIOPErrors> {
191    let start = start_timer!(|| "compute prod(0,x)");
192
193    let input = fx
194        .iter()
195        .zip(gx.iter())
196        .map(|(&fi, &gi)| (fi, gi))
197        .collect::<Vec<_>>();
198    let prod_0x_evals = input.par_iter().map(|(x, y)| *x / *y).collect::<Vec<_>>();
199
200    end_timer!(start);
201    Ok(prod_0x_evals)
202}