From 249e50ebcb6a612b000e800b7670dfcd75302d40 Mon Sep 17 00:00:00 2001
From: Jacqueline Nabaglo <jakub@mirprotocol.org>
Date: Sun, 11 Dec 2022 10:59:14 -0800
Subject: [PATCH] Get/set context (#843)

---
 evm/src/cpu/contextops.rs | 82 +++++++++++++++++++++++++++++++++++++++
 evm/src/cpu/cpu_stark.rs  |  6 ++-
 evm/src/cpu/mod.rs        |  1 +
 evm/src/cpu/stack.rs      | 12 +++++-
 4 files changed, 97 insertions(+), 4 deletions(-)
 create mode 100644 evm/src/cpu/contextops.rs

diff --git a/evm/src/cpu/contextops.rs b/evm/src/cpu/contextops.rs
new file mode 100644
index 00000000..2a05d7c5
--- /dev/null
+++ b/evm/src/cpu/contextops.rs
@@ -0,0 +1,82 @@
+use plonky2::field::extension::Extendable;
+use plonky2::field::packed::PackedField;
+use plonky2::hash::hash_types::RichField;
+use plonky2::iop::ext_target::ExtensionTarget;
+use plonky2::plonk::circuit_builder::CircuitBuilder;
+
+use crate::constraint_consumer::{ConstraintConsumer, RecursiveConstraintConsumer};
+use crate::cpu::columns::CpuColumnsView;
+use crate::cpu::membus::NUM_GP_CHANNELS;
+
+fn eval_packed_get<P: PackedField>(
+    lv: &CpuColumnsView<P>,
+    yield_constr: &mut ConstraintConsumer<P>,
+) {
+    let filter = lv.op.get_context;
+    let push_channel = lv.mem_channels[NUM_GP_CHANNELS - 1];
+    yield_constr.constraint(filter * (push_channel.value[0] - lv.context));
+    for &limb in &push_channel.value[1..] {
+        yield_constr.constraint(filter * limb);
+    }
+}
+
+fn eval_ext_circuit_get<F: RichField + Extendable<D>, const D: usize>(
+    builder: &mut CircuitBuilder<F, D>,
+    lv: &CpuColumnsView<ExtensionTarget<D>>,
+    yield_constr: &mut RecursiveConstraintConsumer<F, D>,
+) {
+    let filter = lv.op.get_context;
+    let push_channel = lv.mem_channels[NUM_GP_CHANNELS - 1];
+    {
+        let diff = builder.sub_extension(push_channel.value[0], lv.context);
+        let constr = builder.mul_extension(filter, diff);
+        yield_constr.constraint(builder, constr);
+    }
+    for &limb in &push_channel.value[1..] {
+        let constr = builder.mul_extension(filter, limb);
+        yield_constr.constraint(builder, constr);
+    }
+}
+
+fn eval_packed_set<P: PackedField>(
+    lv: &CpuColumnsView<P>,
+    nv: &CpuColumnsView<P>,
+    yield_constr: &mut ConstraintConsumer<P>,
+) {
+    let filter = lv.op.set_context;
+    let pop_channel = lv.mem_channels[0];
+    yield_constr.constraint_transition(filter * (pop_channel.value[0] - nv.context));
+}
+
+fn eval_ext_circuit_set<F: RichField + Extendable<D>, const D: usize>(
+    builder: &mut CircuitBuilder<F, D>,
+    lv: &CpuColumnsView<ExtensionTarget<D>>,
+    nv: &CpuColumnsView<ExtensionTarget<D>>,
+    yield_constr: &mut RecursiveConstraintConsumer<F, D>,
+) {
+    let filter = lv.op.set_context;
+    let pop_channel = lv.mem_channels[0];
+
+    let diff = builder.sub_extension(pop_channel.value[0], nv.context);
+    let constr = builder.mul_extension(filter, diff);
+    yield_constr.constraint_transition(builder, constr);
+}
+
+pub fn eval_packed<P: PackedField>(
+    lv: &CpuColumnsView<P>,
+    nv: &CpuColumnsView<P>,
+    yield_constr: &mut ConstraintConsumer<P>,
+) {
+    eval_packed_get(lv, yield_constr);
+    eval_packed_set(lv, nv, yield_constr);
+}
+
+pub fn eval_ext_circuit<F: RichField + Extendable<D>, const D: usize>(
+    builder: &mut CircuitBuilder<F, D>,
+    lv: &CpuColumnsView<ExtensionTarget<D>>,
+    nv: &CpuColumnsView<ExtensionTarget<D>>,
+    yield_constr: &mut RecursiveConstraintConsumer<F, D>,
+) {
+    eval_ext_circuit_get(builder, lv, yield_constr);
+    eval_ext_circuit_set(builder, lv, nv, yield_constr);
+}
diff --git a/evm/src/cpu/cpu_stark.rs b/evm/src/cpu/cpu_stark.rs
index 1c478a3d..cd154e54 100644
--- a/evm/src/cpu/cpu_stark.rs
+++ b/evm/src/cpu/cpu_stark.rs
@@ -11,8 +11,8 @@ 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, pc, shift,
-    simple_logic, stack, stack_bounds, syscalls,
+    bootstrap_kernel, contextops, control_flow, decode, dup_swap, jumps, membus, memio, modfp254,
+    pc, shift, simple_logic, stack, stack_bounds, syscalls,
 };
 use crate::cross_table_lookup::Column;
 use crate::memory::segments::Segment;
@@ -141,6 +141,7 @@ impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for CpuStark<F, D
         // TODO: Some failing constraints temporarily disabled by using this dummy consumer.
         let mut dummy_yield_constr = ConstraintConsumer::new(vec![], P::ZEROS, P::ZEROS, P::ZEROS);
         bootstrap_kernel::eval_bootstrap_kernel(vars, yield_constr);
+        contextops::eval_packed(local_values, next_values, yield_constr);
         control_flow::eval_packed_generic(local_values, next_values, yield_constr);
         decode::eval_packed_generic(local_values, &mut dummy_yield_constr);
         dup_swap::eval_packed(local_values, yield_constr);
@@ -169,6 +170,7 @@ impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for CpuStark<F, D
         let mut dummy_yield_constr =
             RecursiveConstraintConsumer::new(zero, vec![], zero, zero, zero);
         bootstrap_kernel::eval_bootstrap_kernel_circuit(builder, vars, yield_constr);
+        contextops::eval_ext_circuit(builder, local_values, next_values, yield_constr);
         control_flow::eval_ext_circuit(builder, local_values, next_values, yield_constr);
         decode::eval_ext_circuit(builder, local_values, &mut dummy_yield_constr);
         dup_swap::eval_ext_circuit(builder, local_values, yield_constr);
diff --git a/evm/src/cpu/mod.rs b/evm/src/cpu/mod.rs
index 35a2c86a..466cf112 100644
--- a/evm/src/cpu/mod.rs
+++ b/evm/src/cpu/mod.rs
@@ -1,5 +1,6 @@
 pub(crate) mod bootstrap_kernel;
 pub(crate) mod columns;
+mod contextops;
 pub(crate) mod control_flow;
 pub mod cpu_stark;
 pub(crate) mod decode;
diff --git a/evm/src/cpu/stack.rs b/evm/src/cpu/stack.rs
index 8d3b9452..5c3d1576 100644
--- a/evm/src/cpu/stack.rs
+++ b/evm/src/cpu/stack.rs
@@ -76,8 +76,16 @@ const STACK_BEHAVIORS: OpsColumnsView<Option<StackBehavior>> = OpsColumnsView {
     push: None,     // TODO
     dup: None,
     swap: None,
-    get_context: None, // TODO
-    set_context: None, // TODO
+    get_context: Some(StackBehavior {
+        num_pops: 0,
+        pushes: true,
+        disable_other_channels: true,
+    }),
+    set_context: Some(StackBehavior {
+        num_pops: 1,
+        pushes: false,
+        disable_other_channels: true,
+    }),
     consume_gas: None, // TODO
     exit_kernel: None, // TODO
     mload_general: Some(StackBehavior {