From cad0473e1d76dad91aa8c1ae215d928c481b4930 Mon Sep 17 00:00:00 2001 From: Daniel Lubarov Date: Sun, 9 Oct 2022 21:37:46 -0700 Subject: [PATCH] More MPT insert logic --- evm/src/cpu/kernel/aggregator.rs | 2 + evm/src/cpu/kernel/asm/mpt/insert.asm | 47 +---- .../cpu/kernel/asm/mpt/insert_extension.asm | 5 + evm/src/cpu/kernel/asm/mpt/insert_leaf.asm | 167 ++++++++++++++++++ evm/src/cpu/kernel/asm/mpt/util.asm | 10 +- evm/src/cpu/kernel/asm/util/basic_macros.asm | 6 + evm/src/cpu/kernel/tests/mpt/insert.rs | 24 ++- 7 files changed, 212 insertions(+), 49 deletions(-) create mode 100644 evm/src/cpu/kernel/asm/mpt/insert_extension.asm create mode 100644 evm/src/cpu/kernel/asm/mpt/insert_leaf.asm diff --git a/evm/src/cpu/kernel/aggregator.rs b/evm/src/cpu/kernel/aggregator.rs index 6fb2231e..032338b7 100644 --- a/evm/src/cpu/kernel/aggregator.rs +++ b/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"), diff --git a/evm/src/cpu/kernel/asm/mpt/insert.asm b/evm/src/cpu/kernel/asm/mpt/insert.asm index 65c18428..2830d376 100644 --- a/evm/src/cpu/kernel/asm/mpt/insert.asm +++ b/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 diff --git a/evm/src/cpu/kernel/asm/mpt/insert_extension.asm b/evm/src/cpu/kernel/asm/mpt/insert_extension.asm new file mode 100644 index 00000000..36458165 --- /dev/null +++ b/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 diff --git a/evm/src/cpu/kernel/asm/mpt/insert_leaf.asm b/evm/src/cpu/kernel/asm/mpt/insert_leaf.asm new file mode 100644 index 00000000..b82653f2 --- /dev/null +++ b/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 diff --git a/evm/src/cpu/kernel/asm/mpt/util.asm b/evm/src/cpu/kernel/asm/mpt/util.asm index a3ff7f38..0faa72f4 100644 --- a/evm/src/cpu/kernel/asm/mpt/util.asm +++ b/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 diff --git a/evm/src/cpu/kernel/asm/util/basic_macros.asm b/evm/src/cpu/kernel/asm/util/basic_macros.asm index bba2a2c1..02a2c807 100644 --- a/evm/src/cpu/kernel/asm/util/basic_macros.asm +++ b/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 diff --git a/evm/src/cpu/kernel/tests/mpt/insert.rs b/evm/src/cpu/kernel/tests/mpt/insert.rs index 11927d52..1ac7974f 100644 --- a/evm/src/cpu/kernel/tests/mpt/insert.rs +++ b/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(())