Hash MPT extension nodes

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

@@ -113,46 +113,6 @@ global encode_node_hash:
     %stack (hash, encode_value, retdest) -> (retdest, hash, 32)
     JUMP
 
-// Part of the encode_node_branch function. Encodes the i'th child.
-// Stores the result in SEGMENT_KERNEL_GENERAL[i], and its length in
-// SEGMENT_KERNEL_GENERAL_2[i].
-%macro encode_child(i)
-    // stack: node_payload_ptr, encode_value, retdest
-    PUSH %%after_encode
-    DUP3 DUP3
-    // stack: node_payload_ptr, encode_value, %%after_encode, node_payload_ptr, encode_value, retdest
-    %add_const($i) %mload_trie_data
-    // stack: child_i_ptr, encode_value, %%after_encode, node_payload_ptr, encode_value, retdest
-    %jump(encode_or_hash_node)
-%%after_encode:
-    // stack: result, result_len, node_payload_ptr, encode_value, retdest
-    %mstore_kernel_general($i)
-    %mstore_kernel_general_2($i)
-    // stack: node_payload_ptr, encode_value, retdest
-%endmacro
-
-// Part of the encode_node_branch function. Appends the i'th child's RLP.
-%macro append_child(i)
-    // stack: rlp_pos, node_payload_ptr, encode_value, retdest
-    %mload_kernel_general($i) // load result
-    %mload_kernel_general_2($i) // load result_len
-    // stack: result_len, result, rlp_pos, node_payload_ptr, encode_value, retdest
-    // If result_len != 32, result is raw RLP, with an appropriate RLP prefix already.
-    DUP1 %eq_const(32) ISZERO %jumpi(%%unpack)
-    // Otherwise, result is a hash, and we need to add the prefix 0x80 + 32 = 160.
-    // stack: result_len, result, rlp_pos, node_payload_ptr, encode_value, retdest
-    PUSH 160
-    DUP4 // rlp_pos
-    %mstore_rlp
-    SWAP2 %increment SWAP2 // rlp_pos += 1
-%%unpack:
-    %stack (result_len, result, rlp_pos, node_payload_ptr, encode_value, retdest)
-        -> (rlp_pos, result, result_len, %%after_unpacking, node_payload_ptr, encode_value, retdest)
-    %jump(mstore_unpacking_rlp)
-%%after_unpacking:
-    // stack: rlp_pos', node_payload_ptr, encode_value, retdest
-%endmacro
-
 encode_node_branch:
     // stack: node_type, node_payload_ptr, encode_value, retdest
     POP
@@ -202,11 +162,83 @@ encode_node_branch_prepend_prefix:
     %stack (start_pos, rlp_len, retdest) -> (retdest, start_pos, rlp_len)
     JUMP
 
+// Part of the encode_node_branch function. Encodes the i'th child.
+// Stores the result in SEGMENT_KERNEL_GENERAL[i], and its length in
+// SEGMENT_KERNEL_GENERAL_2[i].
+%macro encode_child(i)
+    // stack: node_payload_ptr, encode_value, retdest
+    PUSH %%after_encode
+    DUP3 DUP3
+    // stack: node_payload_ptr, encode_value, %%after_encode, node_payload_ptr, encode_value, retdest
+    %add_const($i) %mload_trie_data
+    // stack: child_i_ptr, encode_value, %%after_encode, node_payload_ptr, encode_value, retdest
+    %jump(encode_or_hash_node)
+%%after_encode:
+    // stack: result, result_len, node_payload_ptr, encode_value, retdest
+    %mstore_kernel_general($i)
+    %mstore_kernel_general_2($i)
+    // stack: node_payload_ptr, encode_value, retdest
+%endmacro
+
+// Part of the encode_node_branch function. Appends the i'th child's RLP.
+%macro append_child(i)
+    // stack: rlp_pos, node_payload_ptr, encode_value, retdest
+    %mload_kernel_general($i) // load result
+    %mload_kernel_general_2($i) // load result_len
+    // stack: result_len, result, rlp_pos, node_payload_ptr, encode_value, retdest
+    // If result_len != 32, result is raw RLP, with an appropriate RLP prefix already.
+    DUP1 %sub_const(32) %jumpi(%%unpack)
+    // Otherwise, result is a hash, and we need to add the prefix 0x80 + 32 = 160.
+    // stack: result_len, result, rlp_pos, node_payload_ptr, encode_value, retdest
+    PUSH 160
+    DUP4 // rlp_pos
+    %mstore_rlp
+    SWAP2 %increment SWAP2 // rlp_pos += 1
+%%unpack:
+    %stack (result_len, result, rlp_pos, node_payload_ptr, encode_value, retdest)
+        -> (rlp_pos, result, result_len, %%after_unpacking, node_payload_ptr, encode_value, retdest)
+    %jump(mstore_unpacking_rlp)
+%%after_unpacking:
+    // stack: rlp_pos', node_payload_ptr, encode_value, retdest
+%endmacro
+
 encode_node_extension:
     // stack: node_type, node_payload_ptr, encode_value, retdest
-    POP
-    // stack: node_payload_ptr, encode_value, retdest
-    PANIC // TODO
+    %stack (node_type, node_payload_ptr, encode_value)
+        -> (node_payload_ptr, encode_value, encode_node_extension_after_encode_child, node_payload_ptr)
+    %add_const(2) %mload_trie_data
+    // stack: child_ptr, encode_value, encode_node_extension_after_encode_child, node_payload_ptr, retdest
+    %jump(encode_or_hash_node)
+encode_node_extension_after_encode_child:
+    // stack: result, result_len, node_payload_ptr, retdest
+    PUSH encode_node_extension_after_hex_prefix // retdest
+    PUSH 0 // terminated
+    // stack: terminated, encode_node_extension_after_hex_prefix, result, result_len, node_payload_ptr, retdest
+    DUP5 %add_const(1) %mload_trie_data // Load the packed_nibbles field, which is at index 1.
+    // stack: packed_nibbles, terminated, encode_node_extension_after_hex_prefix, result, result_len, node_payload_ptr, retdest
+    DUP6 %mload_trie_data // Load the num_nibbles field, which is at index 0.
+    // stack: num_nibbles, packed_nibbles, terminated, encode_node_extension_after_hex_prefix, result, result_len, node_payload_ptr, retdest
+    PUSH 9 // We start at 9 to leave room to prepend the largest possible RLP list header.
+    // stack: rlp_start, num_nibbles, packed_nibbles, terminated, encode_node_extension_after_hex_prefix, result, result_len, node_payload_ptr, retdest
+    %jump(hex_prefix_rlp)
+encode_node_extension_after_hex_prefix:
+    // stack: rlp_pos, result, result_len, node_payload_ptr, retdest
+    // If result_len != 32, result is raw RLP, with an appropriate RLP prefix already.
+    DUP3 %sub_const(32) %jumpi(encode_node_extension_unpack)
+    // Otherwise, result is a hash, and we need to add the prefix 0x80 + 32 = 160.
+    PUSH 160
+    DUP2 // rlp_pos
+    %mstore_rlp
+    %increment // rlp_pos += 1
+encode_node_extension_unpack:
+    %stack (rlp_pos, result, result_len, node_payload_ptr)
+        -> (rlp_pos, result, result_len, encode_node_extension_after_unpacking)
+    %jump(mstore_unpacking_rlp)
+encode_node_extension_after_unpacking:
+    // stack: rlp_end_pos, retdest
+    %prepend_rlp_list_prefix
+    %stack (rlp_start_pos, rlp_len, retdest) -> (retdest, rlp_start_pos, rlp_len)
+    JUMP
 
 encode_node_leaf:
     // stack: node_type, node_payload_ptr, encode_value, retdest

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

@@ -4,12 +4,10 @@ use ethereum_types::{BigEndianHash, H256, U256};
 
 use crate::cpu::kernel::aggregator::KERNEL;
 use crate::cpu::kernel::interpreter::Interpreter;
+use crate::cpu::kernel::tests::mpt::extension_to_leaf;
 use crate::generation::mpt::{all_mpt_prover_inputs_reversed, AccountRlp};
 use crate::generation::TrieInputs;
 
-// TODO: Try this more "advanced" trie.
-// let state_trie = state_trie_ext_to_account_leaf(account_rlp.to_vec());
-
 // TODO: Test with short leaf. Might need to be a storage trie.
 
 #[test]
@@ -53,7 +51,29 @@ fn mpt_hash_leaf() -> Result<()> {
 }
 
 #[test]
-fn mpt_hash_branch_to_account_leaf() -> Result<()> {
+fn mpt_hash_extension_to_leaf() -> Result<()> {
+    let account = AccountRlp {
+        nonce: U256::from(1111),
+        balance: U256::from(2222),
+        storage_root: H256::from_uint(&U256::from(3333)),
+        code_hash: H256::from_uint(&U256::from(4444)),
+    };
+    let account_rlp = rlp::encode(&account);
+
+    let state_trie = extension_to_leaf(account_rlp.to_vec());
+
+    let trie_inputs = TrieInputs {
+        state_trie,
+        transactions_trie: Default::default(),
+        receipts_trie: Default::default(),
+        storage_tries: vec![],
+    };
+
+    test_state_trie(trie_inputs)
+}
+
+#[test]
+fn mpt_hash_branch_to_leaf() -> Result<()> {
     let account = AccountRlp {
         nonce: U256::from(1111),
         balance: U256::from(2222),

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

@@ -5,7 +5,7 @@ use crate::cpu::kernel::aggregator::KERNEL;
 use crate::cpu::kernel::constants::global_metadata::GlobalMetadata;
 use crate::cpu::kernel::constants::trie_type::PartialTrieType;
 use crate::cpu::kernel::interpreter::Interpreter;
-use crate::cpu::kernel::tests::mpt::state_trie_ext_to_account_leaf;
+use crate::cpu::kernel::tests::mpt::extension_to_leaf;
 use crate::generation::mpt::{all_mpt_prover_inputs_reversed, AccountRlp};
 use crate::generation::TrieInputs;
 
@@ -20,7 +20,7 @@ fn load_all_mpts() -> Result<()> {
     let account_rlp = rlp::encode(&account);
 
     let trie_inputs = TrieInputs {
-        state_trie: state_trie_ext_to_account_leaf(account_rlp.to_vec()),
+        state_trie: extension_to_leaf(account_rlp.to_vec()),
         transactions_trie: Default::default(),
         receipts_trie: Default::default(),
         storage_tries: vec![],

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

@@ -6,7 +6,7 @@ mod load;
 mod read;
 
 /// A `PartialTrie` where an extension node leads to a leaf node containing an account.
-pub(crate) fn state_trie_ext_to_account_leaf(value: Vec<u8>) -> PartialTrie {
+pub(crate) fn extension_to_leaf(value: Vec<u8>) -> PartialTrie {
     PartialTrie::Extension {
         nibbles: Nibbles {
             count: 3,

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

@@ -4,7 +4,7 @@ use ethereum_types::{BigEndianHash, H256, U256};
 use crate::cpu::kernel::aggregator::KERNEL;
 use crate::cpu::kernel::constants::global_metadata::GlobalMetadata;
 use crate::cpu::kernel::interpreter::Interpreter;
-use crate::cpu::kernel::tests::mpt::state_trie_ext_to_account_leaf;
+use crate::cpu::kernel::tests::mpt::extension_to_leaf;
 use crate::generation::mpt::{all_mpt_prover_inputs_reversed, AccountRlp};
 use crate::generation::TrieInputs;
 
@@ -19,7 +19,7 @@ fn mpt_read() -> Result<()> {
     let account_rlp = rlp::encode(&account);
 
     let trie_inputs = TrieInputs {
-        state_trie: state_trie_ext_to_account_leaf(account_rlp.to_vec()),
+        state_trie: extension_to_leaf(account_rlp.to_vec()),
         transactions_trie: Default::default(),
         receipts_trie: Default::default(),
         storage_tries: vec![],