From 5795f41738648b4a8e12e8ef44e27ebf641a8562 Mon Sep 17 00:00:00 2001
From: Ayush Shukla <shuklaayush247@gmail.com>
Date: Sun, 18 Feb 2024 10:05:27 +0100
Subject: [PATCH 1/8] fix: remove duplicate constraint

---
 evm_arithmetization/src/keccak/keccak_stark.rs | 13 ++-----------
 1 file changed, 2 insertions(+), 11 deletions(-)

diff --git a/evm_arithmetization/src/keccak/keccak_stark.rs b/evm_arithmetization/src/keccak/keccak_stark.rs
index 969a0357f..b8e804049 100644
--- a/evm_arithmetization/src/keccak/keccak_stark.rs
+++ b/evm_arithmetization/src/keccak/keccak_stark.rs
@@ -275,14 +275,10 @@ impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for KeccakStark<F
         let local_values = vars.get_local_values();
         let next_values = vars.get_next_values();
 
-        // The filter must be 0 or 1.
-        let filter = local_values[reg_step(NUM_ROUNDS - 1)];
-        yield_constr.constraint(filter * (filter - P::ONES));
-
         // If this is not the final step, the filter must be off.
         let final_step = local_values[reg_step(NUM_ROUNDS - 1)];
         let not_final_step = P::ONES - final_step;
-        yield_constr.constraint(not_final_step * filter);
+        yield_constr.constraint(not_final_step * final_step);
 
         // If this is not the final step or a padding row,
         // the local and next timestamps must match.
@@ -446,15 +442,10 @@ impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for KeccakStark<F
         let local_values = vars.get_local_values();
         let next_values = vars.get_next_values();
 
-        // The filter must be 0 or 1.
-        let filter = local_values[reg_step(NUM_ROUNDS - 1)];
-        let constraint = builder.mul_sub_extension(filter, filter, filter);
-        yield_constr.constraint(builder, constraint);
-
         // If this is not the final step, the filter must be off.
         let final_step = local_values[reg_step(NUM_ROUNDS - 1)];
         let not_final_step = builder.sub_extension(one_ext, final_step);
-        let constraint = builder.mul_extension(not_final_step, filter);
+        let constraint = builder.mul_extension(not_final_step, final_step);
         yield_constr.constraint(builder, constraint);
 
         // If this is not the final step or a padding row,

From fe3beaeac2b8206ca0ae562c2daea2aa991fafcf Mon Sep 17 00:00:00 2001
From: Ayush Shukla <shuklaayush247@gmail.com>
Date: Sun, 18 Feb 2024 21:52:51 +0100
Subject: [PATCH 2/8] fix: add proper constraints for padding rows

---
 evm_arithmetization/src/keccak/round_flags.rs | 26 +++++++++++++------
 1 file changed, 18 insertions(+), 8 deletions(-)

diff --git a/evm_arithmetization/src/keccak/round_flags.rs b/evm_arithmetization/src/keccak/round_flags.rs
index 5e76b2ec9..f629bced2 100644
--- a/evm_arithmetization/src/keccak/round_flags.rs
+++ b/evm_arithmetization/src/keccak/round_flags.rs
@@ -25,17 +25,21 @@ pub(crate) fn eval_round_flags<F: Field, P: PackedField<Scalar = F>>(
     }
 
     // Flags should circularly increment, or be all zero for padding rows.
+    let current_any_flag = (0..NUM_ROUNDS)
+        .map(|i| local_values[reg_step(i)])
+        .sum::<P>();
     let next_any_flag = (0..NUM_ROUNDS).map(|i| next_values[reg_step(i)]).sum::<P>();
+    let last_row_flag = local_values[reg_step(NUM_ROUNDS - 1)];
     for i in 0..NUM_ROUNDS {
         let current_round_flag = local_values[reg_step(i)];
         let next_round_flag = next_values[reg_step((i + 1) % NUM_ROUNDS)];
-        yield_constr.constraint_transition(next_any_flag * (next_round_flag - current_round_flag));
+        yield_constr.constraint_transition(
+            next_any_flag * (next_round_flag - current_round_flag)
+                + (next_any_flag - F::ONE) * current_any_flag * (last_row_flag - F::ONE),
+        );
     }
 
     // Padding rows should always be followed by padding rows.
-    let current_any_flag = (0..NUM_ROUNDS)
-        .map(|i| local_values[reg_step(i)])
-        .sum::<P>();
     yield_constr.constraint_transition(next_any_flag * (current_any_flag - F::ONE));
 }
 
@@ -56,19 +60,25 @@ pub(crate) fn eval_round_flags_recursively<F: RichField + Extendable<D>, const D
     }
 
     // Flags should circularly increment, or be all zero for padding rows.
+    let current_any_flag =
+        builder.add_many_extension((0..NUM_ROUNDS).map(|i| local_values[reg_step(i)]));
     let next_any_flag =
         builder.add_many_extension((0..NUM_ROUNDS).map(|i| next_values[reg_step(i)]));
+    let last_row_flag = local_values[reg_step(NUM_ROUNDS - 1)];
     for i in 0..NUM_ROUNDS {
         let current_round_flag = local_values[reg_step(i)];
         let next_round_flag = next_values[reg_step((i + 1) % NUM_ROUNDS)];
-        let diff = builder.sub_extension(next_round_flag, current_round_flag);
-        let constraint = builder.mul_extension(next_any_flag, diff);
+        let diff1 = builder.sub_extension(next_round_flag, current_round_flag);
+        let constraint1 = builder.mul_extension(next_any_flag, diff1);
+        let diff2 = builder.sub_extension(next_any_flag, one);
+        let diff3 = builder.sub_extension(last_row_flag, one);
+        let prod = builder.mul_extension(diff2, diff3);
+        let constraint2 = builder.mul_extension(current_any_flag, prod);
+        let constraint = builder.add_extension(constraint1, constraint2);
         yield_constr.constraint_transition(builder, constraint);
     }
 
     // Padding rows should always be followed by padding rows.
-    let current_any_flag =
-        builder.add_many_extension((0..NUM_ROUNDS).map(|i| local_values[reg_step(i)]));
     let constraint = builder.mul_sub_extension(next_any_flag, current_any_flag, next_any_flag);
     yield_constr.constraint_transition(builder, constraint);
 }

From 7c9c498c34d9feadf639214ba06c6eb706f41240 Mon Sep 17 00:00:00 2001
From: Ayush Shukla <shuklaayush247@gmail.com>
Date: Sun, 18 Feb 2024 21:55:03 +0100
Subject: [PATCH 3/8] fix: remove extra constraints

---
 evm_arithmetization/src/keccak/keccak_stark.rs | 9 ++-------
 1 file changed, 2 insertions(+), 7 deletions(-)

diff --git a/evm_arithmetization/src/keccak/keccak_stark.rs b/evm_arithmetization/src/keccak/keccak_stark.rs
index b8e804049..e1a0484d6 100644
--- a/evm_arithmetization/src/keccak/keccak_stark.rs
+++ b/evm_arithmetization/src/keccak/keccak_stark.rs
@@ -276,9 +276,7 @@ impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for KeccakStark<F
         let next_values = vars.get_next_values();
 
         // If this is not the final step, the filter must be off.
-        let final_step = local_values[reg_step(NUM_ROUNDS - 1)];
-        let not_final_step = P::ONES - final_step;
-        yield_constr.constraint(not_final_step * final_step);
+        let not_final_step = P::ONES - local_values[reg_step(NUM_ROUNDS - 1)];
 
         // If this is not the final step or a padding row,
         // the local and next timestamps must match.
@@ -443,10 +441,7 @@ impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for KeccakStark<F
         let next_values = vars.get_next_values();
 
         // If this is not the final step, the filter must be off.
-        let final_step = local_values[reg_step(NUM_ROUNDS - 1)];
-        let not_final_step = builder.sub_extension(one_ext, final_step);
-        let constraint = builder.mul_extension(not_final_step, final_step);
-        yield_constr.constraint(builder, constraint);
+        let not_final_step = builder.sub_extension(one_ext, local_values[reg_step(NUM_ROUNDS - 1)]);
 
         // If this is not the final step or a padding row,
         // the local and next timestamps must match.

From 54ba1c9e5f7adea0956f1abdc739c2d8f81a7230 Mon Sep 17 00:00:00 2001
From: Ayush Shukla <shuklaayush247@gmail.com>
Date: Sun, 18 Feb 2024 22:54:00 +0100
Subject: [PATCH 4/8] fix: simplify constraints using mul_sub gate

---
 evm_arithmetization/src/keccak/round_flags.rs | 18 +++++++++---------
 1 file changed, 9 insertions(+), 9 deletions(-)

diff --git a/evm_arithmetization/src/keccak/round_flags.rs b/evm_arithmetization/src/keccak/round_flags.rs
index f629bced2..dec87ad07 100644
--- a/evm_arithmetization/src/keccak/round_flags.rs
+++ b/evm_arithmetization/src/keccak/round_flags.rs
@@ -29,13 +29,13 @@ pub(crate) fn eval_round_flags<F: Field, P: PackedField<Scalar = F>>(
         .map(|i| local_values[reg_step(i)])
         .sum::<P>();
     let next_any_flag = (0..NUM_ROUNDS).map(|i| next_values[reg_step(i)]).sum::<P>();
-    let last_row_flag = local_values[reg_step(NUM_ROUNDS - 1)];
+    let last_round_flag = local_values[reg_step(NUM_ROUNDS - 1)];
     for i in 0..NUM_ROUNDS {
         let current_round_flag = local_values[reg_step(i)];
         let next_round_flag = next_values[reg_step((i + 1) % NUM_ROUNDS)];
         yield_constr.constraint_transition(
             next_any_flag * (next_round_flag - current_round_flag)
-                + (next_any_flag - F::ONE) * current_any_flag * (last_row_flag - F::ONE),
+                + (next_any_flag - F::ONE) * current_any_flag * (last_round_flag - F::ONE),
         );
     }
 
@@ -64,16 +64,16 @@ pub(crate) fn eval_round_flags_recursively<F: RichField + Extendable<D>, const D
         builder.add_many_extension((0..NUM_ROUNDS).map(|i| local_values[reg_step(i)]));
     let next_any_flag =
         builder.add_many_extension((0..NUM_ROUNDS).map(|i| next_values[reg_step(i)]));
-    let last_row_flag = local_values[reg_step(NUM_ROUNDS - 1)];
+    let last_round_flag = local_values[reg_step(NUM_ROUNDS - 1)];
     for i in 0..NUM_ROUNDS {
         let current_round_flag = local_values[reg_step(i)];
         let next_round_flag = next_values[reg_step((i + 1) % NUM_ROUNDS)];
-        let diff1 = builder.sub_extension(next_round_flag, current_round_flag);
-        let constraint1 = builder.mul_extension(next_any_flag, diff1);
-        let diff2 = builder.sub_extension(next_any_flag, one);
-        let diff3 = builder.sub_extension(last_row_flag, one);
-        let prod = builder.mul_extension(diff2, diff3);
-        let constraint2 = builder.mul_extension(current_any_flag, prod);
+        let flag_diff = builder.sub_extension(next_round_flag, current_round_flag);
+        let constraint1 = builder.mul_extension(next_any_flag, flag_diff);
+        let constraint2 = {
+            let tmp = builder.mul_sub_extension(current_any_flag, next_any_flag, current_any_flag);
+            builder.mul_sub_extension(tmp, last_round_flag, tmp)
+        };
         let constraint = builder.add_extension(constraint1, constraint2);
         yield_constr.constraint_transition(builder, constraint);
     }

From 4a003d110eb9d17f193201396d768b41234fb57d Mon Sep 17 00:00:00 2001
From: Ayush Shukla <shuklaayush247@gmail.com>
Date: Sun, 18 Feb 2024 23:15:10 +0100
Subject: [PATCH 5/8] fix: move constraint outside loop

---
 evm_arithmetization/src/keccak/round_flags.rs | 16 ++++++++--------
 1 file changed, 8 insertions(+), 8 deletions(-)

diff --git a/evm_arithmetization/src/keccak/round_flags.rs b/evm_arithmetization/src/keccak/round_flags.rs
index dec87ad07..aca3e936d 100644
--- a/evm_arithmetization/src/keccak/round_flags.rs
+++ b/evm_arithmetization/src/keccak/round_flags.rs
@@ -30,12 +30,13 @@ pub(crate) fn eval_round_flags<F: Field, P: PackedField<Scalar = F>>(
         .sum::<P>();
     let next_any_flag = (0..NUM_ROUNDS).map(|i| next_values[reg_step(i)]).sum::<P>();
     let last_round_flag = local_values[reg_step(NUM_ROUNDS - 1)];
+    let padding_constraint =
+        (next_any_flag - F::ONE) * current_any_flag * (last_round_flag - F::ONE);
     for i in 0..NUM_ROUNDS {
         let current_round_flag = local_values[reg_step(i)];
         let next_round_flag = next_values[reg_step((i + 1) % NUM_ROUNDS)];
         yield_constr.constraint_transition(
-            next_any_flag * (next_round_flag - current_round_flag)
-                + (next_any_flag - F::ONE) * current_any_flag * (last_round_flag - F::ONE),
+            next_any_flag * (next_round_flag - current_round_flag) + padding_constraint,
         );
     }
 
@@ -65,16 +66,15 @@ pub(crate) fn eval_round_flags_recursively<F: RichField + Extendable<D>, const D
     let next_any_flag =
         builder.add_many_extension((0..NUM_ROUNDS).map(|i| next_values[reg_step(i)]));
     let last_round_flag = local_values[reg_step(NUM_ROUNDS - 1)];
+    let padding_constraint = {
+        let tmp = builder.mul_sub_extension(current_any_flag, next_any_flag, current_any_flag);
+        builder.mul_sub_extension(tmp, last_round_flag, tmp)
+    };
     for i in 0..NUM_ROUNDS {
         let current_round_flag = local_values[reg_step(i)];
         let next_round_flag = next_values[reg_step((i + 1) % NUM_ROUNDS)];
         let flag_diff = builder.sub_extension(next_round_flag, current_round_flag);
-        let constraint1 = builder.mul_extension(next_any_flag, flag_diff);
-        let constraint2 = {
-            let tmp = builder.mul_sub_extension(current_any_flag, next_any_flag, current_any_flag);
-            builder.mul_sub_extension(tmp, last_round_flag, tmp)
-        };
-        let constraint = builder.add_extension(constraint1, constraint2);
+        let constraint = builder.mul_add_extension(next_any_flag, flag_diff, padding_constraint);
         yield_constr.constraint_transition(builder, constraint);
     }
 

From 01a78b9ea2f8f14047800b33b48b1cdd926f5bcf Mon Sep 17 00:00:00 2001
From: Ayush Shukla <shuklaayush247@gmail.com>
Date: Sun, 18 Feb 2024 23:35:52 +0100
Subject: [PATCH 6/8] chore: rename

---
 evm_arithmetization/src/keccak/round_flags.rs | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/evm_arithmetization/src/keccak/round_flags.rs b/evm_arithmetization/src/keccak/round_flags.rs
index aca3e936d..fbb162c85 100644
--- a/evm_arithmetization/src/keccak/round_flags.rs
+++ b/evm_arithmetization/src/keccak/round_flags.rs
@@ -29,9 +29,8 @@ pub(crate) fn eval_round_flags<F: Field, P: PackedField<Scalar = F>>(
         .map(|i| local_values[reg_step(i)])
         .sum::<P>();
     let next_any_flag = (0..NUM_ROUNDS).map(|i| next_values[reg_step(i)]).sum::<P>();
-    let last_round_flag = local_values[reg_step(NUM_ROUNDS - 1)];
-    let padding_constraint =
-        (next_any_flag - F::ONE) * current_any_flag * (last_round_flag - F::ONE);
+    let not_final_step = P::ONES - local_values[reg_step(NUM_ROUNDS - 1)];
+    let padding_constraint = (next_any_flag - F::ONE) * current_any_flag * not_final_step;
     for i in 0..NUM_ROUNDS {
         let current_round_flag = local_values[reg_step(i)];
         let next_round_flag = next_values[reg_step((i + 1) % NUM_ROUNDS)];
@@ -65,10 +64,10 @@ pub(crate) fn eval_round_flags_recursively<F: RichField + Extendable<D>, const D
         builder.add_many_extension((0..NUM_ROUNDS).map(|i| local_values[reg_step(i)]));
     let next_any_flag =
         builder.add_many_extension((0..NUM_ROUNDS).map(|i| next_values[reg_step(i)]));
-    let last_round_flag = local_values[reg_step(NUM_ROUNDS - 1)];
+    let not_final_step = builder.sub_extension(one, local_values[reg_step(NUM_ROUNDS - 1)]);
     let padding_constraint = {
         let tmp = builder.mul_sub_extension(current_any_flag, next_any_flag, current_any_flag);
-        builder.mul_sub_extension(tmp, last_round_flag, tmp)
+        builder.mul_extension(tmp, not_final_step)
     };
     for i in 0..NUM_ROUNDS {
         let current_round_flag = local_values[reg_step(i)];

From 8d238c4df3442cbf76febf89aae655cd1162d540 Mon Sep 17 00:00:00 2001
From: Ayush Shukla <shuklaayush247@gmail.com>
Date: Sun, 18 Feb 2024 23:43:05 +0100
Subject: [PATCH 7/8] chore: add comment explaining constrain

---
 evm_arithmetization/src/keccak/round_flags.rs | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/evm_arithmetization/src/keccak/round_flags.rs b/evm_arithmetization/src/keccak/round_flags.rs
index fbb162c85..db7829b24 100644
--- a/evm_arithmetization/src/keccak/round_flags.rs
+++ b/evm_arithmetization/src/keccak/round_flags.rs
@@ -29,6 +29,7 @@ pub(crate) fn eval_round_flags<F: Field, P: PackedField<Scalar = F>>(
         .map(|i| local_values[reg_step(i)])
         .sum::<P>();
     let next_any_flag = (0..NUM_ROUNDS).map(|i| next_values[reg_step(i)]).sum::<P>();
+    // Padding row should only start after the last round row.
     let not_final_step = P::ONES - local_values[reg_step(NUM_ROUNDS - 1)];
     let padding_constraint = (next_any_flag - F::ONE) * current_any_flag * not_final_step;
     for i in 0..NUM_ROUNDS {
@@ -64,6 +65,7 @@ pub(crate) fn eval_round_flags_recursively<F: RichField + Extendable<D>, const D
         builder.add_many_extension((0..NUM_ROUNDS).map(|i| local_values[reg_step(i)]));
     let next_any_flag =
         builder.add_many_extension((0..NUM_ROUNDS).map(|i| next_values[reg_step(i)]));
+    // Padding row should only start after the last round row.
     let not_final_step = builder.sub_extension(one, local_values[reg_step(NUM_ROUNDS - 1)]);
     let padding_constraint = {
         let tmp = builder.mul_sub_extension(current_any_flag, next_any_flag, current_any_flag);

From 72955b4cca25c83327a856ac3f63b740f8d80ad0 Mon Sep 17 00:00:00 2001
From: Ayush Shukla <shuklaayush247@gmail.com>
Date: Mon, 19 Feb 2024 00:38:06 +0100
Subject: [PATCH 8/8] fix: remove extra sub constraint

---
 evm_arithmetization/src/keccak/round_flags.rs | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/evm_arithmetization/src/keccak/round_flags.rs b/evm_arithmetization/src/keccak/round_flags.rs
index db7829b24..7b562118e 100644
--- a/evm_arithmetization/src/keccak/round_flags.rs
+++ b/evm_arithmetization/src/keccak/round_flags.rs
@@ -30,8 +30,9 @@ pub(crate) fn eval_round_flags<F: Field, P: PackedField<Scalar = F>>(
         .sum::<P>();
     let next_any_flag = (0..NUM_ROUNDS).map(|i| next_values[reg_step(i)]).sum::<P>();
     // Padding row should only start after the last round row.
-    let not_final_step = P::ONES - local_values[reg_step(NUM_ROUNDS - 1)];
-    let padding_constraint = (next_any_flag - F::ONE) * current_any_flag * not_final_step;
+    let last_round_flag = local_values[reg_step(NUM_ROUNDS - 1)];
+    let padding_constraint =
+        (next_any_flag - F::ONE) * current_any_flag * (last_round_flag - F::ONE);
     for i in 0..NUM_ROUNDS {
         let current_round_flag = local_values[reg_step(i)];
         let next_round_flag = next_values[reg_step((i + 1) % NUM_ROUNDS)];
@@ -66,10 +67,10 @@ pub(crate) fn eval_round_flags_recursively<F: RichField + Extendable<D>, const D
     let next_any_flag =
         builder.add_many_extension((0..NUM_ROUNDS).map(|i| next_values[reg_step(i)]));
     // Padding row should only start after the last round row.
-    let not_final_step = builder.sub_extension(one, local_values[reg_step(NUM_ROUNDS - 1)]);
+    let last_round_flag = local_values[reg_step(NUM_ROUNDS - 1)];
     let padding_constraint = {
         let tmp = builder.mul_sub_extension(current_any_flag, next_any_flag, current_any_flag);
-        builder.mul_extension(tmp, not_final_step)
+        builder.mul_sub_extension(tmp, last_round_flag, tmp)
     };
     for i in 0..NUM_ROUNDS {
         let current_round_flag = local_values[reg_step(i)];