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