From bd02117cdcf56ca09c4ef9bd3978ca3b7073616e Mon Sep 17 00:00:00 2001
From: wborgeaud <williamborgeaud@gmail.com>
Date: Tue, 9 Jan 2024 11:15:53 +0100
Subject: [PATCH] Fix `after_mpt_delete_extension_branch` (#1449)

* Fix after_mpt_delete_extension_branch

* Rename test

* PR feedback
---
 .../asm/mpt/delete/delete_extension.asm       | 14 ++-------
 evm/src/cpu/kernel/tests/mpt/delete.rs        | 29 ++++++++++++++++++-
 2 files changed, 31 insertions(+), 12 deletions(-)

diff --git a/evm/src/cpu/kernel/asm/mpt/delete/delete_extension.asm b/evm/src/cpu/kernel/asm/mpt/delete/delete_extension.asm
index 149b971d76..0627fcba6a 100644
--- a/evm/src/cpu/kernel/asm/mpt/delete/delete_extension.asm
+++ b/evm/src/cpu/kernel/asm/mpt/delete/delete_extension.asm
@@ -37,18 +37,10 @@ after_mpt_delete_extension_branch:
     // stack: child_type, updated_child_node_ptr, node_payload_ptr, node_len, node_key, retdest
     POP
     // stack: updated_child_node_ptr, node_payload_ptr, node_len, node_key, retdest
-    SWAP1
-    // stack: extension_ptr, updated_child_node_ptr, node_len, node_key, retdest
-    PUSH @MPT_NODE_EXTENSION DUP2 %mstore_trie_data
-    // stack: extension_ptr, updated_child_node_ptr, node_len, node_key, retdest
-    DUP3 DUP2 %mstore_trie_data // Append node_len to our node
-    // stack: extension_ptr, updated_child_node_ptr, node_len, node_key, retdest
-    DUP4 DUP2 %mstore_trie_data // Append node_key to our node
-    // stack: extension_ptr, updated_child_node_ptr, node_len, node_key, retdest
-    SWAP1 DUP2 %mstore_trie_data // Append updated_child_node_ptr to our node
-    // stack: extension_ptr, node_len, node_key, retdest
+    DUP2 %add_const(2) %mstore_trie_data
+    // stack: node_payload_ptr, node_len, node_key, retdest
+    %decrement
     %stack (extension_ptr, node_len, node_key, retdest) -> (retdest, extension_ptr)
-    // stack: extension_ptr, retdest
     JUMP
 
 after_mpt_delete_extension_extension:
diff --git a/evm/src/cpu/kernel/tests/mpt/delete.rs b/evm/src/cpu/kernel/tests/mpt/delete.rs
index acb75b7b21..38342dc99c 100644
--- a/evm/src/cpu/kernel/tests/mpt/delete.rs
+++ b/evm/src/cpu/kernel/tests/mpt/delete.rs
@@ -1,7 +1,8 @@
 use anyhow::{anyhow, Result};
 use eth_trie_utils::nibbles::Nibbles;
 use eth_trie_utils::partial_trie::{HashedPartialTrie, PartialTrie};
-use ethereum_types::{BigEndianHash, H256};
+use ethereum_types::{BigEndianHash, H256, U512};
+use rand::random;
 
 use crate::cpu::kernel::aggregator::KERNEL;
 use crate::cpu::kernel::constants::global_metadata::GlobalMetadata;
@@ -48,6 +49,32 @@ fn mpt_delete_branch_into_hash() -> Result<()> {
     test_state_trie(state_trie, nibbles_64(0xADE), test_account_2())
 }
 
+#[test]
+fn test_after_mpt_delete_extension_branch() -> Result<()> {
+    let hash = Node::Hash(H256::random());
+    let branch = Node::Branch {
+        children: std::array::from_fn(|i| {
+            if i == 0 {
+                Node::Empty.into()
+            } else {
+                hash.clone().into()
+            }
+        }),
+        value: vec![],
+    };
+    let nibbles = Nibbles::from_bytes_be(&random::<[u8; 5]>()).unwrap();
+    let state_trie = Node::Extension {
+        nibbles,
+        child: branch.into(),
+    }
+    .into();
+    let key = nibbles.merge_nibbles(&Nibbles {
+        packed: U512::zero(),
+        count: 64 - nibbles.count,
+    });
+    test_state_trie(state_trie, key, test_account_2())
+}
+
 /// Note: The account's storage_root is ignored, as we can't insert a new storage_root without the
 /// accompanying trie data. An empty trie's storage_root is used instead.
 fn test_state_trie(