Refactor linked lists initial hashing

evm_arithmetization/src/cpu/kernel/aggregator.rs

@@ -141,8 +141,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"),

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,

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 == virt == 0
     // stack: result_addr, result_len, new_len, retdest
     KECCAK_GENERAL
     // stack: hash, new_len, retdest
@@ -293,3 +293,4 @@ encode_node_leaf_after_encode_value:
     %stack (rlp_prefix_start_pos, rlp_len, cur_len, retdest)
         -> (retdest, rlp_prefix_start_pos, rlp_len, cur_len)
     JUMP
+

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

evm_arithmetization/src/cpu/kernel/asm/mpt/linked_list/linked_list.asm

@@ -91,6 +91,8 @@ loop_store_initial_accounts:
 store_initial_accounts_end:
     %pop2
     // 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
 
@@ -360,6 +362,8 @@ loop_store_initial_slots:
 store_initial_slots_end:
     POP
     // 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 +898,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

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 leafs we return the pointer
     JUMP

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

@@ -190,36 +190,30 @@ pub(crate) fn prepare_interpreter<F: RichField>(
         interpreter.stack()
     );
 
-    // Set initial tries.
+    // Now, set the payload.
     interpreter
         .push(0xDEADBEEFu32.into())
         .expect("The stack should not overflow");
     interpreter
-        .push(0.into()) // Initial nibbles
-        .expect("The stack should not overflow");
-    interpreter
-        .push(0.into()) // Initial number of nibbles
-        .expect("The stack should not overflow");
-    interpreter
-        .push((Segment::StorageLinkedList as usize + 8).into())
-        .expect("The stack should not overflow");
-    interpreter
-        .push((Segment::AccountsLinkedList as usize + 6).into())
+        .push((Segment::StorageLinkedList as usize + 5).into())
         .expect("The stack should not overflow");
     interpreter
         .push(interpreter.get_global_metadata_field(GlobalMetadata::StateTrieRoot))
+        .unwrap();
+    interpreter
+        .push((Segment::AccountsLinkedList as usize + 4).into())
         .expect("The stack should not overflow");
 
     // Now, set the payload.
     interpreter.generation_state.registers.program_counter =
-        KERNEL.global_labels["mpt_set_payload"];
+        KERNEL.global_labels["insert_all_initial_accounts"];
 
     interpreter.run()?;
 
-    let acc_ptr = interpreter.pop().expect("The stack should not be empty") - 2;
-    let storage_ptr = interpreter.pop().expect("The stack should not be empty") - 3;
-    interpreter.set_global_metadata_field(GlobalMetadata::InitialAccountsLinkedListLen, acc_ptr);
-    interpreter.set_global_metadata_field(GlobalMetadata::InitialStorageLinkedListLen, storage_ptr);
+    assert_eq!(interpreter.stack_len(), 1);
+
+    let state_root = interpreter.pop().expect("The stack should not be empty");
+    interpreter.set_global_metadata_field(GlobalMetadata::StateTrieRoot, state_root);
 
     // Now, execute `mpt_hash_state_trie`.
     state_trie.insert(k, rlp::encode(account).to_vec())?;
@@ -410,43 +404,30 @@ fn prepare_interpreter_all_accounts<F: RichField>(
         KERNEL.global_labels["store_initial_slots"];
     interpreter.run()?;
 
-    // Set the pointers to the initial payloads.
+    // Now, set the payload.
     interpreter
         .push(0xDEADBEEFu32.into())
         .expect("The stack should not overflow");
     interpreter
-        .push(0.into()) // Initial nibbles
-        .expect("The stack should not overflow");
-    interpreter
-        .push(0.into()) // Initial number of nibbles
-        .expect("The stack should not overflow");
-    interpreter
-        .push((Segment::StorageLinkedList as usize + 8).into())
-        .expect("The stack should not overflow");
-    interpreter
-        .push((Segment::AccountsLinkedList as usize + 6).into())
+        .push((Segment::StorageLinkedList as usize + 5).into())
         .expect("The stack should not overflow");
     interpreter
         .push(interpreter.get_global_metadata_field(GlobalMetadata::StateTrieRoot))
+        .unwrap();
+    interpreter
+        .push((Segment::AccountsLinkedList as usize + 4).into())
         .expect("The stack should not overflow");
 
-    // Now, set the payloads in the state trie leaves.
+    // Now, set the payload.
     interpreter.generation_state.registers.program_counter =
-        KERNEL.global_labels["mpt_set_payload"];
+        KERNEL.global_labels["insert_all_initial_accounts"];
 
     interpreter.run()?;
 
-    assert_eq!(
-        interpreter.stack().len(),
-        2,
-        "Expected 2 items on stack after setting the initial trie payloads, found {:?}",
-        interpreter.stack()
-    );
+    assert_eq!(interpreter.stack_len(), 1);
 
-    let acc_ptr = interpreter.pop().expect("The stack should not be empty") - 2;
-    let storage_ptr = interpreter.pop().expect("The stack should not be empty") - 3;
-    interpreter.set_global_metadata_field(GlobalMetadata::InitialAccountsLinkedListLen, acc_ptr);
-    interpreter.set_global_metadata_field(GlobalMetadata::InitialStorageLinkedListLen, storage_ptr);
+    let state_root = interpreter.pop().expect("The stack should not be empty");
+    interpreter.set_global_metadata_field(GlobalMetadata::StateTrieRoot, state_root);
 
     // Switch context and initialize memory with the data we need for the tests.
     interpreter.generation_state.registers.program_counter = 0;

evm_arithmetization/src/cpu/kernel/tests/mpt/delete.rs

@@ -121,33 +121,25 @@ fn test_state_trie(
         .push(0xDEADBEEFu32.into())
         .expect("The stack should not overflow");
     interpreter
-        .push(0.into()) // Initial nibbles
-        .expect("The stack should not overflow");
-    interpreter
-        .push(0.into()) // Initial number of nibbles
-        .expect("The stack should not overflow");
-    interpreter
-        .push((Segment::StorageLinkedList as usize + 8).into())
-        .expect("The stack should not overflow");
-    interpreter
-        .push((Segment::AccountsLinkedList as usize + 6).into())
+        .push((Segment::StorageLinkedList as usize + 5).into())
         .expect("The stack should not overflow");
     interpreter
         .push(interpreter.get_global_metadata_field(GlobalMetadata::StateTrieRoot))
         .unwrap();
+    interpreter
+        .push((Segment::AccountsLinkedList as usize + 4).into())
+        .expect("The stack should not overflow");
 
     // Now, set the payload.
     interpreter.generation_state.registers.program_counter =
-        KERNEL.global_labels["mpt_set_payload"];
+        KERNEL.global_labels["insert_all_initial_accounts"];
 
     interpreter.run()?;
 
-    assert_eq!(interpreter.stack_len(), 2);
+    assert_eq!(interpreter.stack_len(), 1);
 
-    let acc_ptr = interpreter.pop().expect("The stack should not be empty") - 2;
-    let storage_ptr = interpreter.pop().expect("The stack should not be empty") - 3;
-    interpreter.set_global_metadata_field(GlobalMetadata::InitialAccountsLinkedListLen, acc_ptr);
-    interpreter.set_global_metadata_field(GlobalMetadata::InitialStorageLinkedListLen, storage_ptr);
+    let state_root = interpreter.pop().expect("The stack should not be empty");
+    interpreter.set_global_metadata_field(GlobalMetadata::StateTrieRoot, state_root);
 
     // Next, execute mpt_insert_state_trie.
     interpreter.generation_state.registers.program_counter = mpt_insert_state_trie;

evm_arithmetization/src/cpu/kernel/tests/mpt/insert.rs

@@ -190,36 +190,29 @@ fn test_state_trie(
     interpreter.generation_state.registers.program_counter = KERNEL.global_labels["store_initial"];
     interpreter.run()?;
 
-    // Set initial tries.
     interpreter
         .push(0xDEADBEEFu32.into())
         .expect("The stack should not overflow");
     interpreter
-        .push(0.into()) // Initial nibbles
-        .expect("The stack should not overflow");
-    interpreter
-        .push(0.into()) // Initial number of nibbles
-        .expect("The stack should not overflow");
-    interpreter
-        .push((Segment::StorageLinkedList as usize + 8).into())
-        .expect("The stack should not overflow");
-    interpreter
-        .push((Segment::AccountsLinkedList as usize + 6).into())
+        .push((Segment::StorageLinkedList as usize + 5).into())
         .expect("The stack should not overflow");
     interpreter
         .push(interpreter.get_global_metadata_field(GlobalMetadata::StateTrieRoot))
         .unwrap();
+    interpreter
+        .push((Segment::AccountsLinkedList as usize + 4).into())
+        .expect("The stack should not overflow");
 
     // Now, set the payload.
     interpreter.generation_state.registers.program_counter =
-        KERNEL.global_labels["mpt_set_payload"];
+        KERNEL.global_labels["insert_all_initial_accounts"];
 
     interpreter.run()?;
 
-    let acc_ptr = interpreter.pop().expect("The stack should not be empty") - 2;
-    let storage_ptr = interpreter.pop().expect("The stack should not be empty") - 3;
-    interpreter.set_global_metadata_field(GlobalMetadata::InitialAccountsLinkedListLen, acc_ptr);
-    interpreter.set_global_metadata_field(GlobalMetadata::InitialStorageLinkedListLen, storage_ptr);
+    assert_eq!(interpreter.stack_len(), 1);
+
+    let state_root = interpreter.pop().expect("The stack should not be empty");
+    interpreter.set_global_metadata_field(GlobalMetadata::StateTrieRoot, state_root);
 
     // Next, execute mpt_insert_state_trie.
     interpreter.generation_state.registers.program_counter = mpt_insert_state_trie;
@@ -267,10 +260,6 @@ fn test_state_trie(
     );
 
     interpreter.generation_state.registers.program_counter = check_state_trie;
-    interpreter
-        .halt_offsets
-        .push(KERNEL.global_labels["check_txn_trie"]);
-
     interpreter
         .push(interpreter.get_global_metadata_field(GlobalMetadata::TrieDataSize)) // Initial trie data segment size, unused.
         .expect("The stack should not overflow");

evm_arithmetization/src/cpu/kernel/tests/mpt/read.rs

@@ -42,7 +42,9 @@ fn mpt_read() -> Result<()> {
     interpreter.run()?;
 
     assert_eq!(interpreter.stack().len(), 1);
-    let result_ptr = interpreter.stack()[0].as_usize();
+    // mp_read returns a pointer to the accounts pointer
+    let result_ptr_ptr = interpreter.stack()[0].as_usize();
+    let result_ptr = interpreter.get_trie_data()[result_ptr_ptr..][..4][0].as_usize();
     let result = &interpreter.get_trie_data()[result_ptr..][..4];
     assert_eq!(result[0], test_account_1().nonce);
     assert_eq!(result[1], test_account_1().balance);