From 29644e5111e1bd052115c8ec3dfe7aee454098c9 Mon Sep 17 00:00:00 2001
From: Jacqueline Nabaglo <jakub@mirprotocol.org>
Date: Sun, 11 Dec 2022 10:41:32 -0800
Subject: [PATCH] Implement `PC` instruction (#847)

* Implement `PC` instruction

* lints
---
 evm/src/cpu/cpu_stark.rs      |  4 +++-
 evm/src/cpu/mod.rs            |  1 +
 evm/src/cpu/pc.rs             | 38 +++++++++++++++++++++++++++++++++++
 evm/src/cpu/stack.rs          | 12 +++++++----
 evm/src/witness/operation.rs  | 10 +++++++++
 evm/src/witness/transition.rs |  2 +-
 6 files changed, 61 insertions(+), 6 deletions(-)
 create mode 100644 evm/src/cpu/pc.rs

diff --git a/evm/src/cpu/cpu_stark.rs b/evm/src/cpu/cpu_stark.rs
index aa48bb37..1c478a3d 100644
--- a/evm/src/cpu/cpu_stark.rs
+++ b/evm/src/cpu/cpu_stark.rs
@@ -11,7 +11,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, control_flow, decode, dup_swap, jumps, membus, memio, modfp254, shift,
+    bootstrap_kernel, control_flow, decode, dup_swap, jumps, membus, memio, modfp254, pc, shift,
     simple_logic, stack, stack_bounds, syscalls,
 };
 use crate::cross_table_lookup::Column;
@@ -148,6 +148,7 @@ impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for CpuStark<F, D
         membus::eval_packed(local_values, yield_constr);
         memio::eval_packed(local_values, yield_constr);
         modfp254::eval_packed(local_values, yield_constr);
+        pc::eval_packed(local_values, yield_constr);
         shift::eval_packed(local_values, yield_constr);
         simple_logic::eval_packed(local_values, yield_constr);
         stack::eval_packed(local_values, &mut dummy_yield_constr);
@@ -175,6 +176,7 @@ impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for CpuStark<F, D
         membus::eval_ext_circuit(builder, local_values, yield_constr);
         memio::eval_ext_circuit(builder, local_values, yield_constr);
         modfp254::eval_ext_circuit(builder, local_values, yield_constr);
+        pc::eval_ext_circuit(builder, local_values, yield_constr);
         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);
diff --git a/evm/src/cpu/mod.rs b/evm/src/cpu/mod.rs
index d4ff95c8..35a2c86a 100644
--- a/evm/src/cpu/mod.rs
+++ b/evm/src/cpu/mod.rs
@@ -9,6 +9,7 @@ pub mod kernel;
 pub(crate) mod membus;
 mod memio;
 mod modfp254;
+mod pc;
 mod shift;
 pub(crate) mod simple_logic;
 mod stack;
diff --git a/evm/src/cpu/pc.rs b/evm/src/cpu/pc.rs
new file mode 100644
index 00000000..26731c92
--- /dev/null
+++ b/evm/src/cpu/pc.rs
@@ -0,0 +1,38 @@
+use plonky2::field::extension::Extendable;
+use plonky2::field::packed::PackedField;
+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::membus::NUM_GP_CHANNELS;
+
+pub fn eval_packed<P: PackedField>(
+    lv: &CpuColumnsView<P>,
+    yield_constr: &mut ConstraintConsumer<P>,
+) {
+    let filter = lv.op.pc;
+    let push_value = lv.mem_channels[NUM_GP_CHANNELS - 1].value;
+    yield_constr.constraint(filter * (push_value[0] - lv.program_counter));
+    for &limb in &push_value[1..] {
+        yield_constr.constraint(filter * limb);
+    }
+}
+
+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>>,
+    yield_constr: &mut RecursiveConstraintConsumer<F, D>,
+) {
+    let filter = lv.op.pc;
+    let push_value = lv.mem_channels[NUM_GP_CHANNELS - 1].value;
+    {
+        let diff = builder.sub_extension(push_value[0], lv.program_counter);
+        let constr = builder.mul_extension(filter, diff);
+        yield_constr.constraint(builder, constr);
+    }
+    for &limb in &push_value[1..] {
+        let constr = builder.mul_extension(filter, limb);
+        yield_constr.constraint(builder, constr);
+    }
+}
diff --git a/evm/src/cpu/stack.rs b/evm/src/cpu/stack.rs
index 3c358d24..8d3b9452 100644
--- a/evm/src/cpu/stack.rs
+++ b/evm/src/cpu/stack.rs
@@ -66,10 +66,14 @@ const STACK_BEHAVIORS: OpsColumnsView<Option<StackBehavior>> = OpsColumnsView {
     pop: None,            // TODO
     jump: None,           // TODO
     jumpi: None,          // TODO
-    pc: None,             // TODO
-    gas: None,            // TODO
-    jumpdest: None,       // TODO
-    push: None,           // TODO
+    pc: Some(StackBehavior {
+        num_pops: 0,
+        pushes: true,
+        disable_other_channels: true,
+    }),
+    gas: None,      // TODO
+    jumpdest: None, // TODO
+    push: None,     // TODO
     dup: None,
     swap: None,
     get_context: None, // TODO
diff --git a/evm/src/witness/operation.rs b/evm/src/witness/operation.rs
index 5c00b9a5..fa9d0fac 100644
--- a/evm/src/witness/operation.rs
+++ b/evm/src/witness/operation.rs
@@ -215,6 +215,16 @@ pub(crate) fn generate_jumpi<F: Field>(
     Ok(())
 }
 
+pub(crate) fn generate_pc<F: Field>(
+    state: &mut GenerationState<F>,
+    mut row: CpuColumnsView<F>,
+) -> Result<(), ProgramError> {
+    let write = stack_push_log_and_fill(state, &mut row, state.registers.program_counter.into())?;
+    state.traces.push_memory(write);
+    state.traces.push_cpu(row);
+    Ok(())
+}
+
 pub(crate) fn generate_jumpdest<F: Field>(
     state: &mut GenerationState<F>,
     row: CpuColumnsView<F>,
diff --git a/evm/src/witness/transition.rs b/evm/src/witness/transition.rs
index 69be1f83..930028b9 100644
--- a/evm/src/witness/transition.rs
+++ b/evm/src/witness/transition.rs
@@ -208,7 +208,7 @@ fn perform_op<F: Field>(
         Operation::Pop => generate_pop(state, row)?,
         Operation::Jump => generate_jump(state, row)?,
         Operation::Jumpi => generate_jumpi(state, row)?,
-        Operation::Pc => todo!(),
+        Operation::Pc => generate_pc(state, row)?,
         Operation::Gas => todo!(),
         Operation::Jumpdest => generate_jumpdest(state, row)?,
         Operation::GetContext => generate_get_context(state, row)?,