From ed2aac3af392bb661b80167dd910fbb0776f9489 Mon Sep 17 00:00:00 2001 From: Daniel Lubarov Date: Thu, 6 Oct 2022 16:05:28 -0700 Subject: [PATCH] MPT fixes --- evm/src/cpu/kernel/asm/memory/core.asm | 2 +- evm/src/cpu/kernel/asm/mpt/hash.asm | 24 +++++++-- evm/src/cpu/kernel/asm/mpt/load.asm | 25 +++++---- evm/src/cpu/kernel/tests/mpt/hash.rs | 72 +++++++++++++++++++++++--- evm/src/generation/mpt.rs | 11 ++-- 5 files changed, 106 insertions(+), 28 deletions(-) diff --git a/evm/src/cpu/kernel/asm/memory/core.asm b/evm/src/cpu/kernel/asm/memory/core.asm index 12665e55..c2c19811 100644 --- a/evm/src/cpu/kernel/asm/memory/core.asm +++ b/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 diff --git a/evm/src/cpu/kernel/asm/mpt/hash.asm b/evm/src/cpu/kernel/asm/mpt/hash.asm index 840fb429..aefac0da 100644 --- a/evm/src/cpu/kernel/asm/mpt/hash.asm +++ b/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: diff --git a/evm/src/cpu/kernel/asm/mpt/load.asm b/evm/src/cpu/kernel/asm/mpt/load.asm index 2e218ab4..f072f202 100644 --- a/evm/src/cpu/kernel/asm/mpt/load.asm +++ b/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: diff --git a/evm/src/cpu/kernel/tests/mpt/hash.rs b/evm/src/cpu/kernel/tests/mpt/hash.rs index 4a79a5b6..db46e194 100644 --- a/evm/src/cpu/kernel/tests/mpt/hash.rs +++ b/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(()) } diff --git a/evm/src/generation/mpt.rs b/evm/src/generation/mpt.rs index d02e2d59..f6bc630d 100644 --- a/evm/src/generation/mpt.rs +++ b/evm/src/generation/mpt.rs @@ -69,12 +69,17 @@ pub(crate) fn mpt_prover_inputs( PartialTrie::Empty => {} PartialTrie::Hash(h) => prover_inputs.push(U256::from_big_endian(h.as_bytes())), PartialTrie::Branch { children, value } => { + 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); } - let leaf = parse_leaf(value); - prover_inputs.push(leaf.len().into()); - prover_inputs.extend(leaf); } PartialTrie::Extension { nibbles, child } => { prover_inputs.push(nibbles.count.into());