From c7a75b2f94191ca0d868b942c96766b8202a5d00 Mon Sep 17 00:00:00 2001
From: Robin Salen <salenrobin@gmail.com>
Date: Wed, 5 Jun 2024 15:53:09 -0400
Subject: [PATCH 1/2] Reduce overhead on memory copy

---
 .../src/cpu/kernel/interpreter.rs             | 24 ++++++++++++++++---
 evm_arithmetization/src/generation/state.rs   | 21 +++++++---------
 2 files changed, 29 insertions(+), 16 deletions(-)

diff --git a/evm_arithmetization/src/cpu/kernel/interpreter.rs b/evm_arithmetization/src/cpu/kernel/interpreter.rs
index 5501e10d6..b36c5390a 100644
--- a/evm_arithmetization/src/cpu/kernel/interpreter.rs
+++ b/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) {
diff --git a/evm_arithmetization/src/generation/state.rs b/evm_arithmetization/src/generation/state.rs
index d3d585fa6..9879a0cee 100644
--- a/evm_arithmetization/src/generation/state.rs
+++ b/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));

From ca1f51a579c233d18bde37cb895c5d20ee96766d Mon Sep 17 00:00:00 2001
From: Robin Salen <30937548+Nashtare@users.noreply.github.com>
Date: Thu, 6 Jun 2024 06:15:04 +0900
Subject: [PATCH 2/2] Apply suggestion

Co-authored-by: Hamy Ratoanina <hamy.ratoanina@toposware.com>
---
 evm_arithmetization/src/generation/state.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/evm_arithmetization/src/generation/state.rs b/evm_arithmetization/src/generation/state.rs
index 9879a0cee..cabbb60f6 100644
--- a/evm_arithmetization/src/generation/state.rs
+++ b/evm_arithmetization/src/generation/state.rs
@@ -211,7 +211,7 @@ 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, final_mem));
+                        return Ok((final_registers, None));
                     }
                 } else {
                     if !running {