MPT logic to hash branch nodes

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

@@ -257,13 +257,20 @@
     // stack: value
 %endmacro
 
-// Load a single byte from kernel general memory.
+// Load a single value from kernel general memory.
 %macro mload_kernel_general
     // stack: offset
     %mload_kernel(@SEGMENT_KERNEL_GENERAL)
     // stack: value
 %endmacro
 
+// Load a single value from kernel general memory.
+%macro mload_kernel_general(offset)
+    PUSH $offset
+    %mload_kernel(@SEGMENT_KERNEL_GENERAL)
+    // stack: value
+%endmacro
+
 // Load a big-endian u32, consisting of 4 bytes (c_3, c_2, c_1, c_0),
 // from kernel general memory.
 %macro mload_kernel_general_u32
@@ -318,7 +325,7 @@
 %macro mstore_kernel_general
     // stack: offset, value 
     %mstore_kernel(@SEGMENT_KERNEL_GENERAL)
-    // stack: 
+    // stack: (empty)
 %endmacro
 
 %macro mstore_kernel_general(offset)
@@ -326,7 +333,7 @@
     PUSH $offset
     // stack: offset, value 
     %mstore_kernel_general
-    // stack: 
+    // stack: (empty)
 %endmacro
 
 // Store a big-endian u32, consisting of 4 bytes (c_3, c_2, c_1, c_0),
@@ -343,5 +350,33 @@
     // stack: x, i
     SWAP1
     %mstore_kernel_general
-    // stack:
+    // stack: (empty)
+%endmacro
+
+// Load a single value from kernel general 2 memory.
+%macro mload_kernel_general_2
+    // stack: offset
+    %mload_kernel(@SEGMENT_KERNEL_GENERAL_2)
+    // stack: value
+%endmacro
+
+// Load a single value from kernel general memory.
+%macro mload_kernel_general_2(offset)
+    PUSH $offset
+    %mload_kernel(@SEGMENT_KERNEL_GENERAL)
+    // stack: value
+%endmacro
+
+%macro mstore_kernel_general_2
+    // stack: offset, value
+    %mstore_kernel(@SEGMENT_KERNEL_GENERAL_2)
+    // stack: (empty)
+%endmacro
+
+%macro mstore_kernel_general_2(offset)
+    // stack:         value
+    PUSH $offset
+    // stack: offset, value
+    %mstore_kernel_general_2
+    // stack: (empty)
 %endmacro

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

@@ -107,18 +107,84 @@ 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_i
+    %mload_kernel_general_2($i) // load result_i_len
+    %stack (result, result_len, 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
     // stack: node_payload_ptr, encode_value, retdest
+
+    // We will call encode_or_hash_node on each child. For the i'th child, we
+    // will store the result in SEGMENT_KERNEL_GENERAL[i], and its length in
+    // SEGMENT_KERNEL_GENERAL_2[i].
+    %encode_child(0)  %encode_child(1)  %encode_child(2)  %encode_child(3)
+    %encode_child(4)  %encode_child(5)  %encode_child(6)  %encode_child(7)
+    %encode_child(8)  %encode_child(9)  %encode_child(10) %encode_child(11)
+    %encode_child(12) %encode_child(13) %encode_child(14) %encode_child(15)
+    // stack: node_payload_ptr, encode_value, retdest
+
+    // Now, append each child to our RLP tape.
     PUSH 9 // rlp_pos; we start at 9 to leave room to prepend a list prefix
-    %rep 16
-        // stack: rlp_pos, node_child_ptr, encode_value, retdest
-        // TODO: Append encode_or_hash_node(child) to our RLP. Do all encode_or_hash_node calls first to avoid clobbering.
-        SWAP1 %increment SWAP1 // node_child_ptr += 1
-    %endrep
-    // stack: node_value_ptr, encode_value, retdest
-    PANIC // TODO
+    %append_child(0)  %append_child(1)  %append_child(2)  %append_child(3)
+    %append_child(4)  %append_child(5)  %append_child(6)  %append_child(7)
+    %append_child(8)  %append_child(9)  %append_child(10) %append_child(11)
+    %append_child(12) %append_child(13) %append_child(14) %append_child(15)
+
+    // stack: rlp_pos', node_payload_ptr, encode_value, retdest
+    SWAP1
+    %add_const(16)
+    // stack: value_len_ptr, rlp_pos', encode_value, retdest
+    DUP1 %mload_trie_data
+    // stack: value_len, value_len_ptr, rlp_pos', encode_value, retdest
+    %jumpi(encode_node_branch_with_value)
+    // No value; append the empty string (0x80).
+    // stack: value_len_ptr, rlp_pos', encode_value, retdest
+    %stack (value_len_ptr, rlp_pos, encode_value) -> (rlp_pos, 0x80, rlp_pos)
+    %mstore_rlp
+    // stack: rlp_pos', retdest
+    %increment
+    // stack: rlp_pos'', retdest
+    %jump(encode_node_branch_prepend_prefix)
+encode_node_branch_with_value:
+    // stack: value_len_ptr, rlp_pos', encode_value, retdest
+    %increment
+    // stack: value_ptr, rlp_pos', encode_value, retdest
+    %stack (value_ptr, rlp_pos, encode_value)
+        -> (encode_value, rlp_pos, value_ptr, encode_node_branch_prepend_prefix)
+    JUMP // call encode_value
+encode_node_branch_prepend_prefix:
+    // stack: rlp_pos'', retdest
+    %prepend_rlp_list_prefix
+    %stack (start_pos, rlp_len, retdest) -> (retdest, start_pos, rlp_len)
+    JUMP
 
 encode_node_extension:
     // stack: node_type, node_payload_ptr, encode_value, retdest
@@ -143,7 +209,7 @@ encode_node_leaf:
 encode_node_leaf_after_hex_prefix:
     // stack: rlp_pos, node_payload_ptr, encode_value, retdest
     SWAP1
-    %add_const(2) // The value starts at index 2.
+    %add_const(3) // The value starts at index 3, after num_nibbles, packed_nibbles, and value_len.
     // stack: value_ptr, rlp_pos, encode_value, retdest
     %stack (value_ptr, rlp_pos, encode_value, retdest)
         -> (encode_value, rlp_pos, value_ptr, encode_node_leaf_after_encode_value, retdest)

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

@@ -159,10 +159,11 @@ load_mpt_digest:
 
 // Load a leaf from prover input, and append it to trie data.
 %macro load_leaf_value
-    // TODO: Need to store leaf len, or at least a has_leaf flag for branch nodes.
     // stack: (empty)
     PROVER_INPUT(mpt)
     // stack: leaf_len
+    DUP1 %append_to_trie_data
+    // stack: leaf_len
 %%loop:
     DUP1 ISZERO
     // stack: leaf_len == 0, leaf_len

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

@@ -75,7 +75,16 @@ mpt_read_branch_end_of_key:
     %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
+    // stack: value_len_ptr, retdest
+    DUP1 %mload_trie_data
+    // stack: value_len, value_len_ptr, retdest
+    %jumpi(mpt_read_branch_found_value)
+    // This branch node contains no value, so return null.
+    %stack (value_len_ptr, retdest) -> (retdest, 0)
+mpt_read_branch_found_value:
+    // stack: value_len_ptr, retdest
+    %increment
+    // stack: value_ptr, retdest
     SWAP1
     JUMP
 
@@ -138,7 +147,7 @@ mpt_read_leaf:
     JUMP
 mpt_read_leaf_found:
     // stack: node_payload_ptr, retdest
-    %add_const(2) // The leaf data is located after num_nibbles and the key.
+    %add_const(3) // The value is located after num_nibbles, the key, and the value length.
     // stack: value_ptr, retdest
     SWAP1
     JUMP

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

@@ -1,7 +1,7 @@
 global ripemd_storage: // starts by initializing buffer
     // stack: i [init: 64]
     %store_zeros(64, ripemd_storage)
-    // stack:
+    // stack: (empty)
     %jump(store_size)
 
 store_size:

evm/src/cpu/kernel/asm/rlp/encode.asm

@@ -248,7 +248,7 @@ prepend_rlp_list_prefix_big_done_writing_len:
 
 // Convenience macro to call prepend_rlp_list_prefix and return where we left off.
 %macro prepend_rlp_list_prefix
-    %stack (start_pos) -> (start_pos, %%after)
+    %stack (end_pos) -> (end_pos, %%after)
     %jump(prepend_rlp_list_prefix)
 %%after:
 %endmacro
@@ -299,7 +299,7 @@ prepend_rlp_list_prefix_big_done_writing_len:
 // Like mstore_unpacking, but specifically for the RLP segment.
 // Pre stack: offset, value, len, retdest
 // Post stack: offset'
-mstore_unpacking_rlp:
+global mstore_unpacking_rlp:
     // stack: offset, value, len, retdest
     PUSH @SEGMENT_RLP_RAW
     PUSH 0 // context

evm/src/cpu/kernel/optimizer.rs

@@ -22,6 +22,7 @@ pub(crate) fn optimize_asm(code: &mut Vec<Item>) {
 /// A single optimization pass.
 fn optimize_asm_once(code: &mut Vec<Item>) {
     constant_propagation(code);
+    identity_operations(code);
     no_op_jumps(code);
     remove_swapped_pushes(code);
     remove_swaps_commutative(code);
@@ -75,6 +76,25 @@ fn constant_propagation(code: &mut Vec<Item>) {
     });
 }
 
+/// Remove identity operations, e.g. `[PUSH 1, MUL] -> []`.
+fn identity_operations(code: &mut Vec<Item>) {
+    let zero = U256::zero();
+    let one = U256::one();
+    replace_windows(code, |window| {
+        if let [Push(Literal(x)), StandardOp(op)] = window {
+            match op.as_str() {
+                "ADD" => (x == zero).then_some(vec![]),
+                "MUL" => (x == one).then_some(vec![]),
+                "OR" => (x == zero).then_some(vec![]),
+                "XOR" => (x == zero).then_some(vec![]),
+                _ => None,
+            }
+        } else {
+            None
+        }
+    })
+}
+
 /// Remove no-op jumps: `[PUSH label, JUMP, label:] -> [label:]`.
 fn no_op_jumps(code: &mut Vec<Item>) {
     replace_windows(code, |window| {

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

@@ -48,12 +48,13 @@ fn load_all_mpts() -> Result<()> {
             type_leaf,
             3.into(),     // 3 nibbles
             0xDEF.into(), // key part
+            4.into(),     // value length
             account.nonce,
             account.balance,
             account.storage_root.into_uint(),
             account.code_hash.into_uint(),
-            type_empty,
-            type_empty,
+            type_empty, // txn trie
+            type_empty, // receipt trie
         ]
     );
 

evm/src/memory/segments.rs

@@ -18,28 +18,30 @@ pub(crate) enum Segment {
     /// General purpose kernel memory, used by various kernel functions.
     /// In general, calling a helper function can result in this memory being clobbered.
     KernelGeneral = 7,
+    /// Another segment for general purpose kernel use.
+    KernelGeneral2 = 8,
     /// Contains normalized transaction fields; see `NormalizedTxnField`.
-    TxnFields = 8,
+    TxnFields = 9,
     /// Contains the data field of a transaction.
-    TxnData = 9,
+    TxnData = 10,
     /// A buffer used to hold raw RLP data.
-    RlpRaw = 10,
+    RlpRaw = 11,
     /// Contains all trie data. Tries are stored as immutable, copy-on-write trees, so this is an
     /// append-only buffer. It is owned by the kernel, so it only lives on context 0.
-    TrieData = 11,
+    TrieData = 12,
     /// The account address associated with the `i`th storage trie. Only lives on context 0.
-    StorageTrieAddresses = 12,
+    StorageTrieAddresses = 13,
     /// A pointer to the `i`th storage trie within the `TrieData` buffer. Only lives on context 0.
-    StorageTriePointers = 13,
+    StorageTriePointers = 14,
     /// Like `StorageTriePointers`, except that these pointers correspond to the version of each
     /// trie at the creation of a given context. This lets us easily revert a context by replacing
     /// `StorageTriePointers` with `StorageTrieCheckpointPointers`.
     /// See also `StateTrieCheckpointPointer`.
-    StorageTrieCheckpointPointers = 14,
+    StorageTrieCheckpointPointers = 15,
 }
 
 impl Segment {
-    pub(crate) const COUNT: usize = 15;
+    pub(crate) const COUNT: usize = 16;
 
     pub(crate) fn all() -> [Self; Self::COUNT] {
         [
@@ -51,6 +53,7 @@ impl Segment {
             Self::GlobalMetadata,
             Self::ContextMetadata,
             Self::KernelGeneral,
+            Self::KernelGeneral2,
             Self::TxnFields,
             Self::TxnData,
             Self::RlpRaw,
@@ -72,6 +75,7 @@ impl Segment {
             Segment::GlobalMetadata => "SEGMENT_GLOBAL_METADATA",
             Segment::ContextMetadata => "SEGMENT_CONTEXT_METADATA",
             Segment::KernelGeneral => "SEGMENT_KERNEL_GENERAL",
+            Segment::KernelGeneral2 => "SEGMENT_KERNEL_GENERAL_2",
             Segment::TxnFields => "SEGMENT_NORMALIZED_TXN",
             Segment::TxnData => "SEGMENT_TXN_DATA",
             Segment::RlpRaw => "SEGMENT_RLP_RAW",
@@ -93,6 +97,7 @@ impl Segment {
             Segment::GlobalMetadata => 256,
             Segment::ContextMetadata => 256,
             Segment::KernelGeneral => 256,
+            Segment::KernelGeneral2 => 256,
             Segment::TxnFields => 256,
             Segment::TxnData => 256,
             Segment::RlpRaw => 8,