Fixes

evm/src/arithmetic/mod.rs

@@ -23,12 +23,49 @@ pub(crate) enum BinaryOperator {
     Shr,
 }
 
+impl BinaryOperator {
+    pub(crate) fn result(&self, input0: U256, input1: U256) -> U256 {
+        match self {
+            BinaryOperator::Add => input0 + input1,
+            BinaryOperator::Mul => input0 * input1,
+            BinaryOperator::Sub => input0 - input1,
+            BinaryOperator::Div => input0 / input1,
+            BinaryOperator::Mod => input0 % input1,
+            BinaryOperator::Lt => {
+                if input0 < input1 {
+                    U256::one()
+                } else {
+                    U256::zero()
+                }
+            }
+            BinaryOperator::Gt => {
+                if input0 > input1 {
+                    U256::one()
+                } else {
+                    U256::zero()
+                }
+            }
+            BinaryOperator::Shl => input0 << input1,
+            BinaryOperator::Shr => input0 >> input1,
+        }
+    }
+}
+
 #[derive(Clone, Copy, Debug, Eq, PartialEq)]
 pub(crate) enum TernaryOperator {
     AddMod,
     MulMod,
 }
 
+impl TernaryOperator {
+    pub(crate) fn result(&self, input0: U256, input1: U256, input2: U256) -> U256 {
+        match self {
+            TernaryOperator::AddMod => (input0 + input1) % input2,
+            TernaryOperator::MulMod => (input0 * input1) % input2,
+        }
+    }
+}
+
 #[derive(Debug)]
 pub(crate) enum Operation {
     BinaryOperation {
@@ -45,3 +82,38 @@ pub(crate) enum Operation {
         result: U256,
     },
 }
+
+impl Operation {
+    pub(crate) fn binary(operator: BinaryOperator, input0: U256, input1: U256) -> Self {
+        let result = operator.result(input0, input1);
+        Self::BinaryOperation {
+            operator,
+            input0,
+            input1,
+            result,
+        }
+    }
+
+    pub(crate) fn ternary(
+        operator: TernaryOperator,
+        input0: U256,
+        input1: U256,
+        input2: U256,
+    ) -> Self {
+        let result = operator.result(input0, input1, input2);
+        Self::TernaryOperation {
+            operator,
+            input0,
+            input1,
+            input2,
+            result,
+        }
+    }
+
+    pub(crate) fn result(&self) -> U256 {
+        match self {
+            Operation::BinaryOperation { result, .. } => *result,
+            Operation::TernaryOperation { result, .. } => *result,
+        }
+    }
+}

evm/src/generation/mod.rs

@@ -83,7 +83,7 @@ pub(crate) fn generate_traces<F: RichField + Extendable<D>, const D: usize>(
             break;
         }
 
-        registers_state = transition(registers_state, &memory_state, &mut traces);
+        registers_state = transition(registers_state, &mut memory_state, &mut traces);
     }
 
     let read_metadata = |field| {

evm/src/witness/memory.rs

@@ -98,6 +98,17 @@ impl MemoryState {
         Self { contents }
     }
 
+    pub fn apply_ops(&mut self, ops: &[MemoryOp]) {
+        for &op in ops {
+            let MemoryOp {
+                address, op, value, ..
+            } = op;
+            if op == MemoryOpKind::Write {
+                self.set(address, value);
+            }
+        }
+    }
+
     pub fn get(&self, address: MemoryAddress) -> U256 {
         self.contents
             .get(&address)

evm/src/witness/operation.rs

@@ -42,10 +42,11 @@ pub(crate) fn generate_binary_logic_op<F: Field>(
 ) -> Result<RegistersState, ProgramError> {
     let [(in0, log_in0), (in1, log_in1)] =
         stack_pop_with_log_and_fill::<2, _>(&mut registers_state, memory_state, traces, &mut row)?;
-    let result = op.result(in0, in1);
-    let log_out = stack_push_log_and_fill(&mut registers_state, traces, &mut row, result)?;
+    let operation = logic::Operation::new(op, in0, in1);
+    let log_out =
+        stack_push_log_and_fill(&mut registers_state, traces, &mut row, operation.result)?;
 
-    traces.push_logic(logic::Operation::new(op, in0, in1));
+    traces.push_logic(operation);
     traces.push_memory(log_in0);
     traces.push_memory(log_in1);
     traces.push_memory(log_out);
@@ -53,6 +54,49 @@ pub(crate) fn generate_binary_logic_op<F: Field>(
     Ok(registers_state)
 }
 
+pub(crate) fn generate_binary_arithmetic_op<F: Field>(
+    operator: arithmetic::BinaryOperator,
+    mut registers_state: RegistersState,
+    memory_state: &MemoryState,
+    traces: &mut Traces<F>,
+    mut row: CpuColumnsView<F>,
+) -> Result<RegistersState, ProgramError> {
+    let [(input0, log_in0), (input1, log_in1)] =
+        stack_pop_with_log_and_fill::<2, _>(&mut registers_state, memory_state, traces, &mut row)?;
+    let operation = arithmetic::Operation::binary(operator, input0, input1);
+    let log_out =
+        stack_push_log_and_fill(&mut registers_state, traces, &mut row, operation.result())?;
+
+    traces.push_arithmetic(operation);
+    traces.push_memory(log_in0);
+    traces.push_memory(log_in1);
+    traces.push_memory(log_out);
+    traces.push_cpu(row);
+    Ok(registers_state)
+}
+
+pub(crate) fn generate_ternary_arithmetic_op<F: Field>(
+    operator: arithmetic::TernaryOperator,
+    mut registers_state: RegistersState,
+    memory_state: &MemoryState,
+    traces: &mut Traces<F>,
+    mut row: CpuColumnsView<F>,
+) -> Result<RegistersState, ProgramError> {
+    let [(input0, log_in0), (input1, log_in1), (input2, log_in2)] =
+        stack_pop_with_log_and_fill::<3, _>(&mut registers_state, memory_state, traces, &mut row)?;
+    let operation = arithmetic::Operation::ternary(operator, input0, input1, input2);
+    let log_out =
+        stack_push_log_and_fill(&mut registers_state, traces, &mut row, operation.result())?;
+
+    traces.push_arithmetic(operation);
+    traces.push_memory(log_in0);
+    traces.push_memory(log_in1);
+    traces.push_memory(log_in2);
+    traces.push_memory(log_out);
+    traces.push_cpu(row);
+    Ok(registers_state)
+}
+
 pub(crate) fn generate_push<F: Field>(
     n: u8,
     mut registers_state: RegistersState,

evm/src/witness/traces.rs

@@ -4,7 +4,6 @@ use plonky2::hash::hash_types::RichField;
 use plonky2::util::timing::TimingTree;
 
 use crate::all_stark::{AllStark, NUM_TABLES};
-use crate::arithmetic::columns::NUM_ARITH_COLUMNS;
 use crate::config::StarkConfig;
 use crate::cpu::columns::CpuColumnsView;
 use crate::keccak_memory::keccak_memory_stark::KeccakMemoryOp;
@@ -63,6 +62,10 @@ impl<T: Copy> Traces<T> {
         // TODO others
     }
 
+    pub fn mem_ops_since(&self, checkpoint: TraceCheckpoint) -> &[MemoryOp] {
+        &self.memory_ops[checkpoint.memory_len..]
+    }
+
     pub fn push_cpu(&mut self, val: CpuColumnsView<T>) {
         self.cpu.push(val);
     }

evm/src/witness/transition.rs

@@ -5,8 +5,9 @@ use crate::memory::segments::Segment;
 use crate::witness::errors::ProgramError;
 use crate::witness::memory::{MemoryAddress, MemoryState};
 use crate::witness::operation::{
-    generate_binary_logic_op, generate_dup, generate_eq, generate_exit_kernel, generate_iszero,
-    generate_not, generate_push, generate_swap, generate_syscall, Operation,
+    generate_binary_arithmetic_op, generate_binary_logic_op, generate_dup, generate_eq,
+    generate_exit_kernel, generate_iszero, generate_not, generate_push, generate_swap,
+    generate_syscall, generate_ternary_arithmetic_op, Operation,
 };
 use crate::witness::state::RegistersState;
 use crate::witness::traces::Traces;
@@ -188,6 +189,12 @@ fn perform_op<F: Field>(
         Operation::BinaryLogic(binary_logic_op) => {
             generate_binary_logic_op(binary_logic_op, registers_state, memory_state, traces, row)?
         }
+        Operation::BinaryArithmetic(op) => {
+            generate_binary_arithmetic_op(op, registers_state, memory_state, traces, row)?
+        }
+        Operation::TernaryArithmetic(op) => {
+            generate_ternary_arithmetic_op(op, registers_state, memory_state, traces, row)?
+        }
         _ => panic!("operation not implemented: {:?}", op),
     };
 
@@ -233,11 +240,12 @@ fn handle_error<F: Field>(
 
 pub(crate) fn transition<F: Field>(
     registers_state: RegistersState,
-    memory_state: &MemoryState,
+    memory_state: &mut MemoryState,
     traces: &mut Traces<F>,
 ) -> RegistersState {
     let checkpoint = traces.checkpoint();
     let result = try_perform_instruction(registers_state, memory_state, traces);
+    memory_state.apply_ops(traces.mem_ops_since(checkpoint));
 
     match result {
         Ok(new_registers_state) => new_registers_state,