perf(continuations): Only copy non-stale contexts

evm_arithmetization/src/cpu/kernel/interpreter.rs

@@ -36,7 +36,7 @@ use crate::prover::GenerationSegmentData;
 use crate::util::h2u;
 use crate::witness::errors::ProgramError;
 use crate::witness::memory::{
-    MemoryAddress, MemoryOp, MemoryOpKind, MemorySegmentState, MemoryState,
+    MemoryAddress, MemoryContextState, MemoryOp, MemoryOpKind, MemorySegmentState, MemoryState,
 };
 use crate::witness::operation::Operation;
 use crate::witness::state::RegistersState;
@@ -632,8 +632,26 @@ impl<F: Field> State<F> for Interpreter<F> {
         self.halt_offsets.clone()
     }
 
-    fn get_full_memory(&self) -> Option<MemoryState> {
-        Some(self.generation_state.memory.clone())
+    fn get_active_memory(&self) -> Option<MemoryState> {
+        let mut memory_state = MemoryState {
+            contexts: vec![
+                MemoryContextState::default();
+                self.generation_state.memory.contexts.len()
+            ],
+            ..self.generation_state.memory.clone()
+        };
+
+        // Only copy memory from non-stale contexts
+        for (ctx_idx, ctx) in self.generation_state.memory.contexts.iter().enumerate() {
+            if !self
+                .get_generation_state()
+                .stale_contexts
+                .contains(&ctx_idx)
+            {
+                memory_state.contexts[ctx_idx] = ctx.clone();
+            }
+        }
+        Some(memory_state)
     }
 
     fn update_interpreter_final_registers(&mut self, final_registers: RegistersState) {

evm_arithmetization/src/generation/state.rs

@@ -159,12 +159,15 @@ pub(crate) trait State<F: Field> {
     /// Applies a `State`'s operations since a checkpoint.
     fn apply_ops(&mut self, checkpoint: GenerationStateCheckpoint);
 
-    /// Return the offsets at which execution must halt
+    /// Returns the offsets at which execution must halt
     fn get_halt_offsets(&self) -> Vec<usize>;
 
     fn update_interpreter_final_registers(&mut self, _final_registers: RegistersState) {}
 
-    fn get_full_memory(&self) -> Option<MemoryState> {
+    /// Returns all the memory from non-stale contexts.
+    /// This is only necessary during segment data generation, hence the blanket
+    /// impl returns a dummy value.
+    fn get_active_memory(&self) -> Option<MemoryState> {
         None
     }
 
@@ -181,7 +184,7 @@ pub(crate) trait State<F: Field> {
         let halt_offsets = self.get_halt_offsets();
 
         let mut final_registers = RegistersState::default();
-        let final_mem = self.get_generation_state().memory.clone();
+        let final_mem = self.get_active_memory();
         let mut running = true;
         let mut final_clock = 0;
         loop {
@@ -208,21 +211,13 @@ pub(crate) trait State<F: Field> {
                 if let Some(halt_context) = opt_halt_context {
                     if self.get_context() == halt_context {
                         // Only happens during jumpdest analysis, we don't care about the output.
-                        return Ok((final_registers, Some(final_mem)));
+                        return Ok((final_registers, final_mem));
                     }
                 } else {
                     if !running {
                         debug_assert!(self.get_clock() - final_clock == NUM_EXTRA_CYCLES_AFTER - 1);
                     }
-                    let final_mem = if let Some(mut mem) = self.get_full_memory() {
-                        // Clear memory we will not use again.
-                        for &ctx in &self.get_generation_state().stale_contexts {
-                            mem.contexts[ctx] = MemoryContextState::default();
-                        }
-                        Some(mem)
-                    } else {
-                        None
-                    };
+                    let final_mem = self.get_active_memory();
                     #[cfg(not(test))]
                     self.log_info(format!("CPU halted after {} cycles", self.get_clock()));
                     return Ok((final_registers, final_mem));