1//! This module implements useful functions for the product check protocol.
  2
  3use crate::poly_iop::{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 std::sync::Arc;
  9use transcript::IOPTranscript;
 10
 11/// Compute multilinear fractional polynomial s.t. frac(x) = f1(x) * ... * fk(x)
 12/// / (g1(x) * ... * gk(x)) for all x \in {0,1}^n
 13///
 14/// The caller needs to sanity-check that the number of polynomials and
 15/// variables match in fxs and gxs; and gi(x) has no zero entries.
 16pub(super) fn compute_frac_poly<F: PrimeField>(
 17    fxs: &[Arc<DenseMultilinearExtension<F>>],
 18    gxs: &[Arc<DenseMultilinearExtension<F>>],
 19) -> Result<Arc<DenseMultilinearExtension<F>>, PolyIOPErrors> {
 20    let start = start_timer!(|| "compute frac(x)");
 21
 22    let mut f_evals = vec![F::one(); 1 << fxs[0].num_vars];
 23    for fx in fxs.iter() {
 24        for (f_eval, fi) in f_evals.iter_mut().zip(fx.iter()) {
 25            *f_eval *= fi;
 26        }
 27    }
 28    let mut g_evals = vec![F::one(); 1 << gxs[0].num_vars];
 29    for gx in gxs.iter() {
 30        for (g_eval, gi) in g_evals.iter_mut().zip(gx.iter()) {
 31            *g_eval *= gi;
 32        }
 33    }
 34    for (f_eval, g_eval) in f_evals.iter_mut().zip(g_evals.iter()) {
 35        if *g_eval == F::zero() {
 36            return Err(PolyIOPErrors::InvalidParameters(
 37                "gxs has zero entries in the boolean hypercube".to_string(),
 38            ));
 39        }
 40        *f_eval /= g_eval;
 41    }
 42
 43    end_timer!(start);
 44    Ok(Arc::new(DenseMultilinearExtension::from_evaluations_vec(
 45        fxs[0].num_vars,
 46        f_evals,
 47    )))
 48}
 49
 50/// Compute the product polynomial `prod(x)` such that
 51/// `prod(x) = [(1-x1)*frac(x2, ..., xn, 0) + x1*prod(x2, ..., xn, 0)] *
 52/// [(1-x1)*frac(x2, ..., xn, 1) + x1*prod(x2, ..., xn, 1)]` on the boolean
 53/// hypercube {0,1}^n
 54///
 55/// The caller needs to check num_vars matches in f and g
 56/// Cost: linear in N.
 57pub(super) fn compute_product_poly<F: PrimeField>(
 58    frac_poly: &Arc<DenseMultilinearExtension<F>>,
 59) -> Result<Arc<DenseMultilinearExtension<F>>, PolyIOPErrors> {
 60    let start = start_timer!(|| "compute evaluations of prod polynomial");
 61    let num_vars = frac_poly.num_vars;
 62    let frac_evals = &frac_poly.evaluations;
 63
 64    // ===================================
 65    // prod(x)
 66    // ===================================
 67    //
 68    // `prod(x)` can be computed via recursing the following formula for 2^n-1
 69    // times
 70    //
 71    // `prod(x_1, ..., x_n) :=
 72    //      [(1-x1)*frac(x2, ..., xn, 0) + x1*prod(x2, ..., xn, 0)] *
 73    //      [(1-x1)*frac(x2, ..., xn, 1) + x1*prod(x2, ..., xn, 1)]`
 74    //
 75    // At any given step, the right hand side of the equation
 76    // is available via either frac_x or the current view of prod_x
 77    let mut prod_x_evals = vec![];
 78    for x in 0..(1 << num_vars) - 1 {
 79        // sign will decide if the evaluation should be looked up from frac_x or
 80        // prod_x; x_zero_index is the index for the evaluation (x_2, ..., x_n,
 81        // 0); x_one_index is the index for the evaluation (x_2, ..., x_n, 1);
 82        let (x_zero_index, x_one_index, sign) = get_index(x, num_vars);
 83        if !sign {
 84            prod_x_evals.push(frac_evals[x_zero_index] * frac_evals[x_one_index]);
 85        } else {
 86            // sanity check: if we are trying to look up from the prod_x_evals table,
 87            // then the target index must already exist
 88            if x_zero_index >= prod_x_evals.len() || x_one_index >= prod_x_evals.len() {
 89                return Err(PolyIOPErrors::ShouldNotArrive);
 90            }
 91            prod_x_evals.push(prod_x_evals[x_zero_index] * prod_x_evals[x_one_index]);
 92        }
 93    }
 94
 95    // prod(1, 1, ..., 1) := 0
 96    prod_x_evals.push(F::zero());
 97    end_timer!(start);
 98
 99    Ok(Arc::new(DenseMultilinearExtension::from_evaluations_vec(
100        num_vars,
101        prod_x_evals,
102    )))
103}
104
105/// generate the zerocheck proof for the virtual polynomial
106///    prod(x) - p1(x) * p2(x) + alpha * [frac(x) * g1(x) * ... * gk(x) - f1(x)
107/// * ... * fk(x)] where p1(x) = (1-x1) * frac(x2, ..., xn, 0) + x1 * prod(x2,
108///   ..., xn, 0), p2(x) = (1-x1) * frac(x2, ..., xn, 1) + x1 * prod(x2, ...,
109///   xn, 1)
110/// Returns proof.
111///
112/// Cost: O(N)
113pub(super) fn prove_zero_check<F: PrimeField>(
114    fxs: &[Arc<DenseMultilinearExtension<F>>],
115    gxs: &[Arc<DenseMultilinearExtension<F>>],
116    frac_poly: &Arc<DenseMultilinearExtension<F>>,
117    prod_x: &Arc<DenseMultilinearExtension<F>>,
118    alpha: &F,
119    transcript: &mut IOPTranscript<F>,
120) -> Result<(IOPProof<F>, VirtualPolynomial<F>), PolyIOPErrors> {
121    let start = start_timer!(|| "zerocheck in product check");
122    let num_vars = frac_poly.num_vars;
123
124    // compute p1(x) = (1-x1) * frac(x2, ..., xn, 0) + x1 * prod(x2, ..., xn, 0)
125    // compute p2(x) = (1-x1) * frac(x2, ..., xn, 1) + x1 * prod(x2, ..., xn, 1)
126    let mut p1_evals = vec![F::zero(); 1 << num_vars];
127    let mut p2_evals = vec![F::zero(); 1 << num_vars];
128    for x in 0..1 << num_vars {
129        let (x0, x1, sign) = get_index(x, num_vars);
130        if !sign {
131            p1_evals[x] = frac_poly.evaluations[x0];
132            p2_evals[x] = frac_poly.evaluations[x1];
133        } else {
134            p1_evals[x] = prod_x.evaluations[x0];
135            p2_evals[x] = prod_x.evaluations[x1];
136        }
137    }
138    let p1 = Arc::new(DenseMultilinearExtension::from_evaluations_vec(
139        num_vars, p1_evals,
140    ));
141    let p2 = Arc::new(DenseMultilinearExtension::from_evaluations_vec(
142        num_vars, p2_evals,
143    ));
144
145    // compute Q(x)
146    // prod(x)
147    let mut q_x = VirtualPolynomial::new_from_mle(prod_x, F::one());
148
149    //   prod(x)
150    // - p1(x) * p2(x)
151    q_x.add_mle_list([p1, p2], -F::one())?;
152
153    //   prod(x)
154    // - p1(x) * p2(x)
155    // + alpha * frac(x) * g1(x) * ... * gk(x)
156    let mut mle_list = gxs.to_vec();
157    mle_list.push(frac_poly.clone());
158    q_x.add_mle_list(mle_list, *alpha)?;
159
160    //   prod(x)
161    // - p1(x) * p2(x)
162    // + alpha * frac(x) * g1(x) * ... * gk(x)
163    // - alpha * f1(x) * ... * fk(x)]
164    q_x.add_mle_list(fxs.to_vec(), -*alpha)?;
165
166    let iop_proof = <PolyIOP<F> as ZeroCheck<F>>::prove(&q_x, transcript)?;
167
168    end_timer!(start);
169    Ok((iop_proof, q_x))
170}