From 8dcb29e5ad6e15380a7155f030548815319020a6 Mon Sep 17 00:00:00 2001
From: Robin Salen <salenrobin@gmail.com>
Date: Wed, 6 Sep 2023 12:31:17 -0400
Subject: [PATCH 1/2] Display actual trace lengths instead of number of ops

---
 evm/src/generation/mod.rs |  2 +-
 evm/src/witness/traces.rs | 44 ++++++++++++++++++++++++++++++++++++++-
 2 files changed, 44 insertions(+), 2 deletions(-)

diff --git a/evm/src/generation/mod.rs b/evm/src/generation/mod.rs
index 5beded4e..be86b4b6 100644
--- a/evm/src/generation/mod.rs
+++ b/evm/src/generation/mod.rs
@@ -159,7 +159,7 @@ pub fn generate_traces<F: RichField + Extendable<D>, const D: usize>(
 
     log::info!(
         "Trace lengths (before padding): {:?}",
-        state.traces.checkpoint()
+        state.traces.get_lengths()
     );
 
     let outputs = get_outputs(&mut state);
diff --git a/evm/src/witness/traces.rs b/evm/src/witness/traces.rs
index 4a1c8d85..c380ec64 100644
--- a/evm/src/witness/traces.rs
+++ b/evm/src/witness/traces.rs
@@ -8,13 +8,14 @@ use plonky2::timed;
 use plonky2::util::timing::TimingTree;
 
 use crate::all_stark::{AllStark, NUM_TABLES};
+use crate::arithmetic::{BinaryOperator, Operation};
 use crate::config::StarkConfig;
 use crate::cpu::columns::CpuColumnsView;
 use crate::keccak_sponge::columns::KECCAK_WIDTH_BYTES;
 use crate::keccak_sponge::keccak_sponge_stark::KeccakSpongeOp;
 use crate::util::trace_rows_to_poly_values;
 use crate::witness::memory::MemoryOp;
-use crate::{arithmetic, keccak, logic};
+use crate::{arithmetic, keccak, keccak_sponge, logic};
 
 #[derive(Clone, Copy, Debug)]
 pub struct TraceCheckpoint {
@@ -48,6 +49,47 @@ impl<T: Copy> Traces<T> {
         }
     }
 
+    /// Returns the actual trace lengths for each STARK module.
+    //  Uses a `TraceCheckPoint` as return object for convenience.
+    pub fn get_lengths(&self) -> TraceCheckpoint {
+        TraceCheckpoint {
+            arithmetic_len: self
+                .arithmetic_ops
+                .iter()
+                .map(|op| match op {
+                    Operation::TernaryOperation {
+                        operator: _,
+                        input0: _,
+                        input1: _,
+                        input2: _,
+                        result: _,
+                    } => 2,
+                    Operation::BinaryOperation {
+                        operator,
+                        input0: _,
+                        input1: _,
+                        result: _,
+                    } => match operator {
+                        BinaryOperator::Div | BinaryOperator::Mod => 2,
+                        _ => 1,
+                    },
+                })
+                .sum(),
+            cpu_len: self.cpu.len(),
+            keccak_len: self.keccak_inputs.len() * keccak::keccak_stark::NUM_ROUNDS,
+            keccak_sponge_len: self
+                .keccak_sponge_ops
+                .iter()
+                .map(|op| op.input.len() / keccak_sponge::columns::KECCAK_RATE_BYTES + 1)
+                .sum(),
+            logic_len: self.logic_ops.len(),
+            // This is technically a lower-bound, as we may fill gaps,
+            // but this gives a relatively good estimate.
+            memory_len: self.memory_ops.len(),
+        }
+    }
+
+    /// Returns the number of operations for each STARK module.
     pub fn checkpoint(&self) -> TraceCheckpoint {
         TraceCheckpoint {
             arithmetic_len: self.arithmetic_ops.len(),

From d0379e94287c9d85da167c3359152260f18edcaf Mon Sep 17 00:00:00 2001
From: Robin Salen <salenrobin@gmail.com>
Date: Thu, 7 Sep 2023 15:35:34 -0400
Subject: [PATCH 2/2] Apply Nick's comment

---
 evm/src/witness/traces.rs | 15 ++-------------
 1 file changed, 2 insertions(+), 13 deletions(-)

diff --git a/evm/src/witness/traces.rs b/evm/src/witness/traces.rs
index c380ec64..2cc1c500 100644
--- a/evm/src/witness/traces.rs
+++ b/evm/src/witness/traces.rs
@@ -57,19 +57,8 @@ impl<T: Copy> Traces<T> {
                 .arithmetic_ops
                 .iter()
                 .map(|op| match op {
-                    Operation::TernaryOperation {
-                        operator: _,
-                        input0: _,
-                        input1: _,
-                        input2: _,
-                        result: _,
-                    } => 2,
-                    Operation::BinaryOperation {
-                        operator,
-                        input0: _,
-                        input1: _,
-                        result: _,
-                    } => match operator {
+                    Operation::TernaryOperation { .. } => 2,
+                    Operation::BinaryOperation { operator, .. } => match operator {
                         BinaryOperator::Div | BinaryOperator::Mod => 2,
                         _ => 1,
                     },