Merge pull request #738 from mir-protocol/mpt_read

Finish MPT read logic

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

@@ -3,18 +3,18 @@
 // Arguments:
 // - the virtual address of the trie to search in
 // - the key, as a U256
-// - the number of nibbles in the key
+// - the number of nibbles in the key (should start at 64)
 //
 // This function returns a pointer to the leaf, or 0 if the key is not found.
 
 global mpt_read:
-    // stack: node_ptr, key, nibbles, retdest
+    // stack: node_ptr, num_nibbles, key, retdest
     DUP1
     %mload_trie_data
-    // stack: node_type, node_ptr, key, nibbles, retdest
+    // stack: node_type, node_ptr, num_nibbles, key, retdest
     // Increment node_ptr, so it points to the node payload instead of its type.
     SWAP1 %add_const(1) SWAP1
-    // stack: node_type, node_payload_ptr, key, nibbles, retdest
+    // stack: node_type, node_payload_ptr, num_nibbles, key, retdest
 
     DUP1 %eq_const(@MPT_NODE_EMPTY)     %jumpi(mpt_read_empty)
     DUP1 %eq_const(@MPT_NODE_BRANCH)    %jumpi(mpt_read_branch)
@@ -27,43 +27,34 @@ global mpt_read:
 
 mpt_read_empty:
     // Return 0 to indicate that the value was not found.
-    %stack (node_type, node_payload_ptr, key, nibbles, retdest)
+    %stack (node_type, node_payload_ptr, num_nibbles, key, retdest)
         -> (retdest, 0)
     JUMP
 
 mpt_read_branch:
-    // stack: node_type, node_payload_ptr, key, nibbles, retdest
+    // stack: node_type, node_payload_ptr, num_nibbles, key, retdest
     POP
-    // stack: node_payload_ptr, key, nibbles, retdest
-    DUP3 // nibbles
+    // stack: node_payload_ptr, num_nibbles, key, retdest
+    DUP2 // num_nibbles
     ISZERO
-    // stack: nibbles == 0, node_payload_ptr, key, nibbles, retdest
+    // stack: num_nibbles == 0, node_payload_ptr, num_nibbles, key, retdest
     %jumpi(mpt_read_branch_end_of_key)
 
-    // stack: node_payload_ptr, key, nibbles, retdest
     // We have not reached the end of the key, so we descend to one of our children.
-    // Decrement nibbles, then compute current_nibble = (key >> (nibbles * 4)) & 0xF.
-    SWAP2
-    %sub_const(1)
-    // stack: nibbles, key, node_payload_ptr, retdest
-    DUP2 DUP2
-    // stack: nibbles, key, nibbles, key, node_payload_ptr, retdest
-    %mul_const(4)
-    // stack: nibbles * 4, key, nibbles, key, node_payload_ptr, retdest
-    SHR
-    // stack: key >> (nibbles * 4), nibbles, key, node_payload_ptr, retdest
-    %and_const(0xF)
-    // stack: current_nibble, nibbles, key, node_payload_ptr, retdest
-    %stack (current_nibble, nibbles, key, node_payload_ptr, retdest)
-        -> (current_nibble, node_payload_ptr, key, nibbles, retdest)
-    // child_ptr = load(node_payload_ptr + current_nibble)
-    ADD
-    %mload_trie_data
-    // stack: child_ptr, key, nibbles, retdest
+    // stack: node_payload_ptr, num_nibbles, key, retdest
+    %stack (node_payload_ptr, num_nibbles, key)
+        -> (num_nibbles, key, node_payload_ptr)
+    // stack: num_nibbles, key, node_payload_ptr, retdest
+    %split_first_nibble
+    %stack (first_nibble, num_nibbles, key, node_payload_ptr)
+        -> (node_payload_ptr, first_nibble, num_nibbles, key)
+    // child_ptr = load(node_payload_ptr + first_nibble)
+    ADD %mload_trie_data
+    // stack: child_ptr, num_nibbles, key, retdest
     %jump(mpt_read) // recurse
 
 mpt_read_branch_end_of_key:
-    %stack (node_payload_ptr, key, nibbles, retdest) -> (node_payload_ptr, retdest)
+    %stack (node_payload_ptr, num_nibbles, key, retdest) -> (node_payload_ptr, retdest)
     // stack: node_payload_ptr, retdest
     %add_const(16) // skip over the 16 child nodes
     // stack: leaf_ptr, retdest
@@ -71,11 +62,33 @@ mpt_read_branch_end_of_key:
     JUMP
 
 mpt_read_extension:
-    // stack: node_type, node_payload_ptr, key, nibbles, retdest
+    // stack: node_type, node_payload_ptr, num_nibbles, key, retdest
     POP
-    // stack: node_payload_ptr, key, nibbles, retdest
+    // stack: node_payload_ptr, num_nibbles, key, retdest
+    // TODO
 
 mpt_read_leaf:
-    // stack: node_type, node_payload_ptr, key, nibbles, retdest
+    // stack: node_type, node_payload_ptr, key, num_nibbles, retdest
     POP
-    // stack: node_payload_ptr, key, nibbles, retdest
+    // stack: node_payload_ptr, key, num_nibbles, retdest
+    DUP1 %mload_trie_data
+    // stack: node_num_nibbles, node_payload_ptr, key, num_nibbles, retdest
+    DUP2 %add_const(1) %mload_trie_data
+    // stack: node_key, node_num_nibbles, node_payload_ptr, key, num_nibbles, retdest
+    SWAP4
+    // stack: num_nibbles, node_num_nibbles, node_payload_ptr, key, node_key, retdest
+    EQ
+    %stack (num_nibbles_match, node_payload_ptr, key, node_key)
+        -> (key, node_key, num_nibbles_match, node_payload_ptr)
+    EQ
+    AND
+    // stack: keys_match && num_nibbles_match, node_payload_ptr, retdest
+    %jumpi(mpt_read_leaf_found)
+    %stack (node_payload_ptr, retdest) -> (retdest, 0)
+    JUMP
+mpt_read_leaf_found:
+    // stack: node_payload_ptr, retdest
+    %add_const(2) // The leaf data is located after num_nibbles and the key.
+    // stack: value_ptr, retdest
+    SWAP1
+    JUMP

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

@@ -35,3 +35,40 @@
     %mstore_trie_data
     // stack: (empty)
 %endmacro
+
+// Split off the first nibble from a key part. Roughly equivalent to
+// def split_first_nibble(num_nibbles, key):
+//     num_nibbles -= 1
+//     num_nibbles_x4 = num_nibbles * 4
+//     first_nibble = (key >> num_nibbles_x4) & 0xF
+//     key -= (first_nibble << num_nibbles_x4)
+//     return (first_nibble, num_nibbles, key)
+%macro split_first_nibble
+    // stack: num_nibbles, key
+    %sub_const(1) // num_nibbles -= 1
+    // stack: num_nibbles, key
+    DUP2
+    // stack: key, num_nibbles, key
+    DUP2 %mul_const(4)
+    // stack: num_nibbles_x4, key, num_nibbles, key
+    SHR
+    // stack: key >> num_nibbles_x4, num_nibbles, key
+    %and_const(0xF)
+    // stack: first_nibble, num_nibbles, key
+    DUP1
+    // stack: first_nibble, first_nibble, num_nibbles, key
+    DUP3 %mul_const(4)
+    // stack: num_nibbles_x4, first_nibble, first_nibble, num_nibbles, key
+    SHL
+    // stack: first_nibble << num_nibbles_x4, first_nibble, num_nibbles, key
+    DUP1
+    // stack: junk, first_nibble << num_nibbles_x4, first_nibble, num_nibbles, key
+    SWAP4
+    // stack: key, first_nibble << num_nibbles_x4, first_nibble, num_nibbles, junk
+    SUB
+    // stack: key, first_nibble, num_nibbles, junk
+    SWAP3
+    // stack: junk, first_nibble, num_nibbles, key
+    POP
+    // stack: first_nibble, num_nibbles, key
+%endmacro

evm/src/cpu/kernel/interpreter.rs

@@ -58,7 +58,7 @@ impl InterpreterMemory {
 pub struct Interpreter<'a> {
     kernel_mode: bool,
     jumpdests: Vec<usize>,
-    offset: usize,
+    pub(crate) offset: usize,
     context: usize,
     pub(crate) memory: InterpreterMemory,
     pub(crate) generation_state: GenerationState<F>,
@@ -117,11 +117,12 @@ impl<'a> Interpreter<'a> {
             prover_inputs_map: prover_inputs,
             context: 0,
             halt_offsets: vec![DEFAULT_HALT_OFFSET],
-            running: true,
+            running: false,
         }
     }
 
     pub(crate) fn run(&mut self) -> anyhow::Result<()> {
+        self.running = true;
         while self.running {
             self.run_opcode()?;
         }
@@ -185,7 +186,7 @@ impl<'a> Interpreter<'a> {
         &mut self.memory.context_memory[self.context].segments[Segment::Stack as usize].content
     }
 
-    fn push(&mut self, x: U256) {
+    pub(crate) fn push(&mut self, x: U256) {
         self.stack_mut().push(x);
     }
 

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

@@ -16,7 +16,7 @@ fn load_all_mpts() -> Result<()> {
     let storage_root = U256::from(3333);
     let code_hash = U256::from(4444);
 
-    let value_rlp = rlp::encode_list(&[nonce, balance, storage_root, code_hash]);
+    let account_rlp = rlp::encode_list(&[nonce, balance, storage_root, code_hash]);
 
     let trie_inputs = TrieInputs {
         state_trie: PartialTrie::Leaf {
@@ -24,7 +24,7 @@ fn load_all_mpts() -> Result<()> {
                 count: 2,
                 packed: 123.into(),
             },
-            value: value_rlp.to_vec(),
+            value: account_rlp.to_vec(),
         },
         transactions_trie: Default::default(),
         receipts_trie: Default::default(),
@@ -33,7 +33,6 @@ fn load_all_mpts() -> Result<()> {
 
     let load_all_mpts = KERNEL.global_labels["load_all_mpts"];
 
-    // Contract creation transaction.
     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);

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

@@ -1 +1,2 @@
 mod load;
+mod read;

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

@@ -0,0 +1,57 @@
+use anyhow::Result;
+use eth_trie_utils::partial_trie::{Nibbles, PartialTrie};
+use ethereum_types::U256;
+
+use crate::cpu::kernel::aggregator::KERNEL;
+use crate::cpu::kernel::global_metadata::GlobalMetadata;
+use crate::cpu::kernel::interpreter::Interpreter;
+use crate::generation::mpt::all_mpt_prover_inputs_reversed;
+use crate::generation::TrieInputs;
+
+#[test]
+fn mpt_read() -> Result<()> {
+    let nonce = U256::from(1111);
+    let balance = U256::from(2222);
+    let storage_root = U256::from(3333);
+    let code_hash = U256::from(4444);
+
+    let account = &[nonce, balance, storage_root, code_hash];
+    let account_rlp = rlp::encode_list(account);
+
+    let trie_inputs = TrieInputs {
+        state_trie: PartialTrie::Leaf {
+            nibbles: Nibbles {
+                count: 2,
+                packed: 123.into(),
+            },
+            value: account_rlp.to_vec(),
+        },
+        transactions_trie: Default::default(),
+        receipts_trie: Default::default(),
+        storage_tries: vec![],
+    };
+
+    let load_all_mpts = KERNEL.global_labels["load_all_mpts"];
+    let mpt_read = KERNEL.global_labels["mpt_read"];
+
+    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()?;
+    assert_eq!(interpreter.stack(), vec![]);
+
+    // Now, execute mpt_read on the state trie.
+    interpreter.offset = mpt_read;
+    interpreter.push(0xdeadbeefu32.into());
+    interpreter.push(2.into());
+    interpreter.push(123.into());
+    interpreter.push(interpreter.get_global_metadata_field(GlobalMetadata::StateTrieRoot));
+    interpreter.run()?;
+
+    assert_eq!(interpreter.stack().len(), 1);
+    let result_ptr = interpreter.stack()[0].as_usize();
+    let result = &interpreter.get_trie_data()[result_ptr..result_ptr + 4];
+    assert_eq!(result, account);
+
+    Ok(())
+}