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