Merge pull request #1736 from o1-labs/zkvm/keccak/24rounds

Remove `inverse_round` and alternate definition of `is_round()`

kimchi/src/circuits/polynomials/keccak/gadget.rs

@@ -38,8 +38,7 @@ impl<F: PrimeField + SquareRootField> CircuitGate<F> {
                 pad,
                 extra_bytes,
             ));
-            // Start by 1 to avoid the dummy entry
-            for round in 1..=ROUNDS {
+            for round in 0..ROUNDS {
                 gates.push(Self::create_keccak_round(new_row + gates.len(), round));
             }
         }

kimchi/src/circuits/polynomials/keccak/mod.rs

@@ -47,9 +47,8 @@ pub const OFF: [[u64; DIM]; DIM] = [
     [18, 2, 61, 56, 14],
 ];
 
-/// Contains the 24 round constants for Keccak and an initial dummy entry
-pub const RC: [u64; ROUNDS + 1] = [
-    0x0000000000000000,
+/// Contains the 24 round constants for Keccak
+pub const RC: [u64; ROUNDS] = [
     0x0000000000000001,
     0x0000000000008082,
     0x800000000000808a,

kimchi/src/circuits/polynomials/keccak/witness.rs

@@ -430,8 +430,7 @@ pub fn extend_keccak_witness<F: PrimeField>(witness: &mut [Vec<F>; KECCAK_COLS],
 
         let mut ini_state = xor_state.clone();
 
-        // Start by 1 to avoid dummy entry
-        for round in 1..=ROUNDS {
+        for round in 0..ROUNDS {
             // Theta
             let theta = Theta::create(&ini_state);
 

optimism/src/keccak/column.rs

@@ -13,20 +13,19 @@ use rayon::iter::{FromParallelIterator, IntoParallelIterator, ParallelIterator};
 use super::{ZKVM_KECCAK_COLS_CURR, ZKVM_KECCAK_COLS_NEXT};
 
 const ZKVM_KECCAK_COLS_LENGTH: usize =
-    ZKVM_KECCAK_COLS_CURR + ZKVM_KECCAK_COLS_NEXT + QUARTERS + RATE_IN_BYTES + 15;
+    ZKVM_KECCAK_COLS_CURR + ZKVM_KECCAK_COLS_NEXT + QUARTERS + RATE_IN_BYTES + 14;
 
 #[derive(Clone, Copy, Debug, Eq, PartialEq)]
 pub enum KeccakColumn {
     HashIndex,
     StepIndex,
-    FlagRound,              // Coeff Round = 0 | 1 .. 24
+    FlagRound,              // Coeff Round = [0..24)
     FlagAbsorb,             // Coeff Absorb = 0 | 1
     FlagSqueeze,            // Coeff Squeeze = 0 | 1
     FlagRoot,               // Coeff Root = 0 | 1
     FlagPad,                // Coeff Pad = 0 | 1
-    FlagLength,             // Coeff Length 0 | 1 .. 136
+    FlagLength,             // Coeff Length 0 | 1 ..=136
     TwoToPad,               // 2^PadLength
-    InverseRound,           // Round^-1
     FlagsBytes(usize),      // 136 boolean values
     PadSuffix(usize),       // 5 values with padding suffix
     RoundConstants(usize),  // Round constants
@@ -55,14 +54,13 @@ pub enum KeccakColumn {
 pub struct KeccakColumns<T> {
     pub hash_index: T,
     pub step_index: T,
-    pub flag_round: T,                    // Coeff Round = 0 | 1 .. 24
+    pub flag_round: T,                    // Coeff Round = [0..24)
     pub flag_absorb: T,                   // Coeff Absorb = 0 | 1
     pub flag_squeeze: T,                  // Coeff Squeeze = 0 | 1
     pub flag_root: T,                     // Coeff Root = 0 | 1
     pub flag_pad: T,                      // Coeff Pad = 0 | 1
-    pub flag_length: T,                   // Coeff Length 0 | 1 .. 136
+    pub flag_length: T,                   // Coeff Length 0 | 1 ..=136
     pub two_to_pad: T,                    // 2^PadLength
-    pub inverse_round: T,                 // Round^-1
     pub flags_bytes: [T; RATE_IN_BYTES],  // 136 boolean values
     pub pad_suffix: [T; 5],               // 5 values with padding suffix
     pub round_constants: [T; QUARTERS],   // Round constants
@@ -88,10 +86,9 @@ impl<T: Zero + One + Clone> Default for KeccakColumns<T> {
             flag_pad: T::zero(),
             flag_length: T::zero(),
             two_to_pad: T::one(), // So that default 2^0 is in the table
-            inverse_round: T::zero(),
             flags_bytes: std::array::from_fn(|_| T::zero()),
             pad_suffix: std::array::from_fn(|_| T::zero()),
-            round_constants: std::array::from_fn(|_| T::zero()), // RC[0] is set to be all zeros
+            round_constants: std::array::from_fn(|_| T::zero()), // default zeros, but lookup only if is round
             curr: std::array::from_fn(|_| T::zero()),
             next: std::array::from_fn(|_| T::zero()),
         }
@@ -112,7 +109,6 @@ impl<T: Clone> Index<KeccakColumn> for KeccakColumns<T> {
             KeccakColumn::FlagPad => &self.flag_pad,
             KeccakColumn::FlagLength => &self.flag_length,
             KeccakColumn::TwoToPad => &self.two_to_pad,
-            KeccakColumn::InverseRound => &self.inverse_round,
             KeccakColumn::FlagsBytes(idx) => &self.flags_bytes[idx],
             KeccakColumn::PadSuffix(idx) => &self.pad_suffix[idx],
             KeccakColumn::RoundConstants(idx) => &self.round_constants[idx],
@@ -151,7 +147,6 @@ impl<T: Clone> IndexMut<KeccakColumn> for KeccakColumns<T> {
             KeccakColumn::FlagPad => &mut self.flag_pad,
             KeccakColumn::FlagLength => &mut self.flag_length,
             KeccakColumn::TwoToPad => &mut self.two_to_pad,
-            KeccakColumn::InverseRound => &mut self.inverse_round,
             KeccakColumn::FlagsBytes(idx) => &mut self.flags_bytes[idx],
             KeccakColumn::PadSuffix(idx) => &mut self.pad_suffix[idx],
             KeccakColumn::RoundConstants(idx) => &mut self.round_constants[idx],
@@ -193,7 +188,6 @@ impl<F> IntoIterator for KeccakColumns<F> {
         iter_contents.push(self.flag_pad);
         iter_contents.push(self.flag_length);
         iter_contents.push(self.two_to_pad);
-        iter_contents.push(self.inverse_round);
         iter_contents.extend(self.flags_bytes);
         iter_contents.extend(self.pad_suffix);
         iter_contents.extend(self.round_constants);
@@ -221,7 +215,6 @@ where
         iter_contents.push(self.flag_pad);
         iter_contents.push(self.flag_length);
         iter_contents.push(self.two_to_pad);
-        iter_contents.push(self.inverse_round);
         iter_contents.extend(self.flags_bytes);
         iter_contents.extend(self.pad_suffix);
         iter_contents.extend(self.round_constants);
@@ -262,7 +255,6 @@ impl<G: Send + std::fmt::Debug> FromParallelIterator<G> for KeccakColumns<G> {
             .collect::<Vec<G>>()
             .try_into()
             .unwrap();
-        let inverse_round = iter_contents.pop().unwrap();
         let two_to_pad = iter_contents.pop().unwrap();
         let flag_length = iter_contents.pop().unwrap();
         let flag_pad = iter_contents.pop().unwrap();
@@ -282,7 +274,6 @@ impl<G: Send + std::fmt::Debug> FromParallelIterator<G> for KeccakColumns<G> {
             flag_pad,
             flag_length,
             two_to_pad,
-            inverse_round,
             flags_bytes,
             pad_suffix,
             round_constants,
@@ -310,7 +301,6 @@ where
         iter_contents.push(&self.flag_pad);
         iter_contents.push(&self.flag_length);
         iter_contents.push(&self.two_to_pad);
-        iter_contents.push(&self.inverse_round);
         iter_contents.extend(&self.flags_bytes);
         iter_contents.extend(&self.pad_suffix);
         iter_contents.extend(&self.round_constants);
@@ -338,7 +328,6 @@ where
         iter_contents.push(&mut self.flag_pad);
         iter_contents.push(&mut self.flag_length);
         iter_contents.push(&mut self.two_to_pad);
-        iter_contents.push(&mut self.inverse_round);
         iter_contents.extend(&mut self.flags_bytes);
         iter_contents.extend(&mut self.pad_suffix);
         iter_contents.extend(&mut self.round_constants);

optimism/src/keccak/constraints.rs

@@ -77,10 +77,6 @@ impl<Fp: Field> Constraints for KeccakEnv<Fp> {
                 self.constrain(Self::either_false(self.is_round(), self.is_root()));
                 // Absorb and Squeeze cannot happen at the same time
                 self.constrain(Self::either_false(self.is_absorb(), self.is_squeeze()));
-                // Only one of Round and Sponge can be zero
-                // This means either Sponge is true or Round is nonzero -> has an inverse
-                self.constrain(self.is_sponge() * self.round());
-                self.constrain(self.is_round() * Self::is_one(self.round() * self.inverse_round()));
                 // Trivially, is_sponge and is_round are mutually exclusive
             }
         }

optimism/src/keccak/environment.rs

@@ -100,12 +100,12 @@ impl<Fp: Field> KeccakEnv<Fp> {
         match self.keccak_step {
             Some(step) => match step {
                 KeccakStep::Sponge(sponge) => match sponge {
-                    Sponge::Absorb(_) => self.keccak_step = Some(KeccakStep::Round(1)),
+                    Sponge::Absorb(_) => self.keccak_step = Some(KeccakStep::Round(0)),
 
                     Sponge::Squeeze => self.keccak_step = None,
                 },
                 KeccakStep::Round(round) => {
-                    if round < ROUNDS as u64 {
+                    if round < ROUNDS as u64 - 1 {
                         self.keccak_step = Some(KeccakStep::Round(round + 1));
                     } else {
                         self.blocks_left_to_absorb -= 1;
@@ -219,8 +219,6 @@ pub(crate) trait KeccakEnvironment {
 
     fn round(&self) -> Self::Variable;
 
-    fn inverse_round(&self) -> Self::Variable;
-
     fn length(&self) -> Self::Variable;
 
     fn two_to_pad(&self) -> Self::Variable;
@@ -397,10 +395,6 @@ impl<Fp: Field> KeccakEnvironment for KeccakEnv<Fp> {
         self.variable(KeccakColumn::FlagRound)
     }
 
-    fn inverse_round(&self) -> Self::Variable {
-        self.variable(KeccakColumn::InverseRound)
-    }
-
     fn length(&self) -> Self::Variable {
         self.variable(KeccakColumn::FlagLength)
     }

optimism/src/keccak/proof.rs

@@ -38,7 +38,6 @@ impl<G: KimchiCurve> Default for KeccakProofInputs<G> {
                 flag_pad: (0..DOMAIN_SIZE).map(|_| G::ScalarField::zero()).collect(),
                 flag_length: (0..DOMAIN_SIZE).map(|_| G::ScalarField::zero()).collect(),
                 two_to_pad: (0..DOMAIN_SIZE).map(|_| G::ScalarField::one()).collect(),
-                inverse_round: (0..DOMAIN_SIZE).map(|_| G::ScalarField::zero()).collect(),
                 flags_bytes: std::array::from_fn(|_| {
                     (0..DOMAIN_SIZE).map(|_| G::ScalarField::zero()).collect()
                 }),
@@ -149,7 +148,6 @@ where
             flag_pad: eval_col(evaluations.flag_pad),
             flag_length: eval_col(evaluations.flag_length),
             two_to_pad: eval_col(evaluations.two_to_pad),
-            inverse_round: eval_col(evaluations.inverse_round),
             flags_bytes: eval_array_col(&evaluations.flags_bytes).try_into().unwrap(),
             pad_suffix: eval_array_col(&evaluations.pad_suffix).try_into().unwrap(),
             round_constants: eval_array_col(&evaluations.round_constants)
@@ -174,7 +172,6 @@ where
             flag_pad: comm(&polys.flag_pad),
             flag_length: comm(&polys.flag_length),
             two_to_pad: comm(&polys.two_to_pad),
-            inverse_round: comm(&polys.inverse_round),
             flags_bytes: comm_array(&polys.flags_bytes).try_into().unwrap(),
             pad_suffix: comm_array(&polys.pad_suffix).try_into().unwrap(),
             round_constants: comm_array(&polys.round_constants).try_into().unwrap(),
@@ -209,7 +206,6 @@ where
             flag_pad: comm(&polys.flag_pad),
             flag_length: comm(&polys.flag_length),
             two_to_pad: comm(&polys.two_to_pad),
-            inverse_round: comm(&polys.inverse_round),
             flags_bytes: comm_array(&polys.flags_bytes).try_into().unwrap(),
             pad_suffix: comm_array(&polys.pad_suffix).try_into().unwrap(),
             round_constants: comm_array(&polys.round_constants).try_into().unwrap(),
@@ -389,7 +385,6 @@ fn test_keccak_prover() {
                 flag_pad: (0..DOMAIN_SIZE).map(|_| Fp::rand(rng)).collect::<Vec<_>>(),
                 flag_length: (0..DOMAIN_SIZE).map(|_| Fp::rand(rng)).collect::<Vec<_>>(),
                 two_to_pad: (0..DOMAIN_SIZE).map(|_| Fp::rand(rng)).collect::<Vec<_>>(),
-                inverse_round: (0..DOMAIN_SIZE).map(|_| Fp::rand(rng)).collect::<Vec<_>>(),
                 flags_bytes: std::array::from_fn(|_| {
                     (0..DOMAIN_SIZE).map(|_| Fp::rand(rng)).collect::<Vec<_>>()
                 }),

optimism/src/keccak/witness.rs

@@ -91,15 +91,8 @@ impl<Fp: Field> KeccakInterpreter for KeccakEnv<Fp> {
     }
 
     fn set_flag_round(&mut self, round: u64) {
-        assert!(round <= ROUNDS as u64);
-        // Values between 0 (dummy, for sponges) and 24
+        assert!(round < ROUNDS as u64);
         self.write_column(KeccakColumn::FlagRound, round);
-        if round != 0 {
-            self.write_column_field(
-                KeccakColumn::InverseRound,
-                Fp::from(round).inverse().unwrap(),
-            );
-        }
     }
 
     fn run_sponge(&mut self, sponge: Sponge) {

optimism/src/main.rs

@@ -117,7 +117,6 @@ pub fn main() -> ExitCode {
             keccak_columns.flag_pad.clear();
             keccak_columns.flag_length.clear();
             keccak_columns.two_to_pad.clear();
-            keccak_columns.inverse_round.clear();
             keccak_columns.flags_bytes.iter_mut().for_each(Vec::clear);
             keccak_columns.pad_suffix.iter_mut().for_each(Vec::clear);
             keccak_columns
@@ -139,7 +138,6 @@ pub fn main() -> ExitCode {
             flag_pad: Vec::with_capacity(domain_size),
             flag_length: Vec::with_capacity(domain_size),
             two_to_pad: Vec::with_capacity(domain_size),
-            inverse_round: Vec::with_capacity(domain_size),
             flags_bytes: std::array::from_fn(|_| Vec::with_capacity(domain_size)),
             pad_suffix: std::array::from_fn(|_| Vec::with_capacity(domain_size)),
             round_constants: std::array::from_fn(|_| Vec::with_capacity(domain_size)),
@@ -185,9 +183,6 @@ pub fn main() -> ExitCode {
             keccak_current_pre_folding_witness
                 .two_to_pad
                 .push(keccak_env.keccak_witness.two_to_pad);
-            keccak_current_pre_folding_witness
-                .inverse_round
-                .push(keccak_env.keccak_witness.inverse_round);
             for (env_wit, pre_fold_wit) in keccak_env
                 .keccak_witness
                 .flags_bytes