Refactor linked lists initial hashing (#581)

* Deattach hashed nodes from initial trie * Merge mpt_set_payload with mpt_hash * Add missing files * Fix errors in the stack * Fix hash mismatch on initial trie * Refactor final state hash * Fix final hash * [WIP] debugging erc721 * Fix erc721 * [WIP] debugging block 28 * Fix b28 * Fix block 28 * [WIP] Debugging block 20240058 * Fix extension nodes bug * Fix block 20240058 * [WIP] benchmarking * Set inital trie with insertions * Remove hash nodes * Fix missing segment number * Clean code * Fix unit tests * Fix erc20 * [WIP] Fixing blocks * Clean code * Clean and fmt * Address reviews * Apply suggestions from code review Co-authored-by: Hamy Ratoanina <hratoanina@polygon.technology> * [WIP] Fixing trie data length * Fix trie_data_length mismatch * Check correctness of inital next node ptr and check strict keys monotonicity * Address review comment * Apply suggestions from code review Co-authored-by: Hamy Ratoanina <hratoanina@polygon.technology> * Minor * Fix circuit sizes * Apply suggestions from code review Co-authored-by: Robin Salen <30937548+Nashtare@users.noreply.github.com> * Add missing stack comment --------- Co-authored-by: Hamy Ratoanina <hratoanina@polygon.technology> Co-authored-by: Robin Salen <salenrobin@gmail.com> Co-authored-by: Robin Salen <30937548+Nashtare@users.noreply.github.com>
0xPolygonZero · Sep 18, 2024 · 6ef8827 · 6ef8827
1 parent eafbcc6
commit 6ef8827
Show file tree

Hide file tree

Showing 13 changed files with 325 additions and 399 deletions.
diff --git a/evm_arithmetization/src/cpu/kernel/aggregator.rs b/evm_arithmetization/src/cpu/kernel/aggregator.rs
@@ -140,8 +140,8 @@ pub static KERNEL_FILES: [&str; NUMBER_KERNEL_FILES] = [
     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/linked_list/initial_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/main.asm b/evm_arithmetization/src/cpu/kernel/asm/main.asm
@@ -215,14 +215,16 @@ global check_state_trie:
     // `GLOBAL_METADATA_TRIE_DATA_SIZE` is correct.
     %get_trie_data_size
     // stack: trie_data_len
-    PROVER_INPUT(trie_ptr::state)
+    PROVER_INPUT(trie_ptr::initial_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
+    // stack: trie_data_len
+    %set_initial_state_trie
+    // stack: trie_data_len
 
     PUSH @INITIAL_RLP_ADDR
     // stack: rlp_start, trie_data_len
@@ -233,7 +235,7 @@ global check_state_trie:
     %mload_global_metadata(@GLOBAL_METADATA_STATE_TRIE_DIGEST_BEFORE)
     %assert_eq
     // Check that the stored trie data length is correct.
-    %mload_global_metadata(@GLOBAL_METADATA_TRIE_DATA_SIZE) 
+    %mload_global_metadata(@GLOBAL_METADATA_TRIE_DATA_SIZE)
     %assert_eq
 
     // We set a dummy value as an initial trie data length,

diff --git a/evm_arithmetization/src/cpu/kernel/asm/mpt/hash/hash.asm b/evm_arithmetization/src/cpu/kernel/asm/mpt/hash/hash.asm
@@ -29,12 +29,12 @@ mpt_hash_hash_if_rlp:
 mpt_hash_hash_rlp:
     // stack: result, result_len, new_len, retdest
     %stack (result, result_len, new_len)
-        -> (@SEGMENT_RLP_RAW, result, result_len, mpt_hash_hash_rlp_after_unpacking, result_len, new_len)
+        -> (@INITIAL_RLP_ADDR, result, result_len, mpt_hash_hash_rlp_after_unpacking, result_len, new_len)
     // stack: addr, result, result_len, mpt_hash_hash_rlp_after_unpacking, result_len, new_len
     %jump(mstore_unpacking)
 mpt_hash_hash_rlp_after_unpacking:
     // stack: result_addr, result_len, new_len, retdest
-    POP PUSH @SEGMENT_RLP_RAW // ctx == virt == 0
+    POP PUSH @INITIAL_RLP_ADDR // ctx == 0, virt == 1
     // stack: result_addr, result_len, new_len, retdest
     KECCAK_GENERAL
     // stack: hash, new_len, retdest

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
@@ -209,21 +209,17 @@ after_mpt_delete_slot:
 
 global set_final_tries:
     PUSH set_final_tries_after
-    PUSH @SEGMENT_STORAGE_LINKED_LIST
-    %add_const(@STORAGE_LINKED_LISTS_NODE_SIZE) // Skip the first node.
+    %first_initial_slot // Skip the first node.
     %mload_global_metadata(@GLOBAL_METADATA_STATE_TRIE_ROOT)
-    PUSH @SEGMENT_ACCOUNTS_LINKED_LIST
-    %add_const(@ACCOUNTS_LINKED_LISTS_NODE_SIZE) // Skip the first node.
+    %first_initial_account // Skip the first node.
     %jump(delete_removed_accounts)
 set_final_tries_after:
     // stack: new_state_root
     PUSH set_final_tries_after_after SWAP1
     // stack: new_state_root, set_final_tries_after_after
-    PUSH @SEGMENT_STORAGE_LINKED_LIST
-    %next_slot
+    %first_slot
     SWAP1
-    PUSH @SEGMENT_ACCOUNTS_LINKED_LIST
-    %next_account
+    %first_account
     %jump(insert_all_accounts)
 set_final_tries_after_after:
     //stack: new_state_root

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
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
@@ -38,22 +38,36 @@
 /// the accounts, storing a pointer to the copied account in the node.
 /// Computes the length of `SEGMENT_ACCOUNTS_LINKED_LIST` and 
 /// stores it in `GLOBAL_METADATA_ACCOUNTS_LINKED_LIST_NEXT_AVAILABLE`.
+/// It also checks that the next node address is current address + 4
+/// and that all keys are strictly increasing.
+/// NOTE: It may be more efficient to check that the next node addres != U256_MAX
+/// (i.e. node was not deleted) and ensure that no node with repeated key
+/// is ever read.
 global store_initial_accounts:
     // stack: retdest
     PUSH @ACCOUNTS_LINKED_LISTS_NODE_SIZE
     PUSH @SEGMENT_ACCOUNTS_LINKED_LIST
     ADD
     // stack: cur_len, retdest
     PUSH @SEGMENT_ACCOUNTS_LINKED_LIST
+    // stack: current_node_ptr, cur_len, retdest
+    DUP1
+    MLOAD_GENERAL
+    // stack: current_addr_key, current_node_ptr, cur_len', retdest
+    %assert_eq_const(@U256_MAX)
+    DUP1
     %next_account
+    // stack: next_node_ptr, current_node_ptr, cur_len', retdest
+    DUP1
+    SWAP2
+    %next_initial_account
+    %assert_eq(store_initial_accounts_end) // next_node_ptr ==  current_node_ptr + node_size
+    // stack: next_node_ptr, cur_len', retdest
+    
 loop_store_initial_accounts:
     // stack: current_node_ptr, cur_len, retdest
     %get_trie_data_size
-    DUP2
-    MLOAD_GENERAL
-    // stack: current_addr_key, cpy_ptr, current_node_ptr, cur_len, retdest
-    %eq_const(@U256_MAX)
-    %jumpi(store_initial_accounts_end)
+    // stack: cpy_ptr, current_node_ptr, cur_len, retdest
     DUP2
     %increment
     MLOAD_GENERAL
@@ -84,13 +98,35 @@ loop_store_initial_accounts:
     SWAP1 PUSH @ACCOUNTS_LINKED_LISTS_NODE_SIZE 
     ADD
     SWAP1
-    // stack: current_node_ptr, cur_len', retdest
+    // Check next node ptr validity and strict keys monotonicity
+    DUP1
+    MLOAD_GENERAL
+    // stack: current_addr_key, current_node_ptr, cur_len', retdest
+    SWAP1
+    DUP1
     %next_account
+    // stack: next_node_ptr, current_node_ptr, current_addr_key, cur_len', retdest
+    DUP1
+    SWAP2
+    %next_initial_account
+    %assert_eq(store_initial_accounts_end_pop_key) // next_node_ptr ==  current_node_ptr + node_size
+    // stack: next_node_ptr, current_addr_key, cur_len', retdest
+    SWAP1
+    DUP2
+    MLOAD_GENERAL
+    %assert_gt // next_addr_key > current_addr_key
+    // stack: next_node_ptr, cur_len', retdest
     %jump(loop_store_initial_accounts)
 
+store_initial_accounts_end_pop_key:
+    // stack: next_node_ptr, current_addr_key, cur_len', retdest
+    SWAP1 POP
 store_initial_accounts_end:
-    %pop2
+    // stack: next_node_ptr, cur_len', retdest
+    %assert_eq_const(@SEGMENT_ACCOUNTS_LINKED_LIST)
     // stack: cur_len, retdest
+    DUP1
+    %mstore_global_metadata(@GLOBAL_METADATA_INITIAL_ACCOUNTS_LINKED_LIST_LEN)
     %mstore_global_metadata(@GLOBAL_METADATA_ACCOUNTS_LINKED_LIST_NEXT_AVAILABLE)
     JUMP
 
@@ -324,22 +360,35 @@ global remove_account:
 /// the accounts, storing a pointer to the copied account in the node.
 /// Computes the length of `SEGMENT_STORAGE_LINKED_LIST` and 
 /// checks against `GLOBAL_METADATA_STORAGE_LINKED_LIST_NEXT_AVAILABLE`.
+/// It also checks that the next node address is current address + 5
+/// and that all keys are strictly increasing.
+/// NOTE: It may be more efficient to check that the next node addres != U256_MAX
+/// (i.e. node was not deleted) and ensure that no node with repeated key
+/// is ever read.
 global store_initial_slots:
     // stack: retdest
     PUSH @STORAGE_LINKED_LISTS_NODE_SIZE
     PUSH @SEGMENT_STORAGE_LINKED_LIST
     ADD
     // stack: cur_len, retdest
     PUSH @SEGMENT_STORAGE_LINKED_LIST
-    %next_slot
-
-loop_store_initial_slots:
-    // stack: current_node_ptr, cur_len, retdest
     DUP1
     MLOAD_GENERAL
     // stack: current_addr_key, current_node_ptr, cur_len, retdest
-    %eq_const(@U256_MAX)
-    %jumpi(store_initial_slots_end)
+    %assert_eq_const(@U256_MAX)
+
+    // stack: current_node_ptr, cur_len', retdest
+    DUP1
+    %next_slot
+    // stack: next_node_ptr, current_node_ptr, cur_len, retdest
+    DUP1
+    SWAP2
+    %next_initial_slot
+    %assert_eq(store_initial_slots_end) // next_node_ptr == current_node_ptr + node_size
+    // stack: next_node_ptr, cur_len', retdest
+ 
+loop_store_initial_slots:
+    // stack: current_node_ptr, cur_len, retdest
     DUP1
     %add_const(2)
     MLOAD_GENERAL
@@ -353,13 +402,65 @@ loop_store_initial_slots:
     SWAP1 PUSH @STORAGE_LINKED_LISTS_NODE_SIZE
     ADD
     SWAP1
-    // stack: current_node_ptr, cur_len', retdest
+    // Check correctness of next node ptr and strict key monotonicity.
+    DUP1
+    MLOAD_GENERAL
+    // stack: current_addr_key, current_node_ptr, cur_len', retdest
+    SWAP1
+    DUP1
+    %increment
+    MLOAD_GENERAL
+    // stack: current_slot_key, current_node_ptr, current_addr_key, cur_len', retdest
+    SWAP1
+    DUP1
     %next_slot
+    // stack: next_node_ptr, current_node_ptr, current_slot_key, current_addr_key, cur_len', retdest
+    DUP1
+    SWAP2
+    %next_initial_slot
+    %assert_eq(store_initial_slots_end_pop_keys) // next_node_ptr == current_node_ptr + node_size
+    // stack: next_node_ptr, current_slot_key, current_addr_key, cur_len', retdest
+    DUP1
+    DUP1
+    %increment
+    MLOAD_GENERAL
+    // stack: next_node_slot_key, next_node_ptr, next_node_ptr, current_slot_key, current_addr_key, cur_len', retdest
+    SWAP1
+    MLOAD_GENERAL
+    // stack: next_node_addr_key, next_node_slot_key, next_node_ptr, current_slot_key, current_addr_key, cur_len', retdest
+    SWAP3
+    LT
+    // stack: current_slot_key > next_node_slot_key, next_node_ptr, next_node_addr_key, current_addr_key, cur_len', retdest
+    SWAP2
+    SWAP1
+    SWAP3
+    // stack: current_addr_key, next_node_addr_key, current_slot_key > next_node_slot_key, next_node_ptr, cur_len', retdest
+    DUP2
+    DUP2
+    EQ
+    // stack: current_addr_key == next_node_addr_key, current_addr_key, next_node_addr_key, current_slot_key > next_node_slot_key, next_node_ptr, cur_len', retdest
+    SWAP1
+    SWAP3
+    MUL // AND
+    // stack  current_slot_key > next_node_slot_key AND current_addr_key == next_node_addr_key, next_node_addr_key, current_addr_key, next_node_ptr, cur_len', retdest
+    SWAP2
+    LT
+    ADD // OR
+    %assert_nonzero
     %jump(loop_store_initial_slots)
 
+store_initial_slots_end_pop_keys:
+    // stack: next_node_ptr, current_slot_key, current_addr_key, cur_len', retdest
+    SWAP2
+    %pop2
+
 store_initial_slots_end:
-    POP
+    // stack: next_node_ptr, cur_len', retdest
+    %assert_eq_const(@SEGMENT_STORAGE_LINKED_LIST)
+    
     // stack: cur_len, retdest
+    DUP1
+    %mstore_global_metadata(@GLOBAL_METADATA_INITIAL_STORAGE_LINKED_LIST_LEN)
     %mstore_global_metadata(@GLOBAL_METADATA_STORAGE_LINKED_LIST_NEXT_AVAILABLE)
     JUMP
 
@@ -894,22 +995,62 @@ remove_all_slots_end:
     %next_account
 %endmacro
 
+%macro first_initial_account
+    // stack: empty
+    PUSH @SEGMENT_ACCOUNTS_LINKED_LIST
+    %next_initial_account
+%endmacro
+
 %macro next_account
     // stack: node_ptr
     %add_const(@ACCOUNTS_NEXT_NODE_PTR)
     MLOAD_GENERAL
     // stack: next_node_ptr
 %endmacro
 
+%macro next_initial_account
+    // stack: node_ptr
+    %add_const(@ACCOUNTS_LINKED_LISTS_NODE_SIZE)
+    // stack: next_node_ptr
+%endmacro
+
 %macro first_slot
     // stack: empty
     PUSH @SEGMENT_STORAGE_LINKED_LIST
     %next_slot
 %endmacro
 
+%macro first_initial_slot
+    // stack: empty
+    PUSH @SEGMENT_STORAGE_LINKED_LIST
+    %next_initial_slot
+%endmacro
+
 %macro next_slot
     // stack: node_ptr
     %add_const(@STORAGE_NEXT_NODE_PTR)
     MLOAD_GENERAL
     // stack: next_node_ptr
 %endmacro
+
+%macro next_initial_slot
+    // stack: node_ptr
+    %add_const(@STORAGE_LINKED_LISTS_NODE_SIZE)
+    // stack: next_node_ptr
+%endmacro
+
+%macro next_hash_node
+    // stack: hash_node_ptr
+    %add_const(4)
+    // stack: next_hash_node_ptr
+%endmacro
+
+// Skip over the the first three words (number of nibbles and keys)
+// and load the hash from memory.
+%macro get_hash
+    // stack: hash_node_ptr
+    %add_const(3)
+    // stack: next_ptr
+    MLOAD_GENERAL
+    // stack: hash
+%endmacro
diff --git a/evm_arithmetization/src/cpu/kernel/asm/mpt/read.asm b/evm_arithmetization/src/cpu/kernel/asm/mpt/read.asm
@@ -24,7 +24,8 @@ global mpt_read_state_trie:
 // - the key, as a U256
 // - return destination
 //
-// This function returns a pointer to the value, or 0 if the key is not found.
+// This function returns a pointer to the value, or 0 if the key is not found. If the key
+// is a leaf, it returns a pointer to a pointer.
 global mpt_read:
     // stack: node_ptr, num_nibbles, key, retdest
     DUP1
@@ -145,7 +146,6 @@ global mpt_read_leaf_found:
     // stack: node_payload_ptr, retdest
     %add_const(2) // The value pointer is located after num_nibbles and the key.
     // stack: value_ptr_ptr, retdest
-    %mload_trie_data
-    // stack: value_ptr, retdest
     SWAP1
+    // For leaves, we return the pointer
     JUMP