From 938e3bd5dab45e7b6b49d9b9e368b85350637be1 Mon Sep 17 00:00:00 2001
From: Hamy Ratoanina <hamy.ratoanina@toposware.com>
Date: Tue, 11 Apr 2023 15:13:55 -0400
Subject: [PATCH] Set stack_len_bounds_aux properly

---
 evm/src/cpu/cpu_stark.rs      |  4 ++--
 evm/src/cpu/stack_bounds.rs   |  4 ++--
 evm/src/witness/transition.rs | 19 +++++++++++++++++++
 3 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/evm/src/cpu/cpu_stark.rs b/evm/src/cpu/cpu_stark.rs
index 09c7d576..21bf75cc 100644
--- a/evm/src/cpu/cpu_stark.rs
+++ b/evm/src/cpu/cpu_stark.rs
@@ -157,7 +157,7 @@ impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for CpuStark<F, D
         shift::eval_packed(local_values, yield_constr);
         simple_logic::eval_packed(local_values, yield_constr);
         stack::eval_packed(local_values, &mut dummy_yield_constr);
-        stack_bounds::eval_packed(local_values, &mut dummy_yield_constr);
+        stack_bounds::eval_packed(local_values, yield_constr);
         syscalls::eval_packed(local_values, next_values, yield_constr);
     }
 
@@ -187,7 +187,7 @@ impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for CpuStark<F, D
         shift::eval_ext_circuit(builder, local_values, yield_constr);
         simple_logic::eval_ext_circuit(builder, local_values, yield_constr);
         stack::eval_ext_circuit(builder, local_values, &mut dummy_yield_constr);
-        stack_bounds::eval_ext_circuit(builder, local_values, &mut dummy_yield_constr);
+        stack_bounds::eval_ext_circuit(builder, local_values, yield_constr);
         syscalls::eval_ext_circuit(builder, local_values, next_values, yield_constr);
     }
 
diff --git a/evm/src/cpu/stack_bounds.rs b/evm/src/cpu/stack_bounds.rs
index 627411ea..4c4995f4 100644
--- a/evm/src/cpu/stack_bounds.rs
+++ b/evm/src/cpu/stack_bounds.rs
@@ -26,13 +26,13 @@ pub const MAX_USER_STACK_SIZE: usize = 1024;
 //   Other operations that have a minimum stack size (e.g. `MULMOD`, which has three inputs) read
 // all their inputs from memory. On underflow, the cross-table lookup fails, as -1, ..., -17 are
 // invalid memory addresses.
-const DECREMENTING_FLAGS: [usize; 1] = [COL_MAP.op.pop];
+pub(crate) const DECREMENTING_FLAGS: [usize; 1] = [COL_MAP.op.pop];
 
 // Operations that increase the stack length by 1, but excluding:
 //  - privileged (kernel-only) operations (superfluous; doesn't affect correctness),
 //  - operations that from userspace to the kernel (required for correctness).
 // TODO: This list is incomplete.
-const INCREMENTING_FLAGS: [usize; 2] = [COL_MAP.op.pc, COL_MAP.op.dup];
+pub(crate) const INCREMENTING_FLAGS: [usize; 2] = [COL_MAP.op.pc, COL_MAP.op.dup];
 
 /// Calculates `lv.stack_len_bounds_aux`. Note that this must be run after decode.
 pub fn generate<F: Field>(lv: &mut CpuColumnsView<F>) {
diff --git a/evm/src/witness/transition.rs b/evm/src/witness/transition.rs
index 65b7ddcf..0184e183 100644
--- a/evm/src/witness/transition.rs
+++ b/evm/src/witness/transition.rs
@@ -4,6 +4,7 @@ use plonky2::field::types::Field;
 
 use crate::cpu::columns::CpuColumnsView;
 use crate::cpu::kernel::aggregator::KERNEL;
+use crate::cpu::stack_bounds::{DECREMENTING_FLAGS, INCREMENTING_FLAGS, MAX_USER_STACK_SIZE};
 use crate::generation::state::GenerationState;
 use crate::memory::segments::Segment;
 use crate::witness::errors::ProgramError;
@@ -258,6 +259,24 @@ fn try_perform_instruction<F: Field>(state: &mut GenerationState<F>) -> Result<(
 
     fill_op_flag(op, &mut row);
 
+    let check_underflow: F = DECREMENTING_FLAGS.map(|i| row[i]).into_iter().sum();
+    let check_overflow: F = INCREMENTING_FLAGS.map(|i| row[i]).into_iter().sum();
+    let no_check = F::ONE - (check_underflow + check_overflow);
+
+    let disallowed_len = check_overflow * F::from_canonical_usize(MAX_USER_STACK_SIZE) - no_check;
+    let diff = row.stack_len - disallowed_len;
+
+    let user_mode = F::ONE - row.is_kernel_mode;
+    let rhs = user_mode + check_underflow;
+
+    row.stack_len_bounds_aux = match diff.try_inverse() {
+        Some(diff_inv) => diff_inv * rhs, // `rhs` may be a value other than 1 or 0
+        None => {
+            assert_eq!(rhs, F::ZERO);
+            F::ZERO
+        }
+    };
+
     perform_op(state, op, row)
 }