fix (#738)

0xPolygonZero · Oct 20, 2024 · f8c5f9c · f8c5f9c
1 parent ae6c3c4
commit f8c5f9c
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 18 deletions.
diff --git a/evm_arithmetization/src/generation/mod.rs b/evm_arithmetization/src/generation/mod.rs
@@ -607,7 +607,7 @@ pub fn generate_traces<F: RichField + Extendable<D>, const D: usize>(
         table_in_use[*Table::Poseidon] = false;
     }
 
-    let tables = timed!(
+    let (tables, final_len) = timed!(
         timing,
         "convert trace data to tables",
         state.traces.into_tables(
@@ -620,9 +620,7 @@ pub fn generate_traces<F: RichField + Extendable<D>, const D: usize>(
         )
     );
 
-    let is_last_segment =
-        segment_data.registers_after.program_counter == KERNEL.global_labels["halt"];
-    if is_last_segment && OPTIONAL_TABLE_INDICES.contains(&MemAfter) {
+    if final_len == 0 && OPTIONAL_TABLE_INDICES.contains(&MemAfter) {
         log::debug!("MemAfter table not in use");
         table_in_use[*MemAfter] = false;
     }

diff --git a/evm_arithmetization/src/witness/traces.rs b/evm_arithmetization/src/witness/traces.rs
@@ -144,7 +144,7 @@ impl<T: Copy + Field> Traces<T> {
         mut trace_lengths: TraceCheckpoint,
         config: &StarkConfig,
         timing: &mut TimingTree,
-    ) -> [Vec<PolynomialValues<T>>; NUM_TABLES]
+    ) -> ([Vec<PolynomialValues<T>>; NUM_TABLES], usize)
     where
         T: RichField + Extendable<D>,
     {
@@ -240,19 +240,22 @@ impl<T: Copy + Field> Traces<T> {
             final_values.len()
         );
 
-        [
-            arithmetic_trace,
-            byte_packing_trace,
-            cpu_trace,
-            keccak_trace,
-            keccak_sponge_trace,
-            logic_trace,
-            memory_trace,
-            mem_before_trace,
-            mem_after_trace,
-            #[cfg(feature = "cdk_erigon")]
-            poseidon_trace,
-        ]
+        (
+            [
+                arithmetic_trace,
+                byte_packing_trace,
+                cpu_trace,
+                keccak_trace,
+                keccak_sponge_trace,
+                logic_trace,
+                memory_trace,
+                mem_before_trace,
+                mem_after_trace,
+                #[cfg(feature = "cdk_erigon")]
+                poseidon_trace,
+            ],
+            final_values.len(),
+        )
     }
 }