From 50002df8e49cd1675ae13d20b126b507d81196d6 Mon Sep 17 00:00:00 2001
From: Daniel Lubarov <daniel@lubarov.com>
Date: Mon, 10 Oct 2022 10:42:02 -0700
Subject: [PATCH] MPT insert into leaf, overlapping keys case

---
 evm/src/cpu/kernel/asm/mpt/insert_leaf.asm | 16 +++++++++++++--
 evm/src/cpu/kernel/tests/mpt/insert.rs     | 24 ++++++++++++++++++++--
 2 files changed, 36 insertions(+), 4 deletions(-)

diff --git a/evm/src/cpu/kernel/asm/mpt/insert_leaf.asm b/evm/src/cpu/kernel/asm/mpt/insert_leaf.asm
index b82653f2..eeb7612a 100644
--- a/evm/src/cpu/kernel/asm/mpt/insert_leaf.asm
+++ b/evm/src/cpu/kernel/asm/mpt/insert_leaf.asm
@@ -21,7 +21,7 @@
 //     branch[17] = insert_value
 //
 // if common_len > 0:
-//     return Extension[common_key, branch]
+//     return Extension[common_len, common_key, branch]
 // else:
 //     return branch
 
@@ -104,7 +104,19 @@ finished_processing_insert_value:
 
 extension_for_common_key:
     // stack: branch_ptr, common_len, common_key, node_len, node_key, insert_len, insert_key, node_value_ptr, insert_value_ptr, retdest
-    PANIC // TODO
+    // return Extension[common_len, common_key, branch]
+    %get_trie_data_size
+    // stack: extension_ptr, branch_ptr, common_len, common_key, ...
+    PUSH @MPT_NODE_EXTENSION %append_to_trie_data
+    SWAP2 %append_to_trie_data // Append common_len to our node
+    SWAP2 %append_to_trie_data // Append common_key to our node
+    SWAP1 %append_to_trie_data // Append branch_ptr to our node
+    // stack: extension_ptr, node_len, node_key, insert_len, insert_key, node_value_ptr, insert_value_ptr, retdest
+    SWAP6
+    %pop6
+    // stack: extension_ptr, retdest
+    SWAP1
+    JUMP
 
 node_key_continues:
     // stack: branch_ptr, common_len, common_key, node_len, node_key, insert_len, insert_key, node_value_ptr, insert_value_ptr, retdest
diff --git a/evm/src/cpu/kernel/tests/mpt/insert.rs b/evm/src/cpu/kernel/tests/mpt/insert.rs
index 1ac7974f..872ef4af 100644
--- a/evm/src/cpu/kernel/tests/mpt/insert.rs
+++ b/evm/src/cpu/kernel/tests/mpt/insert.rs
@@ -23,7 +23,7 @@ fn mpt_insert_empty() -> Result<()> {
 }
 
 #[test]
-fn mpt_insert_leaf_same_key() -> Result<()> {
+fn mpt_insert_leaf_identical_keys() -> Result<()> {
     let key = Nibbles {
         count: 3,
         packed: 0xABC.into(),
@@ -41,7 +41,7 @@ fn mpt_insert_leaf_same_key() -> Result<()> {
 }
 
 #[test]
-fn mpt_insert_leaf_nonoverlapping_key() -> Result<()> {
+fn mpt_insert_leaf_nonoverlapping_keys() -> Result<()> {
     let state_trie = PartialTrie::Leaf {
         nibbles: Nibbles {
             count: 3,
@@ -60,6 +60,26 @@ fn mpt_insert_leaf_nonoverlapping_key() -> Result<()> {
     test_state_trie(state_trie, insert)
 }
 
+#[test]
+fn mpt_insert_leaf_overlapping_keys() -> Result<()> {
+    let state_trie = PartialTrie::Leaf {
+        nibbles: Nibbles {
+            count: 3,
+            packed: 0xABC.into(),
+        },
+        value: test_account_1_rlp(),
+    };
+    let insert = InsertEntry {
+        nibbles: Nibbles {
+            count: 3,
+            packed: 0xADE.into(),
+        },
+        v: test_account_2_rlp(),
+    };
+
+    test_state_trie(state_trie, insert)
+}
+
 #[test]
 fn mpt_insert_branch_replacing_empty_child() -> Result<()> {
     let children = std::array::from_fn(|_| Box::new(PartialTrie::Empty));