Fix storage write for beacons root contract (#102)

* Fix storage write for beacons root contract

* Fix test

* Add comment about panicking for invalid timestamps

* Add trie read check for missing key

evm_arithmetization/src/cpu/kernel/asm/beacon_roots.asm

@@ -1,5 +1,7 @@
 /// EIP-4788: Beacon block root in the EVM
 /// <https://eips.ethereum.org/EIPS/eip-4788#pseudocode>
+///
+/// *NOTE*: This will panic if one of the provided timestamps is zero.
 
 global set_beacon_root:
     PUSH set_global_exit_roots
@@ -17,9 +19,12 @@ global set_beacon_root:
     %add_const(@HISTORY_BUFFER_LENGTH)
     // stack: root_idx, calldata, write_beacon_roots_to_storage, timestamp_idx, timestamp, retdest
 
+    // If the calldata is zero, delete the slot from the storage trie.
+    DUP2 ISZERO %jumpi(delete_root_idx_slot)
+
 write_beacon_roots_to_storage:
     // stack: slot, value, retdest
-     // First we write the value to MPT data, and get a pointer to it.
+    // 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
@@ -46,3 +51,45 @@ after_beacon_roots_storage_insert:
     // stack: account_storage_root_ptr_ptr, new_storage_root_ptr, retdest
     %mstore_trie_data
     JUMP
+
+delete_root_idx_slot:
+    // stack: root_idx, 0, write_beacon_roots_to_storage, timestamp_idx, timestamp, retdest
+    PUSH after_root_idx_slot_delete
+    SWAP2 POP
+    // stack: root_idx, after_root_idx_slot_delete, write_beacon_roots_to_storage, timestamp_idx, timestamp, retdest
+    %slot_to_storage_key
+    // stack: storage_key, after_root_idx_slot_delete, write_beacon_roots_to_storage, timestamp_idx, timestamp, retdest
+    PUSH 64 // storage_key has 64 nibbles
+    %get_storage_trie(@BEACON_ROOTS_ADDRESS)
+    // stack: storage_root_ptr, 64, storage_key, after_root_idx_slot_delete, write_beacon_roots_to_storage, timestamp_idx, timestamp, retdest
+
+    // If the slot is empty (i.e. ptr defaulting to 0), skip the deletion.
+    DUP1 ISZERO %jumpi(skip_empty_slot)
+
+    // stack: storage_root_ptr, 64, storage_key, after_root_idx_slot_delete, write_beacon_roots_to_storage, timestamp_idx, timestamp, retdest
+    %stack (storage_root_ptr, nibbles, storage_key) -> (storage_root_ptr, nibbles, storage_key, checkpoint_delete_root_idx, storage_root_ptr, nibbles, storage_key)
+    %jump(mpt_read)
+checkpoint_delete_root_idx:
+    // stack: value_ptr, storage_root_ptr, 64, storage_key, after_root_idx_slot_delete, write_beacon_roots_to_storage, timestamp_idx, timestamp, retdest
+    // If the the storage key is not found (i.e. ptr defaulting to 0), skip the deletion.
+    ISZERO %jumpi(skip_empty_slot)
+
+    // stack: storage_root_ptr, 64, storage_key, after_root_idx_slot_delete, write_beacon_roots_to_storage, timestamp_idx, timestamp, retdest
+    %jump(mpt_delete)
+
+after_root_idx_slot_delete:
+    // stack: new_storage_root_ptr, write_beacon_roots_to_storage, timestamp_idx, timestamp, retdest
+    %get_account_data(@BEACON_ROOTS_ADDRESS)
+    // stack: account_ptr, new_storage_root_ptr, write_beacon_roots_to_storage, timestamp_idx, timestamp, retdest
+
+    // Update the copied account with our new storage root pointer.
+    %add_const(2)
+    // stack: account_storage_root_ptr_ptr, new_storage_root_ptr, write_beacon_roots_to_storage, timestamp_idx, timestamp, retdest
+    %mstore_trie_data
+    // stack: write_beacon_roots_to_storage, timestamp_idx, timestamp, retdest
+    JUMP
+
+skip_empty_slot:
+    // stack: 0, 64, storage_key, after_root_idx_slot_delete, write_beacon_roots_to_storage, timestamp_idx, timestamp, retdest
+    %pop4
+    JUMP

evm_arithmetization/src/cpu/kernel/asm/global_exit_root.asm

@@ -2,6 +2,8 @@
 /// Global exit roots (GER) are of the form `(timestamp, root)` and are loaded from prover inputs.
 /// The timestamp is written to the storage of address `ADDRESS_GLOBAL_EXIT_ROOT_MANAGER_L2` in the slot `keccak256(abi.encodePacked(root, GLOBAL_EXIT_ROOT_STORAGE_POS))`.
 /// See https://github.com/0xPolygonHermez/cdk-erigon/blob/zkevm/zk/utils/global_exit_root.go for reference.
+///
+/// *NOTE*: This will panic if one of the provided timestamps is zero.
 
 global set_global_exit_roots:
     // stack: (empty)

evm_arithmetization/src/testing_utils.rs

@@ -40,8 +40,12 @@ fn insert_storage(trie: &mut HashedPartialTrie, slot: U256, value: U256) {
     slot.to_big_endian(&mut bytes);
     let key = keccak(bytes);
     let nibbles = Nibbles::from_bytes_be(key.as_bytes()).unwrap();
-    let r = rlp::encode(&value);
-    trie.insert(nibbles, r.freeze().to_vec());
+    if value.is_zero() {
+        trie.delete(nibbles);
+    } else {
+        let r = rlp::encode(&value);
+        trie.insert(nibbles, r.freeze().to_vec());
+    }
 }
 
 /// Creates a storage trie for an account, given a list of `(slot, value)`
@@ -81,10 +85,8 @@ pub fn update_beacon_roots_account_storage(
 /// Returns the beacon roots contract account from its provided storage trie.
 pub fn beacon_roots_contract_from_storage(storage_trie: &HashedPartialTrie) -> AccountRlp {
     AccountRlp {
-        nonce: 0.into(),
-        balance: 0.into(),
         storage_root: storage_trie.hash(),
-        code_hash: H256(BEACON_ROOTS_CONTRACT_CODE_HASH),
+        ..BEACON_ROOTS_ACCOUNT
     }
 }
 

evm_arithmetization/tests/empty_txn_list.rs

@@ -35,6 +35,7 @@ fn test_empty_txn_list() -> anyhow::Result<()> {
 
     let block_metadata = BlockMetadata {
         block_number: 1.into(),
+        block_timestamp: 1.into(),
         parent_beacon_block_root: H256(hex!(
             "44e2566c06c03b132e0ede3e90af477ebca74393b89dd6cb29f9c79cbcb6e963"
         )),

evm_arithmetization/tests/global_exit_root.rs

@@ -31,7 +31,10 @@ fn test_global_exit_root() -> anyhow::Result<()> {
     let all_stark = AllStark::<F, D>::default();
     let config = StarkConfig::standard_fast_config();
 
-    let block_metadata = BlockMetadata::default();
+    let block_metadata = BlockMetadata {
+        block_timestamp: 1.into(),
+        ..BlockMetadata::default()
+    };
 
     let (state_trie_before, storage_tries) = preinitialized_state_and_storage_tries();
     let mut beacon_roots_account_storage = storage_tries[0].1.clone();
@@ -48,7 +51,7 @@ fn test_global_exit_root() -> anyhow::Result<()> {
         let mut trie = HashedPartialTrie::from(Node::Empty);
         update_beacon_roots_account_storage(
             &mut beacon_roots_account_storage,
-            0.into(),
+            block_metadata.block_timestamp,
             block_metadata.parent_beacon_block_root,
         );
         let beacon_roots_account =

evm_arithmetization/tests/withdrawals.rs

@@ -33,7 +33,10 @@ fn test_withdrawals() -> anyhow::Result<()> {
     let all_stark = AllStark::<F, D>::default();
     let config = StarkConfig::standard_fast_config();
 
-    let block_metadata = BlockMetadata::default();
+    let block_metadata = BlockMetadata {
+        block_timestamp: 1.into(),
+        ..BlockMetadata::default()
+    };
 
     let (state_trie_before, storage_tries) = preinitialized_state_and_storage_tries();
     let mut beacon_roots_account_storage = storage_tries[0].1.clone();
@@ -50,7 +53,7 @@ fn test_withdrawals() -> anyhow::Result<()> {
         let mut trie = HashedPartialTrie::from(Node::Empty);
         update_beacon_roots_account_storage(
             &mut beacon_roots_account_storage,
-            0.into(),
+            block_metadata.block_timestamp,
             block_metadata.parent_beacon_block_root,
         );
         let beacon_roots_account =