Move stack_len_bounds_aux to general columns (0xPolygonZero#1360)

* Move stack_len_bounds_aux to general columns

* Update specs

* Apply comments

* Apply comment

evm/spec/cpulogic.tex

@@ -114,6 +114,7 @@ \subsection{Privileged instructions}
 
 
 \subsection{Stack handling}
+\label{stackhandling}
 
 \subsubsection{Top of the stack}
 
@@ -188,6 +189,9 @@ \subsubsection{Stack length checking}
 $$\texttt{(1 - is\_kernel\_mode) - stack\_len * stack\_len\_bounds\_aux}.$$
 The flag is 1 if \texttt{stack\_len = 1025} and we're in user mode, and 0 otherwise.
 
+Because \texttt{stack\_len\_bounds\_aux} is a shared general column, we only check this constraint after an instruction that can actually trigger an overflow,
+i.e. a pushing, non-popping instruction.
+
 \subsection{Gas handling}
 
 \subsubsection{Out of gas errors}

evm/spec/tables/cpu.tex

@@ -31,7 +31,6 @@ \subsubsection{CPU columns}
     \item \texttt{code\_context}: Indicates in which context the code to execute resides. It's equal to \texttt{context} in user mode, but is always 0 in kernel mode.
     \item \texttt{program\_counter}: The address of the instruction to be read and executed.
     \item \texttt{stack\_len}: The current length of the stack.
-    \item \texttt{stack\_len\_bounds\_aux}: Helper column used to check that the stack doesn't overflow in user mode.
     \item \texttt{is\_kernel\_mode}: Boolean indicating whether we are in kernel (i.e. privileged) mode. This means we are executing kernel code, and we have access to
 privileged instructions.
     \item \texttt{gas}: The current amount of gas used in the current context. It is eventually checked to be below the current gas limit. Must fit in 32 bits.
@@ -73,6 +72,7 @@ \subsubsection{CPU columns}
 of the sum of the seven high limbs, and is used to check it's non-zero like the previous cases.
 Contrary to the logic operations, we do not need to check limbs individually: each limb has been range-checked to 32 bits, meaning that it's not possible for the sum to
 overflow and be zero if some of the limbs are non-zero.
-    \item  \texttt{Stack}: The last three columns are used by popping-only (resp. pushing-only) instructions to check if the stack is empty after (resp. was empty
-before) the instruction. We use the last columns to prevent conflicts with the other general columns. More details are provided in the stack handling section.
+    \item  \texttt{Stack}: \texttt{stack\_inv}, \texttt{stack\_inv\_aux} and \texttt{stack\_inv\_aux\_2} are used by popping-only (resp. pushing-only) instructions to check if the stack is empty after (resp. was empty
+before) the instruction. \texttt{stack\_len\_bounds\_ aux} is used to check that the stack doesn't overflow in user mode. We use the last four columns to prevent conflicts with the other general columns.
+See \ref{stackhandling} for more details.
 \end{itemize}

evm/src/cpu/columns/general.rs

@@ -135,18 +135,20 @@ pub(crate) struct CpuShiftView<T: Copy> {
     pub(crate) high_limb_sum_inv: T,
 }
 
-/// View of the last three `CpuGeneralColumns` storing the stack length pseudoinverse `stack_inv`,
-/// stack_len * stack_inv and filter * stack_inv_aux when needed.
+/// View of the last four `CpuGeneralColumns` storing stack-related variables. The first three are used
+/// for conditionally enabling and disabling channels when reading the next `stack_top`, and the fourth one
+/// is used to check for stack overflow.
 #[derive(Copy, Clone)]
 pub(crate) struct CpuStackView<T: Copy> {
-    // Used for conditionally enabling and disabling channels when reading the next `stack_top`.
-    _unused: [T; 5],
-    /// Pseudoinverse of the stack len.
+    _unused: [T; 4],
+    /// Pseudoinverse of `stack_len - num_pops`.
     pub(crate) stack_inv: T,
     /// stack_inv * stack_len.
     pub(crate) stack_inv_aux: T,
-    /// Holds filter * stack_inv_aux when necessary, to reduce the degree of stack constraints.
+    /// Used to reduce the degree of stack constraints when needed.
     pub(crate) stack_inv_aux_2: T,
+    /// Pseudoinverse of `nv.stack_len - (MAX_USER_STACK_SIZE + 1)` to check for stack overflow.
+    pub(crate) stack_len_bounds_aux: T,
 }
 
 /// Number of columns shared by all the views of `CpuGeneralColumnsView`.

evm/src/cpu/columns/mod.rs

@@ -68,10 +68,6 @@ pub(crate) struct CpuColumnsView<T: Copy> {
     /// If CPU cycle: The stack length.
     pub stack_len: T,
 
-    /// If CPU cycle: A prover-provided value needed to show that the instruction does not cause the
-    /// stack to underflow or overflow.
-    pub stack_len_bounds_aux: T,
-
     /// If CPU cycle: We're in kernel (privileged) mode.
     pub is_kernel_mode: T,
 

evm/src/cpu/cpu_stark.rs

@@ -17,8 +17,7 @@ use crate::constraint_consumer::{ConstraintConsumer, RecursiveConstraintConsumer
 use crate::cpu::columns::{COL_MAP, NUM_CPU_COLUMNS};
 use crate::cpu::{
     bootstrap_kernel, byte_unpacking, clock, contextops, control_flow, decode, dup_swap, gas,
-    jumps, membus, memio, modfp254, pc, push0, shift, simple_logic, stack, stack_bounds,
-    syscalls_exceptions,
+    jumps, membus, memio, modfp254, pc, push0, shift, simple_logic, stack, syscalls_exceptions,
 };
 use crate::cross_table_lookup::{Column, TableWithColumns};
 use crate::evaluation_frame::{StarkEvaluationFrame, StarkFrame};
@@ -314,7 +313,6 @@ 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, next_values, yield_constr);
         stack::eval_packed(local_values, next_values, yield_constr);
-        stack_bounds::eval_packed(local_values, yield_constr);
         syscalls_exceptions::eval_packed(local_values, next_values, yield_constr);
     }
 
@@ -356,7 +354,6 @@ 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, next_values, yield_constr);
         stack::eval_ext_circuit(builder, local_values, next_values, yield_constr);
-        stack_bounds::eval_ext_circuit(builder, local_values, yield_constr);
         syscalls_exceptions::eval_ext_circuit(builder, local_values, next_values, yield_constr);
     }
 

evm/src/cpu/kernel/interpreter.rs

@@ -14,7 +14,7 @@ use crate::cpu::kernel::aggregator::KERNEL;
 use crate::cpu::kernel::constants::context_metadata::ContextMetadata;
 use crate::cpu::kernel::constants::global_metadata::GlobalMetadata;
 use crate::cpu::kernel::constants::txn_fields::NormalizedTxnField;
-use crate::cpu::stack_bounds::MAX_USER_STACK_SIZE;
+use crate::cpu::stack::MAX_USER_STACK_SIZE;
 use crate::extension_tower::BN_BASE;
 use crate::generation::prover_input::ProverInputFn;
 use crate::generation::state::GenerationState;

evm/src/cpu/mod.rs

@@ -19,5 +19,4 @@ mod push0;
 mod shift;
 pub(crate) mod simple_logic;
 pub(crate) mod stack;
-pub(crate) mod stack_bounds;
 mod syscalls_exceptions;

evm/src/cpu/stack.rs

@@ -13,6 +13,35 @@ use crate::cpu::columns::CpuColumnsView;
 use crate::cpu::membus::NUM_GP_CHANNELS;
 use crate::memory::segments::Segment;
 
+pub(crate) const MAX_USER_STACK_SIZE: usize = 1024;
+
+// We check for stack overflows here. An overflow occurs when the stack length is 1025 in user mode,
+// which can happen after a non-kernel-only, non-popping, pushing instruction/syscall.
+// The check uses `stack_len_bounds_aux`, which is either 0 if next row's `stack_len` is 1025 or
+// next row is in kernel mode, or the inverse of `nv.stack_len - 1025` otherwise.
+pub(crate) const MIGHT_OVERFLOW: OpsColumnsView<bool> = OpsColumnsView {
+    binary_op: false,
+    ternary_op: false,
+    fp254_op: false,
+    eq_iszero: false,
+    logic_op: false,
+    not_pop: false,
+    shift: false,
+    jumpdest_keccak_general: false,
+    prover_input: false,
+    jumps: false,
+    pc_push0: true,
+    push: true,
+    dup_swap: true,
+    context_op: false,
+    mload_32bytes: false,
+    mstore_32bytes: false,
+    exit_kernel: true, // Doesn't directly push, but the syscall it's returning from might.
+    m_op_general: false,
+    syscall: false,
+    exception: false,
+};
+
 /// Structure to represent opcodes stack behaviours:
 /// - number of pops
 /// - whether the opcode(s) push
@@ -270,10 +299,22 @@ pub(crate) fn eval_packed<P: PackedField>(
     nv: &CpuColumnsView<P>,
     yield_constr: &mut ConstraintConsumer<P>,
 ) {
-    for (op, stack_behavior) in izip!(lv.op.into_iter(), STACK_BEHAVIORS.into_iter()) {
+    for (op, stack_behavior, might_overflow) in izip!(
+        lv.op.into_iter(),
+        STACK_BEHAVIORS.into_iter(),
+        MIGHT_OVERFLOW.into_iter()
+    ) {
         if let Some(stack_behavior) = stack_behavior {
             eval_packed_one(lv, nv, op, stack_behavior, yield_constr);
         }
+
+        if might_overflow {
+            // Check for stack overflow in the next row.
+            let diff = nv.stack_len - P::Scalar::from_canonical_usize(MAX_USER_STACK_SIZE + 1);
+            let lhs = diff * lv.general.stack().stack_len_bounds_aux;
+            let rhs = P::ONES - nv.is_kernel_mode;
+            yield_constr.constraint_transition(op * (lhs - rhs));
+        }
     }
 
     // Constrain stack for JUMPDEST.
@@ -549,18 +590,38 @@ pub(crate) fn eval_ext_circuit_one<F: RichField + Extendable<D>, const D: usize>
     yield_constr.constraint_transition(builder, constr);
 }
 
-/// Circuti version of `eval_packed`.
+/// Circuit version of `eval_packed`.
 /// Evaluates constraints for all opcodes' `StackBehavior`s.
 pub(crate) fn eval_ext_circuit<F: RichField + Extendable<D>, const D: usize>(
     builder: &mut plonky2::plonk::circuit_builder::CircuitBuilder<F, D>,
     lv: &CpuColumnsView<ExtensionTarget<D>>,
     nv: &CpuColumnsView<ExtensionTarget<D>>,
     yield_constr: &mut RecursiveConstraintConsumer<F, D>,
 ) {
-    for (op, stack_behavior) in izip!(lv.op.into_iter(), STACK_BEHAVIORS.into_iter()) {
+    for (op, stack_behavior, might_overflow) in izip!(
+        lv.op.into_iter(),
+        STACK_BEHAVIORS.into_iter(),
+        MIGHT_OVERFLOW.into_iter()
+    ) {
         if let Some(stack_behavior) = stack_behavior {
             eval_ext_circuit_one(builder, lv, nv, op, stack_behavior, yield_constr);
         }
+
+        if might_overflow {
+            // Check for stack overflow in the next row.
+            let diff = builder.add_const_extension(
+                nv.stack_len,
+                -F::from_canonical_usize(MAX_USER_STACK_SIZE + 1),
+            );
+            let prod = builder.mul_add_extension(
+                diff,
+                lv.general.stack().stack_len_bounds_aux,
+                nv.is_kernel_mode,
+            );
+            let rhs = builder.add_const_extension(prod, -F::ONE);
+            let constr = builder.mul_extension(op, rhs);
+            yield_constr.constraint_transition(builder, constr);
+        }
     }
 
     // Constrain stack for JUMPDEST.

evm/src/witness/operation.rs

@@ -14,7 +14,7 @@ use crate::cpu::kernel::assembler::BYTES_PER_OFFSET;
 use crate::cpu::kernel::constants::context_metadata::ContextMetadata;
 use crate::cpu::membus::NUM_GP_CHANNELS;
 use crate::cpu::simple_logic::eq_iszero::generate_pinv_diff;
-use crate::cpu::stack_bounds::MAX_USER_STACK_SIZE;
+use crate::cpu::stack::MAX_USER_STACK_SIZE;
 use crate::extension_tower::BN_BASE;
 use crate::generation::state::GenerationState;
 use crate::memory::segments::Segment;
@@ -24,6 +24,7 @@ use crate::witness::errors::ProgramError;
 use crate::witness::errors::ProgramError::MemoryError;
 use crate::witness::memory::{MemoryAddress, MemoryChannel, MemoryOp, MemoryOpKind};
 use crate::witness::operation::MemoryChannel::GeneralPurpose;
+use crate::witness::transition::fill_stack_fields;
 use crate::witness::util::{
     keccak_sponge_log, mem_read_gp_with_log_and_fill, mem_write_gp_log_and_fill,
     stack_pop_with_log_and_fill,
@@ -950,44 +951,12 @@ pub(crate) fn generate_exception<F: Field>(
 
     row.op.exception = F::ONE;
 
-    let disallowed_len = F::from_canonical_usize(MAX_USER_STACK_SIZE + 1);
-    let diff = row.stack_len - disallowed_len;
-    if let Some(inv) = diff.try_inverse() {
-        row.stack_len_bounds_aux = inv;
-    } else {
-        // This is a stack overflow that should have been caught earlier.
-        return Err(ProgramError::InterpreterError);
-    }
-
     if let Some(inv) = row.stack_len.try_inverse() {
         row.general.stack_mut().stack_inv = inv;
         row.general.stack_mut().stack_inv_aux = F::ONE;
     }
 
-    if state.registers.is_stack_top_read {
-        let channel = &mut row.mem_channels[0];
-        channel.used = F::ONE;
-        channel.is_read = F::ONE;
-        channel.addr_context = F::from_canonical_usize(state.registers.context);
-        channel.addr_segment = F::from_canonical_usize(Segment::Stack as usize);
-        channel.addr_virtual = F::from_canonical_usize(state.registers.stack_len - 1);
-
-        let address = MemoryAddress {
-            context: state.registers.context,
-            segment: Segment::Stack as usize,
-            virt: state.registers.stack_len - 1,
-        };
-
-        let mem_op = MemoryOp::new(
-            GeneralPurpose(0),
-            state.traces.clock(),
-            address,
-            MemoryOpKind::Read,
-            state.registers.stack_top,
-        );
-        state.traces.push_memory(mem_op);
-        state.registers.is_stack_top_read = false;
-    }
+    fill_stack_fields(state, &mut row);
 
     row.general.exception_mut().exc_code_bits = [
         F::from_bool(exc_code & 1 != 0),

evm/src/witness/state.rs

@@ -12,6 +12,9 @@ pub struct RegistersState {
     pub stack_top: U256,
     // Indicates if you read the new stack_top from memory to set the channel accordingly.
     pub is_stack_top_read: bool,
+    // Indicates if the previous operation might have caused an overflow, and we must check
+    // if it's the case.
+    pub check_overflow: bool,
     pub context: usize,
     pub gas_used: u64,
 }
@@ -34,6 +37,7 @@ impl Default for RegistersState {
             stack_len: 0,
             stack_top: U256::zero(),
             is_stack_top_read: false,
+            check_overflow: false,
             context: 0,
             gas_used: 0,
         }