Small storage fixes

evm/src/cpu/kernel/asm/mpt/insert/insert_trie_specific.asm

@@ -2,9 +2,10 @@
 
 // Mutate the state trie, inserting the given key-value pair.
 global mpt_insert_state_trie:
-    // stack: num_nibbles, key, value_ptr, retdest
-    %stack (num_nibbles, key, value_ptr)
-        -> (num_nibbles, key, value_ptr, mpt_insert_state_trie_save)
+    // stack: key, value_ptr, retdest
+    %stack (key, value_ptr)
+        -> (key, value_ptr, mpt_insert_state_trie_save)
+    PUSH 64 // num_nibbles
     %mload_global_metadata(@GLOBAL_METADATA_STATE_TRIE_ROOT)
     // stack: state_root_ptr, num_nibbles, key, value_ptr, mpt_insert_state_trie_save, retdest
     %jump(mpt_insert)

evm/src/cpu/kernel/asm/mpt/storage/storage_read.asm

@@ -17,12 +17,14 @@ after_storage_read:
     // stack: value_ptr, retdest
     DUP1 %jumpi(storage_key_exists)
 
-    // Storage key not found; return default value of 0.
+    // Storage key not found. Return default value_ptr = 0,
+    // which derefs to 0 since @SEGMENT_TRIE_DATA[0] = 0.
     %stack (value_ptr, retdest) -> (retdest, 0)
     JUMP
 
 storage_key_exists:
     // stack: value_ptr, retdest
-    %mload_trie_data // TODO: If we end up not using value pointers in storage tries, remove this.
+    %mload_trie_data
+    // stack: value, retdest
     SWAP1
     JUMP

evm/src/cpu/kernel/asm/mpt/storage/storage_write.asm

@@ -4,15 +4,24 @@
 // Post stack: (empty)
 
 global storage_write:
+    // TODO: If value = 0, delete the key instead of inserting 0.
     // stack: slot, value, retdest
-    // TODO: If value = 0, delete the key instead of inserting 0?
-    // TODO: Do we need to write value to MPT data and insert value_ptr? Currently some logic assumes all values are pointers, but could be relaxed so a value is any single word.
-    %stack (slot, value) -> (slot, value, after_storage_insert)
+
+    // First we write the value to MPT data, and get a pointer to it.
+    %get_trie_data_size
+    // stack: value_ptr, slot, value, retdest
+    SWAP2
+    // stack: value, slot, value_ptr, retdest
+    %append_to_trie_data
+    // stack: slot, value_ptr, retdest
+
+    // Next, call mpt_insert on the current account's storage root.
+    %stack (slot, value_ptr) -> (slot, value_ptr, after_storage_insert)
     %slot_to_storage_key
-    // stack: storage_key, value, after_storage_write, retdest
+    // stack: storage_key, value_ptr, after_storage_write, retdest
     PUSH 64 // storage_key has 64 nibbles
     %current_storage_trie
-    // stack: storage_root_ptr, 64, storage_key, value, after_storage_insert, retdest
+    // stack: storage_root_ptr, 64, storage_key, value_ptr, after_storage_insert, retdest
     %jump(mpt_insert)
 
 after_storage_insert:
@@ -29,5 +38,7 @@ after_storage_insert:
     %mstore_trie_data
     // stack: new_account_ptr, retdest
 
-    SWAP1
-    JUMP
+    // Save this updated account to the state trie.
+    ADDRESS %addr_to_state_key
+    // stack: state_key, new_account_ptr, retdest
+    %jump(mpt_insert_state_trie)

evm/src/cpu/kernel/interpreter.rs

@@ -328,6 +328,8 @@ impl<'a> Interpreter<'a> {
 
         if self.debug_offsets.contains(&self.offset) {
             println!("At {}, stack={:?}", self.offset_name(), self.stack());
+        } else if let Some(label) = self.offset_label() {
+            println!("At {}", label);
         }
 
         Ok(())
@@ -335,12 +337,16 @@ impl<'a> Interpreter<'a> {
 
     /// Get a string representation of the current offset for debugging purposes.
     fn offset_name(&self) -> String {
+        self.offset_label()
+            .unwrap_or_else(|| self.offset.to_string())
+    }
+
+    fn offset_label(&self) -> Option<String> {
         // TODO: Not sure we should use KERNEL? Interpreter is more general in other places.
-        let label = KERNEL
+        KERNEL
             .global_labels
             .iter()
-            .find_map(|(k, v)| (*v == self.offset).then(|| k.clone()));
-        label.unwrap_or_else(|| self.offset.to_string())
+            .find_map(|(k, v)| (*v == self.offset).then(|| k.clone()))
     }
 
     fn run_stop(&mut self) {

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

@@ -6,18 +6,20 @@ use super::nibbles;
 use crate::cpu::kernel::aggregator::KERNEL;
 use crate::cpu::kernel::constants::global_metadata::GlobalMetadata;
 use crate::cpu::kernel::interpreter::Interpreter;
-use crate::cpu::kernel::tests::mpt::{test_account_1_rlp, test_account_2};
+use crate::cpu::kernel::tests::mpt::{
+    nibbles_64, nibbles_count, test_account_1_rlp, test_account_2,
+};
 use crate::generation::mpt::{all_mpt_prover_inputs_reversed, AccountRlp};
 use crate::generation::TrieInputs;
 
 #[test]
 fn mpt_insert_empty() -> Result<()> {
-    test_state_trie(Default::default(), nibbles(0xABC), test_account_2())
+    test_state_trie(Default::default(), nibbles_64(0xABC), test_account_2())
 }
 
 #[test]
 fn mpt_insert_leaf_identical_keys() -> Result<()> {
-    let key = nibbles(0xABC);
+    let key = nibbles_64(0xABC);
     let state_trie = PartialTrie::Leaf {
         nibbles: key,
         value: test_account_1_rlp(),
@@ -28,37 +30,39 @@ fn mpt_insert_leaf_identical_keys() -> Result<()> {
 #[test]
 fn mpt_insert_leaf_nonoverlapping_keys() -> Result<()> {
     let state_trie = PartialTrie::Leaf {
-        nibbles: nibbles(0xABC),
+        nibbles: nibbles_64(0xABC),
         value: test_account_1_rlp(),
     };
-    test_state_trie(state_trie, nibbles(0x123), test_account_2())
+    test_state_trie(state_trie, nibbles_64(0x123), test_account_2())
 }
 
 #[test]
 fn mpt_insert_leaf_overlapping_keys() -> Result<()> {
     let state_trie = PartialTrie::Leaf {
-        nibbles: nibbles(0xABC),
+        nibbles: nibbles_64(0xABC),
         value: test_account_1_rlp(),
     };
-    test_state_trie(state_trie, nibbles(0xADE), test_account_2())
+    test_state_trie(state_trie, nibbles_64(0xADE), test_account_2())
 }
 
 #[test]
+#[ignore] // TODO: Not valid for state trie, all keys have same len.
 fn mpt_insert_leaf_insert_key_extends_leaf_key() -> Result<()> {
     let state_trie = PartialTrie::Leaf {
         nibbles: nibbles(0xABC),
         value: test_account_1_rlp(),
     };
-    test_state_trie(state_trie, nibbles(0xABCDE), test_account_2())
+    test_state_trie(state_trie, nibbles_64(0xABCDE), test_account_2())
 }
 
 #[test]
+#[ignore] // TODO: Not valid for state trie, all keys have same len.
 fn mpt_insert_leaf_leaf_key_extends_insert_key() -> Result<()> {
     let state_trie = PartialTrie::Leaf {
         nibbles: nibbles(0xABCDE),
         value: test_account_1_rlp(),
     };
-    test_state_trie(state_trie, nibbles(0xABC), test_account_2())
+    test_state_trie(state_trie, nibbles_64(0xABC), test_account_2())
 }
 
 #[test]
@@ -69,7 +73,7 @@ fn mpt_insert_branch_replacing_empty_child() -> Result<()> {
         value: vec![],
     };
 
-    test_state_trie(state_trie, nibbles(0xABC), test_account_2())
+    test_state_trie(state_trie, nibbles_64(0xABC), test_account_2())
 }
 
 #[test]
@@ -92,7 +96,7 @@ fn mpt_insert_extension_nonoverlapping_keys() -> Result<()> {
         }
         .into(),
     };
-    test_state_trie(state_trie, nibbles(0x12345), test_account_2())
+    test_state_trie(state_trie, nibbles_64(0x12345), test_account_2())
 }
 
 #[test]
@@ -115,29 +119,33 @@ fn mpt_insert_extension_insert_key_extends_node_key() -> Result<()> {
         }
         .into(),
     };
-    test_state_trie(state_trie, nibbles(0xABCDEF), test_account_2())
+    test_state_trie(state_trie, nibbles_64(0xABCDEF), test_account_2())
 }
 
 #[test]
 fn mpt_insert_branch_to_leaf_same_key() -> Result<()> {
     let leaf = PartialTrie::Leaf {
-        nibbles: nibbles(0xBCD),
+        nibbles: nibbles_count(0xBCD, 63),
         value: test_account_1_rlp(),
     }
     .into();
+
     let mut children = std::array::from_fn(|_| PartialTrie::Empty.into());
-    children[0xA] = leaf;
+    children[0] = leaf;
     let state_trie = PartialTrie::Branch {
         children,
         value: vec![],
     };
 
-    test_state_trie(state_trie, nibbles(0xABCD), test_account_2())
+    test_state_trie(state_trie, nibbles_64(0xABCD), test_account_2())
 }
 
 /// Note: The account's storage_root is ignored, as we can't insert a new storage_root without the
 /// accompanying trie data. An empty trie's storage_root is used instead.
 fn test_state_trie(state_trie: PartialTrie, k: Nibbles, mut account: AccountRlp) -> Result<()> {
+    assert_eq!(k.count, 64);
+
+    // Ignore any storage_root; see documentation note.
     account.storage_root = PartialTrie::Empty.calc_hash();
 
     let trie_inputs = TrieInputs {
@@ -177,7 +185,6 @@ fn test_state_trie(state_trie: PartialTrie, k: Nibbles, mut account: AccountRlp)
     interpreter.push(0xDEADBEEFu32.into());
     interpreter.push(value_ptr.into()); // value_ptr
     interpreter.push(k.packed); // key
-    interpreter.push(k.count.into()); // num_nibbles
 
     interpreter.run()?;
     assert_eq!(

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

@@ -13,13 +13,22 @@ mod read;
 /// Note that this preserves all nibbles (eg. `0x123` is not interpreted as `0x0123`).
 pub(crate) fn nibbles<T: Into<U256>>(v: T) -> Nibbles {
     let packed = v.into();
-
     Nibbles {
         count: Nibbles::get_num_nibbles_in_key(&packed),
         packed,
     }
 }
 
+pub(crate) fn nibbles_64<T: Into<U256>>(v: T) -> Nibbles {
+    let packed = v.into();
+    Nibbles { count: 64, packed }
+}
+
+pub(crate) fn nibbles_count<T: Into<U256>>(v: T, count: usize) -> Nibbles {
+    let packed = v.into();
+    Nibbles { count, packed }
+}
+
 pub(crate) fn test_account_1() -> AccountRlp {
     AccountRlp {
         nonce: U256::from(1111),