From 301aedf072720ecd6dfb9af18d1635ef97ce5c81 Mon Sep 17 00:00:00 2001
From: Ayush Shukla <shuklaayush247@gmail.com>
Date: Sun, 27 Aug 2023 22:32:44 +0200
Subject: [PATCH] fix: constrain higher bits of reg_preimage

---
 evm/src/keccak/keccak_stark.rs | 26 ++++++++++++++++++++------
 1 file changed, 20 insertions(+), 6 deletions(-)

diff --git a/evm/src/keccak/keccak_stark.rs b/evm/src/keccak/keccak_stark.rs
index c8fe8086..7c077da1 100644
--- a/evm/src/keccak/keccak_stark.rs
+++ b/evm/src/keccak/keccak_stark.rs
@@ -266,9 +266,14 @@ impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for KeccakStark<F
         // If this is not the final step, the local and next preimages must match.
         for x in 0..5 {
             for y in 0..5 {
-                let preimage = reg_preimage(x, y);
-                let diff = vars.local_values[preimage] - vars.next_values[preimage];
-                yield_constr.constraint_transition(not_final_step * diff);
+                let reg_preimage_lo = reg_preimage(x, y);
+                let reg_preimage_hi = reg_preimage_lo + 1;
+                let diff_lo =
+                    vars.local_values[reg_preimage_lo] - vars.next_values[reg_preimage_lo];
+                let diff_hi =
+                    vars.local_values[reg_preimage_hi] - vars.next_values[reg_preimage_hi];
+                yield_constr.constraint_transition(not_final_step * diff_lo);
+                yield_constr.constraint_transition(not_final_step * diff_hi);
             }
         }
 
@@ -436,9 +441,18 @@ impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for KeccakStark<F
         // If this is not the final step, the local and next preimages must match.
         for x in 0..5 {
             for y in 0..5 {
-                let preimage = reg_preimage(x, y);
-                let diff =
-                    builder.sub_extension(vars.local_values[preimage], vars.next_values[preimage]);
+                let reg_preimage_lo = reg_preimage(x, y);
+                let reg_preimage_hi = reg_preimage_lo + 1;
+                let diff = builder.sub_extension(
+                    vars.local_values[reg_preimage_lo],
+                    vars.next_values[reg_preimage_lo],
+                );
+                let constraint = builder.mul_extension(not_final_step, diff);
+                yield_constr.constraint_transition(builder, constraint);
+                let diff = builder.sub_extension(
+                    vars.local_values[reg_preimage_hi],
+                    vars.next_values[reg_preimage_hi],
+                );
                 let constraint = builder.mul_extension(not_final_step, diff);
                 yield_constr.constraint_transition(builder, constraint);
             }