Merge pull request #2684 from o1-labs/dw/arrabiata-challenges-express…

…ions

Arrabiata: define correctly the different challenges used in the protocol

arrabiata/Cargo.toml

@@ -31,5 +31,6 @@ once_cell.workspace = true
 poly-commitment.workspace = true
 rand.workspace = true
 rayon.workspace = true
+serde.workspace = true
 strum.workspace = true
 strum_macros.workspace = true

arrabiata/src/columns.rs

@@ -1,8 +1,11 @@
-use kimchi::circuits::{
-    berkeley_columns::BerkeleyChallengeTerm,
-    expr::{CacheId, ConstantExpr, Expr, FormattedOutput},
+use ark_ff::Field;
+use kimchi::circuits::expr::{AlphaChallengeTerm, CacheId, ConstantExpr, Expr, FormattedOutput};
+use serde::{Deserialize, Serialize};
+use std::{
+    collections::HashMap,
+    fmt::{Display, Formatter, Result},
+    ops::Index,
 };
-use std::collections::HashMap;
 use strum_macros::{EnumCount as EnumCountMacro, EnumIter};
 
 /// This enum represents the different gadgets that can be used in the circuit.
@@ -48,8 +51,64 @@ pub enum Column {
     X(usize),
 }
 
-// FIXME: We should use something different than BerkeleyChallengeTerm here
-pub type E<Fp> = Expr<ConstantExpr<Fp, BerkeleyChallengeTerm>, Column>;
+pub struct Challenges<F: Field> {
+    /// Challenge used to aggregate the constraints
+    pub alpha: F,
+
+    /// Both challenges used in the permutation argument
+    pub beta: F,
+    pub gamma: F,
+
+    /// Challenge to homogenize the constraints
+    pub homogenous_challenge: F,
+
+    /// Random coin used to aggregate witnesses while folding
+    pub r: F,
+}
+
+#[derive(Copy, Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
+pub enum ChallengeTerm {
+    /// Challenge used to aggregate the constraints
+    Alpha,
+    /// Both challenges used in the permutation argument
+    Beta,
+    Gamma,
+    /// Challenge to homogenize the constraints
+    HomogenousChallenge,
+    /// Random coin used to aggregate witnesses while folding
+    R,
+}
+
+impl Display for ChallengeTerm {
+    fn fmt(&self, f: &mut Formatter) -> Result {
+        match self {
+            ChallengeTerm::Alpha => write!(f, "alpha"),
+            ChallengeTerm::Beta => write!(f, "beta"),
+            ChallengeTerm::Gamma => write!(f, "gamma"),
+            ChallengeTerm::HomogenousChallenge => write!(f, "u"),
+            ChallengeTerm::R => write!(f, "r"),
+        }
+    }
+}
+impl<F: Field> Index<ChallengeTerm> for Challenges<F> {
+    type Output = F;
+
+    fn index(&self, term: ChallengeTerm) -> &Self::Output {
+        match term {
+            ChallengeTerm::Alpha => &self.alpha,
+            ChallengeTerm::Beta => &self.beta,
+            ChallengeTerm::Gamma => &self.gamma,
+            ChallengeTerm::HomogenousChallenge => &self.homogenous_challenge,
+            ChallengeTerm::R => &self.r,
+        }
+    }
+}
+
+impl<'a> AlphaChallengeTerm<'a> for ChallengeTerm {
+    const ALPHA: Self = Self::Alpha;
+}
+
+pub type E<Fp> = Expr<ConstantExpr<Fp, ChallengeTerm>, Column>;
 
 // Code to allow for pretty printing of the expressions
 impl FormattedOutput for Column {

kimchi/src/circuits/berkeley_columns.rs

@@ -1,29 +1,52 @@
+//! This module defines the particular form of the expressions used in the Mina
+//! Berkeley hardfork. You can find more information in [this blog
+//! article](https://www.o1labs.org/blog/reintroducing-kimchi).
+//! This module is also a good starting point if you want to implement your own
+//! variant of Kimchi using the expression framework.
+//!
+//! The module uses the generic expression framework defined in the
+//! [crate::circuits::expr] module.
+//! The expressions define the polynomials that can be used to describe the
+//! constraints.
+//! It starts by defining the different challenges used by the PLONK IOP in
+//! [BerkeleyChallengeTerm] and [BerkeleyChallenges].
+//! It then defines the [Column] type which represents the different variables
+//! the polynomials are defined over.
+//!
+//! Two "environments" are after that defined: one for the lookup argument
+//! [LookupEnvironment], and one for the main argument [Environment], which
+//! contains the former.
+//! The trait [ColumnEnvironment] is then defined to provide the necessary
+//! primitives used to evaluate the quotient polynomial.
+
 use crate::{
     circuits::{
         domains::EvaluationDomains,
-        expr::{CacheId, ColumnEvaluations, ConstantExpr, ConstantTerm, Expr, ExprError},
+        expr::{
+            CacheId, ColumnEnvironment, ColumnEvaluations, ConstantExpr, ConstantTerm, Constants,
+            Domain, Expr, ExprError, FormattedOutput,
+        },
         gate::{CurrOrNext, GateType},
         lookup::{index::LookupSelectors, lookups::LookupPattern},
+        wires::COLUMNS,
     },
     proof::{PointEvaluations, ProofEvaluations},
 };
-use serde::{Deserialize, Serialize};
-
 use ark_ff::FftField;
 use ark_poly::{Evaluations, Radix2EvaluationDomain as D};
-
-use crate::circuits::expr::{ColumnEnvironment, Constants, Domain, FormattedOutput};
-
-use crate::circuits::wires::COLUMNS;
-
+use serde::{Deserialize, Serialize};
 use std::collections::HashMap;
 
-/// The challenge terms used in Berkeley
+/// The challenge terms used in Berkeley.
 #[derive(Copy, Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
 pub enum BerkeleyChallengeTerm {
+    /// Used to combine constraints
     Alpha,
+    /// The first challenge used in the permutation argument
     Beta,
+    /// The second challenge used in the permutation argument
     Gamma,
+    /// A challenge used to columns of a lookup table
     JointCombiner,
 }
 
@@ -45,13 +68,14 @@ impl<'a> super::expr::AlphaChallengeTerm<'a> for BerkeleyChallengeTerm {
 }
 
 pub struct BerkeleyChallenges<F> {
-    /// The challenge alpha from the PLONK IOP.
+    /// The challenge α from the PLONK IOP.
     pub alpha: F,
-    /// The challenge beta from the PLONK IOP.
+    /// The challenge β from the PLONK IOP.
     pub beta: F,
-    /// The challenge gamma from the PLONK IOP.
+    /// The challenge γ from the PLONK IOP.
     pub gamma: F,
-    /// The challenge joint_combiner which is used to combine joint lookup tables.
+    /// The challenge joint_combiner which is used to combine joint lookup
+    /// tables.
     pub joint_combiner: F,
 }
 
@@ -68,9 +92,16 @@ impl<F: ark_ff::Field> std::ops::Index<BerkeleyChallengeTerm> for BerkeleyChalle
     }
 }
 
+/// A type representing the variables involved in the constraints of the
+/// Berkeley hardfork.
+///
+/// In Berkeley, the constraints are defined over the following variables:
+/// - The [COLUMNS] witness columns.
+/// - The permutation polynomial, Z.
+/// - The public coefficients, `Coefficients`, which can be used for public
+/// values. For instance, it is used for the Poseidon round constants.
+/// - ...
 #[derive(Copy, Clone, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
-/// A type representing one of the polynomials involved in the PLONK IOP, use in
-/// the Berkeley hardfork.
 pub enum Column {
     Witness(usize),
     Z,

kimchi/src/circuits/domains.rs

@@ -19,9 +19,10 @@ pub struct EvaluationDomains<F: FftField> {
 }
 
 impl<F: FftField> EvaluationDomains<F> {
-    /// Creates 4 evaluation domains `d1` (of size `n`), `d2` (of size `2n`), `d4` (of size `4n`),
-    /// and `d8` (of size `8n`). If generator of `d8` is `g`, the generator
-    /// of `d4` is `g^2`, the generator of `d2` is `g^4`, and the generator of `d1` is `g^8`.
+    /// Creates 4 evaluation domains `d1` (of size `n`), `d2` (of size `2n`),
+    /// `d4` (of size `4n`), and `d8` (of size `8n`). If generator of `d8` is
+    /// `g`, the generator of `d4` is `g^2`, the generator of `d2` is `g^4`, and
+    /// the generator of `d1` is `g^8`.
     pub fn create(n: usize) -> Result<Self, DomainCreationError> {
         let n = Domain::<F>::compute_size_of_domain(n)
             .ok_or(DomainCreationError::DomainSizeFailed(n))?;
@@ -33,7 +34,8 @@ impl<F: FftField> EvaluationDomains<F> {
 
         // we also create domains of larger sizes
         // to efficiently operate on polynomials in evaluation form.
-        // (in evaluation form, the domain needs to grow as the degree of a polynomial grows)
+        // (in evaluation form, the domain needs to grow as the degree of a
+        // polynomial grows)
         let d2 = Domain::<F>::new(2 * n).ok_or(DomainCreationError::DomainConstructionFailed(
             "d2".to_string(),
             2 * n,
@@ -47,32 +49,12 @@ impl<F: FftField> EvaluationDomains<F> {
             8 * n,
         ))?;
 
-        // ensure the relationship between the three domains in case the library's behavior changes
+        // ensure the relationship between the three domains in case the
+        // library's behavior changes
         assert_eq!(d2.group_gen.square(), d1.group_gen);
         assert_eq!(d4.group_gen.square(), d2.group_gen);
         assert_eq!(d8.group_gen.square(), d4.group_gen);
 
         Ok(EvaluationDomains { d1, d2, d4, d8 })
     }
 }
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-    use ark_ff::Field;
-    use mina_curves::pasta::Fp;
-
-    #[test]
-    #[ignore] // TODO(mimoo): wait for fix upstream (https://github.com/arkworks-rs/algebra/pull/307)
-    fn test_create_domain() {
-        if let Ok(d) = EvaluationDomains::<Fp>::create(usize::MAX) {
-            assert!(d.d4.group_gen.pow([4]) == d.d1.group_gen);
-            assert!(d.d8.group_gen.pow([2]) == d.d4.group_gen);
-            println!("d8 = {:?}", d.d8.group_gen);
-            println!("d8^2 = {:?}", d.d8.group_gen.pow([2]));
-            println!("d4 = {:?}", d.d4.group_gen);
-            println!("d4 = {:?}", d.d4.group_gen.pow([4]));
-            println!("d1 = {:?}", d.d1.group_gen);
-        }
-    }
-}

kimchi/src/circuits/expr.rs

@@ -146,18 +146,6 @@ pub struct Variable<Column> {
     pub row: CurrOrNext,
 }
 
-/// Define challenges the verifier coins during the interactive protocol.
-/// It has been defined initially to handle the PLONK IOP, hence:
-/// - `alpha` for the quotient polynomial
-/// - `beta` and `gamma` for the permutation challenges.
-/// The joint combiner is to handle vector lookups, initially designed to be
-/// used with PLOOKUP.
-/// The terms have no built-in semantic in the expression framework, and can be
-/// used for any other four challenges the verifier coins in other polynomial
-/// interactive protocol.
-/// TODO: we should generalize the expression type over challenges and constants.
-/// See <https://github.com/MinaProtocol/mina/issues/15287>
-
 /// Define the constant terms an expression can use.
 /// It can be any constant term (`Literal`), a matrix (`Mds` - used by the
 /// permutation used by Poseidon for instance), or endomorphism coefficients
@@ -855,28 +843,6 @@ impl<Column: Copy> Variable<Column> {
     }
 }
 
-impl<Column: FormattedOutput + Debug> Variable<Column> {
-    pub fn ocaml(&self) -> String {
-        format!("var({:?}, {:?})", self.col, self.row)
-    }
-
-    pub fn latex(&self) -> String {
-        let col = self.col.latex(&mut HashMap::new());
-        match self.row {
-            Curr => col,
-            Next => format!("\\tilde{{{col}}}"),
-        }
-    }
-
-    pub fn text(&self) -> String {
-        let col = self.col.text(&mut HashMap::new());
-        match self.row {
-            Curr => format!("Curr({col})"),
-            Next => format!("Next({col})"),
-        }
-    }
-}
-
 impl<F: FftField, Column: Copy, ChallengeTerm: Copy> PolishToken<F, Column, ChallengeTerm> {
     /// Evaluate an RPN expression to a field element.
     pub fn evaluate<Evaluations: ColumnEvaluations<F, Column = Column>>(
@@ -2882,6 +2848,32 @@ where
     }
 }
 
+impl<Column: FormattedOutput + Debug> FormattedOutput for Variable<Column> {
+    fn is_alpha(&self) -> bool {
+        false
+    }
+
+    fn ocaml(&self, _cache: &mut HashMap<CacheId, Self>) -> String {
+        format!("var({:?}, {:?})", self.col, self.row)
+    }
+
+    fn latex(&self, _cache: &mut HashMap<CacheId, Self>) -> String {
+        let col = self.col.latex(&mut HashMap::new());
+        match self.row {
+            Curr => col,
+            Next => format!("\\tilde{{{col}}}"),
+        }
+    }
+
+    fn text(&self, _cache: &mut HashMap<CacheId, Self>) -> String {
+        let col = self.col.text(&mut HashMap::new());
+        match self.row {
+            Curr => format!("Curr({col})"),
+            Next => format!("Next({col})"),
+        }
+    }
+}
+
 impl<T: FormattedOutput + Clone> FormattedOutput for Operations<T> {
     fn is_alpha(&self) -> bool {
         match self {
@@ -3023,7 +3015,7 @@ where
                 });
                 res
             }
-            Atom(Cell(v)) => format!("cell({})", v.ocaml()),
+            Atom(Cell(v)) => format!("cell({})", v.ocaml(&mut HashMap::new())),
             Atom(UnnormalizedLagrangeBasis(i)) => {
                 format!("unnormalized_lagrange_basis({}, {})", i.zk_rows, i.offset)
             }
@@ -3087,7 +3079,7 @@ where
                 });
                 res
             }
-            Atom(Cell(v)) => v.latex(),
+            Atom(Cell(v)) => v.latex(&mut HashMap::new()),
             Atom(UnnormalizedLagrangeBasis(RowOffset {
                 zk_rows: true,
                 offset: i,
@@ -3134,7 +3126,7 @@ where
                 });
                 res
             }
-            Atom(Cell(v)) => v.text(),
+            Atom(Cell(v)) => v.text(&mut HashMap::new()),
             Atom(UnnormalizedLagrangeBasis(RowOffset {
                 zk_rows: true,
                 offset: i,

kimchi/src/circuits/gate.rs

@@ -587,14 +587,6 @@ mod tests {
     use proptest::prelude::*;
     use rand::SeedableRng as _;
 
-    // TODO: move to mina-curves
-    prop_compose! {
-        pub fn arb_fp()(seed: [u8; 32]) -> Fp {
-            let rng = &mut rand::rngs::StdRng::from_seed(seed);
-            Fp::rand(rng)
-        }
-    }
-
     prop_compose! {
         fn arb_fp_vec(max: usize)(seed: [u8; 32], num in 0..max) -> Vec<Fp> {
             let rng = &mut rand::rngs::StdRng::from_seed(seed);

kimchi/tests/test_domain.rs

@@ -0,0 +1,16 @@
+use ark_ff::Field;
+use kimchi::circuits::domains::EvaluationDomains;
+use mina_curves::pasta::Fp;
+
+#[test]
+fn test_create_domain() {
+    if let Ok(d) = EvaluationDomains::<Fp>::create(usize::MAX) {
+        assert!(d.d4.group_gen.pow([4]) == d.d1.group_gen);
+        assert!(d.d8.group_gen.pow([2]) == d.d4.group_gen);
+        println!("d8 = {:?}", d.d8.group_gen);
+        println!("d8^2 = {:?}", d.d8.group_gen.pow([2]));
+        println!("d4 = {:?}", d.d4.group_gen);
+        println!("d4 = {:?}", d.d4.group_gen.pow([4]));
+        println!("d1 = {:?}", d.d1.group_gen);
+    }
+}