Arrabiata: prepare constraints for cross-terms computation

arrabiata/src/columns.rs

@@ -8,6 +8,8 @@ use std::{
 };
 use strum_macros::{EnumCount as EnumCountMacro, EnumIter};
 
+use crate::NUMBER_OF_COLUMNS;
+
 /// This enum represents the different gadgets that can be used in the circuit.
 /// The selectors are defined at setup time, can take only the values `0` or
 /// `1` and are public.
@@ -36,6 +38,27 @@ pub enum Column {
     X(usize),
 }
 
+/// Convert a column to a usize. This is used by the library [mvpoly] when we
+/// need to compute the cross-terms.
+/// For now, only the private inputs and the public inputs are converted,
+/// because there might not need to treat the selectors in the polynomial while
+/// computing the cross-terms (FIXME: check this later, but pretty sure it's the
+/// case).
+///
+/// Also, the [mvpoly::monomials] implementation of the trait [mvpoly::MVPoly]
+/// will be used, and the mapping here is consistent with the one expected by
+/// this implementation, i.e. we simply map to an increasing number starting at
+/// 0, without any gap.
+impl From<Column> for usize {
+    fn from(val: Column) -> usize {
+        match val {
+            Column::X(i) => i,
+            Column::PublicInput(i) => NUMBER_OF_COLUMNS + i,
+            Column::Selector(_) => unimplemented!("Selectors are not supported. This method is supposed to be called only to compute the cross-term and an optimisation is in progress to avoid the inclusion of the selectors in the multi-variate polynomial."),
+        }
+    }
+}
+
 pub struct Challenges<F: Field> {
     /// Challenge used to aggregate the constraints
     pub alpha: F,

arrabiata/src/main.rs

@@ -101,6 +101,12 @@ pub fn main() {
         // FIXME:
         // Compute the accumulator for the permutation argument
 
+        // FIXME:
+        // Commit to the accumulator and absorb the commitment
+
+        // FIXME:
+        // Coin challenge α for combining the constraints
+
         // FIXME:
         // Compute the cross-terms
 

arrabiata/tests/constraints.rs

@@ -1,13 +1,15 @@
-use num_bigint::BigInt;
-use std::collections::HashMap;
-
+use ark_ff::UniformRand;
 use arrabiata::{
-    columns::Gadget,
+    columns::{ChallengeTerm, Column, Gadget},
     constraints,
     interpreter::{self, Instruction},
-    poseidon_3_60_0_5_5_fp, poseidon_3_60_0_5_5_fq,
+    poseidon_3_60_0_5_5_fp, poseidon_3_60_0_5_5_fq, MAX_DEGREE, NUMBER_OF_COLUMNS,
+    NUMBER_OF_PUBLIC_INPUTS,
 };
 use mina_curves::pasta::fields::{Fp, Fq};
+use mvpoly::{monomials::Sparse, MVPoly};
+use num_bigint::BigInt;
+use std::collections::HashMap;
 
 fn helper_compute_constraints_gadget(instr: Instruction, exp_constraints: usize) {
     let mut constraints_fp = {
@@ -164,3 +166,46 @@ fn test_gadget_elliptic_curve_scaling() {
 
     helper_check_gadget_activated(instr, Gadget::EllipticCurveScaling);
 }
+
+// This test is mostly an example to show to compute the cross-terms when we do
+// have the expressions, and some evaluations.
+// It doesn't test anything in particular. It is mostly an "integration" test.
+#[test]
+fn test_integration_with_mvpoly_to_compute_cross_terms() {
+    let constraints_fp = {
+        let poseidon_mds = poseidon_3_60_0_5_5_fp::static_params().mds.clone();
+        constraints::Env::<Fp>::new(poseidon_mds.to_vec(), BigInt::from(0_usize))
+    };
+
+    let constraints = constraints_fp.get_all_constraints_for_ivc();
+    let mut rng = o1_utils::tests::make_test_rng(None);
+
+    // Simulating two homogenising values
+    let u1 = Fp::rand(&mut rng);
+    let u2 = Fp::rand(&mut rng);
+
+    // Simulating two constraint combiners
+    let alpha_1 = Fp::rand(&mut rng);
+    let alpha_2 = Fp::rand(&mut rng);
+
+    // Simulating some row evaluations. Only 15 columns for now.
+    let eval1: [Fp; NUMBER_OF_COLUMNS + NUMBER_OF_PUBLIC_INPUTS] =
+        std::array::from_fn(|_| Fp::rand(&mut rng));
+    let eval2: [Fp; NUMBER_OF_COLUMNS + NUMBER_OF_PUBLIC_INPUTS] =
+        std::array::from_fn(|_| Fp::rand(&mut rng));
+
+    let polys = constraints
+        .iter()
+        .map(|c| {
+            println!("Converting into mvpoly the constraint: {:?}", c);
+            // Adding one to the maximum degree to account for the variable α.
+            Sparse::<
+                    Fp,
+                    { NUMBER_OF_COLUMNS + NUMBER_OF_PUBLIC_INPUTS },
+                    { MAX_DEGREE as usize + 1 },
+                >::from_expr::<Column, ChallengeTerm>(c.clone())
+        })
+        .collect();
+    let _cross_terms =
+        mvpoly::compute_combined_cross_terms(polys, alpha_1, alpha_2, eval1, eval2, u1, u2);
+}

mvpoly/src/lib.rs

@@ -261,3 +261,70 @@ pub trait MVPoly<F: PrimeField, const N: usize, const D: usize>:
     /// variable in each monomial is of maximum degree 1.
     fn is_multilinear(&self) -> bool;
 }
+
+/// Compute the cross terms of a list of polynomials. The polynomials are
+/// linearly combined using the power of a combiner, often called `α`.
+/// Powers of the combiner are considered as an additional variable of the
+/// multi-variate polynomials, therefore increasing the degree of the sum of all
+/// polynomials by one. In other words, if we combine (M + 1) polynomials
+/// `P_1(X_1, ..., X_N)`, ..., `P_{M + 1}(X_1, ..., X_N)`, we will compute the
+/// cross-terms of the polynomial:
+///
+/// ```text
+/// P(X_1, ..., X_N, α_1, ..., α_M) =     P_1(X_1, ..., X_N)
+///                                 + α_1 P_2(X_1, ..., X_N)
+///                                 + ...
+///                                 + α_M P_{M + 1}(X_1, ..., X_N)
+/// ```
+///
+/// We notice that we simply do have a shifted linear combination of the
+/// polynomials. Based on the property that the cross-terms of a sum of
+/// polynomials is the sum of the cross-terms of the individual polynomials, and
+/// that shifting a polynomial with a new variable simply requires to shift the
+/// power of cross-terms by one and multiply by the evaluation at the new point,
+/// we get a simple algorithm.
+///
+/// All polynomials are supposed to be given as multivariate polynomials of the
+/// same degree and the same number of variables, even if they do not have the
+/// same number of variables nor the same degree.
+/// Therefore, `N` is the maximum number of variables of the polynomials and `D - 1`
+/// is the maximum degree of the polynomials, `D` being the the degree of the
+/// combined polynomial when including the combiner.
+///
+/// The number of keys in the output is always `D`.
+pub fn compute_combined_cross_terms<
+    F: PrimeField,
+    const N: usize,
+    const D: usize,
+    T: MVPoly<F, N, D>,
+>(
+    polys: Vec<T>,
+    combiner1: F,
+    combiner2: F,
+    eval1: [F; N],
+    eval2: [F; N],
+    u1: F,
+    u2: F,
+) -> HashMap<usize, F> {
+    let mut cross_terms = HashMap::new();
+    polys.into_iter().enumerate().for_each(|(i, poly)| {
+        let cross_terms_poly: HashMap<usize, F> = poly.compute_cross_terms(&eval1, &eval2, u1, u2);
+        cross_terms_poly.into_iter().for_each(|(p, value)| {
+            // Computing α^i * value, and adding it to the same power of r
+            // Handling 0^0 = 1
+            let entry = cross_terms.entry(p).or_insert(F::zero());
+            if combiner1 != F::zero() {
+                let alpha_i = combiner1.pow([i as u64]);
+                *entry += alpha_i * value;
+            };
+            // Computing α'^i * value, and adding it to the next power of r
+            // Handling 0^0 = 1
+            let entry_prime = cross_terms.entry(p + 1).or_insert_with(F::zero);
+            if combiner2 != F::zero() {
+                let alpha_i = combiner2.pow([i as u64]);
+                *entry_prime += alpha_i * value;
+            };
+        });
+    });
+    cross_terms
+}