MPT fixes

evm/src/cpu/kernel/asm/memory/core.asm

@@ -363,7 +363,7 @@
 // Load a single value from kernel general memory.
 %macro mload_kernel_general_2(offset)
     PUSH $offset
-    %mload_kernel(@SEGMENT_KERNEL_GENERAL)
+    %mload_kernel(@SEGMENT_KERNEL_GENERAL_2)
     // stack: value
 %endmacro
 

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

@@ -64,7 +64,13 @@ maybe_hash_node:
     JUMP
 pack_small_rlp:
     // stack: result_ptr, result_len, retdest
-    PANIC // TODO: Return packed RLP
+    %stack (result_ptr, result_len)
+        -> (0, @SEGMENT_RLP_RAW, result_ptr, result_len,
+            after_packed_small_rlp, result_len)
+    %jump(mload_packing)
+after_packed_small_rlp:
+    %stack (result, result_len, retdest) -> (retdest, result, result_len)
+    JUMP
 
 // RLP encode the given trie node, and return an (pointer, length) pair
 // indicating where the data lives within @SEGMENT_RLP_RAW.
@@ -128,9 +134,19 @@ global encode_node_hash:
 // 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_i
-    %mload_kernel_general_2($i) // load result_i_len
-    %stack (result, result_len, 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:

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

@@ -70,6 +70,9 @@ load_mpt_branch:
     // stack: node_type, retdest
     POP
     // stack: retdest
+    // Save the offset of our 16 child pointers so we can write them later.
+    // Then advance out current trie pointer beyond them, so we can load the
+    // value and have it placed after our child pointers.
     %get_trie_data_size
     // stack: ptr_children, retdest
     DUP1 %add_const(16)
@@ -78,24 +81,20 @@ load_mpt_branch:
     // stack: ptr_children, retdest
     %load_leaf_value
 
-    // Save the current trie_data_size (which now points to the end of the leaf)
-    // for later, then have it point to the start of our 16 child pointers.
-    %get_trie_data_size
-    // stack: ptr_end_of_leaf, ptr_children, retdest
-    SWAP1
-    %set_trie_data_size
-    // stack: ptr_end_of_leaf, retdest
-
     // Load the 16 children.
     %rep 16
         %load_mpt_and_return_root_ptr
-        // stack: child_ptr, ptr_end_of_leaf, retdest
-        %append_to_trie_data
-        // stack: ptr_end_of_leaf, retdest
+        // stack: child_ptr, ptr_next_child, retdest
+        DUP2
+        // stack: ptr_next_child, child_ptr, ptr_next_child, retdest
+        %mstore_trie_data
+        // stack: ptr_next_child, retdest
+        %increment
+        // stack: ptr_next_child, retdest
     %endrep
 
-    %set_trie_data_size
-    // stack: retdest
+    // stack: ptr_next_child, retdest
+    POP
     JUMP
 
 load_mpt_extension:

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

@@ -7,8 +7,25 @@ use crate::cpu::kernel::interpreter::Interpreter;
 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]
-fn mpt_hash() -> Result<()> {
+fn mpt_hash_empty() -> Result<()> {
+    let trie_inputs = TrieInputs {
+        state_trie: Default::default(),
+        transactions_trie: Default::default(),
+        receipts_trie: Default::default(),
+        storage_tries: vec![],
+    };
+
+    test_state_trie(trie_inputs)
+}
+
+#[test]
+fn mpt_hash_leaf() -> Result<()> {
     let account = AccountRlp {
         nonce: U256::from(1111),
         balance: U256::from(2222),
@@ -17,8 +34,6 @@ fn mpt_hash() -> Result<()> {
     };
     let account_rlp = rlp::encode(&account);
 
-    // TODO: Try this more "advanced" trie.
-    // let state_trie = state_trie_ext_to_account_leaf(account_rlp.to_vec());
     let state_trie = PartialTrie::Leaf {
         nibbles: Nibbles {
             count: 3,
@@ -26,7 +41,6 @@ fn mpt_hash() -> Result<()> {
         },
         value: account_rlp.to_vec(),
     };
-    let state_trie_hash = state_trie.calc_hash();
 
     let trie_inputs = TrieInputs {
         state_trie,
@@ -35,10 +49,48 @@ fn mpt_hash() -> Result<()> {
         storage_tries: vec![],
     };
 
+    test_state_trie(trie_inputs)
+}
+
+#[test]
+fn mpt_hash_branch_to_account_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 leaf = PartialTrie::Leaf {
+        nibbles: Nibbles {
+            count: 3,
+            packed: 0xABC.into(),
+        },
+        value: account_rlp.to_vec(),
+    };
+    let mut children = std::array::from_fn(|_| Box::new(PartialTrie::Empty));
+    children[0] = Box::new(leaf);
+    let state_trie = PartialTrie::Branch {
+        children,
+        value: vec![],
+    };
+
+    let trie_inputs = TrieInputs {
+        state_trie,
+        transactions_trie: Default::default(),
+        receipts_trie: Default::default(),
+        storage_tries: vec![],
+    };
+
+    test_state_trie(trie_inputs)
+}
+
+fn test_state_trie(trie_inputs: TrieInputs) -> Result<()> {
     let load_all_mpts = KERNEL.global_labels["load_all_mpts"];
     let mpt_hash_state_trie = KERNEL.global_labels["mpt_hash_state_trie"];
 
-    let initial_stack = vec![0xdeadbeefu32.into()];
+    let initial_stack = vec![0xDEADBEEFu32.into()];
     let mut interpreter = Interpreter::new_with_kernel(load_all_mpts, initial_stack);
     interpreter.generation_state.mpt_prover_inputs = all_mpt_prover_inputs_reversed(&trie_inputs);
     interpreter.run()?;
@@ -49,9 +101,15 @@ fn mpt_hash() -> Result<()> {
     interpreter.push(0xDEADBEEFu32.into());
     interpreter.run()?;
 
-    assert_eq!(interpreter.stack().len(), 1);
+    assert_eq!(
+        interpreter.stack().len(),
+        1,
+        "Expected 1 item on stack, found {:?}",
+        interpreter.stack()
+    );
     let hash = H256::from_uint(&interpreter.stack()[0]);
-    assert_eq!(hash, state_trie_hash);
+    let expected_state_trie_hash = trie_inputs.state_trie.calc_hash();
+    assert_eq!(hash, expected_state_trie_hash);
 
     Ok(())
 }

evm/src/generation/mpt.rs

@@ -69,13 +69,18 @@ pub(crate) fn mpt_prover_inputs<F>(
         PartialTrie::Empty => {}
         PartialTrie::Hash(h) => prover_inputs.push(U256::from_big_endian(h.as_bytes())),
         PartialTrie::Branch { children, value } => {
-            for child in children {
-                mpt_prover_inputs(child, prover_inputs, parse_leaf);
-            }
+            if value.is_empty() {
+                // There's no value, so length=0.
+                prover_inputs.push(U256::zero());
+            } else {
                 let leaf = parse_leaf(value);
                 prover_inputs.push(leaf.len().into());
                 prover_inputs.extend(leaf);
             }
+            for child in children {
+                mpt_prover_inputs(child, prover_inputs, parse_leaf);
+            }
+        }
         PartialTrie::Extension { nibbles, child } => {
             prover_inputs.push(nibbles.count.into());
             prover_inputs.push(nibbles.packed);