Fix trace lengths (#110)

0xPolygonZero · Mar 18, 2024 · 39bf13b · 39bf13b
1 parent 6c3d507
commit 39bf13b
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 10 deletions.
diff --git a/evm_arithmetization/src/generation/mod.rs b/evm_arithmetization/src/generation/mod.rs
@@ -380,10 +380,7 @@ pub fn generate_traces<F: RichField + Extendable<D>, const D: usize>(
         (RegistersState::default(), None)
     };
 
-    log::info!(
-        "Trace lengths (before padding): {:?}",
-        state.traces.get_lengths()
-    );
+    let mut trace_lengths = state.traces.get_lengths();
 
     let read_metadata = |field| state.memory.read_global_metadata(field);
     let trie_roots_before = TrieRoots {
@@ -428,7 +425,7 @@ pub fn generate_traces<F: RichField + Extendable<D>, const D: usize>(
         "convert trace data to tables",
         state
             .traces
-            .into_tables(all_stark, &memory_before, config, timing)
+            .into_tables(all_stark, &memory_before, trace_lengths, config, timing)
     );
     Ok((tables, public_values, final_values))
 }

diff --git a/evm_arithmetization/src/memory/memory_stark.rs b/evm_arithmetization/src/memory/memory_stark.rs
@@ -186,11 +186,16 @@ pub(crate) fn generate_first_change_flags_and_rc<F: RichField>(
 impl<F: RichField + Extendable<D>, const D: usize> MemoryStark<F, D> {
     /// Generate most of the trace rows. Excludes a few columns like `COUNTER`,
     /// which are generated later, after transposing to column-major form.
-    fn generate_trace_row_major(&self, mut memory_ops: Vec<MemoryOp>) -> Vec<[F; NUM_COLUMNS]> {
+    fn generate_trace_row_major(
+        &self,
+        mut memory_ops: Vec<MemoryOp>,
+    ) -> (Vec<[F; NUM_COLUMNS]>, usize) {
         // fill_gaps expects an ordered list of operations.
         memory_ops.sort_by_key(MemoryOp::sorting_key);
         Self::fill_gaps(&mut memory_ops);
 
+        let unpadded_length = memory_ops.len();
+
         memory_ops.sort_by_key(MemoryOp::sorting_key);
 
         Self::pad_memory_ops(&mut memory_ops);
@@ -204,7 +209,7 @@ impl<F: RichField + Extendable<D>, const D: usize> MemoryStark<F, D> {
             .map(|op| op.into_row())
             .collect::<Vec<_>>();
         generate_first_change_flags_and_rc(trace_rows.as_mut_slice());
-        trace_rows
+        (trace_rows, unpadded_length)
     }
 
     /// Generates the `COUNTER`, `RANGE_CHECK` and `FREQUENCIES` columns, given
@@ -317,7 +322,7 @@ impl<F: RichField + Extendable<D>, const D: usize> MemoryStark<F, D> {
         mut memory_ops: Vec<MemoryOp>,
         mem_before_values: &[(MemoryAddress, U256)],
         timing: &mut TimingTree,
-    ) -> (Vec<PolynomialValues<F>>, Vec<Vec<F>>) {
+    ) -> (Vec<PolynomialValues<F>>, Vec<Vec<F>>, usize) {
         // First, push `mem_before` operations.
         for &(address, value) in mem_before_values {
             memory_ops.push(MemoryOp {
@@ -329,7 +334,7 @@ impl<F: RichField + Extendable<D>, const D: usize> MemoryStark<F, D> {
             });
         }
         // Generate most of the trace in row-major form.
-        let trace_rows = timed!(
+        let (trace_rows, unpadded_length) = timed!(
             timing,
             "generate trace rows",
             self.generate_trace_row_major(memory_ops)
@@ -372,6 +377,7 @@ impl<F: RichField + Extendable<D>, const D: usize> MemoryStark<F, D> {
                 .map(|column| PolynomialValues::new(column))
                 .collect(),
             final_values,
+            unpadded_length,
         )
     }
 }

diff --git a/evm_arithmetization/src/witness/traces.rs b/evm_arithmetization/src/witness/traces.rs
@@ -124,6 +124,7 @@ impl<T: Copy> Traces<T> {
         self,
         all_stark: &AllStark<T, D>,
         mem_before_values: &MemBeforeValues,
+        mut trace_lengths: TraceCheckpoint,
         config: &StarkConfig,
         timing: &mut TimingTree,
     ) -> ([Vec<PolynomialValues<T>>; NUM_TABLES], Vec<Vec<T>>)
@@ -176,13 +177,15 @@ impl<T: Copy> Traces<T> {
                 .logic_stark
                 .generate_trace(logic_ops, cap_elements, timing)
         );
-        let (memory_trace, final_values) = timed!(
+        let (memory_trace, final_values, unpadded_memory_length) = timed!(
             timing,
             "generate memory trace",
             all_stark
                 .memory_stark
                 .generate_trace(memory_ops, mem_before_values, timing)
         );
+        trace_lengths.memory_len = unpadded_memory_length;
+
         let mem_before_trace = timed!(
             timing,
             "generate mem_before trace",
@@ -198,6 +201,13 @@ impl<T: Copy> Traces<T> {
                 .generate_trace(final_values.clone(), timing)
         );
 
+        log::info!(
+            "Trace lengths (before padding): {:?}, mem_before_len: {}, mem_after_len: {}",
+            trace_lengths,
+            mem_before_values.len(),
+            final_values.len()
+        );
+
         (
             [
                 arithmetic_trace,