Merge pull request #771 from mir-protocol/mpt_insert_4

More MPT insert logic

evm/src/cpu/kernel/aggregator.rs

@@ -44,6 +44,8 @@ pub(crate) fn combined_kernel() -> Kernel {
         include_str!("asm/mpt/hash_trie_specific.asm"),
         include_str!("asm/mpt/hex_prefix.asm"),
         include_str!("asm/mpt/insert.asm"),
+        include_str!("asm/mpt/insert_extension.asm"),
+        include_str!("asm/mpt/insert_leaf.asm"),
         include_str!("asm/mpt/insert_trie_specific.asm"),
         include_str!("asm/mpt/load.asm"),
         include_str!("asm/mpt/read.asm"),

evm/src/cpu/kernel/asm/mpt/insert.asm

@@ -98,55 +98,10 @@ mpt_insert_branch_nonterminal:
             child_ptr_ptr, updated_branch_ptr)
     %mload_trie_data // Deref child_ptr_ptr, giving child_ptr
     %jump(mpt_insert)
+
 mpt_insert_branch_nonterminal_after_recursion:
     // stack: updated_child_ptr, child_ptr_ptr, updated_branch_ptr, retdest
     SWAP1 %mstore_trie_data // Store the pointer to the updated child.
     // stack: updated_branch_ptr, retdest
     SWAP1
     JUMP
-
-mpt_insert_extension:
-    // stack: node_type, node_payload_ptr, insert_len, insert_key, value_ptr, retdest
-    POP
-    // stack: node_payload_ptr, insert_len, insert_key, value_ptr, retdest
-    PANIC // TODO
-
-mpt_insert_leaf:
-    // stack: node_type, node_payload_ptr, insert_len, insert_key, value_ptr, retdest
-    POP
-    // stack: node_payload_ptr, insert_len, insert_key, value_ptr, retdest
-    %stack (node_payload_ptr, insert_len, insert_key) -> (insert_len, insert_key, node_payload_ptr)
-    // stack: insert_len, insert_key, node_payload_ptr, value_ptr, retdest
-    DUP3 %increment %mload_trie_data
-    // stack: node_key, insert_len, insert_key, node_payload_ptr, value_ptr, retdest
-    DUP4 %mload_trie_data
-    // stack: node_len, node_key, insert_len, insert_key, node_payload_ptr, value_ptr, retdest
-
-    // If the keys match, i.e. node_len == insert_len && node_key == insert_key,
-    // then we're simply replacing the leaf node's value. Since this is a common
-    // case, it's best to detect it early. Calling %split_common_prefix could be
-    // expensive as leaf keys tend to be long.
-    DUP1 DUP4 EQ // node_len == insert_len
-    DUP3 DUP6 EQ // node_key == insert_key
-    MUL // Cheaper than AND
-    // stack: key_match, node_len, node_key, insert_len, insert_key, node_payload_ptr, value_ptr, retdest
-    %jumpi(mpt_insert_leaf_keys_match)
-
-    %split_common_prefix
-    PANIC // TODO
-
-mpt_insert_leaf_keys_match:
-    // The keys match exactly, so we simply create a new leaf node with the new value.xs
-    // stack: node_len, node_key, insert_len, insert_key, node_payload_ptr, value_ptr, retdest
-    %stack (node_len, node_key, insert_len, insert_key, node_payload_ptr, value_ptr)
-        -> (node_len, node_key, value_ptr)
-    // stack: common_len, common_key, value_ptr, retdest
-    %get_trie_data_size
-    // stack: updated_leaf_ptr, common_len, common_key, value_ptr, retdest
-    PUSH @MPT_NODE_LEAF %append_to_trie_data
-    SWAP1 %append_to_trie_data // append common_len
-    SWAP1 %append_to_trie_data // append common_key
-    SWAP1 %append_to_trie_data // append value_ptr
-    // stack: updated_leaf_ptr, retdestx
-    SWAP1
-    JUMP

evm/src/cpu/kernel/asm/mpt/insert_extension.asm

@@ -0,0 +1,5 @@
+global mpt_insert_extension:
+    // stack: node_type, node_payload_ptr, insert_len, insert_key, value_ptr, retdest
+    POP
+    // stack: node_payload_ptr, insert_len, insert_key, value_ptr, retdest
+    PANIC // TODO

evm/src/cpu/kernel/asm/mpt/insert_leaf.asm

@@ -0,0 +1,167 @@
+// The high-level logic can be expressed with the following pseudocode:
+//
+// if node_len == insert_len && node_key == insert_key:
+//     return Leaf[node_key, insert_value]
+//
+// common_len, common_key, node_len, node_key, insert_len, insert_key =
+//     consume_common_prefix(node_len, node_key, insert_len, insert_key)
+//
+// branch = [MPT_TYPE_BRANCH] + [0] * 17
+//
+// if node_len > 0:
+//     node_key_first, node_len, node_key = split_first_nibble(node_len, node_key)
+//     branch[node_key_first + 1] = Leaf[node_len, node_key, node_value]
+// else:
+//     branch[17] = node_value
+//
+// if insert_len > 0:
+//     insert_key_first, insert_len, insert_key = split_first_nibble(insert_len, insert_key)
+//     branch[insert_key_first + 1] = Leaf[insert_len, insert_key, insert_value]
+// else:
+//     branch[17] = insert_value
+//
+// if common_len > 0:
+//     return Extension[common_key, branch]
+// else:
+//     return branch
+
+global mpt_insert_leaf:
+    // stack: node_type, node_payload_ptr, insert_len, insert_key, insert_value_ptr, retdest
+    POP
+    // stack: node_payload_ptr, insert_len, insert_key, insert_value_ptr, retdest
+    %stack (node_payload_ptr, insert_len, insert_key) -> (insert_len, insert_key, node_payload_ptr)
+    // stack: insert_len, insert_key, node_payload_ptr, insert_value_ptr, retdest
+    DUP3 %increment %mload_trie_data
+    // stack: node_key, insert_len, insert_key, node_payload_ptr, insert_value_ptr, retdest
+    DUP4 %mload_trie_data
+    // stack: node_len, node_key, insert_len, insert_key, node_payload_ptr, insert_value_ptr, retdest
+
+    // If the keys match, i.e. node_len == insert_len && node_key == insert_key,
+    // then we're simply replacing the leaf node's value. Since this is a common
+    // case, it's best to detect it early. Calling %split_common_prefix could be
+    // expensive as leaf keys tend to be long.
+    DUP1 DUP4 EQ // node_len == insert_len
+    DUP3 DUP6 EQ // node_key == insert_key
+    MUL // Cheaper than AND
+    // stack: keys_match, node_len, node_key, insert_len, insert_key, node_payload_ptr, insert_value_ptr, retdest
+    %jumpi(keys_match)
+
+    // Replace node_payload_ptr with node_value, which is node_payload[2].
+    // stack: node_len, node_key, insert_len, insert_key, node_payload_ptr, insert_value_ptr, retdest
+    SWAP4
+    %add_const(2)
+    %mload_trie_data
+    SWAP4
+    // stack: node_len, node_key, insert_len, insert_key, node_value_ptr, insert_value_ptr, retdest
+
+    // Split off any common prefix between the node key and the inserted key.
+    %split_common_prefix
+    // stack: common_len, common_key, node_len, node_key, insert_len, insert_key, node_value_ptr, insert_value_ptr, retdest
+
+    // For the remaining cases, we will need a new branch node since the two keys diverge.
+    // We may also need an extension node above it (if common_len > 0); we will handle that later.
+    // For now, we allocate the branch node, initially with no children or value.
+    %get_trie_data_size
+    PUSH @MPT_NODE_BRANCH %append_to_trie_data
+    %rep 17
+        PUSH 0 %append_to_trie_data
+    %endrep
+    // stack: branch_ptr, common_len, common_key, node_len, node_key, insert_len, insert_key, node_value_ptr, insert_value_ptr, retdest
+
+    // Here, we branch based on whether each key continues beyond the common
+    // prefix, starting with the node key.
+    DUP4 // node_len
+    %jumpi(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
+    // branch[17] = node_value_ptr
+    DUP8 // node_value_ptr
+    DUP2 // branch_ptr
+    %add_const(17)
+    %mstore_trie_data
+
+finished_processing_node_value:
+    DUP6 // insert_len
+    %jumpi(insert_key_continues)
+
+    // stack: branch_ptr, common_len, common_key, node_len, node_key, insert_len, insert_key, node_value_ptr, insert_value_ptr, retdest
+    // branch[17] = insert_value_ptr
+    DUP9 // insert_value_ptr
+    DUP2 // branch_ptr
+    %add_const(17)
+    %mstore_trie_data
+
+finished_processing_insert_value:
+    // stack: branch_ptr, common_len, common_key, node_len, node_key, insert_len, insert_key, node_value_ptr, insert_value_ptr, retdest
+    // If common_len > 0, we need to add an extension node.
+    DUP2 %jumpi(extension_for_common_key)
+    // Otherwise, we simply return our branch node.
+    SWAP8
+    %pop8
+    // stack: branch_ptr, retdest
+    SWAP1
+    JUMP
+
+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
+
+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
+    // branch[node_key_first + 1] = Leaf[node_len, node_key, node_value]
+    DUP5 DUP5
+    // stack: node_len, node_key, branch_ptr, ...
+    %split_first_nibble
+    // stack: node_key_first, node_len, node_key, branch_ptr, ...
+    %get_trie_data_size
+    // stack: leaf_ptr, node_key_first, node_len, node_key, branch_ptr, ...
+    SWAP1
+    DUP5 // branch_ptr
+    %increment // Skip over node type field
+    ADD // Add node_key_first
+    %mstore_trie_data
+    // stack: node_len, node_key, branch_ptr, ...
+    PUSH @MPT_NODE_LEAF %append_to_trie_data
+    %append_to_trie_data // Append node_len to our leaf node
+    %append_to_trie_data // Append node_key to our leaf node
+    // stack: branch_ptr, common_len, common_key, node_len, node_key, insert_len, insert_key, node_value_ptr, insert_value_ptr, retdest
+    DUP8 %append_to_trie_data // Append node_value_ptr to our leaf node
+    %jump(finished_processing_node_value)
+
+insert_key_continues:
+    // stack: branch_ptr, common_len, common_key, node_len, node_key, insert_len, insert_key, node_value_ptr, insert_value_ptr, retdest
+    // branch[insert_key_first + 1] = Leaf[insert_len, insert_key, insert_value]
+    DUP7 DUP7
+    // stack: insert_len, insert_key, branch_ptr, ...
+    %split_first_nibble
+    // stack: insert_key_first, insert_len, insert_key, branch_ptr, ...
+    %get_trie_data_size
+    // stack: leaf_ptr, insert_key_first, insert_len, insert_key, branch_ptr, ...
+    SWAP1
+    DUP5 // branch_ptr
+    %increment // Skip over node type field
+    ADD // Add insert_key_first
+    %mstore_trie_data
+    // stack: insert_len, insert_key, branch_ptr, ...
+    PUSH @MPT_NODE_LEAF %append_to_trie_data
+    %append_to_trie_data // Append insert_len to our leaf node
+    %append_to_trie_data // Append insert_key to our leaf node
+    // stack: branch_ptr, common_len, common_key, node_len, node_key, insert_len, insert_key, node_value_ptr, insert_value_ptr, retdest
+    DUP9 %append_to_trie_data // Append insert_value_ptr to our leaf node
+    %jump(finished_processing_insert_value)
+
+keys_match:
+    // The keys match exactly, so we simply create a new leaf node with the new value.xs
+    // stack: node_len, node_key, insert_len, insert_key, node_payload_ptr, insert_value_ptr, retdest
+    %stack (node_len, node_key, insert_len, insert_key, node_payload_ptr, insert_value_ptr)
+        -> (node_len, node_key, insert_value_ptr)
+    // stack: common_len, common_key, insert_value_ptr, retdest
+    %get_trie_data_size
+    // stack: updated_leaf_ptr, common_len, common_key, insert_value_ptr, retdest
+    PUSH @MPT_NODE_LEAF %append_to_trie_data
+    SWAP1 %append_to_trie_data // Append common_len to our leaf node
+    SWAP1 %append_to_trie_data // Append common_key to our leaf node
+    SWAP1 %append_to_trie_data // Append insert_value_ptr to our leaf node
+    // stack: updated_leaf_ptr, retdestx
+    SWAP1
+    JUMP

evm/src/cpu/kernel/asm/mpt/util.asm

@@ -73,7 +73,12 @@
     // stack: first_nibble, num_nibbles, key
 %endmacro
 
-// Split off the common prefix among two key parts. Roughly equivalent to
+// Split off the common prefix among two key parts.
+//
+// Pre stack: len_1, key_1, len_2, key_2
+// Post stack: len_common, key_common, len_1, key_1, len_2, key_2
+//
+// Roughly equivalent to
 // def split_common_prefix(len_1, key_1, len_2, key_2):
 //     bits_1 = len_1 * 4
 //     bits_2 = len_2 * 4
@@ -155,5 +160,8 @@
     // stack: first_nib_1, first_nib_2, len_common, key_common, bits_1, key_1, bits_2, key_2
     %pop2
 %%return:
+    // stack: len_common, key_common, bits_1, key_1, bits_2, key_2
+    SWAP2 %div_const(4) SWAP2 // bits_1 -> len_1 (in nibbles)
+    SWAP4 %div_const(4) SWAP4 // bits_2 -> len_2 (in nibbles)
     // stack: len_common, key_common, len_1, key_1, len_2, key_2
 %endmacro

evm/src/cpu/kernel/asm/util/basic_macros.asm

@@ -44,6 +44,12 @@
     %endrep
 %endmacro
 
+%macro pop8
+    %rep 8
+        POP
+    %endrep
+%endmacro
+
 %macro and_const(c)
     // stack: input, ...
     PUSH $c

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

@@ -40,6 +40,26 @@ fn mpt_insert_leaf_same_key() -> Result<()> {
     test_state_trie(state_trie, insert)
 }
 
+#[test]
+fn mpt_insert_leaf_nonoverlapping_key() -> 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: 0x123.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));
@@ -76,7 +96,6 @@ fn mpt_insert_extension_to_leaf_same_key() -> Result<()> {
 }
 
 #[test]
-#[ignore] // TODO: Enable when mpt_insert_leaf is done.
 fn mpt_insert_branch_to_leaf_same_key() -> Result<()> {
     let leaf = PartialTrie::Leaf {
         nibbles: Nibbles {
@@ -161,7 +180,8 @@ fn test_state_trie(state_trie: PartialTrie, insert: InsertEntry) -> Result<()> {
     );
     let hash = H256::from_uint(&interpreter.stack()[0]);
 
-    let expected_state_trie_hash = apply_insert(state_trie, insert).calc_hash();
+    let updated_trie = apply_insert(state_trie, insert);
+    let expected_state_trie_hash = updated_trie.calc_hash();
     assert_eq!(hash, expected_state_trie_hash);
 
     Ok(())