Check initial and final state tries at the end (#476)

0xPolygonZero · Aug 8, 2024 · 5905cf7 · 5905cf7
1 parent 21d22c6
commit 5905cf7
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 13 deletions.
diff --git a/evm_arithmetization/src/cpu/kernel/asm/main.asm b/evm_arithmetization/src/cpu/kernel/asm/main.asm
@@ -157,6 +157,14 @@ global perform_final_checks:
     %mload_global_metadata(@GLOBAL_METADATA_TXN_NUMBER_AFTER) %assert_eq
     %pop3
 
+    PUSH 1 // initial trie data length
+    
+global check_txn_trie:
+    %mpt_hash_txn_trie     %mload_global_metadata(@GLOBAL_METADATA_TXN_TRIE_DIGEST_AFTER)       %assert_eq
+global check_receipt_trie:
+    %mpt_hash_receipt_trie %mload_global_metadata(@GLOBAL_METADATA_RECEIPT_TRIE_DIGEST_AFTER)   %assert_eq
+global check_state_trie:
+    // First, check initial trie.
     PROVER_INPUT(trie_ptr::state)
 
     %mstore_global_metadata(@GLOBAL_METADATA_STATE_TRIE_ROOT)
@@ -172,15 +180,9 @@ global perform_final_checks:
     %mload_global_metadata(@GLOBAL_METADATA_STATE_TRIE_DIGEST_BEFORE)
     %assert_eq
 
-    PUSH 1 // initial trie data length
-    
-global check_state_trie:
+global check_final_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
-global check_receipt_trie:
-    %mpt_hash_receipt_trie %mload_global_metadata(@GLOBAL_METADATA_RECEIPT_TRIE_DIGEST_AFTER)   %assert_eq
     // We don't need the trie data length here.
     POP
 

diff --git a/evm_arithmetization/src/cpu/kernel/tests/account_code.rs b/evm_arithmetization/src/cpu/kernel/tests/account_code.rs
@@ -139,7 +139,7 @@ pub(crate) fn prepare_interpreter<F: Field>(
     account: &AccountRlp,
 ) -> Result<()> {
     let mpt_insert_state_trie = KERNEL.global_labels["mpt_insert_state_trie"];
-    let check_state_trie = KERNEL.global_labels["check_state_trie"];
+    let check_state_trie = KERNEL.global_labels["check_final_state_trie"];
     let mut state_trie: HashedPartialTrie = HashedPartialTrie::from(Node::Empty);
     let trie_inputs = TrieInputs {
         state_trie: HashedPartialTrie::from(Node::Empty),
@@ -527,11 +527,9 @@ fn sstore() -> Result<()> {
         GlobalMetadata::StateTrieRootDigestAfter,
         h2u(expected_state_trie_hash),
     );
-    interpreter
-        .halt_offsets
-        .push(KERNEL.global_labels["check_txn_trie"]);
+
     // Now, execute `mpt_hash_state_trie` and check that the hash is correct.
-    let mpt_hash_state_trie = KERNEL.global_labels["check_state_trie"];
+    let mpt_hash_state_trie = KERNEL.global_labels["check_final_state_trie"];
     interpreter.generation_state.registers.program_counter = mpt_hash_state_trie;
     interpreter.set_is_kernel(true);
     interpreter.set_context(0);

diff --git a/evm_arithmetization/src/cpu/kernel/tests/mpt/insert.rs b/evm_arithmetization/src/cpu/kernel/tests/mpt/insert.rs
@@ -175,7 +175,7 @@ fn test_state_trie(
         storage_tries: vec![],
     };
     let mpt_insert_state_trie = KERNEL.global_labels["mpt_insert_state_trie"];
-    let check_state_trie = KERNEL.global_labels["check_state_trie"];
+    let check_state_trie = KERNEL.global_labels["check_final_state_trie"];
 
     let initial_stack = vec![];
     let mut interpreter: Interpreter<F> = Interpreter::new(0, initial_stack, None);