From bffdc553dc4c5bf1dcb816b0f4d8b580e0e820cc Mon Sep 17 00:00:00 2001
From: Daniel Lubarov <daniel@lubarov.com>
Date: Sat, 3 Dec 2022 22:44:54 -0800
Subject: [PATCH] fixes

---
 evm/src/generation/mod.rs    |  2 ++
 evm/src/witness/operation.rs | 25 ++++++++++++++-----------
 evm/src/witness/traces.rs    |  9 +++++++--
 3 files changed, 23 insertions(+), 13 deletions(-)

diff --git a/evm/src/generation/mod.rs b/evm/src/generation/mod.rs
index bff2a6c9..e2384770 100644
--- a/evm/src/generation/mod.rs
+++ b/evm/src/generation/mod.rs
@@ -74,6 +74,8 @@ pub(crate) fn generate_traces<F: RichField + Extendable<D>, const D: usize>(
 
     timed!(timing, "simulate CPU", simulate_cpu(&mut state));
 
+    log::info!("Trace lengths (before padding): {:?}", state.traces.checkpoint());
+
     let read_metadata = |field| {
         state.memory.get(MemoryAddress::new(
             0,
diff --git a/evm/src/witness/operation.rs b/evm/src/witness/operation.rs
index fcb96d8f..6d65f16c 100644
--- a/evm/src/witness/operation.rs
+++ b/evm/src/witness/operation.rs
@@ -75,17 +75,20 @@ pub(crate) fn generate_binary_arithmetic_op<F: Field>(
 
     let log_out = stack_push_log_and_fill(state, &mut row, operation.result())?;
 
-    const LOOKUP_CHANNEL: usize = 2;
-    let lookup_addr = MemoryAddress::new(0, Segment::ShiftTable, input0.low_u32() as usize);
-    if input0.bits() <= 32 {
-        let (_, read) = mem_read_gp_with_log_and_fill(LOOKUP_CHANNEL, lookup_addr, state, &mut row);
-        state.traces.push_memory(read);
-    } else {
-        // The shift constraints still expect the address to be set, even though no read will occur.
-        let mut channel = &mut row.mem_channels[LOOKUP_CHANNEL];
-        channel.addr_context = F::from_canonical_usize(lookup_addr.context);
-        channel.addr_segment = F::from_canonical_usize(lookup_addr.segment);
-        channel.addr_virtual = F::from_canonical_usize(lookup_addr.virt);
+    if operator == arithmetic::BinaryOperator::Shl || operator == arithmetic::BinaryOperator::Shr {
+        const LOOKUP_CHANNEL: usize = 2;
+        let lookup_addr = MemoryAddress::new(0, Segment::ShiftTable, input0.low_u32() as usize);
+        if input0.bits() <= 32 {
+            let (_, read) =
+                mem_read_gp_with_log_and_fill(LOOKUP_CHANNEL, lookup_addr, state, &mut row);
+            state.traces.push_memory(read);
+        } else {
+            // The shift constraints still expect the address to be set, even though no read will occur.
+            let mut channel = &mut row.mem_channels[LOOKUP_CHANNEL];
+            channel.addr_context = F::from_canonical_usize(lookup_addr.context);
+            channel.addr_segment = F::from_canonical_usize(lookup_addr.segment);
+            channel.addr_virtual = F::from_canonical_usize(lookup_addr.virt);
+        }
     }
 
     state.traces.push_arithmetic(operation);
diff --git a/evm/src/witness/traces.rs b/evm/src/witness/traces.rs
index 60e0d8af..41b654fb 100644
--- a/evm/src/witness/traces.rs
+++ b/evm/src/witness/traces.rs
@@ -18,6 +18,8 @@ use crate::{arithmetic, keccak, logic};
 #[derive(Clone, Copy, Debug)]
 pub struct TraceCheckpoint {
     pub(self) cpu_len: usize,
+    pub(self) keccak_len: usize,
+    pub(self) keccak_sponge_len: usize,
     pub(self) logic_len: usize,
     pub(self) arithmetic_len: usize,
     pub(self) memory_len: usize,
@@ -48,19 +50,22 @@ impl<T: Copy> Traces<T> {
     pub fn checkpoint(&self) -> TraceCheckpoint {
         TraceCheckpoint {
             cpu_len: self.cpu.len(),
+            keccak_len: self.keccak_inputs.len(),
+            keccak_sponge_len: self.keccak_sponge_ops.len(),
             logic_len: self.logic_ops.len(),
             arithmetic_len: self.arithmetic.len(),
             memory_len: self.memory_ops.len(),
-            // TODO others
         }
     }
 
     pub fn rollback(&mut self, checkpoint: TraceCheckpoint) {
         self.cpu.truncate(checkpoint.cpu_len);
+        self.keccak_inputs.truncate(checkpoint.keccak_len);
+        self.keccak_sponge_ops
+            .truncate(checkpoint.keccak_sponge_len);
         self.logic_ops.truncate(checkpoint.logic_len);
         self.arithmetic.truncate(checkpoint.arithmetic_len);
         self.memory_ops.truncate(checkpoint.memory_len);
-        // TODO others
     }
 
     pub fn mem_ops_since(&self, checkpoint: TraceCheckpoint) -> &[MemoryOp] {