MPT hashing logic, part 3

evm/src/cpu/columns/ops.rs

@@ -34,6 +34,7 @@ pub struct OpsColumnsView<T> {
     pub shr: T,
     pub sar: T,
     pub keccak256: T,
+    pub keccak_general: T,
     pub address: T,
     pub balance: T,
     pub origin: T,

evm/src/cpu/decode.rs

@@ -22,7 +22,7 @@ use crate::cpu::columns::{CpuColumnsView, COL_MAP};
 /// behavior.
 /// Note: invalid opcodes are not represented here. _Any_ opcode is permitted to decode to
 /// `is_invalid`. The kernel then verifies that the opcode was _actually_ invalid.
-const OPCODES: [(u8, usize, bool, usize); 92] = [
+const OPCODES: [(u8, usize, bool, usize); 93] = [
     // (start index of block, number of top bits to check (log2), kernel-only, flag column)
     (0x00, 0, false, COL_MAP.op.stop),
     (0x01, 0, false, COL_MAP.op.add),
@@ -51,6 +51,7 @@ const OPCODES: [(u8, usize, bool, usize); 92] = [
     (0x1c, 0, false, COL_MAP.op.shr),
     (0x1d, 0, false, COL_MAP.op.sar),
     (0x20, 0, false, COL_MAP.op.keccak256),
+    (0x21, 0, true, COL_MAP.op.keccak_general),
     (0x30, 0, false, COL_MAP.op.address),
     (0x31, 0, false, COL_MAP.op.balance),
     (0x32, 0, false, COL_MAP.op.origin),

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

@@ -1,3 +1,46 @@
+global mpt_hash_state_trie:
+    // stack: retdest
+    %mload_global_metadata(@GLOBAL_METADATA_STATE_TRIE_ROOT)
+    // stack: node_ptr, retdest
+    %mpt_hash(encode_account)
+
+encode_account:
+    // stack: rlp_pos, value_ptr, retdest
+    // First, we compute the length of the RLP data we're about to write.
+    // The nonce and balance fields are variable-length, so we need to load them
+    // to determine their contribution, while the other two fields are fixed
+    // 32-bytes integers.
+    DUP2 %mload_trie_data // nonce = value[0]
+    %scalar_rlp_len
+    // stack: nonce_rlp_len, rlp_pos, value_ptr, retdest
+    DUP3 %add_const(1) %mload_trie_data // balance = value[1]
+    %scalar_rlp_len
+    // stack: balance_rlp_lenm, nonce_rlp_len, rlp_pos, value_ptr, retdest
+    PUSH 66 // storage_root and code_hash fields each take 1 + 32 bytes
+    ADD ADD
+    // stack: payload_len, rlp_pos, value_ptr, retdest
+    SWAP1
+    %encode_rlp_list_prefix
+    // stack: rlp_pos', value_ptr, retdest
+    DUP2 %mload_trie_data // nonce = value[0]
+    // stack: nonce, rlp_pos', value_ptr, retdest
+    SWAP1 %encode_rlp_scalar
+    // stack: rlp_pos'', value_ptr, retdest
+    DUP2 %add_const(1) %mload_trie_data // balance = value[1]
+    // stack: balance, rlp_pos'', value_ptr, retdest
+    SWAP1 %encode_rlp_scalar
+    // stack: rlp_pos''', value_ptr, retdest
+    DUP2 %add_const(2) %mload_trie_data // storage_root = value[2]
+    // stack: storage_root, rlp_pos''', value_ptr, retdest
+    SWAP1 %encode_rlp_256
+    // stack: rlp_pos'''', value_ptr, retdest
+    SWAP1 %add_const(3) %mload_trie_data // code_hash = value[3]
+    // stack: code_hash, rlp_pos'''', retdest
+    SWAP1 %encode_rlp_256
+    // stack: rlp_pos''''', retdest
+    SWAP1
+    JUMP
+
 // Computes the Merkle root of the given trie node.
 //
 // The encode_value function should take as input
@@ -47,7 +90,7 @@
     // stack: num_nibbles, packed_nibbles, terminated, %%mpt_hash_leaf_after_hex_prefix, 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, %%mpt_hash_leaf_after_hex_prefix, node_payload_ptr, retdest
-    %jump(hex_prefix)
+    %jump(hex_prefix_rlp)
 %%mpt_hash_leaf_after_hex_prefix:
     // stack: rlp_pos, node_payload_ptr, retdest
     SWAP1
@@ -59,12 +102,12 @@
     // stack: rlp_end_pos, retdest
     %prepend_rlp_list_prefix
     // stack: rlp_start_pos, rlp_len, retdest
-    PUSH $SEGMENT_RLP
+    PUSH @SEGMENT_RLP_RAW
     PUSH 0 // kernel context
     // stack: rlp_start_addr: 3, rlp_len, retdest
     KECCAK_GENERAL
     // stack: hash, retdest
-    SWAP
+    SWAP1
     JUMP
 %endmacro
 

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

@@ -30,7 +30,7 @@ global hex_prefix_rlp:
 
 rlp_header_medium:
     // stack: hp_len, i, rlp_pos, num_nibbles, packed_nibbles, terminated, retdest
-    DUP1 // value = hp_len
+    DUP1 %add_const(0x80) // value = 0x80 + hp_len
     DUP4 // offset = rlp_pos
     %mstore_rlp
 
@@ -44,7 +44,7 @@ rlp_header_large:
     // In practice hex-prefix length will never exceed 256, so the length of the
     // length will always be 1 byte in this case.
 
-    PUSH 1 // value = len_of_len = 1
+    PUSH 0xb8 // value = 0xb7 + len_of_len = 0xb8
     DUP4 // offset = rlp_pos
     %mstore_rlp
 

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

@@ -9,8 +9,7 @@ global load_all_mpts:
     PUSH 1
     %set_trie_data_size
 
-    %load_mpt_and_return_root_ptr
-    %mstore_global_metadata(@GLOBAL_METADATA_STATE_TRIE_ROOT)
+    %load_mpt_and_return_root_ptr %mstore_global_metadata(@GLOBAL_METADATA_STATE_TRIE_ROOT)
     %load_mpt_and_return_root_ptr %mstore_global_metadata(@GLOBAL_METADATA_TXN_TRIE_ROOT)
     %load_mpt_and_return_root_ptr %mstore_global_metadata(@GLOBAL_METADATA_RECEIPT_TRIE_ROOT)
 

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

@@ -79,7 +79,7 @@ encode_rlp_fixed:
     %add_const(1) // increment pos
     // stack: pos, len, string, retdest
     %stack (pos, len, string) -> (@SEGMENT_RLP_RAW, pos, string, len, encode_rlp_fixed_finish, pos, len)
-    GET_CONTEXT
+    PUSH 0 // context
     // stack: context, segment, pos, string, len, encode_rlp_fixed, pos, retdest
     %jump(mstore_unpacking)
 
@@ -90,6 +90,54 @@ encode_rlp_fixed_finish:
     SWAP1
     JUMP
 
+// Pre stack: pos, payload_len, retdest
+// Post stack: pos'
+global encode_rlp_list_prefix:
+    // stack: pos, payload_len, retdest
+    DUP2 %gt_const(55)
+    %jumpi(encode_rlp_list_prefix_large)
+    // Small case: prefix is just 0xc0 + length.
+    // stack: pos, payload_len, retdest
+    SWAP1
+    %add_const(0xc0)
+    // stack: prefix, pos, retdest
+    DUP2
+    // stack: pos, prefix, pos, retdest
+    %mstore_rlp
+    // stack: pos, retdest
+    %add_const(1)
+    SWAP1
+    JUMP
+encode_rlp_list_prefix_large:
+    // Write 0xf7 + len_of_len.
+    // stack: pos, payload_len, retdest
+    DUP2 %num_bytes
+    // stack: len_of_len, pos, payload_len, retdest
+    DUP1 %add_const(0xf7)
+    // stack: first_byte, len_of_len, pos, payload_len, retdest
+    DUP3 // pos
+    %mstore_rlp
+    // stack: len_of_len, pos, payload_len, retdest
+    SWAP1 %add_const(1)
+    // stack: pos', len_of_len, payload_len, retdest
+    %stack (pos, len_of_len, payload_len, retdest)
+        -> (0, @SEGMENT_RLP_RAW, pos, payload_len, len_of_len,
+            encode_rlp_list_prefix_large_done_writing_len,
+            pos, len_of_len, retdest)
+    %jump(mstore_unpacking)
+encode_rlp_list_prefix_large_done_writing_len:
+    // stack: pos', len_of_len, retdest
+    ADD
+    // stack: pos'', retdest
+    SWAP1
+    JUMP
+
+%macro encode_rlp_list_prefix
+    %stack (pos, payload_len) -> (pos, payload_len, %%after)
+    %jump(encode_rlp_list_prefix)
+%%after:
+%endmacro
+
 // Given an RLP list payload which starts at position 9 and ends at the given
 // position, prepend the appropriate RLP list prefix. Returns the updated start
 // position, as well as the length of the RLP data (including the newly-added
@@ -131,16 +179,15 @@ prepend_rlp_list_prefix_big:
     PUSH 8
     SUB
     // stack: start_pos, len_of_len, payload_len, end_pos, retdest
-    DUP2 DUP2 %mstore_rlp // rlp[start_pos] = len_of_len
+    DUP2 %add_const(0xf7) DUP2 %mstore_rlp // rlp[start_pos] = 0xf7 + len_of_len
     DUP1 %add_const(1) // start_len_pos = start_pos + 1
     %stack (start_len_pos, start_pos, len_of_len, payload_len, end_pos, retdest)
-        -> (len_of_len, start_len_pos, payload_len,
+        -> (0, @SEGMENT_RLP_RAW, start_len_pos, // context, segment, offset
+            payload_len, len_of_len,
             prepend_rlp_list_prefix_big_done_writing_len,
             start_pos, end_pos, retdest)
-    %jump(encode_rlp_fixed)
+    %jump(mstore_unpacking)
 prepend_rlp_list_prefix_big_done_writing_len:
-    // stack: start_payload_pos, start_pos, end_pos, retdest
-    POP
     // stack: start_pos, end_pos, retdest
     DUP1
     SWAP2
@@ -160,7 +207,7 @@ prepend_rlp_list_prefix_big_done_writing_len:
 // Get the number of bytes required to represent the given scalar.
 // The scalar is assumed to be non-zero, as small scalars like zero should
 // have already been handled with the small-scalar encoding.
-num_bytes:
+global num_bytes:
     // stack: x, retdest
     PUSH 0 // i
     // stack: i, x, retdest
@@ -192,3 +239,22 @@ num_bytes_finish:
     %jump(num_bytes)
 %%after:
 %endmacro
+
+// Given some scalar, compute the number of bytes used in its RLP encoding,
+// including any length prefix.
+%macro scalar_rlp_len
+    // stack: scalar
+    // Since the scalar fits in a word, we can't hit the large (>55 byte)
+    // case, so we just check for small vs medium.
+    DUP1 %gt_const(0x7f)
+    // stack: is_medium, scalar
+    %jumpi(%%medium)
+    // Small case; result is 1.
+    %stack (scalar) -> (1)
+%%medium:
+    // stack: scalar
+    %num_bytes
+    // stack: scalar_bytes
+    %add_const(1) // Account for the length prefix.
+    // stack: rlp_len
+%endmacro

evm/src/cpu/kernel/interpreter.rs

@@ -47,10 +47,21 @@ impl InterpreterMemory {
 
 impl InterpreterMemory {
     fn mload_general(&self, context: usize, segment: Segment, offset: usize) -> U256 {
-        self.context_memory[context].segments[segment as usize].get(offset)
+        let value = self.context_memory[context].segments[segment as usize].get(offset);
+        assert!(
+            value.bits() <= segment.bit_range(),
+            "Value read from memory exceeds expected range of {:?} segment",
+            segment
+        );
+        value
     }
 
     fn mstore_general(&mut self, context: usize, segment: Segment, offset: usize, value: U256) {
+        assert!(
+            value.bits() <= segment.bit_range(),
+            "Value written to memory exceeds expected range of {:?} segment",
+            segment
+        );
         self.context_memory[context].segments[segment as usize].set(offset, value)
     }
 }
@@ -162,7 +173,7 @@ impl<'a> Interpreter<'a> {
     }
 
     pub(crate) fn get_rlp_memory(&self) -> Vec<u8> {
-        self.memory.context_memory[self.context].segments[Segment::RlpRaw as usize]
+        self.memory.context_memory[0].segments[Segment::RlpRaw as usize]
             .content
             .iter()
             .map(|x| x.as_u32() as u8)
@@ -170,7 +181,7 @@ impl<'a> Interpreter<'a> {
     }
 
     pub(crate) fn set_rlp_memory(&mut self, rlp: Vec<u8>) {
-        self.memory.context_memory[self.context].segments[Segment::RlpRaw as usize].content =
+        self.memory.context_memory[0].segments[Segment::RlpRaw as usize].content =
             rlp.into_iter().map(U256::from).collect();
     }
 
@@ -229,6 +240,7 @@ impl<'a> Interpreter<'a> {
             0x1c => self.run_shr(),                                     // "SHR",
             0x1d => todo!(),                                            // "SAR",
             0x20 => self.run_keccak256(),                               // "KECCAK256",
+            0x21 => self.run_keccak_general(),                          // "KECCAK_GENERAL",
             0x30 => todo!(),                                            // "ADDRESS",
             0x31 => todo!(),                                            // "BALANCE",
             0x32 => todo!(),                                            // "ORIGIN",
@@ -447,6 +459,18 @@ impl<'a> Interpreter<'a> {
         self.push(U256::from_big_endian(hash.as_bytes()));
     }
 
+    fn run_keccak_general(&mut self) {
+        let context = self.pop().as_usize();
+        let segment = Segment::all()[self.pop().as_usize()];
+        let offset = self.pop().as_usize();
+        let size = self.pop().as_usize();
+        let bytes = (offset..offset + size)
+            .map(|i| self.memory.mload_general(context, segment, i).byte(0))
+            .collect::<Vec<_>>();
+        let hash = keccak(bytes);
+        self.push(U256::from_big_endian(hash.as_bytes()));
+    }
+
     fn run_prover_input(&mut self) -> anyhow::Result<()> {
         let prover_input_fn = self
             .prover_inputs_map
@@ -567,7 +591,6 @@ impl<'a> Interpreter<'a> {
         let segment = Segment::all()[self.pop().as_usize()];
         let offset = self.pop().as_usize();
         let value = self.memory.mload_general(context, segment, offset);
-        assert!(value.bits() <= segment.bit_range());
         self.push(value);
     }
 
@@ -576,7 +599,6 @@ impl<'a> Interpreter<'a> {
         let segment = Segment::all()[self.pop().as_usize()];
         let offset = self.pop().as_usize();
         let value = self.pop();
-        assert!(value.bits() <= segment.bit_range());
         self.memory.mstore_general(context, segment, offset, value);
     }
 }

evm/src/cpu/kernel/opcodes.rs

@@ -35,6 +35,7 @@ pub(crate) fn get_opcode(mnemonic: &str) -> u8 {
         "SHR" => 0x1c,
         "SAR" => 0x1d,
         "KECCAK256" => 0x20,
+        "KECCAK_GENERAL" => 0x21,
         "ADDRESS" => 0x30,
         "BALANCE" => 0x31,
         "ORIGIN" => 0x32,

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

@@ -0,0 +1,62 @@
+use anyhow::Result;
+use eth_trie_utils::partial_trie::{Nibbles, PartialTrie};
+use ethereum_types::{BigEndianHash, H256, U256};
+use hex_literal::hex;
+
+use crate::cpu::kernel::aggregator::KERNEL;
+use crate::cpu::kernel::interpreter::Interpreter;
+use crate::generation::mpt::all_mpt_prover_inputs_reversed;
+use crate::generation::TrieInputs;
+
+#[test]
+fn mpt_hash() -> 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);
+
+    // 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,
+            packed: 0xABC.into(),
+        },
+        value: account_rlp.to_vec(),
+    };
+    // TODO: It seems like calc_hash isn't giving the expected hash yet, so for now, I'm using a
+    // hardcoded hash obtained from py-evm.
+    // let state_trie_hash = state_trie.calc_hash();
+    let state_trie_hash =
+        hex!("e38d6053838fe057c865ec0c74a8f0de21865d74fac222a2d3241fe57c9c3a0f").into();
+
+    let trie_inputs = TrieInputs {
+        state_trie,
+        transactions_trie: Default::default(),
+        receipts_trie: Default::default(),
+        storage_tries: vec![],
+    };
+
+    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 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_hash_state_trie.
+    interpreter.offset = mpt_hash_state_trie;
+    interpreter.push(0xDEADBEEFu32.into());
+    interpreter.run()?;
+
+    assert_eq!(interpreter.stack().len(), 1);
+    let hash = H256::from_uint(&interpreter.stack()[0]);
+    assert_eq!(hash, state_trie_hash);
+
+    Ok(())
+}

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

@@ -20,9 +20,11 @@ fn hex_prefix_even_nonterminated() -> Result<()> {
     assert_eq!(
         interpreter.get_rlp_memory(),
         vec![
-            4, // length
-            0, // neither flag is set
-            0xAB, 0xCD, 0xEF
+            0x80 + 4, // prefix
+            0,        // neither flag is set
+            0xAB,
+            0xCD,
+            0xEF
         ]
     );
 
@@ -46,7 +48,7 @@ fn hex_prefix_odd_terminated() -> Result<()> {
     assert_eq!(
         interpreter.get_rlp_memory(),
         vec![
-            3, // length
+            0x80 + 3, // prefix
             (2 + 1) * 16 + 0xA,
             0xBC,
             0xDE,

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

@@ -1,5 +1,6 @@
 use eth_trie_utils::partial_trie::{Nibbles, PartialTrie};
 
+mod hash;
 mod hex_prefix;
 mod load;
 mod read;

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

@@ -42,10 +42,8 @@ fn mpt_read() -> Result<()> {
     interpreter.push(interpreter.get_global_metadata_field(GlobalMetadata::StateTrieRoot));
     interpreter.run()?;
 
-    dbg!(interpreter.stack());
     assert_eq!(interpreter.stack().len(), 1);
     let result_ptr = interpreter.stack()[0].as_usize();
-    dbg!(result_ptr);
     let result = &interpreter.get_trie_data()[result_ptr..][..account.len()];
     assert_eq!(result, account);
 

evm/src/cpu/kernel/tests/rlp/encode.rs

@@ -108,3 +108,48 @@ fn test_prepend_rlp_list_prefix_small() -> Result<()> {
 
     Ok(())
 }
+
+#[test]
+fn test_prepend_rlp_list_prefix_large() -> Result<()> {
+    let prepend_rlp_list_prefix = KERNEL.global_labels["prepend_rlp_list_prefix"];
+
+    let retdest = 0xDEADBEEFu32.into();
+    let end_pos = (9 + 60).into();
+    let initial_stack = vec![retdest, end_pos];
+    let mut interpreter = Interpreter::new_with_kernel(prepend_rlp_list_prefix, initial_stack);
+
+    #[rustfmt::skip]
+    interpreter.set_rlp_memory(vec![
+        // Nine 0s to leave room for the longest possible RLP list prefix.
+        0, 0, 0, 0, 0, 0, 0, 0, 0,
+        // The actual RLP list payload, consisting of 60 tiny strings.
+        0, 1, 2, 3, 4, 5, 6, 7, 8, 9,
+        10, 11, 12, 13, 14, 15, 16, 17, 18, 19,
+        20, 21, 22, 23, 24, 25, 26, 27, 28, 29,
+        30, 31, 32, 33, 34, 35, 36, 37, 38, 39,
+        40, 41, 42, 43, 44, 45, 46, 47, 48, 49,
+        50, 51, 52, 53, 54, 55, 56, 57, 58, 59,
+    ]);
+
+    interpreter.run()?;
+
+    let expected_rlp_len = 62.into();
+    let expected_start_pos = 7.into();
+    let expected_stack = vec![expected_rlp_len, expected_start_pos];
+
+    #[rustfmt::skip]
+    let expected_rlp = vec![
+        0, 0, 0, 0, 0, 0, 0, 0xf7 + 1, 60,
+        0, 1, 2, 3, 4, 5, 6, 7, 8, 9,
+        10, 11, 12, 13, 14, 15, 16, 17, 18, 19,
+        20, 21, 22, 23, 24, 25, 26, 27, 28, 29,
+        30, 31, 32, 33, 34, 35, 36, 37, 38, 39,
+        40, 41, 42, 43, 44, 45, 46, 47, 48, 49,
+        50, 51, 52, 53, 54, 55, 56, 57, 58, 59,
+    ];
+
+    assert_eq!(interpreter.stack(), expected_stack);
+    assert_eq!(interpreter.get_rlp_memory(), expected_rlp);
+
+    Ok(())
+}

evm/src/cpu/stack.rs

@@ -62,6 +62,7 @@ const STACK_BEHAVIORS: OpsColumnsView<Option<StackBehavior>> = OpsColumnsView {
     shr: BASIC_BINARY_OP,
     sar: BASIC_BINARY_OP,
     keccak256: None,        // TODO
+    keccak_general: None,   // TODO
     address: None,          // TODO
     balance: None,          // TODO
     origin: None,           // TODO