fix(cancun): correct search loop in transient storage (#257)

evm_arithmetization/src/cpu/kernel/asm/journal/transient_storage_change.asm

@@ -9,13 +9,15 @@ global revert_transient_storage_change:
     POP
     %journal_load_3
     // We will always write 0 for deletions as it makes no difference.
-    // stack: address, slot, prev_value, retdest
+    // stack: addr, slot, prev_value, retdest
     %search_transient_storage
+    // stack: found, pos, addr, value, slot, prev_value, retdest
     // The value must have been stored
     %assert_nonzero
-    // stack: pos, addr, value, key, prev_value, retdest
+    // stack: pos, addr, value, slot, prev_value, retdest
     %add_const(2)
     DUP5
+    // stack: prev_value, pos+2, addr, value, slot, prev_value, retdest
     MSTORE_GENERAL
     %pop4
-    JUMP
+    JUMP

evm_arithmetization/src/cpu/kernel/asm/main.asm

@@ -15,6 +15,9 @@ global main:
     // Initialize accessed addresses and storage keys lists
     %init_access_lists
 
+    // Initialize transient storage length
+    %init_transient_storage_len
+
     // Initialize the RLP DATA pointer to its initial position, 
     // skipping over the preinitialized empty node.
     PUSH @INITIAL_TXN_RLP_ADDR

evm_arithmetization/src/cpu/kernel/asm/memory/transient_storage.asm

@@ -8,21 +8,29 @@
 /// If the key isn't found in the array, it is inserted at the end.
 /// TODO: Look into using a more efficient data structure.
 
+/// The initial length, 0, must be scaled by its segment for
+/// comparison with the accumulator when iterating through the list.
+%macro init_transient_storage_len
+    PUSH @SEGMENT_TRANSIENT_STORAGE
+    %mstore_global_metadata(@GLOBAL_METADATA_TRANSIENT_STORAGE_LEN)
+%endmacro
+
 %macro search_transient_storage
     %stack (addr, key) -> (addr, key, %%after)
     %jump(search_transient_storage)
 %%after:
     // stack:    (is_present, pos, addr, key, val)
 %endmacro
 
-/// Looks for an address, key pair into the transient storage. Returns 1 and the position in @SEGMENT_TRANSIENT_STORAGE
-/// if present or 0 and @GLOBAL_METADATA_TRANSIENT_STORAGE_LEN if not.
+/// Looks for an address, key pair into the transient storage.
+/// Returns 1 and the position in @SEGMENT_TRANSIENT_STORAGE if present,
+/// or 0 and @GLOBAL_METADATA_TRANSIENT_STORAGE_LEN if not.
 global search_transient_storage:
     // stack: addr, key, retdest
     %mload_global_metadata(@GLOBAL_METADATA_TRANSIENT_STORAGE_LEN)
     // stack: len, addr, key, retdest
-    PUSH @SEGMENT_TRANSIENT_STORAGE ADD
     PUSH @SEGMENT_TRANSIENT_STORAGE
+    // stack: i = 0, len, addr, key, retdest
 search_transient_storage_loop:
     // `i` and `len` are both scaled by SEGMENT_TRANSIENT_STORAGE
     %stack (i, len, addr, key, retdest) -> (i, len, i, len, addr, key, retdest)
@@ -56,6 +64,7 @@ search_transient_storage_not_found:
     JUMP
 
 search_transient_storage_found:
+    // stack: i, len, addr, key, retdest
     DUP1 %add_const(2)
     MLOAD_GENERAL
     %stack (val, i, len, addr, key, retdest) -> (retdest, 1, i, addr, val, key) // Return 1 to indicate that the address was already present.
@@ -119,6 +128,7 @@ global sys_tstore:
     MSTORE_GENERAL
     %increment DUP1
     DUP6
+    // stack: value, pos'', pos'', addr, original_value, slot, value, kexit_info
     MSTORE_GENERAL
     // stack: pos'', addr, original_value, slot, value, kexit_info
     // If pos'' > @GLOBAL_METADATA_TRANSIENT_STORAGE_LEN we need to also store the new @GLOBAL_METADATA_TRANSIENT_STORAGE_LEN
@@ -133,26 +143,25 @@ sys_tstore_charge_gas:
     %stack 
         (addr, original_value, slot, value, kexit_info) -> 
         (value, original_value, addr, slot, original_value, kexit_info)
-    EQ %jumpi(sstore_noop)
+    EQ %jumpi(tstore_noop)
 
+add_to_journal:
     // stack: addr, slot, original_value, kexit_info
-global debug_journal:
     %journal_add_transient_storage_change
 
     // stack: kexit_info
     EXIT_KERNEL
 
 new_transient_storage_len:
     // Store the new (unscaled) length.
-    // stack: addr, original_value, slot, value, kexit_info
-    PUSH @SEGMENT_TRANSIENT_STORAGE
-    PUSH 1
-    SUB // 1 - seg
-    ADD // new_len = (addr - seg) + 1
+    // stack: pos, addr, original_value, slot, value, kexit_info
+    %increment
+    // stack: pos + 1, addr, original_value, slot, value, kexit_info
     %mstore_global_metadata(@GLOBAL_METADATA_TRANSIENT_STORAGE_LEN)
+    // stack: addr, original_value, slot, value, kexit_info
     %jump(sys_tstore_charge_gas)
 
-sstore_noop:
+tstore_noop:
     // stack: current_value, slot, value, kexit_info
     %pop3
     EXIT_KERNEL

evm_arithmetization/src/cpu/kernel/tests/transient_storage.rs

@@ -2,46 +2,64 @@ use anyhow::Result;
 use ethereum_types::U256;
 use once_cell::sync::Lazy;
 use plonky2::field::goldilocks_field::GoldilocksField as F;
+use plonky2::field::types::Field;
 
 use crate::cpu::kernel::aggregator::{
     combined_kernel_from_files, KERNEL_FILES, NUMBER_KERNEL_FILES,
 };
 use crate::cpu::kernel::assembler::Kernel;
 use crate::cpu::kernel::constants::context_metadata::ContextMetadata;
+use crate::cpu::kernel::constants::global_metadata::GlobalMetadata;
 use crate::cpu::kernel::interpreter::Interpreter;
 use crate::generation::state::GenerationState;
 use crate::memory::segments::Segment;
 use crate::witness::memory::MemoryAddress;
 use crate::GenerationInputs;
 
-#[test]
-fn test_tstore() -> Result<()> {
-    let sys_tstore = crate::cpu::kernel::aggregator::KERNEL.global_labels["sys_tstore"];
-
-    let kexit_info = U256::from(0xdeadbeefu32) + (U256::from(u64::from(true)) << 32);
-
-    let initial_stack = vec![
-        1.into(), // val
-        2.into(), // slot
-        kexit_info,
-    ];
-
-    let mut interpreter: Interpreter<F> = Interpreter::new(sys_tstore, initial_stack);
+fn initialize_interpreter<F: Field>(interpreter: &mut Interpreter<F>, gas_limit: U256) {
     let gas_limit_address = MemoryAddress {
         context: 0,
         segment: Segment::ContextMetadata.unscale(),
         virt: ContextMetadata::GasLimit.unscale(),
     };
+    interpreter
+        .generation_state
+        .memory
+        .set(gas_limit_address, gas_limit);
+
     let addr_addr = MemoryAddress {
         context: 0,
         segment: Segment::ContextMetadata.unscale(),
         virt: ContextMetadata::Address.unscale(),
     };
+    interpreter.generation_state.memory.set(addr_addr, 3.into());
+
+    let initial_len = MemoryAddress {
+        context: 0,
+        segment: Segment::GlobalMetadata.unscale(),
+        virt: GlobalMetadata::TransientStorageLen.unscale(),
+    };
+    // Store 0 but scaled by its segment
     interpreter
         .generation_state
         .memory
-        .set(gas_limit_address, 100.into());
-    interpreter.generation_state.memory.set(addr_addr, 3.into());
+        .set(initial_len, (Segment::TransientStorage as usize).into());
+}
+
+#[test]
+fn test_tstore() -> Result<()> {
+    let sys_tstore = crate::cpu::kernel::aggregator::KERNEL.global_labels["sys_tstore"];
+
+    let kexit_info = U256::from(0xdeadbeefu32) + (U256::from(u64::from(true)) << 32);
+
+    let initial_stack = vec![
+        1.into(), // val
+        2.into(), // slot
+        kexit_info,
+    ];
+
+    let mut interpreter: Interpreter<F> = Interpreter::new(sys_tstore, initial_stack);
+    initialize_interpreter(&mut interpreter, 100.into());
 
     interpreter.run()?;
 
@@ -87,21 +105,7 @@ fn test_tstore_tload() -> Result<()> {
     ];
 
     let mut interpreter: Interpreter<F> = Interpreter::new(sys_tstore, initial_stack);
-    let gas_limit_address = MemoryAddress {
-        context: 0,
-        segment: Segment::ContextMetadata.unscale(),
-        virt: ContextMetadata::GasLimit.unscale(),
-    };
-    let addr_addr = MemoryAddress {
-        context: 0,
-        segment: Segment::ContextMetadata.unscale(),
-        virt: ContextMetadata::Address.unscale(),
-    };
-    interpreter
-        .generation_state
-        .memory
-        .set(gas_limit_address, 200.into());
-    interpreter.generation_state.memory.set(addr_addr, 3.into());
+    initialize_interpreter(&mut interpreter, 200.into());
 
     interpreter.run()?;
 
@@ -150,6 +154,7 @@ fn test_tstore_tload() -> Result<()> {
 #[test]
 fn test_many_tstore_many_tload() -> Result<()> {
     let kexit_info = U256::from(0xdeadbeefu32) + (U256::from(u64::from(true)) << 32);
+    let sys_tstore = crate::cpu::kernel::aggregator::KERNEL.global_labels["sys_tstore"];
 
     let initial_stack = vec![
         1.into(), // val
@@ -158,24 +163,7 @@ fn test_many_tstore_many_tload() -> Result<()> {
     ];
 
     let mut interpreter: Interpreter<F> = Interpreter::new(0, initial_stack);
-    let gas_limit_address = MemoryAddress {
-        context: 0,
-        segment: Segment::ContextMetadata.unscale(),
-        virt: ContextMetadata::GasLimit.unscale(),
-    };
-    let addr_addr = MemoryAddress {
-        context: 0,
-        segment: Segment::ContextMetadata.unscale(),
-        virt: ContextMetadata::Address.unscale(),
-    };
-
-    interpreter
-        .generation_state
-        .memory
-        .set(gas_limit_address, (10 * 200).into());
-    interpreter.generation_state.memory.set(addr_addr, 3.into());
-
-    let sys_tstore = crate::cpu::kernel::aggregator::KERNEL.global_labels["sys_tstore"];
+    initialize_interpreter(&mut interpreter, (10 * 200).into());
 
     for i in 0..10 {
         interpreter.generation_state.registers.program_counter = sys_tstore;
@@ -250,23 +238,7 @@ fn test_revert() -> Result<()> {
     let mut interpreter = Interpreter::<F>::new(sys_tstore, vec![]);
     interpreter.generation_state =
         GenerationState::<F>::new(GenerationInputs::default(), &KERNEL.code).unwrap();
-
-    let gas_limit_address = MemoryAddress {
-        context: 0,
-        segment: Segment::ContextMetadata.unscale(),
-        virt: ContextMetadata::GasLimit.unscale(),
-    };
-    let addr_addr = MemoryAddress {
-        context: 0,
-        segment: Segment::ContextMetadata.unscale(),
-        virt: ContextMetadata::Address.unscale(),
-    };
-
-    interpreter
-        .generation_state
-        .memory
-        .set(gas_limit_address, (20 * 100).into());
-    interpreter.generation_state.memory.set(addr_addr, 3.into());
+    initialize_interpreter(&mut interpreter, (20 * 100).into());
 
     // Store different values at slot 1
     for i in 0..10 {