From 2f3a280b907e83626308eedf3531d8a216146685 Mon Sep 17 00:00:00 2001
From: wborgeaud <williamborgeaud@gmail.com>
Date: Thu, 9 Jun 2022 23:40:58 +0200
Subject: [PATCH] Circuit fix

---
 evm/src/keccak/keccak_stark.rs | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/evm/src/keccak/keccak_stark.rs b/evm/src/keccak/keccak_stark.rs
index 995b6fca..771513a6 100644
--- a/evm/src/keccak/keccak_stark.rs
+++ b/evm/src/keccak/keccak_stark.rs
@@ -497,10 +497,13 @@ impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for KeccakStark<F
                     reduce_with_powers_ext_circuit(builder, &input_bits[0..32], two);
                 let input_bits_combined_hi =
                     reduce_with_powers_ext_circuit(builder, &input_bits[32..64], two);
-                let diff = builder.sub_extension(output_lo, input_bits_combined_lo);
-                yield_constr.constraint_transition(builder, diff);
-                let diff = builder.sub_extension(output_hi, input_bits_combined_hi);
-                yield_constr.constraint_transition(builder, diff);
+                let is_last_round = vars.local_values[reg_step(NUM_ROUNDS - 1)];
+                let diff = builder.sub_extension(input_bits_combined_lo, output_lo);
+                let filtered_diff = builder.mul_sub_extension(is_last_round, diff, diff);
+                yield_constr.constraint_transition(builder, filtered_diff);
+                let diff = builder.sub_extension(input_bits_combined_hi, output_hi);
+                let filtered_diff = builder.mul_sub_extension(is_last_round, diff, diff);
+                yield_constr.constraint_transition(builder, filtered_diff);
             }
         }
     }