diff --git a/evm/src/cpu/kernel/constants/mod.rs b/evm/src/cpu/kernel/constants/mod.rs
index fa25fa87..de4dfd4b 100644
--- a/evm/src/cpu/kernel/constants/mod.rs
+++ b/evm/src/cpu/kernel/constants/mod.rs
@@ -94,7 +94,7 @@ const MISC_CONSTANTS: [(&str, [u8; 32]); 3] = [
         hex!("0000000000000000000000000000000100000000000000000000000000000000"),
     ),
     // Position in SEGMENT_RLP_RAW where the empty node encoding is stored. It is
-    // equal to usize::MAX so that all rlp pointers all much smalled than that
+    // equal to u32::MAX so that all rlp pointers are much smaller than that
     (
         "ENCODED_EMPTY_NODE_POS",
         hex!("00000000000000000000000000000000000000000000000000000000FFFFFFFF"),
diff --git a/evm/src/cpu/kernel/tests/account_code.rs b/evm/src/cpu/kernel/tests/account_code.rs
index f985ebd4..20c98bf9 100644
--- a/evm/src/cpu/kernel/tests/account_code.rs
+++ b/evm/src/cpu/kernel/tests/account_code.rs
@@ -201,36 +201,21 @@ fn test_extcodecopy() -> Result<()> {
     prepare_interpreter(&mut interpreter, address, &account)?;
 
     let context = interpreter.context();
-    interpreter.generation_state.memory.set(
-        MemoryAddress {
-            context,
-            segment: Segment::ContextMetadata as usize,
-            virt: GasLimit as usize,
-        },
-        U256::from(1000000000000u64),
-    );
+    interpreter.generation_state.memory.contexts[context].segments
+        [Segment::ContextMetadata as usize]
+        .set(GasLimit as usize, U256::from(1000000000000u64));
 
     let extcodecopy = KERNEL.global_labels["sys_extcodecopy"];
 
     // Put random data in main memory and the `KernelAccountCode` segment for realism.
     let mut rng = thread_rng();
     for i in 0..2000 {
-        interpreter.generation_state.memory.set(
-            MemoryAddress {
-                context,
-                segment: Segment::MainMemory as usize,
-                virt: i,
-            },
-            U256::from(rng.gen::<u8>()),
-        );
-        interpreter.generation_state.memory.set(
-            MemoryAddress {
-                context,
-                segment: Segment::KernelAccountCode as usize,
-                virt: i,
-            },
-            U256::from(rng.gen::<u8>()),
-        );
+        interpreter.generation_state.memory.contexts[context].segments
+            [Segment::MainMemory as usize]
+            .set(i, U256::from(rng.gen::<u8>()));
+        interpreter.generation_state.memory.contexts[context].segments
+            [Segment::KernelAccountCode as usize]
+            .set(i, U256::from(rng.gen::<u8>()));
     }
 
     // Random inputs
@@ -265,11 +250,9 @@ fn test_extcodecopy() -> Result<()> {
     assert!(interpreter.stack().is_empty());
     // Check that the code was correctly copied to memory.
     for i in 0..size {
-        let memory = interpreter.generation_state.memory.get(MemoryAddress {
-            context,
-            segment: Segment::MainMemory as usize,
-            virt: dest_offset + i,
-        });
+        let memory = interpreter.generation_state.memory.contexts[context].segments
+            [Segment::MainMemory as usize]
+            .get(dest_offset + i);
         assert_eq!(
             memory,
             code.get(offset + i).copied().unwrap_or_default().into()
@@ -294,22 +277,13 @@ fn prepare_interpreter_all_accounts(
     // Switch context and initialize memory with the data we need for the tests.
     interpreter.generation_state.registers.program_counter = 0;
     interpreter.set_code(1, code.to_vec());
-    interpreter.generation_state.memory.set(
-        MemoryAddress {
-            context: 1,
-            segment: Segment::ContextMetadata as usize,
-            virt: ContextMetadata::Address as usize,
-        },
-        U256::from_big_endian(&addr),
-    );
-    interpreter.generation_state.memory.set(
-        MemoryAddress {
-            context: 1,
-            segment: Segment::ContextMetadata as usize,
-            virt: ContextMetadata::GasLimit as usize,
-        },
-        100_000.into(),
-    );
+    interpreter.generation_state.memory.contexts[1].segments[Segment::ContextMetadata as usize]
+        .set(
+            ContextMetadata::Address as usize,
+            U256::from_big_endian(&addr),
+        );
+    interpreter.generation_state.memory.contexts[1].segments[Segment::ContextMetadata as usize]
+        .set(ContextMetadata::GasLimit as usize, 100_000.into());
     interpreter.set_context(1);
     interpreter.set_is_kernel(false);
     interpreter.generation_state.memory.set(
diff --git a/evm/src/cpu/kernel/tests/add11.rs b/evm/src/cpu/kernel/tests/add11.rs
index b85bfe66..9ba65db2 100644
--- a/evm/src/cpu/kernel/tests/add11.rs
+++ b/evm/src/cpu/kernel/tests/add11.rs
@@ -19,7 +19,6 @@ use crate::generation::TrieInputs;
 use crate::memory::segments::Segment;
 use crate::proof::TrieRoots;
 use crate::util::h2u;
-use crate::witness::memory::MemoryAddress;
 
 // Stolen from `tests/mpt/insert.rs`
 // Prepare the interpreter by loading the initial MPTs and
@@ -200,14 +199,8 @@ fn test_add11_yml() {
     let route_txn_label = KERNEL.global_labels["hash_initial_tries"];
     // Switch context and initialize memory with the data we need for the tests.
     interpreter.generation_state.registers.program_counter = route_txn_label;
-    interpreter.generation_state.memory.set(
-        MemoryAddress {
-            context: 0,
-            segment: Segment::ContextMetadata as usize,
-            virt: ContextMetadata::GasLimit as usize,
-        },
-        1_000_000.into(),
-    );
+    interpreter.generation_state.memory.contexts[0].segments[Segment::ContextMetadata as usize]
+        .set(ContextMetadata::GasLimit as usize, 1_000_000.into());
     interpreter.set_is_kernel(true);
     interpreter.run().expect("Proving add11 failed.");
 }