[#1793] MSM FFA: Rewrite existing circuits to use interpreter

msm/src/ffa/constraint.rs

@@ -1,62 +1,116 @@
 use crate::{
     columns::{Column, ColumnIndexer},
     expr::MSMExpr,
-    ffa::columns::FFAColumnIndexer,
-    {Fp, N_LIMBS},
+    ffa::{columns::FFAColumnIndexer, interpreter::FFAInterpreterEnv},
 };
+use ark_ff::PrimeField;
 use kimchi::circuits::{
-    expr::{ConstantExprInner, ExprInner, Operations, Variable},
+    expr::{ConstantExpr, ConstantTerm, Expr, ExprInner, Variable},
     gate::CurrOrNext,
 };
 
-/// Access exprs for addition
-pub fn get_exprs_add() -> Vec<MSMExpr<Fp>> {
-    let mut limb_exprs: Vec<_> = vec![];
-    for i in 0..N_LIMBS {
-        let limb_constraint = {
-            let a_i = MSMExpr::Atom(
-                ExprInner::<Operations<ConstantExprInner<Fp>>, Column>::Cell(Variable {
-                    col: FFAColumnIndexer::A(i).ix_to_column(),
-                    row: CurrOrNext::Curr,
-                }),
-            );
-            let b_i = MSMExpr::Atom(ExprInner::Cell(Variable {
-                col: FFAColumnIndexer::B(i).ix_to_column(),
-                row: CurrOrNext::Curr,
-            }));
-            let c_i = MSMExpr::Atom(ExprInner::Cell(Variable {
-                col: FFAColumnIndexer::C(i).ix_to_column(),
-                row: CurrOrNext::Curr,
-            }));
-            a_i + b_i - c_i
-        };
-        limb_exprs.push(limb_constraint);
+/// Contains constraints for just one row.
+pub struct ConstraintBuilder<F> {
+    pub constraints: Vec<MSMExpr<F>>,
+}
+
+impl<F: PrimeField> FFAInterpreterEnv<F> for ConstraintBuilder<F> {
+    type Position = Column;
+
+    type Variable = MSMExpr<F>;
+
+    fn assert_zero(&mut self, cst: Self::Variable) {
+        self.constraints.push(cst)
+    }
+
+    fn copy(&mut self, x: &Self::Variable, position: Self::Position) -> Self::Variable {
+        let y = Expr::Atom(ExprInner::Cell(Variable {
+            col: position,
+            row: CurrOrNext::Curr,
+        }));
+        self.constraints.push(y.clone() - x.clone());
+        y
+    }
+
+    fn constant(value: F) -> Self::Variable {
+        let cst_expr_inner = ConstantExpr::from(ConstantTerm::Literal(value));
+        Expr::Atom(ExprInner::Constant(cst_expr_inner))
+    }
+
+    // TODO deduplicate, remove this
+    fn column_pos(ix: FFAColumnIndexer) -> Self::Position {
+        ix.ix_to_column()
+    }
+
+    fn read_column(&self, ix: FFAColumnIndexer) -> Self::Variable {
+        Expr::Atom(ExprInner::Cell(Variable {
+            col: ix.ix_to_column(),
+            row: CurrOrNext::Curr,
+        }))
+    }
+
+    fn next_row(&mut self) {
+        panic!("Please don't call this");
     }
-    limb_exprs
 }
 
-/// Get expressions for multiplication
-pub fn get_exprs_mul() -> Vec<MSMExpr<Fp>> {
-    let mut limb_exprs: Vec<_> = vec![];
-    for i in 0..N_LIMBS {
-        let limb_constraint = {
-            let a_i = MSMExpr::Atom(
-                ExprInner::<Operations<ConstantExprInner<Fp>>, Column>::Cell(Variable {
-                    col: FFAColumnIndexer::A(i).ix_to_column(),
-                    row: CurrOrNext::Curr,
-                }),
-            );
-            let b_i = MSMExpr::Atom(ExprInner::Cell(Variable {
-                col: FFAColumnIndexer::B(i).ix_to_column(),
-                row: CurrOrNext::Curr,
-            }));
-            let d_i = MSMExpr::Atom(ExprInner::Cell(Variable {
-                col: FFAColumnIndexer::D(i).ix_to_column(),
-                row: CurrOrNext::Curr,
-            }));
-            a_i * b_i - d_i
-        };
-        limb_exprs.push(limb_constraint);
+impl<F: PrimeField> ConstraintBuilder<F> {
+    pub fn empty() -> Self {
+        ConstraintBuilder {
+            constraints: vec![],
+        }
     }
-    limb_exprs
 }
+
+//
+///// Access exprs for addition
+//pub fn get_exprs_add() -> Vec<MSMExpr<Fp>> {
+//    let mut limb_exprs: Vec<_> = vec![];
+//    for i in 0..N_LIMBS {
+//        let limb_constraint = {
+//            let a_i = MSMExpr::Atom(
+//                ExprInner::<Operations<ConstantExprInner<Fp>>, Column>::Cell(Variable {
+//                    col: FFAColumnIndexer::A(i).ix_to_column(),
+//                    row: CurrOrNext::Curr,
+//                }),
+//            );
+//            let b_i = MSMExpr::Atom(ExprInner::Cell(Variable {
+//                col: FFAColumnIndexer::B(i).ix_to_column(),
+//                row: CurrOrNext::Curr,
+//            }));
+//            let c_i = MSMExpr::Atom(ExprInner::Cell(Variable {
+//                col: FFAColumnIndexer::C(i).ix_to_column(),
+//                row: CurrOrNext::Curr,
+//            }));
+//            a_i + b_i - c_i
+//        };
+//        limb_exprs.push(limb_constraint);
+//    }
+//    limb_exprs
+//}
+//
+///// Get expressions for multiplication
+//pub fn get_exprs_mul() -> Vec<MSMExpr<Fp>> {
+//    let mut limb_exprs: Vec<_> = vec![];
+//    for i in 0..N_LIMBS {
+//        let limb_constraint = {
+//            let a_i = MSMExpr::Atom(
+//                ExprInner::<Operations<ConstantExprInner<Fp>>, Column>::Cell(Variable {
+//                    col: FFAColumnIndexer::A(i).ix_to_column(),
+//                    row: CurrOrNext::Curr,
+//                }),
+//            );
+//            let b_i = MSMExpr::Atom(ExprInner::Cell(Variable {
+//                col: FFAColumnIndexer::B(i).ix_to_column(),
+//                row: CurrOrNext::Curr,
+//            }));
+//            let d_i = MSMExpr::Atom(ExprInner::Cell(Variable {
+//                col: FFAColumnIndexer::D(i).ix_to_column(),
+//                row: CurrOrNext::Curr,
+//            }));
+//            a_i * b_i - d_i
+//        };
+//        limb_exprs.push(limb_constraint);
+//    }
+//    limb_exprs
+//}

msm/src/ffa/interpreter.rs

@@ -1,7 +1,9 @@
-use crate::ffa::columns::FFAColumnIndexer;
-use ark_ff::PrimeField;
+use crate::{ffa::columns::FFAColumnIndexer, Ff1, N_LIMBS};
+use ark_ff::{PrimeField, Zero};
+use num_bigint::BigUint;
+use o1_utils::{field_helpers::FieldHelpers, foreign_field::ForeignElement};
 
-pub trait FFAInterpreterEnv<Fp: PrimeField> {
+pub trait FFAInterpreterEnv<F: PrimeField> {
     type Position;
 
     type Variable: Clone
@@ -10,11 +12,97 @@ pub trait FFAInterpreterEnv<Fp: PrimeField> {
         + std::ops::Mul<Self::Variable, Output = Self::Variable>
         + std::fmt::Debug;
 
-    fn add_constraint(&mut self, cst: Self::Variable);
+    fn assert_zero(&mut self, cst: Self::Variable);
 
     fn copy(&mut self, x: &Self::Variable, position: Self::Position) -> Self::Variable;
 
-    fn constant(value: Fp) -> Self::Variable;
+    fn constant(value: F) -> Self::Variable;
 
-    fn get_column(ix: FFAColumnIndexer) -> Self::Position;
+    fn column_pos(ix: FFAColumnIndexer) -> Self::Position;
+
+    fn read_column(&self, ix: FFAColumnIndexer) -> Self::Variable;
+
+    /// In constraint environment does nothing (?). In witness environment progresses to the next row.
+    fn next_row(&mut self);
+}
+
+// TODO use more foreign_field.rs with from/to bigint conversion
+fn limb_decompose<F: PrimeField>(input: &Ff1) -> [F; N_LIMBS] {
+    let input_bi: BigUint = FieldHelpers::to_biguint(input);
+    let ff_el: ForeignElement<F, N_LIMBS> = ForeignElement::from_biguint(input_bi);
+    ff_el.limbs
+}
+
+/// Reads values from limbs A and B, returns resulting value in C.
+pub fn constrain_multiplication<F: PrimeField, Env: FFAInterpreterEnv<F>>(
+    env: &mut Env,
+) -> [Env::Variable; N_LIMBS] {
+    let a_limbs: [Env::Variable; N_LIMBS] =
+        core::array::from_fn(|i| Env::read_column(env, FFAColumnIndexer::A(i)));
+    let b_limbs: [Env::Variable; N_LIMBS] =
+        core::array::from_fn(|i| Env::read_column(env, FFAColumnIndexer::B(i)));
+    // fix cloning
+    let c_limbs: [Env::Variable; N_LIMBS] =
+        core::array::from_fn(|i| a_limbs[i].clone() * b_limbs[i].clone());
+    c_limbs.iter().enumerate().for_each(|(i, var)| {
+        env.copy(var, Env::column_pos(FFAColumnIndexer::C(i)));
+    });
+    c_limbs
+}
+
+pub fn test_multiplication<F: PrimeField, Env: FFAInterpreterEnv<F>>(
+    env: &mut Env,
+    a: Ff1,
+    b: Ff1,
+) {
+    let a_limbs: [Env::Variable; N_LIMBS] = limb_decompose(&a).map(Env::constant);
+    let b_limbs: [Env::Variable; N_LIMBS] = limb_decompose(&b).map(Env::constant);
+    a_limbs.iter().enumerate().for_each(|(i, var)| {
+        env.copy(var, Env::column_pos(FFAColumnIndexer::A(i)));
+    });
+    b_limbs.iter().enumerate().for_each(|(i, var)| {
+        env.copy(var, Env::column_pos(FFAColumnIndexer::B(i)));
+    });
+
+    let _ = constrain_multiplication(env); // we don't do anything else further with c_limbs
+
+    let d_limbs: [Env::Variable; N_LIMBS] = [Zero::zero(); N_LIMBS].map(Env::constant);
+    d_limbs.iter().enumerate().for_each(|(i, var)| {
+        env.copy(var, Env::column_pos(FFAColumnIndexer::D(i)));
+    });
+}
+
+/// Reads values from limbs A and B, returns resulting value in C.
+pub fn constrain_addition<F: PrimeField, Env: FFAInterpreterEnv<F>>(
+    env: &mut Env,
+) -> [Env::Variable; N_LIMBS] {
+    let a_limbs: [Env::Variable; N_LIMBS] =
+        core::array::from_fn(|i| Env::read_column(env, FFAColumnIndexer::A(i)));
+    let b_limbs: [Env::Variable; N_LIMBS] =
+        core::array::from_fn(|i| Env::read_column(env, FFAColumnIndexer::B(i)));
+    // fix cloning
+    let c_limbs: [Env::Variable; N_LIMBS] =
+        core::array::from_fn(|i| a_limbs[i].clone() + b_limbs[i].clone());
+    c_limbs.iter().enumerate().for_each(|(i, var)| {
+        env.copy(var, Env::column_pos(FFAColumnIndexer::C(i)));
+    });
+    c_limbs
+}
+
+pub fn test_addition<F: PrimeField, Env: FFAInterpreterEnv<F>>(env: &mut Env, a: Ff1, b: Ff1) {
+    let a_limbs: [Env::Variable; N_LIMBS] = limb_decompose(&a).map(Env::constant);
+    let b_limbs: [Env::Variable; N_LIMBS] = limb_decompose(&b).map(Env::constant);
+    a_limbs.iter().enumerate().for_each(|(i, var)| {
+        env.copy(var, Env::column_pos(FFAColumnIndexer::A(i)));
+    });
+    b_limbs.iter().enumerate().for_each(|(i, var)| {
+        env.copy(var, Env::column_pos(FFAColumnIndexer::B(i)));
+    });
+
+    let _ = constrain_addition(env); // we don't do anything else further with c_limbs
+
+    let d_limbs: [Env::Variable; N_LIMBS] = [Zero::zero(); N_LIMBS].map(Env::constant);
+    d_limbs.iter().enumerate().for_each(|(i, var)| {
+        env.copy(var, Env::column_pos(FFAColumnIndexer::D(i)));
+    });
 }

msm/src/ffa/main.rs

@@ -1,34 +1,21 @@
-use rand::{thread_rng, Rng};
+use rand::Rng;
 
 use kimchi::circuits::domains::EvaluationDomains;
 use poly_commitment::pairing_proof::PairingSRS;
 
 use kimchi_msm::columns::Column;
 use kimchi_msm::ffa::{
-    columns::FFA_N_COLUMNS, constraint::get_exprs_add, witness::WitnessBuilder as FFAWitnessBuilder,
+    columns::FFA_N_COLUMNS,
+    constraint::ConstraintBuilder as FFAConstraintBuilder,
+    interpreter::{self as ffa_interpreter, FFAInterpreterEnv},
+    witness::WitnessBuilder as FFAWitnessBuilder,
 };
 use kimchi_msm::lookups::LookupTableIDs;
 use kimchi_msm::precomputed_srs::get_bn254_srs;
 use kimchi_msm::prover::prove;
 use kimchi_msm::verifier::verify;
 use kimchi_msm::{BaseSponge, Ff1, Fp, OpeningProof, ScalarSponge, BN254, DOMAIN_SIZE};
 
-pub fn generate_random_msm_witness() -> FFAWitnessBuilder<Fp> {
-    let mut circuit_env = FFAWitnessBuilder::<Fp>::empty();
-    let mut rng = thread_rng();
-
-    let row_num = DOMAIN_SIZE;
-    assert!(row_num <= DOMAIN_SIZE);
-
-    for _row_i in 0..row_num {
-        let a: Ff1 = From::from(rng.gen_range(0..(1 << 16)));
-        let b: Ff1 = From::from(rng.gen_range(0..(1 << 16)));
-        circuit_env.add_test_addition(a, b);
-    }
-
-    circuit_env
-}
-
 pub fn main() {
     // FIXME: use a proper RNG
     let mut rng = o1_utils::tests::make_test_rng();
@@ -38,9 +25,23 @@ pub fn main() {
 
     let srs: PairingSRS<BN254> = get_bn254_srs(domain);
 
-    let circuit_env = generate_random_msm_witness();
-    let proof_inputs = circuit_env.get_witness();
-    let constraint_exprs = get_exprs_add();
+    let mut witness_env = FFAWitnessBuilder::<Fp>::empty();
+    let mut constraint_env = FFAConstraintBuilder::<Fp>::empty();
+
+    ffa_interpreter::constrain_multiplication(&mut constraint_env);
+
+    let row_num = 10;
+    assert!(row_num <= DOMAIN_SIZE);
+
+    for _row_i in 0..row_num {
+        let a: Ff1 = From::from(rng.gen_range(0..(1 << 16)));
+        let b: Ff1 = From::from(rng.gen_range(0..(1 << 16)));
+        ffa_interpreter::test_multiplication(&mut witness_env, a, b);
+        witness_env.next_row();
+    }
+
+    let proof_inputs = witness_env.get_witness();
+    let constraint_exprs = constraint_env.constraints;
 
     println!("Proof inputs: {:?}", proof_inputs);
     println!("Constraints: {:?}", constraint_exprs);