Changes to proof system to support Keccak gates (#1267, #1312 Revival)

book/src/specs/kimchi.md

@@ -1856,6 +1856,14 @@ pub struct VerifierIndex<
     #[serde(bound = "Option<PolyComm<G>>: Serialize + DeserializeOwned")]
     pub rot_comm: Option<PolyComm<G>>,
 
+    /// Keccak round commitments
+    #[serde(bound = "Option<PolyComm<G>>: Serialize + DeserializeOwned")]
+    pub keccak_round_comm: Option<PolyComm<G>>,
+
+    /// Keccak sponge commitments
+    #[serde(bound = "Option<PolyComm<G>>: Serialize + DeserializeOwned")]
+    pub keccak_sponge_comm: Option<PolyComm<G>>,
+
     /// wire coordinate shifts
     #[serde_as(as = "[o1_utils::serialization::SerdeAs; PERMUTS]")]
     pub shift: [G::ScalarField; PERMUTS],
@@ -2021,6 +2029,10 @@ pub struct ProofEvaluations<Evals, const COLUMNS: usize = KIMCHI_COLS> {
     pub xor_selector: Option<Evals>,
     /// evaluation of the Rot selector polynomial
     pub rot_selector: Option<Evals>,
+    /// evaluation of the KeccakRound selector polynomial
+    pub keccak_round_selector: Option<Evals>,
+    /// evaluation of the KeccakRound selector polynomial
+    pub keccak_sponge_selector: Option<Evals>,
 
     // lookup-related evaluations
     /// evaluation of lookup aggregation polynomial
@@ -2043,6 +2055,10 @@ pub struct ProofEvaluations<Evals, const COLUMNS: usize = KIMCHI_COLS> {
     pub range_check_lookup_selector: Option<Evals>,
     /// evaluation of the ForeignFieldMul range check pattern selector polynomial
     pub foreign_field_mul_lookup_selector: Option<Evals>,
+    /// evaluation of the KeccakRound pattern selector polynomial
+    pub keccak_round_lookup_selector: Option<Evals>,
+    /// evaluation of the KeccakSponge pattern selector polynomial
+    pub keccak_sponge_lookup_selector: Option<Evals>,
 }
 
 /// Commitments linked to the lookup feature

kimchi/src/circuits/berkeley_columns.rs

@@ -152,8 +152,12 @@ impl<F: Copy, const COLUMNS: usize> ColumnEvaluations<F>
             LookupKindIndex(LookupPattern::ForeignFieldMul) => self
                 .foreign_field_mul_lookup_selector
                 .ok_or(ExprError::MissingIndexEvaluation(col)),
-            LookupKindIndex(LookupPattern::KeccakRound) => todo!(),
-            LookupKindIndex(LookupPattern::KeccakSponge) => todo!(),
+            LookupKindIndex(LookupPattern::KeccakRound) => self
+                .keccak_round_lookup_selector
+                .ok_or(ExprError::MissingIndexEvaluation(col)),
+            LookupKindIndex(LookupPattern::KeccakSponge) => self
+                .keccak_sponge_lookup_selector
+                .ok_or(ExprError::MissingIndexEvaluation(col)),
             LookupRuntimeSelector => self
                 .runtime_lookup_table_selector
                 .ok_or(ExprError::MissingIndexEvaluation(col)),

kimchi/src/circuits/constraints.rs

@@ -46,6 +46,8 @@ pub struct FeatureFlags {
     pub xor: bool,
     /// ROT gate
     pub rot: bool,
+    /// Keccak gates
+    pub keccak: bool,
     /// Lookup features
     pub lookup_features: LookupFeatures,
 }
@@ -131,6 +133,14 @@ pub struct ColumnEvaluations<F: PrimeField, const COLUMNS: usize = KIMCHI_COLS>
     /// Rot gate selector over domain d8
     #[serde_as(as = "Option<o1_utils::serialization::SerdeAs>")]
     pub rot_selector8: Option<E<F, D<F>>>,
+
+    /// Keccak round gate selector over domain d8
+    #[serde_as(as = "Option<o1_utils::serialization::SerdeAs>")]
+    pub keccak_round_selector8: Option<E<F, D<F>>>,
+
+    /// Keccak sponge gate selector over domain d8
+    #[serde_as(as = "Option<o1_utils::serialization::SerdeAs>")]
+    pub keccak_sponge_selector8: Option<E<F, D<F>>>,
 }
 
 #[serde_as]
@@ -592,6 +602,34 @@ impl<F: PrimeField + SquareRootField> ConstraintSystem<F> {
             }
         };
 
+        let keccak_round_selector8 = {
+            if !self.feature_flags.keccak {
+                None
+            } else {
+                Some(selector_polynomial(
+                    GateType::KeccakRound,
+                    &self.gates,
+                    &self.domain,
+                    &self.domain.d8,
+                    self.disable_gates_checks,
+                ))
+            }
+        };
+
+        let keccak_sponge_selector8 = {
+            if !self.feature_flags.keccak {
+                None
+            } else {
+                Some(selector_polynomial(
+                    GateType::KeccakSponge,
+                    &self.gates,
+                    &self.domain,
+                    &self.domain.d8,
+                    self.disable_gates_checks,
+                ))
+            }
+        };
+
         // TODO: This doesn't need to be degree 8 but that would require some changes in expr
         let coefficients8 = array::from_fn(|i| {
             evaluated_column_coefficients.coefficients[i]
@@ -613,6 +651,8 @@ impl<F: PrimeField + SquareRootField> ConstraintSystem<F> {
             foreign_field_mul_selector8,
             xor_selector8,
             rot_selector8,
+            keccak_round_selector8,
+            keccak_sponge_selector8,
         }
     }
 }
@@ -804,6 +844,7 @@ impl<F: PrimeField + SquareRootField> Builder<F> {
             foreign_field_mul: false,
             xor: false,
             rot: false,
+            keccak: false,
         };
 
         for gate in &gates {
@@ -814,6 +855,7 @@ impl<F: PrimeField + SquareRootField> Builder<F> {
                 GateType::ForeignFieldMul => feature_flags.foreign_field_mul = true,
                 GateType::Xor16 => feature_flags.xor = true,
                 GateType::Rot64 => feature_flags.rot = true,
+                GateType::KeccakRound | GateType::KeccakSponge => feature_flags.keccak = true,
                 _ => (),
             }
         }

kimchi/src/circuits/expr.rs

@@ -402,6 +402,7 @@ pub enum FeatureFlag {
     ForeignFieldMul,
     Xor,
     Rot,
+    Keccak,
     LookupTables,
     RuntimeLookupTables,
     LookupPattern(LookupPattern),
@@ -560,6 +561,7 @@ impl<C: Zero + One + Neg<Output = C> + PartialEq + Clone, Column: Clone + Partia
                         ForeignFieldMul => features.foreign_field_mul,
                         Xor => features.xor,
                         Rot => features.rot,
+                        Keccak => features.keccak,
                         LookupTables => {
                             features.lookup_features.patterns != LookupPatterns::default()
                         }

kimchi/src/circuits/lookup/lookups.rs

@@ -733,6 +733,8 @@ impl LookupPattern {
             }
             (ForeignFieldMul, Curr | Next) => Some(LookupPattern::ForeignFieldMul),
             (Xor16, Curr) => Some(LookupPattern::Xor),
+            (KeccakRound, Curr) => Some(LookupPattern::KeccakRound),
+            (KeccakSponge, Curr) => Some(LookupPattern::KeccakSponge),
             _ => None,
         }
     }

kimchi/src/linearization.rs

@@ -1,5 +1,7 @@
 //! This module implements the linearization.
 
+use std::cmp::max;
+
 use crate::alphas::Alphas;
 use crate::circuits::argument::{Argument, ArgumentType};
 use crate::circuits::expr;
@@ -14,7 +16,9 @@ use crate::circuits::polynomials::{
     endosclmul::EndosclMul,
     foreign_field_add::circuitgates::ForeignFieldAdd,
     foreign_field_mul::circuitgates::ForeignFieldMul,
-    generic, permutation,
+    generic,
+    keccak::circuitgates::{KeccakRound, KeccakSponge},
+    permutation,
     poseidon::Poseidon,
     range_check::circuitgates::{RangeCheck0, RangeCheck1},
     rot,
@@ -44,7 +48,15 @@ pub fn constraints_expr<F: PrimeField + SquareRootField, const COLUMNS: usize>(
 
     // Set up powers of alpha. Only the max number of constraints matters.
     // The gate type argument can just be the zero gate.
-    let max_exponents = VarbaseMul::<F>::CONSTRAINTS;
+    let mut max_exponents = VarbaseMul::<F>::CONSTRAINTS;
+    if let Some(feature_flags) = feature_flags {
+        if feature_flags.keccak {
+            max_exponents = max(
+                KeccakRound::<F>::CONSTRAINTS,
+                KeccakSponge::<F>::CONSTRAINTS,
+            );
+        }
+    }
     powers_of_alpha.register(ArgumentType::Gate(GateType::Zero), max_exponents);
 
     let mut cache = expr::Cache::default();
@@ -151,6 +163,38 @@ pub fn constraints_expr<F: PrimeField + SquareRootField, const COLUMNS: usize>(
         }
     }
 
+    {
+        let mut keccak_round_expr =
+            || KeccakRound::combined_constraints(&powers_of_alpha, &mut cache);
+        if let Some(feature_flags) = feature_flags {
+            if feature_flags.keccak {
+                expr += keccak_round_expr();
+            }
+        } else {
+            expr += Expr::IfFeature(
+                FeatureFlag::Keccak,
+                Box::new(keccak_round_expr()),
+                Box::new(Expr::zero()),
+            );
+        }
+    }
+
+    {
+        let mut keccak_sponge_expr =
+            || KeccakSponge::combined_constraints(&powers_of_alpha, &mut cache);
+        if let Some(feature_flags) = feature_flags {
+            if feature_flags.keccak {
+                expr += keccak_sponge_expr();
+            }
+        } else {
+            expr += Expr::IfFeature(
+                FeatureFlag::Keccak,
+                Box::new(keccak_sponge_expr()),
+                Box::new(Expr::zero()),
+            );
+        }
+    }
+
     if generic {
         expr += generic::Generic::combined_constraints(&powers_of_alpha, &mut cache);
     }
@@ -252,14 +296,15 @@ pub fn linearization_columns<F: FftField + SquareRootField, const COLUMNS: usize
                 foreign_field_mul: true,
                 xor: true,
                 rot: true,
+                keccak: true,
                 lookup_features: LookupFeatures {
                     patterns: LookupPatterns {
                         xor: true,
                         lookup: true,
                         range_check: true,
                         foreign_field_mul: true,
-                        keccak_round: false,
-                        keccak_sponge: false,
+                        keccak_round: true,
+                        keccak_sponge: true,
                     },
                     joint_lookup_used: true,
                     uses_runtime_tables: true,
@@ -319,13 +364,17 @@ pub fn linearization_columns<F: FftField + SquareRootField, const COLUMNS: usize
     h.insert(Index(GateType::ForeignFieldMul));
     h.insert(Index(GateType::Xor16));
     h.insert(Index(GateType::Rot64));
+    h.insert(Index(GateType::KeccakRound));
+    h.insert(Index(GateType::KeccakSponge));
 
     // lookup selectors
     h.insert(LookupRuntimeSelector);
     h.insert(LookupKindIndex(LookupPattern::Xor));
     h.insert(LookupKindIndex(LookupPattern::Lookup));
     h.insert(LookupKindIndex(LookupPattern::RangeCheck));
     h.insert(LookupKindIndex(LookupPattern::ForeignFieldMul));
+    h.insert(LookupKindIndex(LookupPattern::KeccakRound));
+    h.insert(LookupKindIndex(LookupPattern::KeccakSponge));
 
     h
 }

kimchi/src/plonk_sponge.rs

@@ -78,6 +78,8 @@ impl<Fr: PrimeField, const COLUMNS: usize> FrSponge<Fr, COLUMNS> for DefaultFrSp
             foreign_field_mul_selector,
             xor_selector,
             rot_selector,
+            keccak_round_selector,
+            keccak_sponge_selector,
             lookup_aggregation,
             lookup_table,
             lookup_sorted,
@@ -87,6 +89,8 @@ impl<Fr: PrimeField, const COLUMNS: usize> FrSponge<Fr, COLUMNS> for DefaultFrSp
             lookup_gate_lookup_selector,
             range_check_lookup_selector,
             foreign_field_mul_lookup_selector,
+            keccak_round_lookup_selector,
+            keccak_sponge_lookup_selector,
         } = e;
 
         let mut points = vec![
@@ -122,6 +126,12 @@ impl<Fr: PrimeField, const COLUMNS: usize> FrSponge<Fr, COLUMNS> for DefaultFrSp
         if let Some(rot_selector) = rot_selector.as_ref() {
             points.push(rot_selector)
         }
+        if let Some(keccak_round_selector) = keccak_round_selector.as_ref() {
+            points.push(keccak_round_selector)
+        }
+        if let Some(keccak_sponge_selector) = keccak_sponge_selector.as_ref() {
+            points.push(keccak_sponge_selector)
+        }
         if let Some(lookup_aggregation) = lookup_aggregation.as_ref() {
             points.push(lookup_aggregation)
         }
@@ -152,6 +162,12 @@ impl<Fr: PrimeField, const COLUMNS: usize> FrSponge<Fr, COLUMNS> for DefaultFrSp
         {
             points.push(foreign_field_mul_lookup_selector)
         }
+        if let Some(keccak_round_lookup_selector) = keccak_round_lookup_selector.as_ref() {
+            points.push(keccak_round_lookup_selector)
+        }
+        if let Some(keccak_sponge_lookup_selector) = keccak_sponge_lookup_selector.as_ref() {
+            points.push(keccak_sponge_lookup_selector)
+        }
 
         points.into_iter().for_each(|p| {
             self.sponge.absorb(&p.zeta);