Working smt hashing

0xPolygonZero · Oct 18, 2023 · d95c430 · d95c430
1 parent fd09915
commit d95c430
Show file tree

Hide file tree

Showing 18 changed files with 1,403 additions and 780 deletions.
diff --git a/evm/src/cpu/kernel/aggregator.rs b/evm/src/cpu/kernel/aggregator.rs
@@ -110,6 +110,8 @@ pub(crate) fn combined_kernel() -> Kernel {
         include_str!("asm/memory/packing.asm"),
         include_str!("asm/memory/syscalls.asm"),
         include_str!("asm/memory/txn_fields.asm"),
+        include_str!("asm/smt/load.asm"),
+        include_str!("asm/smt/hash.asm"),
         include_str!("asm/mpt/accounts.asm"),
         include_str!("asm/mpt/delete/delete.asm"),
         include_str!("asm/mpt/delete/delete_branch.asm"),

diff --git a/evm/src/cpu/kernel/asm/mpt/load/load.asm b/evm/src/cpu/kernel/asm/mpt/load/load.asm
@@ -3,10 +3,10 @@ global load_all_mpts:
     // stack: retdest
     // First set @GLOBAL_METADATA_TRIE_DATA_SIZE = 1.
     // We don't want it to start at 0, as we use 0 as a null pointer.
-    PUSH 1
-    %set_trie_data_size
+    // PUSH 1
+    // %set_trie_data_size
 
-    %load_mpt(mpt_load_state_trie_value)   %mstore_global_metadata(@GLOBAL_METADATA_STATE_TRIE_ROOT)
+    %load_state_smt
     %load_mpt(mpt_load_txn_trie_value)     %mstore_global_metadata(@GLOBAL_METADATA_TXN_TRIE_ROOT)
     %load_mpt(mpt_load_receipt_trie_value) %mstore_global_metadata(@GLOBAL_METADATA_RECEIPT_TRIE_ROOT)
 

diff --git a/evm/src/cpu/kernel/asm/smt/hash.asm b/evm/src/cpu/kernel/asm/smt/hash.asm
@@ -0,0 +1,166 @@
+%macro smt_hash_state
+    PUSH %%after %jump(smt_hash_state)
+%%after:
+%endmacro
+
+global smt_hash_state:
+    // stack: retdest
+    PUSH 0 %mstore_kernel_general(13371337)
+    %mload_global_metadata(@GLOBAL_METADATA_STATE_TRIE_ROOT)
+smt_hash:
+    // stack: ptr, retdest
+    DUP1
+    %mload_trie_data
+    // stack: node, node_ptr, retdest
+global smt_hash_node:
+    DUP1 %eq_const(@SMT_NODE_HASH) %jumpi(smt_hash_hash)
+    DUP1 %eq_const(@SMT_NODE_INTERNAL) %jumpi(smt_hash_internal)
+    DUP1 %eq_const(@SMT_NODE_LEAF) %jumpi(smt_hash_leaf)
+global wtf:
+    // stack: node, node_ptr, retdest
+    %pop2 JUMP
+
+smt_hash_hash:
+    // stack: node, node_ptr, retdest
+    POP
+    // stack: node_ptr, retdest
+    %increment
+    // stack: node_ptr+1, retdest
+    %mload_trie_data
+    // stack: hash, retdest
+global smt_hash_hash_yo:
+    SWAP1 JUMP
+
+smt_hash_internal:
+    // stack: node, node_ptr, retdest
+    POP
+    // stack: node_ptr, retdest
+    %increment
+    // stack: node_ptr+1, retdest
+    DUP1
+    %mload_trie_data
+    %stack (left_child_ptr, node_ptr_plus_1, retdest) -> (left_child_ptr, smt_hash_internal_after_left, node_ptr_plus_1, retdest)
+    %jump(smt_hash)
+smt_hash_internal_after_left:
+    // stack: left_hash, node_ptr+1, retdest
+    SWAP1 %increment
+    // stack: node_ptr+2, left_hash, retdest
+    %mload_trie_data
+    %stack (right_child_ptr, left_hash, retdest) -> (right_child_ptr, smt_hash_internal_after_right, left_hash, retdest)
+    %jump(smt_hash)
+smt_hash_internal_after_right:
+    // stack: right_hash, left_hash, retdest
+    %stack (right_hash) -> (0, @SEGMENT_KERNEL_GENERAL, 33, right_hash, 32)
+    %mstore_unpacking POP
+    %stack (left_hash) -> (0, @SEGMENT_KERNEL_GENERAL, 1, left_hash, 32)
+    %mstore_unpacking POP
+    // stack: retdest
+    PUSH 1 %mstore_kernel_general(0)
+    %stack () -> (0, @SEGMENT_KERNEL_GENERAL, 0, 65)
+    KECCAK_GENERAL
+    // stack: hash, retdest
+global smt_hash_internal_yo:
+    SWAP1 JUMP
+
+smt_hash_leaf:
+    // stack: node, node_ptr, retdest
+    POP
+    // stack: node_ptr, retdest
+    %increment
+    // stack: node_ptr+1, retdest
+    DUP1 %mload_trie_data
+    // stack: key, node_ptr+1, retdest
+    %mload_kernel_general(13371337)
+    // stack: is_value, key, node_ptr+1, retdest
+    %jumpi(smt_hash_leaf_value)
+smt_hash_leaf_account:
+    // stack: key, node_ptr+1, retdest
+    SWAP1
+    // stack: node_ptr+1, key, retdest
+    %increment
+    // stack: node_ptr+2, key, retdest
+    DUP1 %mload_trie_data
+    // stack: nonce, node_ptr+2, key, retdest
+    SWAP1
+    // stack: node_ptr+2, nonce, key, retdest
+    %increment
+    // stack: node_ptr+3, nonce, key, retdest
+    DUP1 %mload_trie_data
+    // stack: balance, node_ptr+3, nonce, key, retdest
+    SWAP1
+    // stack: node_ptr+3, balance, nonce, key, retdest
+    %increment
+    // stack: node_ptr+4, balance, nonce, key, retdest
+    DUP1 %mload_trie_data
+    // stack: storage_root, node_ptr+4, balance, nonce, key, retdest
+    PUSH 1 %mstore_kernel_general(13371337)
+    %stack (storage_root) -> (storage_root, smt_hash_leaf_account_after_storage)
+    %jump(smt_hash)
+smt_hash_leaf_account_after_storage:
+    PUSH 0 %mstore_kernel_general(13371337)
+    // stack: storage_root_hash, node_ptr+4, balance, nonce, key, retdest
+    SWAP1
+    // stack: node_ptr+4, storage_root_hash, balance, nonce, key, retdest
+    %increment
+    // stack: node_ptr+5, storage_root_hash, balance, nonce, key, retdest
+    %mload_trie_data
+    // stack: code_hash, storage_root_hash, balance, nonce, key, retdest
+
+    // 0    | 1----32 | 33----40 | 41----72 | 73-------104 | 105----136
+    // flag |   key   |  nonce   | balance  | storage_root | code_hash
+
+    // 0----7 | 8----39 | 40--------71 | 72----103
+    // nonce  | balance | storage_root | code_hash
+
+    %stack (code_hash) -> (0, @SEGMENT_KERNEL_GENERAL, 72, code_hash, 32)
+    %mstore_unpacking POP
+
+    %stack (storage_root) -> (0, @SEGMENT_KERNEL_GENERAL, 40, storage_root, 32)
+    %mstore_unpacking POP
+
+    %stack (balance) -> (0, @SEGMENT_KERNEL_GENERAL, 8, balance, 32)
+    %mstore_unpacking POP
+
+    %stack (nonce) -> (0, @SEGMENT_KERNEL_GENERAL, 0, nonce)
+    %mstore_unpacking_u64_LE
+
+    // stack: key, retdest
+    %stack () -> (0, @SEGMENT_KERNEL_GENERAL, 0, 104)
+    KECCAK_GENERAL
+    // stack: hash, key, retdest
+
+    // flag
+    PUSH 0 %mstore_kernel_general(0)
+
+    %stack (hash) -> (0, @SEGMENT_KERNEL_GENERAL, 33, hash, 32)
+    %mstore_unpacking POP
+
+    %stack (key) -> (0, @SEGMENT_KERNEL_GENERAL, 1, key, 32)
+    %mstore_unpacking POP
+
+    %stack () -> (0, @SEGMENT_KERNEL_GENERAL, 0, 65)
+    KECCAK_GENERAL
+
+global smt_hash_leaf_account_yo:
+    SWAP1 JUMP
+
+smt_hash_leaf_value:
+    // stack: key, node_ptr+1, retdest
+    SWAP1
+    // stack: node_ptr+1, key, retdest
+    %increment
+    // stack: node_ptr+2, key, retdest
+    %mload_trie_data
+    // stack: value, key, retdest
+    PUSH 0 %mstore_kernel_general(0)
+    %stack (value) -> (0, @SEGMENT_KERNEL_GENERAL, 33, value, 32)
+    %mstore_unpacking POP
+    // stack: key, retdest
+    %stack (key) -> (0, @SEGMENT_KERNEL_GENERAL, 1, key, 32)
+    %mstore_unpacking POP
+    // stack: retdest
+    %stack () -> (0, @SEGMENT_KERNEL_GENERAL, 0, 65)
+    KECCAK_GENERAL
+    // stack: hash, retdest
+global smt_hash_leaf_value_yo:
+    SWAP1 JUMP
diff --git a/evm/src/cpu/kernel/asm/smt/load.asm b/evm/src/cpu/kernel/asm/smt/load.asm
@@ -0,0 +1,29 @@
+%macro load_state_smt
+    PUSH %%after %jump(load_state_smt)
+%%after:
+%endmacro
+
+global load_state_smt:
+    // stack: retdest
+    PROVER_INPUT(smt::state)
+    // stack: len, retdest
+    %get_trie_data_size
+    DUP1 %mstore_global_metadata(@GLOBAL_METADATA_STATE_TRIE_ROOT)
+    // stack: i, len, retdest
+    %stack (i, len) -> (len, i, i)
+    ADD SWAP1
+    // stack: i, len, retdest
+loop:
+    // stack: i, len, retdest
+    DUP2 DUP2 EQ %jumpi(loop_end)
+    // stack: i, len, retdest
+    PROVER_INPUT(smt::state)
+    DUP2
+    // stack: i, x, i, len, retdest
+    %mstore_trie_data
+    // stack: i, len, retdest
+    %increment
+    %jump(loop)
+loop_end:
+    // stack: i, len, retdest
+    %pop2 JUMP
diff --git a/evm/src/cpu/kernel/constants/smt_type.rs b/evm/src/cpu/kernel/constants/smt_type.rs
@@ -16,7 +16,7 @@ impl PartialSmtType {
     pub(crate) fn var_name(&self) -> &'static str {
         match self {
             Self::Hash => "SMT_NODE_HASH",
-            Self::Internal => "SMT_NODE_BRANCH",
+            Self::Internal => "SMT_NODE_INTERNAL",
             Self::Leaf => "SMT_NODE_LEAF",
         }
     }

diff --git a/evm/src/cpu/kernel/interpreter.rs b/evm/src/cpu/kernel/interpreter.rs
@@ -441,7 +441,7 @@ impl<'a> Interpreter<'a> {
         {
             println!("At {}, stack={:?}", self.offset_name(), self.stack());
         } else if let Some(label) = self.offset_label() {
-            println!("At {label}");
+            println!("At {}, stack={:?}", label, self.stack());
         }
 
         Ok(())

diff --git a/evm/src/cpu/kernel/tests/mod.rs b/evm/src/cpu/kernel/tests/mod.rs
@@ -15,6 +15,7 @@ mod packing;
 mod receipt;
 mod rlp;
 mod signed_syscalls;
+mod smt;
 mod transaction_parsing;
 
 use std::str::FromStr;