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