feat(continuations): write storage values directly in linked lists (#433

) * Remove unnecessary linked_lists methods and payload_ptr * Store storage value in linked list * Fix Clippy and cleanup * Apply comments * Only write to TrieData if value nonzero * Fix comments and name changed in merge * Apply comments --------- Co-authored-by: 4l0n50 <[email protected]> Co-authored-by: Alonso Gonzalez <[email protected]> Co-authored-by: Robin Salen <[email protected]>
0xPolygonZero · Aug 2, 2024 · e4b2854 · e4b2854
1 parent b0cda17
commit e4b2854
Show file tree

Hide file tree

Showing 9 changed files with 66 additions and 85 deletions.
diff --git a/evm_arithmetization/src/cpu/kernel/asm/journal/storage_change.asm b/evm_arithmetization/src/cpu/kernel/asm/journal/storage_change.asm
@@ -12,8 +12,6 @@ global revert_storage_change:
     DUP3 ISZERO %jumpi(delete)
     // stack: address, slot, prev_value, retdest
     %insert_slot_with_value
-    // stack: value_ptr
-    POP
     JUMP
 
 delete:

diff --git a/evm_arithmetization/src/cpu/kernel/asm/main.asm b/evm_arithmetization/src/cpu/kernel/asm/main.asm
@@ -160,6 +160,10 @@ global perform_final_checks:
     PROVER_INPUT(trie_ptr::state)
 
     %mstore_global_metadata(@GLOBAL_METADATA_STATE_TRIE_ROOT)
+
+    PROVER_INPUT(trie_ptr::trie_data_size)
+    %mstore_global_metadata(@GLOBAL_METADATA_TRIE_DATA_SIZE)
+
     %set_initial_tries
     %get_trie_data_size
     %mpt_hash_state_trie

diff --git a/evm_arithmetization/src/cpu/kernel/asm/mpt/linked_list/final_tries.asm b/evm_arithmetization/src/cpu/kernel/asm/mpt/linked_list/final_tries.asm
@@ -69,11 +69,16 @@ insert_next_slot:
     DUP2
     %increment
     MLOAD_GENERAL
-    // key, addr, storage_ptr_ptr, root_ptr, retdest
+    // stack: key, addr, storage_ptr_ptr, root_ptr, retdest
     DUP3
     %add_const(2)
     MLOAD_GENERAL
-    // stack: payload_ptr, key, addr, storage_ptr_ptr, root_ptr, retdest
+    // stack: value, key, addr, storage_ptr_ptr, root_ptr, retdest
+    // If the value is 0, then payload_ptr = 0, and we don't need to insert a value in the `TrieData` segment.
+    DUP1 ISZERO %jumpi(insert_with_payload_ptr)
+    %get_trie_data_size // payload_ptr
+    SWAP1 %append_to_trie_data // append the value to the trie data segment
+insert_with_payload_ptr:
     %stack (payload_ptr, key, addr, storage_ptr_ptr, root_ptr) -> (root_ptr, 64, key, payload_ptr, after_insert_slot, storage_ptr_ptr, addr)
     %jump(mpt_insert)
 after_insert_slot:
@@ -112,7 +117,7 @@ global delete_removed_accounts:
     // stack: key, account_ptr_ptr, root_ptr, storage_ptr_ptr, retdest
     DUP2
     %add_const(2)
-    MLOAD_GENERAL // get intitial payload_ptr
+    MLOAD_GENERAL // get initial payload_ptr
     %add_const(2) // storage_root_ptr_ptr = payload_ptr + 2
     %mload_trie_data
     // stack: storage_root_ptr, key, account_ptr_ptr, root_ptr, storage_ptr_ptr, retdest

diff --git a/evm_arithmetization/src/cpu/kernel/asm/mpt/linked_list/initial_tries.asm b/evm_arithmetization/src/cpu/kernel/asm/mpt/linked_list/initial_tries.asm
@@ -34,7 +34,7 @@ skip:
 %macro set_initial_tries
     PUSH %%after
     PUSH @SEGMENT_STORAGE_LINKED_LIST
-    %add_const(8) // The first node is the special node, of size 5, so the first payload is at position 5 + 3.
+    %add_const(8) // The first node is the special node, of size 5, so the first value is at position 5 + 3.
     PUSH @SEGMENT_ACCOUNTS_LINKED_LIST
     %add_const(6) // The first node is the special node, of size 4, so the first payload is at position 4 + 2.
     %mload_global_metadata(@GLOBAL_METADATA_STATE_TRIE_ROOT)
@@ -175,10 +175,19 @@ after_set_storage_payload:
 set_payload_storage_leaf:
     // stack: node_type, after_node_type, storage_ptr_ptr, retdest
     POP
-    // stack:  after_node_type, storage_ptr_ptr, retdest
+    // stack: after_node_type, storage_ptr_ptr, retdest
     %add_const(2) // The value pointer starts at index 3, after num_nibbles and packed_nibbles.
-    DUP2
-    MLOAD_GENERAL
+    // stack: value_ptr_ptr, storage_ptr_ptr, retdest
+    DUP2 MLOAD_GENERAL
+    // stack: value, value_ptr_ptr, storage_ptr_ptr, retdest
+    // If value == 0, then value_ptr = 0, and we don't need to append the value to the `TrieData` segment.
+    DUP1 ISZERO %jumpi(set_payload_storage_leaf_end)
+    %get_trie_data_size
+    // stack: value_ptr, value, value_ptr_ptr, storage_ptr_ptr, retdest
+    SWAP1
+    %append_to_trie_data
+set_payload_storage_leaf_end:
+    // stack: value_ptr, value_ptr_ptr, storage_ptr_ptr, retdest
     SWAP1
     %mstore_trie_data
     // stack: storage_ptr_ptr, retdest

diff --git a/evm_arithmetization/src/cpu/kernel/asm/mpt/linked_list/linked_list.asm b/evm_arithmetization/src/cpu/kernel/asm/mpt/linked_list/linked_list.asm
@@ -335,23 +335,20 @@ global store_initial_slots:
 
 loop_store_initial_slots:
     // stack: current_node_ptr, cur_len, retdest
-    %get_trie_data_size
-    DUP2
+    DUP1
     MLOAD_GENERAL
-    // stack: current_addr_key, cpy_ptr, current_node_ptr, cur_len, retdest
+    // stack: current_addr_key, current_node_ptr, cur_len, retdest
     %eq_const(@U256_MAX)
     %jumpi(store_initial_slots_end)
-    DUP2
+    DUP1
     %add_const(2)
     MLOAD_GENERAL
-    // stack: payload_ptr, cpy_ptr, current_node_ptr, cur_len, retdest
-    %mload_trie_data
-    %append_to_trie_data
-    // stack: cpy_ptr, current_node_ptr, cur_len, retdest
+    // stack: value, current_node_ptr, cur_len, retdest
     DUP2
     %add_const(@STORAGE_COPY_PAYLOAD_PTR)
+    // stack: cpy_value_ptr, value, current_node_ptr, cur_len, retdest
     SWAP1
-    MSTORE_GENERAL // Store cpy_ptr
+    MSTORE_GENERAL // Store cpy_value
     // stack: current_node_ptr, cur_len, retdest
     SWAP1 PUSH @STORAGE_LINKED_LISTS_NODE_SIZE
     ADD
@@ -361,7 +358,7 @@ loop_store_initial_slots:
     %jump(loop_store_initial_slots)
 
 store_initial_slots_end:
-    %pop2
+    POP
     // stack: cur_len, retdest
     %mstore_global_metadata(@GLOBAL_METADATA_STORAGE_LINKED_LIST_NEXT_AVAILABLE)
     JUMP
@@ -376,7 +373,6 @@ store_initial_slots_end:
 
 %macro insert_slot_no_return
     %insert_slot
-    POP
 %endmacro
 
 // Multiplies the value at the top of the stack, denoted by ptr/5, by 5
@@ -402,7 +398,6 @@ store_initial_slots_end:
 
 /// Inserts the pair (address_key, storage_key) and a new payload pointer into the linked list if it is not already present,
 /// or modifies its payload if it was already present.
-/// Returns `new_payload_ptr` if the storage key was inserted, `original_ptr` if it was already present.
 global insert_slot_with_value:
     // stack: addr_key, key, value, retdest
     PROVER_INPUT(linked_list::insert_slot)
@@ -505,43 +500,33 @@ next_node_ok_with_value:
     DUP5
     MSTORE_GENERAL
     // stack: new_ptr + 1, next_ptr, addr_key, key, value, retdest
-    // Append the value to `TrieDataSegment` and write the resulting payload_ptr.
+    // Write the value in the linked list.
     %increment
-    DUP1
-    %get_trie_data_size
-    // stack: new_payload_ptr, new_ptr+2, new_ptr+2, next_ptr, addr_key, key, value, retdest
-    %stack (new_payload_ptr, new_payload_ptr_ptr, new_payload_ptr_ptr, next_ptr, addr_key, key, value, retdest)
-        -> (value, new_payload_ptr, new_payload_ptr_ptr, new_payload_ptr_ptr, next_ptr, new_payload_ptr, retdest)
-    %append_to_trie_data
-    MSTORE_GENERAL
-
-    // stack: new_ptr + 2, next_ptr, new_payload_ptr, retdest
-    // Store the payload ptr copy
-    %increment
-    DUP1
-    DUP4
-    %clone_slot
-    MSTORE_GENERAL
-    // stack: new_ptr + 3, next_ptr, new_payload_ptr, retdest
+    DUP1 %increment
+    // stack: new_ptr+3, new_value_ptr, next_ptr, addr_key, key, value, retdest
+    %stack (new_cloned_value_ptr, new_value_ptr, next_ptr, addr_key, key, value, retdest)
+        -> (value, new_cloned_value_ptr, value, new_value_ptr, new_cloned_value_ptr, next_ptr, retdest)
+    MSTORE_GENERAL // Store copied value.
+    MSTORE_GENERAL // Store value.
+
+    // stack: new_ptr + 3, next_ptr, retdest
     %increment
     DUP1
-    // stack: new_next_ptr, new_next_ptr, next_ptr, new_payload_ptr, retdest
+    // stack: new_next_ptr_ptr, new_next_ptr_ptr, next_ptr, retdest
     SWAP2
     MSTORE_GENERAL
-    // stack: new_next_ptr, new_payload_ptr, retdest
+    // stack: new_next_ptr_ptr, retdest
     %increment
     %mstore_global_metadata(@GLOBAL_METADATA_STORAGE_LINKED_LIST_NEXT_AVAILABLE)
-    // stack: new_payload_ptr, retdest
-    SWAP1
+    // stack: retdest
     JUMP
 
 slot_found_write_value:
     // stack: pred_ptr, addr_key, key, value, retdest
-    %add_const(2) MLOAD_GENERAL
-    %stack (payload_ptr, addr_key, key, value) -> (payload_ptr, value, payload_ptr)
-    %mstore_trie_data
-    // stack: payload_ptr, retdest
-    %stack (payload_ptr, retdest) -> (retdest, payload_ptr)
+    %add_const(2)
+    %stack (payload_ptr, addr_key, key, value) -> (value, payload_ptr)
+    MSTORE_GENERAL
+    // stack: retdest
     JUMP
 
 %macro insert_slot_with_value
@@ -552,7 +537,6 @@ slot_found_write_value:
     %stack (slot_key, addr_key, value) -> (addr_key, slot_key, value, %%after)
     %jump(insert_slot_with_value)
 %%after:
-    // stack: value_ptr
 %endmacro
 
 /// Inserts the pair (address_key, storage_key) and payload pointer into the linked list if it is not already present,
@@ -758,8 +742,8 @@ slot_found_no_write:
     // Load the the payload pointer and access counter
     %add_const(2)
     MLOAD_GENERAL
-    // stack: orig_payload_ptr, addr_key, key, payload_ptr, retdest
-    %stack (orig_payload_ptr, addr_key, key, payload_ptr, retdest) -> (retdest, orig_payload_ptr)
+    // stack: orig_value, addr_key, key, payload_ptr, retdest
+    %stack (orig_value, addr_key, key, payload_ptr, retdest) -> (retdest, orig_value)
     JUMP
 
 

diff --git a/evm_arithmetization/src/cpu/kernel/asm/mpt/storage/storage_write.asm b/evm_arithmetization/src/cpu/kernel/asm/mpt/storage/storage_write.asm
@@ -111,18 +111,10 @@ sstore_after_refund:
     // stack: slot, value, kexit_info
     DUP2 ISZERO %jumpi(sstore_delete)
 
-    // First we write the value to MPT data, and get a pointer to it.
-    %get_trie_data_size
-    // stack: value_ptr, slot, value, kexit_info
-    SWAP2
-    // stack: value, slot, value_ptr, kexit_info
-    %append_to_trie_data
-    // stack: slot, value_ptr, kexit_info
 
-    %slot_to_storage_key
+    // stack: slot, value, kexit_info
     %address
-    %addr_to_state_key
-    %insert_slot_no_return
+    %insert_slot_with_value
 
     EXIT_KERNEL
 

diff --git a/evm_arithmetization/src/cpu/kernel/asm/transactions/common_decoding.asm b/evm_arithmetization/src/cpu/kernel/asm/transactions/common_decoding.asm
@@ -252,12 +252,11 @@ sload_with_addr:
 
     // Storage key not found. Return default value_ptr = 0,
     // which derefs to 0 since @SEGMENT_TRIE_DATA[0] = 0.
-    %stack (value_ptr, retdest) -> (retdest, 0)
+    %stack (value, retdest) -> (retdest, 0)
+    
     JUMP
 
 global storage_key_exists:
-    // stack: value_ptr, retdest
-    %mload_trie_data
     // stack: value, retdest
     SWAP1
     JUMP
diff --git a/evm_arithmetization/src/cpu/kernel/tests/mpt/delete.rs b/evm_arithmetization/src/cpu/kernel/tests/mpt/delete.rs
@@ -115,6 +115,7 @@ fn test_state_trie(
     interpreter.generation_state.registers.program_counter = KERNEL.global_labels["store_initial"];
     interpreter.run();
 
+    assert_eq!(interpreter.stack(), vec![]);
     // Set initial tries.
     interpreter
         .push(0xDEADBEEFu32.into())

diff --git a/evm_arithmetization/src/generation/mpt.rs b/evm_arithmetization/src/generation/mpt.rs
@@ -410,7 +410,6 @@ associated storage trie hash"
                 empty_nibbles(),
                 storage_trie,
                 storage_leaves,
-                trie_data,
                 &parse_storage_value,
             )?;
 
@@ -425,7 +424,6 @@ pub(crate) fn get_storage_leaves<F>(
     key: Nibbles,
     trie: &HashedPartialTrie,
     storage_leaves: &mut Vec<Option<U256>>,
-    trie_data: &mut Vec<Option<U256>>,
     parse_value: &F,
 ) -> Result<(), ProgramError>
 where
@@ -439,29 +437,15 @@ where
                     count: 1,
                     packed: i.into(),
                 });
-                get_storage_leaves(
-                    address,
-                    extended_key,
-                    child,
-                    storage_leaves,
-                    trie_data,
-                    parse_value,
-                )?;
+                get_storage_leaves(address, extended_key, child, storage_leaves, parse_value)?;
             }
 
             Ok(())
         }
 
         Node::Extension { nibbles, child } => {
             let extended_key = key.merge_nibbles(nibbles);
-            get_storage_leaves(
-                address,
-                extended_key,
-                child,
-                storage_leaves,
-                trie_data,
-                parse_value,
-            )?;
+            get_storage_leaves(address, extended_key, child, storage_leaves, parse_value)?;
 
             Ok(())
         }
@@ -481,15 +465,20 @@ where
                     .map_err(|_| ProgramError::IntegerTooLarge)?,
             ));
             // Write `value_ptr_ptr`.
-            storage_leaves.push(Some((trie_data.len()).into()));
+            let leaves = parse_value(value)?
+                .into_iter()
+                .map(Some)
+                .collect::<Vec<_>>();
+            let leaf = match leaves.len() {
+                1 => leaves[0],
+                _ => panic!("Slot can only store exactly one value."),
+            };
+            storage_leaves.push(leaf);
             // Write the counter.
             storage_leaves.push(Some(0.into()));
             // Set the next node as the initial node.
             storage_leaves.push(Some((Segment::StorageLinkedList as usize).into()));
 
-            let leaf = parse_value(value)?.into_iter().map(Some);
-            trie_data.extend(leaf);
-
             Ok(())
         }
         _ => Ok(()),