From ea03e4183f6c7583f6177c96bbda9d7c5c2e5c85 Mon Sep 17 00:00:00 2001
From: Nicholas Ward <npward@berkeley.edu>
Date: Mon, 28 Aug 2023 12:42:01 -0700
Subject: [PATCH] Keccak STARK: constraint preimage to equal A on first round

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

diff --git a/evm/src/keccak/keccak_stark.rs b/evm/src/keccak/keccak_stark.rs
index 7c077da1..6d89c6fd 100644
--- a/evm/src/keccak/keccak_stark.rs
+++ b/evm/src/keccak/keccak_stark.rs
@@ -264,6 +264,8 @@ impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for KeccakStark<F
         yield_constr.constraint(not_final_step * filter);
 
         // If this is not the final step, the local and next preimages must match.
+        // Also, if this is the first step, the preimage must match A.
+        let is_first_step = vars.local_values[reg_step(0)];
         for x in 0..5 {
             for y in 0..5 {
                 let reg_preimage_lo = reg_preimage(x, y);
@@ -274,6 +276,13 @@ impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for KeccakStark<F
                     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);
+
+                let reg_a_lo = reg_a(x, y);
+                let reg_a_hi = reg_a_lo + 1;
+                let diff_lo = vars.local_values[reg_preimage_lo] - vars.local_values[reg_a_lo];
+                let diff_hi = vars.local_values[reg_preimage_hi] - vars.local_values[reg_a_hi];
+                yield_constr.constraint(is_first_step * diff_lo);
+                yield_constr.constraint(is_first_step * diff_hi);
             }
         }
 
@@ -439,6 +448,8 @@ impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for KeccakStark<F
         yield_constr.constraint(builder, constraint);
 
         // If this is not the final step, the local and next preimages must match.
+        // Also, if this is the first step, the preimage must match A.
+        let is_first_step = vars.local_values[reg_step(0)];
         for x in 0..5 {
             for y in 0..5 {
                 let reg_preimage_lo = reg_preimage(x, y);
@@ -455,6 +466,21 @@ impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for KeccakStark<F
                 );
                 let constraint = builder.mul_extension(not_final_step, diff);
                 yield_constr.constraint_transition(builder, constraint);
+
+                let reg_a_lo = reg_a(x, y);
+                let reg_a_hi = reg_a_lo + 1;
+                let diff_lo = builder.sub_extension(
+                    vars.local_values[reg_preimage_lo],
+                    vars.local_values[reg_a_lo],
+                );
+                let constraint = builder.mul_extension(is_first_step, diff_lo);
+                yield_constr.constraint(builder, constraint);
+                let diff_hi = builder.sub_extension(
+                    vars.local_values[reg_preimage_hi],
+                    vars.local_values[reg_a_hi],
+                );
+                let constraint = builder.mul_extension(is_first_step, diff_hi);
+                yield_constr.constraint(builder, constraint);
             }
         }