Inter-row program counter constraints (#639)

* Beginning of control flow support

* Fixes to halt spin loop

evm/src/all_stark.rs

@@ -143,6 +143,7 @@ mod tests {
     use crate::all_stark::AllStark;
     use crate::config::StarkConfig;
     use crate::cpu::cpu_stark::CpuStark;
+    use crate::cpu::kernel::aggregator::KERNEL;
     use crate::cross_table_lookup::testutils::check_ctls;
     use crate::keccak::keccak_stark::{KeccakStark, NUM_INPUTS, NUM_ROUNDS};
     use crate::logic::{self, LogicStark, Operation};
@@ -321,8 +322,27 @@ mod tests {
 
         // Pad to a power of two.
         for _ in cpu_trace_rows.len()..cpu_trace_rows.len().next_power_of_two() {
-            cpu_trace_rows.push([F::ZERO; CpuStark::<F, D>::COLUMNS]);
+            let mut row: cpu::columns::CpuColumnsView<F> =
+                [F::ZERO; CpuStark::<F, D>::COLUMNS].into();
+            row.is_cpu_cycle = F::ONE;
+            cpu_stark.generate(row.borrow_mut());
+            cpu_trace_rows.push(row.into());
         }
+
+        // Ensure we finish in a halted state.
+        {
+            let num_rows = cpu_trace_rows.len();
+            let halt_label = F::from_canonical_usize(KERNEL.global_labels["halt_pc0"]);
+
+            let last_row: &mut cpu::columns::CpuColumnsView<F> =
+                cpu_trace_rows[num_rows - 1].borrow_mut();
+            last_row.program_counter = halt_label;
+
+            let second_last_row: &mut cpu::columns::CpuColumnsView<F> =
+                cpu_trace_rows[num_rows - 2].borrow_mut();
+            second_last_row.next_program_counter = halt_label;
+        }
+
         trace_rows_to_poly_values(cpu_trace_rows)
     }
 

evm/src/cpu/columns.rs

@@ -17,9 +17,12 @@ pub struct CpuColumnsView<T> {
     pub is_bootstrap_contract: T,
 
     /// Filter. 1 if the row corresponds to a cycle of execution and 0 otherwise.
-    /// Lets us re-use decode columns in non-cycle rows.
+    /// Lets us re-use columns in non-cycle rows.
     pub is_cpu_cycle: T,
 
+    /// If CPU cycle: The program counter for the current instruction.
+    pub program_counter: T,
+
     /// If CPU cycle: The opcode being decoded, in {0, ..., 255}.
     pub opcode: T,
 
@@ -139,6 +142,9 @@ pub struct CpuColumnsView<T> {
     /// If CPU cycle: the opcode, broken up into bits in **big-endian** order.
     pub opcode_bits: [T; 8],
 
+    /// If CPU cycle: The program counter for the next instruction.
+    pub next_program_counter: T,
+
     /// Filter. 1 iff a Keccak permutation is computed on this row.
     pub is_keccak: T,
     pub keccak_input_limbs: [T; 50],

evm/src/cpu/control_flow.rs

@@ -0,0 +1,112 @@
+use plonky2::field::extension::Extendable;
+use plonky2::field::packed::PackedField;
+use plonky2::field::types::Field;
+use plonky2::hash::hash_types::RichField;
+use plonky2::iop::ext_target::ExtensionTarget;
+
+use crate::constraint_consumer::{ConstraintConsumer, RecursiveConstraintConsumer};
+use crate::cpu::columns::CpuColumnsView;
+use crate::cpu::kernel::aggregator::KERNEL;
+
+fn get_halt_pcs<F: Field>() -> (F, F) {
+    let halt_pc0 = KERNEL.global_labels["halt_pc0"];
+    let halt_pc1 = KERNEL.global_labels["halt_pc1"];
+
+    (
+        F::from_canonical_usize(halt_pc0),
+        F::from_canonical_usize(halt_pc1),
+    )
+}
+
+pub fn eval_packed_generic<P: PackedField>(
+    lv: &CpuColumnsView<P>,
+    nv: &CpuColumnsView<P>,
+    yield_constr: &mut ConstraintConsumer<P>,
+) {
+    // Once we start executing instructions, then we continue until the end of the table.
+    yield_constr.constraint_transition(lv.is_cpu_cycle * (nv.is_cpu_cycle - P::ONES));
+
+    // If a row is a CPU cycle, then its `next_program_counter` becomes the `program_counter` of the
+    // next row.
+    yield_constr
+        .constraint_transition(lv.is_cpu_cycle * (nv.program_counter - lv.next_program_counter));
+
+    // If a non-CPU cycle row is followed by a CPU cycle row, then the `program_counter` of the CPU
+    // cycle row is 0.
+    yield_constr
+        .constraint_transition((lv.is_cpu_cycle - P::ONES) * nv.is_cpu_cycle * nv.program_counter);
+
+    // The first row has nowhere to continue execution from, so if it's a cycle row, then its
+    // `program_counter` must be 0.
+    // NB: I know the first few rows will be used for initialization and will not be CPU cycle rows.
+    // Once that's done, then this constraint can be removed. Until then, it is needed to ensure
+    // that execution starts at 0 and not at any arbitrary offset.
+    yield_constr.constraint_first_row(lv.is_cpu_cycle * lv.program_counter);
+
+    // The last row must be a CPU cycle row.
+    yield_constr.constraint_last_row(lv.is_cpu_cycle - P::ONES);
+    // Also, the last row's `program_counter` must be inside the `halt` infinite loop. Note that
+    // that loop consists of two instructions, so we must check for `halt` and `halt_inner` labels.
+    let (halt_pc0, halt_pc1) = get_halt_pcs::<P::Scalar>();
+    yield_constr
+        .constraint_last_row((lv.program_counter - halt_pc0) * (lv.program_counter - halt_pc1));
+}
+
+pub 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>,
+) {
+    // Once we start executing instructions, then we continue until the end of the table.
+    {
+        let constr = builder.mul_sub_extension(lv.is_cpu_cycle, nv.is_cpu_cycle, lv.is_cpu_cycle);
+        yield_constr.constraint_transition(builder, constr);
+    }
+
+    // If a row is a CPU cycle, then its `next_program_counter` becomes the `program_counter` of the
+    // next row.
+    {
+        let constr = builder.sub_extension(nv.program_counter, lv.next_program_counter);
+        let constr = builder.mul_extension(lv.is_cpu_cycle, constr);
+        yield_constr.constraint_transition(builder, constr);
+    }
+
+    // If a non-CPU cycle row is followed by a CPU cycle row, then the `program_counter` of the CPU
+    // cycle row is 0.
+    {
+        let constr = builder.mul_extension(nv.is_cpu_cycle, nv.program_counter);
+        let constr = builder.mul_sub_extension(lv.is_cpu_cycle, constr, constr);
+        yield_constr.constraint_transition(builder, constr);
+    }
+
+    // The first row has nowhere to continue execution from, so if it's a cycle row, then its
+    // `program_counter` must be 0.
+    // NB: I know the first few rows will be used for initialization and will not be CPU cycle rows.
+    // Once that's done, then this constraint can be removed. Until then, it is needed to ensure
+    // that execution starts at 0 and not at any arbitrary offset.
+    {
+        let constr = builder.mul_extension(lv.is_cpu_cycle, lv.program_counter);
+        yield_constr.constraint_first_row(builder, constr);
+    }
+
+    // The last row must be a CPU cycle row.
+    {
+        let one = builder.one_extension();
+        let constr = builder.sub_extension(lv.is_cpu_cycle, one);
+        yield_constr.constraint_last_row(builder, constr);
+    }
+    // Also, the last row's `program_counter` must be inside the `halt` infinite loop. Note that
+    // that loop consists of two instructions, so we must check for `halt` and `halt_inner` labels.
+    {
+        let (halt_pc0, halt_pc1) = get_halt_pcs();
+        let halt_pc0_target = builder.constant_extension(halt_pc0);
+        let halt_pc1_target = builder.constant_extension(halt_pc1);
+
+        let halt_pc0_offset = builder.sub_extension(lv.program_counter, halt_pc0_target);
+        let halt_pc1_offset = builder.sub_extension(lv.program_counter, halt_pc1_target);
+        let constr = builder.mul_extension(halt_pc0_offset, halt_pc1_offset);
+
+        yield_constr.constraint_last_row(builder, constr);
+    }
+}

evm/src/cpu/cpu_stark.rs

@@ -9,7 +9,7 @@ use plonky2::hash::hash_types::RichField;
 
 use crate::constraint_consumer::{ConstraintConsumer, RecursiveConstraintConsumer};
 use crate::cpu::columns::{CpuColumnsView, COL_MAP, NUM_CPU_COLUMNS};
-use crate::cpu::{bootstrap_kernel, decode, simple_logic};
+use crate::cpu::{bootstrap_kernel, control_flow, decode, simple_logic};
 use crate::cross_table_lookup::Column;
 use crate::memory::NUM_CHANNELS;
 use crate::stark::Stark;
@@ -88,7 +88,9 @@ impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for CpuStark<F, D
         P: PackedField<Scalar = FE>,
     {
         let local_values = vars.local_values.borrow();
+        let next_values = vars.next_values.borrow();
         bootstrap_kernel::eval_bootstrap_kernel(vars, yield_constr);
+        control_flow::eval_packed_generic(local_values, next_values, yield_constr);
         decode::eval_packed_generic(local_values, yield_constr);
         simple_logic::eval_packed(local_values, yield_constr);
     }
@@ -100,7 +102,9 @@ impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for CpuStark<F, D
         yield_constr: &mut RecursiveConstraintConsumer<F, D>,
     ) {
         let local_values = vars.local_values.borrow();
+        let next_values = vars.next_values.borrow();
         bootstrap_kernel::eval_bootstrap_kernel_circuit(builder, vars, yield_constr);
+        control_flow::eval_ext_circuit(builder, local_values, next_values, yield_constr);
         decode::eval_ext_circuit(builder, local_values, yield_constr);
         simple_logic::eval_ext_circuit(builder, local_values, yield_constr);
     }

evm/src/cpu/kernel/aggregator.rs

@@ -39,6 +39,7 @@ pub(crate) fn combined_kernel() -> Kernel {
         include_str!("asm/exp.asm"),
         include_str!("asm/curve_mul.asm"),
         include_str!("asm/curve_add.asm"),
+        include_str!("asm/halt.asm"),
         include_str!("asm/memory.asm"),
         include_str!("asm/moddiv.asm"),
         include_str!("asm/secp256k1/curve_mul.asm"),

evm/src/cpu/kernel/asm/halt.asm

@@ -0,0 +1,6 @@
+global halt:
+    PUSH halt_pc0
+global halt_pc0:
+    DUP1
+global halt_pc1:
+    JUMP

evm/src/cpu/mod.rs

@@ -1,5 +1,6 @@
 pub(crate) mod bootstrap_kernel;
 pub(crate) mod columns;
+mod control_flow;
 pub mod cpu_stark;
 pub(crate) mod decode;
 pub mod kernel;