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