From 5905cf797cbc274447e7156c5986550eedb16a29 Mon Sep 17 00:00:00 2001
From: Linda Guiga <101227802+LindaGuiga@users.noreply.github.com>
Date: Thu, 8 Aug 2024 14:22:00 +0100
Subject: [PATCH] Check initial and final state tries at the end (#476)

---
 evm_arithmetization/src/cpu/kernel/asm/main.asm  | 16 +++++++++-------
 .../src/cpu/kernel/tests/account_code.rs         |  8 +++-----
 .../src/cpu/kernel/tests/mpt/insert.rs           |  2 +-
 3 files changed, 13 insertions(+), 13 deletions(-)

diff --git a/evm_arithmetization/src/cpu/kernel/asm/main.asm b/evm_arithmetization/src/cpu/kernel/asm/main.asm
index fde4dfeb3..76c32d472 100644
--- 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
index b72eaf7ca..604eda6c9 100644
--- 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
index 28352c69e..8e50fcb38 100644
--- 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);