From 51f66d4dac9970d50f16fc51f413ed5bcf03ab08 Mon Sep 17 00:00:00 2001
From: wborgeaud <williamborgeaud@gmail.com>
Date: Thu, 9 Jun 2022 23:36:43 +0200
Subject: [PATCH] Fix constraint

---
 evm/src/all_stark.rs           | 5 +++--
 evm/src/keccak/keccak_stark.rs | 9 +++++++--
 2 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/evm/src/all_stark.rs b/evm/src/all_stark.rs
index c3c9653c..ef5e1352 100644
--- a/evm/src/all_stark.rs
+++ b/evm/src/all_stark.rs
@@ -87,15 +87,16 @@ mod tests {
         let keccak_stark = KeccakStark::<F, D> {
             f: Default::default(),
         };
-        let keccak_rows = (NUM_ROUNDS + 1).next_power_of_two();
+        let keccak_rows = (2 * NUM_ROUNDS + 1).next_power_of_two();
         let keccak_looked_col = 3;
 
         let mut rng = ChaCha8Rng::seed_from_u64(0x6feb51b7ec230f25);
-        let num_inputs = 1;
+        let num_inputs = 2;
         let keccak_inputs = (0..num_inputs)
             .map(|_| [0u64; INPUT_LIMBS].map(|_| rng.gen()))
             .collect_vec();
         let keccak_trace = keccak_stark.generate_trace(keccak_inputs);
+        dbg!(keccak_trace[0].len(), keccak_rows);
         let column_to_copy: Vec<_> = keccak_trace[keccak_looked_col].values[..].into();
 
         let default = vec![F::ONE; 1];
diff --git a/evm/src/keccak/keccak_stark.rs b/evm/src/keccak/keccak_stark.rs
index cbc2408b..995b6fca 100644
--- a/evm/src/keccak/keccak_stark.rs
+++ b/evm/src/keccak/keccak_stark.rs
@@ -345,8 +345,13 @@ impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for KeccakStark<F
                 let input_bits_combined_hi = (32..64)
                     .rev()
                     .fold(P::ZEROS, |acc, z| acc.doubles() + input_bits[z]);
-                yield_constr.constraint_transition(output_lo - input_bits_combined_lo);
-                yield_constr.constraint_transition(output_hi - input_bits_combined_hi);
+                let is_last_round = vars.local_values[reg_step(NUM_ROUNDS - 1)];
+                yield_constr.constraint_transition(
+                    (P::ONES - is_last_round) * (output_lo - input_bits_combined_lo),
+                );
+                yield_constr.constraint_transition(
+                    (P::ONES - is_last_round) * (output_hi - input_bits_combined_hi),
+                );
             }
         }
     }