Linked lists for the state trie (#402)

0xPolygonZero · Aug 1, 2024 · 5edb969 · 5edb969
1 parent 4fa4b77
commit 5edb969
Show file tree

Hide file tree

Showing 44 changed files with 3,411 additions and 568 deletions.
diff --git a/evm_arithmetization/src/cpu/kernel/aggregator.rs b/evm_arithmetization/src/cpu/kernel/aggregator.rs
@@ -125,6 +125,9 @@ pub(crate) fn combined_kernel() -> Kernel {
         include_str!("asm/mpt/insert/insert_extension.asm"),
         include_str!("asm/mpt/insert/insert_leaf.asm"),
         include_str!("asm/mpt/insert/insert_trie_specific.asm"),
+        include_str!("asm/mpt/linked_list/linked_list.asm"),
+        include_str!("asm/mpt/linked_list/initial_tries.asm"),
+        include_str!("asm/mpt/linked_list/final_tries.asm"),
         include_str!("asm/mpt/read.asm"),
         include_str!("asm/mpt/storage/storage_read.asm"),
         include_str!("asm/mpt/storage/storage_write.asm"),

diff --git a/evm_arithmetization/src/cpu/kernel/asm/core/access_lists.asm b/evm_arithmetization/src/cpu/kernel/asm/core/access_lists.asm
@@ -21,8 +21,7 @@ global init_access_lists:
     // Store @U256_MAX at the beginning of the segment
     PUSH @SEGMENT_ACCESSED_ADDRESSES // ctx == virt == 0
     DUP1
-    PUSH @U256_MAX
-    MSTORE_GENERAL
+    %mstore_u256_max
     // Store @SEGMENT_ACCESSED_ADDRESSES at address 1
     %increment
     DUP1
@@ -38,8 +37,7 @@ global init_access_lists:
     // Store @U256_MAX at the beginning of the segment
     PUSH @SEGMENT_ACCESSED_STORAGE_KEYS // ctx == virt == 0
     DUP1
-    PUSH @U256_MAX
-    MSTORE_GENERAL
+    %mstore_u256_max
     // Store @SEGMENT_ACCESSED_STORAGE_KEYS at address 3
     %add_const(3)
     DUP1
@@ -201,8 +199,7 @@ global remove_accessed_addresses:
     MLOAD_GENERAL
     // stack: next_next_ptr, next_next_ptr_ptr, next_ptr_ptr, addr, retdest
     SWAP1
-    PUSH @U256_MAX
-    MSTORE_GENERAL
+    %mstore_u256_max
     // stack: next_next_ptr, next_ptr_ptr, addr, retdest
     MSTORE_GENERAL
     POP
@@ -385,8 +382,7 @@ global remove_accessed_storage_keys:
     MLOAD_GENERAL
     // stack: next_next_ptr, next_next_ptr_ptr, next_ptr_ptr, addr, key, retdest
     SWAP1
-    PUSH @U256_MAX
-    MSTORE_GENERAL
+    %mstore_u256_max
     // stack: next_next_ptr, next_ptr_ptr, addr, key, retdest
     MSTORE_GENERAL
     %pop2

diff --git a/evm_arithmetization/src/cpu/kernel/asm/core/util.asm b/evm_arithmetization/src/cpu/kernel/asm/core/util.asm
@@ -93,3 +93,9 @@
     SET_CONTEXT
     // stack: (empty)
 %endmacro
+
+%macro mstore_u256_max
+    // stack: addr
+    PUSH @U256_MAX
+    MSTORE_GENERAL
+%endmacro
diff --git a/evm_arithmetization/src/cpu/kernel/asm/journal/account_destroyed.asm b/evm_arithmetization/src/cpu/kernel/asm/journal/account_destroyed.asm
@@ -16,7 +16,10 @@ revert_account_destroyed_contd:
     SWAP1
     // Remove `prev_balance` from `target`'s balance.
     // stack: target, address, prev_balance, retdest
-    %mpt_read_state_trie
+    %read_accounts_linked_list
+    // stack: target_payload_ptr, address, prev_balance, retdest
+    DUP1
+    %assert_nonzero
     %add_const(1)
     // stack: target_balance_ptr, address, prev_balance, retdest
     DUP3
@@ -25,8 +28,12 @@ revert_account_destroyed_contd:
     SUB SWAP1 %mstore_trie_data
     // Set `address`'s balance to `prev_balance`.
     // stack: address, prev_balance, retdest
-    %mpt_read_state_trie
-    %add_const(1)
+    %read_accounts_linked_list
+    // stack: account_payload_ptr, prev_balance, retdest
+    DUP1 
+    %assert_nonzero
+    %increment
+    // stack: account_balance_payload_ptr, prev_balance, retdest
     %mstore_trie_data
     JUMP
 
diff --git a/evm_arithmetization/src/cpu/kernel/asm/journal/code_change.asm b/evm_arithmetization/src/cpu/kernel/asm/journal/code_change.asm
@@ -9,7 +9,9 @@ global revert_code_change:
     POP
     %journal_load_2
     // stack: address, prev_codehash, retdest
-    %mpt_read_state_trie
+    %read_accounts_linked_list
+    // stack: account_ptr, prev_codehash, retdest
+    DUP1 %assert_nonzero
     // stack: account_ptr, prev_codehash, retdest
     %add_const(3)
     // stack: codehash_ptr, prev_codehash, retdest

diff --git a/evm_arithmetization/src/cpu/kernel/asm/journal/nonce_change.asm b/evm_arithmetization/src/cpu/kernel/asm/journal/nonce_change.asm
@@ -9,8 +9,10 @@ global revert_nonce_change:
     POP
     %journal_load_2
     // stack: address, prev_nonce, retdest
-    %mpt_read_state_trie
-    // stack: nonce_ptr, prev_nonce retdest
+    %read_accounts_linked_list
+    // stack: payload_ptr, prev_nonce, retdest
+    DUP1 %assert_nonzero
+    // stack: nonce_ptr, prev_nonce, retdest
     %mstore_trie_data
     // stack: retdest
     JUMP

diff --git a/evm_arithmetization/src/cpu/kernel/asm/journal/storage_change.asm b/evm_arithmetization/src/cpu/kernel/asm/journal/storage_change.asm
@@ -11,47 +11,16 @@ global revert_storage_change:
     // stack: address, slot, prev_value, retdest
     DUP3 ISZERO %jumpi(delete)
     // stack: address, slot, prev_value, retdest
-    SWAP1 %slot_to_storage_key
-    // stack: storage_key, address, prev_value, retdest
-    PUSH 64 // storage_key has 64 nibbles
-    // stack: 64, storage_key, address, prev_value, retdest
-    DUP3 %mpt_read_state_trie
-    DUP1 ISZERO %jumpi(panic)
-    // stack: account_ptr, 64, storage_key, address, prev_value, retdest
-    %add_const(2)
-    // stack: storage_root_ptr_ptr, 64, storage_key, address, prev_value, retdest
-    %mload_trie_data
-    %get_trie_data_size
-    DUP6 %append_to_trie_data
-    %stack (prev_value_ptr, storage_root_ptr, num_nibbles, storage_key, address, prev_value, retdest) ->
-        (storage_root_ptr, num_nibbles, storage_key, prev_value_ptr, new_storage_root, address, retdest)
-    %jump(mpt_insert)
+    %insert_slot_with_value
+    // stack: value_ptr
+    POP
+    JUMP
 
 delete:
     // stack: address, slot, prev_value, retdest
     SWAP2 POP
-    %stack (slot, address, retdest) -> (slot, new_storage_root, address, retdest)
+    // stack: slot, address, retdest
     %slot_to_storage_key
-    // stack: storage_key, new_storage_root, address, retdest
-    PUSH 64 // storage_key has 64 nibbles
-    // stack: 64, storage_key, new_storage_root, address, retdest
-    DUP4 %mpt_read_state_trie
-    DUP1 ISZERO %jumpi(panic)
-    // stack: account_ptr, 64, storage_key, new_storage_root, address, retdest
-    %add_const(2)
-    // stack: storage_root_ptr_ptr, 64, storage_key, new_storage_root, address, retdest
-    %mload_trie_data
-    // stack: storage_root_ptr, 64, storage_key, new_storage_root, address, retdest
-    %jump(mpt_delete)
-
-new_storage_root:
-    // stack: new_storage_root_ptr, address, retdest
-    DUP2 %mpt_read_state_trie
-    // stack: account_ptr, new_storage_root_ptr, address, retdest
-
-    // Update account with our new storage root pointer.
-    %add_const(2)
-    // stack: account_storage_root_ptr_ptr, new_storage_root_ptr, address, retdest
-    %mstore_trie_data
-    // stack: address, retdest
-    POP JUMP
+    SWAP1 %addr_to_state_key
+    // stack: addr_key, slot_key, retdest
+    %jump(remove_slot)
diff --git a/evm_arithmetization/src/cpu/kernel/asm/main.asm b/evm_arithmetization/src/cpu/kernel/asm/main.asm
@@ -67,14 +67,23 @@ global main:
     // Initialize the RLP DATA pointer to its initial position, 
     // skipping over the preinitialized empty node.
     PUSH @INITIAL_TXN_RLP_ADDR
+    %add_const(@MAX_RLP_BLOB_SIZE)
     %mstore_global_metadata(@GLOBAL_METADATA_RLP_DATA_SIZE)
 
     // Encode constant nodes
     %initialize_rlp_segment
+
+    // Initialize trie data size.
+    PROVER_INPUT(trie_ptr::trie_data_size)
+    %mstore_global_metadata(@GLOBAL_METADATA_TRIE_DATA_SIZE)
+
+global store_initial:
+    // Store the initial accounts and slots for hashing later
+    %store_initial_accounts
+    %store_initial_slots
    
-    // Initialize the state, transaction and receipt trie root pointers.
-    PROVER_INPUT(trie_ptr::state)
-    %mstore_global_metadata(@GLOBAL_METADATA_STATE_TRIE_ROOT)
+global after_store_initial:
+    // Initialize the transaction and receipt trie root pointers.
     PROVER_INPUT(trie_ptr::txn)
     %mstore_global_metadata(@GLOBAL_METADATA_TXN_TRIE_ROOT)
     PROVER_INPUT(trie_ptr::receipt)
@@ -83,16 +92,16 @@ global main:
 global hash_initial_tries:
     // We compute the length of the trie data segment in `mpt_hash` so that we
     // can check the value provided by the prover.
-    // We initialize the segment length with 1 because the segment contains 
-    // the null pointer `0` when the tries are empty.
-    PUSH 1
-    %mpt_hash_state_trie  %mload_global_metadata(@GLOBAL_METADATA_STATE_TRIE_DIGEST_BEFORE)    %assert_eq
+    // The trie data segment is already written by the linked lists
+    %get_trie_data_size
+
     // stack: trie_data_len
     %mpt_hash_txn_trie     %mload_global_metadata(@GLOBAL_METADATA_TXN_TRIE_DIGEST_BEFORE)      %assert_eq
     // stack: trie_data_len
     %mpt_hash_receipt_trie %mload_global_metadata(@GLOBAL_METADATA_RECEIPT_TRIE_DIGEST_BEFORE)  %assert_eq
     // stack: trie_data_full_len
-    %mstore_global_metadata(@GLOBAL_METADATA_TRIE_DATA_SIZE)
+
+    %set_trie_data_size
 
 global start_txns:
     // stack: (empty)
@@ -144,10 +153,25 @@ global perform_final_checks:
     // stack: cum_gas, txn_counter, num_nibbles, txn_nb
     // Check that we end up with the correct `cum_gas`, `txn_nb` and bloom filter.
     %mload_global_metadata(@GLOBAL_METADATA_BLOCK_GAS_USED_AFTER) %assert_eq
-    DUP3 %mload_global_metadata(@GLOBAL_METADATA_TXN_NUMBER_AFTER) %assert_eq
+    DUP3
+    %mload_global_metadata(@GLOBAL_METADATA_TXN_NUMBER_AFTER) %assert_eq
     %pop3
+
+    PROVER_INPUT(trie_ptr::state)
+
+    %mstore_global_metadata(@GLOBAL_METADATA_STATE_TRIE_ROOT)
+    %set_initial_tries
+    %get_trie_data_size
+    %mpt_hash_state_trie
+
+    SWAP1 %set_trie_data_size
+    %mload_global_metadata(@GLOBAL_METADATA_STATE_TRIE_DIGEST_BEFORE)
+    %assert_eq
+
     PUSH 1 // initial trie data length
+    
 global check_state_trie:
+    %set_final_tries
     %mpt_hash_state_trie   %mload_global_metadata(@GLOBAL_METADATA_STATE_TRIE_DIGEST_AFTER)     %assert_eq
 global check_txn_trie:
     %mpt_hash_txn_trie     %mload_global_metadata(@GLOBAL_METADATA_TXN_TRIE_DIGEST_AFTER)       %assert_eq

diff --git a/evm_arithmetization/src/cpu/kernel/asm/mpt/accounts.asm b/evm_arithmetization/src/cpu/kernel/asm/mpt/accounts.asm
@@ -19,3 +19,44 @@
     %mload_trie_data
     // stack: storage_root_ptr
 %endmacro
+
+%macro clone_account
+    // stack: account_ptr
+    %get_trie_data_size
+    // stack: cloned_account_ptr, account_ptr
+    SWAP1
+    DUP1
+    // stack: nonce_ptr, account_ptr, cloned_account_ptr
+    // Nonce
+    %mload_trie_data
+    %append_to_trie_data
+    // stack: account_ptr, cloned_account_ptr
+    %increment
+    // Balance
+    DUP1
+    %mload_trie_data
+    %append_to_trie_data
+    // Storage trie root
+    %increment
+    DUP1
+    %mload_trie_data
+    %append_to_trie_data
+    // Codehash 
+    %increment
+    %mload_trie_data
+    %append_to_trie_data
+    // stack: cloned_account_ptr
+%endmacro
+
+// The slot_ptr cannot be 0, because `insert_slot` 
+// is only called in `revert_storage_change` (where the case `slot_ptr = 0` 
+// is dealt with differently), and in `storage_write`, 
+// where writing 0 actually corresponds to a `delete`.
+%macro clone_slot
+    // stack: slot_ptr
+    %get_trie_data_size
+    // stack: cloned_slot_ptr, slot_ptr
+    SWAP1
+    %mload_trie_data
+    %append_to_trie_data
+%endmacro
diff --git a/evm_arithmetization/src/cpu/kernel/asm/mpt/delete/delete.asm b/evm_arithmetization/src/cpu/kernel/asm/mpt/delete/delete.asm
@@ -24,17 +24,13 @@ mpt_delete_leaf:
     SWAP1 JUMP
 
 global delete_account:
-    %stack (address, retdest) -> (address, delete_account_save, retdest)
     %addr_to_state_key
-    // stack: key, delete_account_save, retdest
-    PUSH 64
-    // stack: 64, key, delete_account_save, retdest
-    %mload_global_metadata(@GLOBAL_METADATA_STATE_TRIE_ROOT)
-    // stack: state_root_prt, 64, key, delete_account_save, retdest
-    %jump(mpt_delete)
-delete_account_save:
-    // stack: updated_state_root_ptr, retdest
-    %mstore_global_metadata(@GLOBAL_METADATA_STATE_TRIE_ROOT)
+    DUP1
+    %remove_account_from_linked_list
+    // stack: addr_to_state_key, retdest
+
+    // Now we also need to remove all the storage nodes associated with the deleted account.
+    %remove_all_account_slots
     JUMP
 
 %macro delete_account

diff --git a/evm_arithmetization/src/cpu/kernel/asm/mpt/hash/hash_trie_specific.asm b/evm_arithmetization/src/cpu/kernel/asm/mpt/hash/hash_trie_specific.asm
@@ -101,10 +101,6 @@ global encode_account:
     DUP3 %add_const(2) %mload_trie_data // storage_root_ptr = value[2]
     // stack: storage_root_ptr, cur_len, rlp_pos_5, value_ptr, cur_len, retdest
 
-
-    PUSH debug_after_hash_storage_trie
-    POP
-
     // Hash storage trie.
     %mpt_hash_storage_trie
     // stack: storage_root_digest, new_len, rlp_pos_5, value_ptr, cur_len, retdest

diff --git a/evm_arithmetization/src/cpu/kernel/asm/mpt/insert/insert_trie_specific.asm b/evm_arithmetization/src/cpu/kernel/asm/mpt/insert/insert_trie_specific.asm
@@ -6,15 +6,7 @@
 // TODO: Have this take an address and do %mpt_insert_state_trie? To match mpt_read_state_trie.
 global mpt_insert_state_trie:
     // stack: key, value_ptr, retdest
-    %stack (key, value_ptr)
-        -> (key, value_ptr, mpt_insert_state_trie_save)
-    PUSH 64 // num_nibbles
-    %mload_global_metadata(@GLOBAL_METADATA_STATE_TRIE_ROOT)
-    // stack: state_root_ptr, num_nibbles, key, value_ptr, mpt_insert_state_trie_save, retdest
-    %jump(mpt_insert)
-mpt_insert_state_trie_save:
-    // stack: updated_node_ptr, retdest
-    %mstore_global_metadata(@GLOBAL_METADATA_STATE_TRIE_ROOT)
+    %insert_account_with_overwrite
     JUMP
 
 %macro mpt_insert_state_trie