feat: add pre-state execution for polygon-cdk (#594)

* WIP

* Tweak

* Final tweaks and test

* Fix

* Clippy

* Add feature-gating

* Comment

* Add missing fix from feat/feature_gat_chain

* Apply review

* Forgotten file change

evm_arithmetization/benches/fibonacci_25m_gas.rs

@@ -17,9 +17,8 @@ use evm_arithmetization::generation::{GenerationInputs, TrieInputs};
 use evm_arithmetization::proof::{BlockHashes, BlockMetadata, TrieRoots};
 use evm_arithmetization::prover::testing::simulate_execution;
 use evm_arithmetization::testing_utils::{
-    beacon_roots_account_nibbles, beacon_roots_contract_from_storage, ger_account_nibbles,
+    beacon_roots_account_nibbles, beacon_roots_contract_from_storage,
     preinitialized_state_and_storage_tries, update_beacon_roots_account_storage,
-    GLOBAL_EXIT_ROOT_ACCOUNT,
 };
 use evm_arithmetization::Node;
 use hex_literal::hex;
@@ -148,10 +147,6 @@ fn prepare_setup() -> anyhow::Result<GenerationInputs> {
         beacon_roots_account_nibbles(),
         rlp::encode(&beacon_roots_account).to_vec(),
     )?;
-    expected_state_trie_after.insert(
-        ger_account_nibbles(),
-        rlp::encode(&GLOBAL_EXIT_ROOT_ACCOUNT).to_vec(),
-    )?;
 
     let receipt_0 = LegacyReceiptRlp {
         status: false,
@@ -186,7 +181,7 @@ fn prepare_setup() -> anyhow::Result<GenerationInputs> {
         checkpoint_state_trie_root: H256(hex!(
             "fe07ff6d1ab215df17884b89112ccf2373597285a56c5902150313ad1a53ee57"
         )),
-        global_exit_roots: vec![],
+        ger_data: None,
         block_metadata,
         txn_number_before: 0.into(),
         gas_used_before: 0.into(),

evm_arithmetization/src/cpu/kernel/aggregator.rs

@@ -9,7 +9,11 @@ use super::assembler::{assemble, Kernel};
 use crate::cpu::kernel::constants::evm_constants;
 use crate::cpu::kernel::parser::parse;
 
-pub const NUMBER_KERNEL_FILES: usize = 159;
+pub const NUMBER_KERNEL_FILES: usize = if cfg!(feature = "cdk_erigon") {
+    159
+} else {
+    158
+};
 
 pub static KERNEL_FILES: [&str; NUMBER_KERNEL_FILES] = [
     "global jumped_to_0: PANIC",
@@ -172,7 +176,8 @@ pub static KERNEL_FILES: [&str; NUMBER_KERNEL_FILES] = [
     include_str!("asm/account_code.asm"),
     include_str!("asm/balance.asm"),
     include_str!("asm/bloom_filter.asm"),
-    include_str!("asm/global_exit_root.asm"),
+    #[cfg(feature = "cdk_erigon")]
+    include_str!("asm/cdk_pre_execution.asm"),
 ];
 
 pub static KERNEL: Lazy<Kernel> = Lazy::new(combined_kernel);

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

@@ -4,7 +4,15 @@
 /// *NOTE*: This will panic if one of the provided timestamps is zero.
 
 global set_beacon_root:
-    PUSH set_global_exit_roots
+    #[cfg(feature = cdk_erigon)]
+    {
+        PUSH pre_block_execution
+    }
+    #[cfg(not(feature = cdk_erigon))]
+    {
+        PUSH txn_loop
+    }
+
     %timestamp
     // stack: timestamp, retdest
     PUSH @HISTORY_BUFFER_LENGTH

evm_arithmetization/src/cpu/kernel/asm/cdk_pre_execution.asm

@@ -0,0 +1,121 @@
+/// CDK-Erigon pre-block execution logic.
+/// Reference implementation: `cdk-erigon/core/state/intra_block_state_zkevm.go`.
+/// This currently supports the Etrog upgrade.
+
+/// Pre-stack: (empty)
+/// Post-stack: (empty)
+global pre_block_execution:
+    // stack: (empty)
+    PUSH txn_loop
+    // stack: retdest
+    PUSH @ADDRESS_SCALABLE_L2
+    %is_non_existent
+    %jumpi(create_scalable_l2_account)
+
+global update_scalable_block_number:
+    // stack: retdest
+    %blocknumber
+    PUSH @LAST_BLOCK_STORAGE_POS
+    // stack: last_block_slot, block_number, retdest
+    %write_scalable_storage
+    // stack: retdest
+
+    // Check timestamp
+    PUSH @ADDRESS_SCALABLE_L2_STATE_KEY
+    PUSH @TIMESTAMP_STORAGE_POS
+    %read_storage_linked_list_w_state_key
+    // stack: old_timestamp, retdest
+    %timestamp
+    GT %jumpi(update_scalable_timestamp)
+
+global update_scalable_prev_block_root_hash:
+    // stack: retdest
+    %mload_global_metadata(@GLOBAL_METADATA_STATE_TRIE_DIGEST_BEFORE)
+    // stack: prev_block_root, retdest
+    PUSH @STATE_ROOT_STORAGE_POS
+    PUSH 1 %blocknumber SUB
+    // stack: block_number - 1, STATE_ROOT_STORAGE_POS, prev_block_root, retdest
+    PUSH @SEGMENT_KERNEL_GENERAL
+    // stack: addr, block_number - 1, STATE_ROOT_STORAGE_POS, prev_block_root, retdest
+    MSTORE_32BYTES_32
+    // stack: addr, STATE_ROOT_STORAGE_POS, prev_block_root, retdest
+    MSTORE_32BYTES_32
+    // stack: addr, prev_block_root, retdest
+    POP
+    // stack: prev_block_root, retdest
+    PUSH 64 PUSH @SEGMENT_KERNEL_GENERAL
+    // stack: addr, len, prev_block_root, retdest
+    KECCAK_GENERAL
+    // stack: slot, prev_block_root, retdest
+    %write_scalable_storage
+    // stack: retdest
+
+// Note: We assume that if the l1 info tree has been re-used or the GER does not exist,
+// the payload will not contain any root to store, in which case calling `PROVER_INPUT(ger)`
+// will return `U256::MAX` causing this to return early.
+global update_scalable_l1blockhash:
+    // stack: retdest
+    PROVER_INPUT(ger)
+    // stack: l1blockhash?, retdest
+    DUP1 %eq_const(@U256_MAX) %jumpi(skip_and_exit)
+    // stack: l1blockhash, retdest
+    PUSH @GLOBAL_EXIT_ROOT_STORAGE_POS
+    PROVER_INPUT(ger)
+    // stack: root, GLOBAL_EXIT_ROOT_STORAGE_POS, l1blockhash, retdest
+    PUSH @SEGMENT_KERNEL_GENERAL
+    // stack: addr, root, GLOBAL_EXIT_ROOT_STORAGE_POS, l1blockhash, retdest
+    MSTORE_32BYTES_32
+    // stack: addr, GLOBAL_EXIT_ROOT_STORAGE_POS, l1blockhash, retdest
+    MSTORE_32BYTES_32
+    // stack: addr, l1blockhash, retdest
+    POP
+    // stack: l1blockhash, retdest
+    PUSH 64 PUSH @SEGMENT_KERNEL_GENERAL
+    // stack: addr, len, l1blockhash, retdest
+    KECCAK_GENERAL
+    // stack: slot, l1blockhash, retdest
+    %slot_to_storage_key
+    // stack: storage_key, l1blockhash, retdest
+    PUSH @GLOBAL_EXIT_ROOT_MANAGER_L2_STATE_KEY
+    // stack: state_key, storage_key, l1blockhash, retdest
+    %insert_slot_with_value_from_keys
+    // stack: retdest
+    JUMP
+
+skip_and_exit:
+    // stack: null, retdest
+    POP
+    JUMP
+
+global update_scalable_timestamp:
+    // stack: retdest
+    %timestamp
+    PUSH @TIMESTAMP_STORAGE_POS
+    // stack: timestamp_slot, timestamp, retdest
+    %write_scalable_storage
+    %jump(update_scalable_prev_block_root_hash)
+
+global create_scalable_l2_account:
+    // stack: (empty)
+    PUSH update_scalable_block_number
+    // stack: retdest
+    %get_trie_data_size // pointer to new account we're about to create
+    // stack: new_account_ptr, retdest
+    PUSH 0 %append_to_trie_data // nonce
+    PUSH 0 %append_to_trie_data // balance
+    PUSH 0 %append_to_trie_data // storage root pointer
+    PUSH @EMPTY_STRING_HASH %append_to_trie_data // code hash
+    // stack: new_account_ptr, retdest
+    PUSH @ADDRESS_SCALABLE_L2_STATE_KEY
+    // stack: key, new_account_ptr, retdest
+    %jump(mpt_insert_state_trie)
+
+%macro write_scalable_storage
+    // stack: slot, value
+    %slot_to_storage_key
+    // stack: storage_key, value
+    PUSH @ADDRESS_SCALABLE_L2_STATE_KEY
+    // stack: state_key, storage_key, value
+    %insert_slot_with_value_from_keys
+    // stack: (empty)
+%endmacro

evm_arithmetization/src/cpu/kernel/asm/mpt/linked_list/linked_list.asm

@@ -859,9 +859,10 @@ remove_all_slots_end:
 %macro read_storage_linked_list
     // stack: slot
     %slot_to_storage_key
+    %stack (storage_key) -> (storage_key, 0, %%after)
     %address
     %addr_to_state_key
-    %stack (addr_key, key) -> (addr_key, key, 0, %%after)
+    // stack: addr_key, storage_key, 0, %%after
     %jump(search_slot)
 %%after:
     // stack: slot_value
@@ -870,14 +871,23 @@ remove_all_slots_end:
 %macro read_storage_linked_list_w_addr
     // stack: slot, address
     %slot_to_storage_key
-    SWAP1
+    %stack (storage_key, address) -> (address, storage_key, 0, %%after)
     %addr_to_state_key
-    %stack (addr_key, key) -> (addr_key, key, 0, %%after)
+    // stack: addr_key, storage_key, 0, %%after
     %jump(search_slot)
 %%after:
     // stack: slot_value
 %endmacro
 
+%macro read_storage_linked_list_w_state_key
+    // stack: slot, state_key
+    %slot_to_storage_key
+    %stack (storage_key, state_key) -> (state_key, storage_key, 0, %%after)
+    %jump(search_slot)
+%%after:
+    // stack: slot_ptr
+%endmacro
+
 %macro first_account
     // stack: empty
     PUSH @SEGMENT_ACCOUNTS_LINKED_LIST

evm_arithmetization/src/cpu/kernel/constants/mod.rs

@@ -85,6 +85,30 @@ pub(crate) fn evm_constants() -> HashMap<String, U256> {
         global_exit_root::GLOBAL_EXIT_ROOT_STORAGE_POS.0.into(),
         U256::from(global_exit_root::GLOBAL_EXIT_ROOT_STORAGE_POS.1),
     );
+    c.insert(
+        global_exit_root::ADDRESS_SCALABLE_L2.0.into(),
+        U256::from_big_endian(&global_exit_root::ADDRESS_SCALABLE_L2.1),
+    );
+    c.insert(
+        global_exit_root::ADDRESS_SCALABLE_L2_STATE_KEY.0.into(),
+        U256::from_big_endian(&global_exit_root::ADDRESS_SCALABLE_L2_STATE_KEY.1),
+    );
+    c.insert(
+        global_exit_root::LAST_BLOCK_STORAGE_POS.0.into(),
+        U256::from(global_exit_root::LAST_BLOCK_STORAGE_POS.1),
+    );
+    c.insert(
+        global_exit_root::STATE_ROOT_STORAGE_POS.0.into(),
+        U256::from(global_exit_root::STATE_ROOT_STORAGE_POS.1),
+    );
+    c.insert(
+        global_exit_root::TIMESTAMP_STORAGE_POS.0.into(),
+        U256::from(global_exit_root::TIMESTAMP_STORAGE_POS.1),
+    );
+    c.insert(
+        global_exit_root::BLOCK_INFO_ROOT_STORAGE_POS.0.into(),
+        U256::from(global_exit_root::BLOCK_INFO_ROOT_STORAGE_POS.1),
+    );
 
     for segment in Segment::all() {
         c.insert(segment.var_name().into(), (segment as usize).into());
@@ -397,11 +421,11 @@ pub mod cancun_constants {
         hex!("000000000000000000000000000000001666c54b0a32529503432fcae0181b4bef79de09fc63671fda5ed1ba9bfa07899495346f3d7ac9cd23048ef30d0a154f"), // y_im
     ];
 
-    pub const HISTORY_BUFFER_LENGTH: (&str, u64) = ("HISTORY_BUFFER_LENGTH", 8191);
-
     // Beacon constants
     ///////////////////
 
+    pub const HISTORY_BUFFER_LENGTH: (&str, u64) = ("HISTORY_BUFFER_LENGTH", 8191);
+
     pub const BEACON_ROOTS_CONTRACT_ADDRESS: Address =
         H160(hex!("000F3df6D732807Ef1319fB7B8bB8522d0Beac02"));
 
@@ -440,19 +464,42 @@ pub mod global_exit_root {
     use super::*;
 
     /// Taken from https://github.com/0xPolygonHermez/cdk-erigon/blob/61f0b6912055c73f6879ea7e9b5bac22ea5fc85c/zk/utils/global_exit_root.go#L16.
-    pub const GLOBAL_EXIT_ROOT_MANAGER_L2_STATE_KEY: (&str, [u8; 20]) = (
-        "GLOBAL_EXIT_ROOT_MANAGER_L2_STATE_KEY",
+    pub const GLOBAL_EXIT_ROOT_MANAGER_L2: (&str, [u8; 20]) = (
+        "GLOBAL_EXIT_ROOT_MANAGER_L2",
         hex!("a40D5f56745a118D0906a34E69aeC8C0Db1cB8fA"),
     );
+    pub const GLOBAL_EXIT_ROOT_ADDRESS_HASHED: H256 = H256(hex!(
+        "1d5e9c22b4b1a781d0ef63e9c1293c2a45fee966809019aa9804b5e7148b0ca9"
+    ));
+    pub const GLOBAL_EXIT_ROOT_MANAGER_L2_STATE_KEY: (&str, [u8; 32]) = (
+        "GLOBAL_EXIT_ROOT_MANAGER_L2_STATE_KEY",
+        GLOBAL_EXIT_ROOT_ADDRESS_HASHED.to_fixed_bytes(),
+    );
+    /// Taken from https://github.com/0xPolygonHermez/cdk-erigon/blob/dc3cbcc59a95769626056c7bc70aade501e7741d/core/state/intra_block_state_zkevm.go#L20.
+    pub const ADDRESS_SCALABLE_L2: (&str, [u8; 20]) = (
+        "ADDRESS_SCALABLE_L2",
+        hex!("000000000000000000000000000000005ca1ab1e"),
+    );
+    pub const ADDRESS_SCALABLE_L2_ADDRESS_HASHED: H256 = H256(hex!(
+        "4bff39c4f33dafcd90c45c9a0a1f2e72949b200788587b306eda5dd84aa87577"
+    ));
+    pub const ADDRESS_SCALABLE_L2_STATE_KEY: (&str, [u8; 32]) = (
+        "ADDRESS_SCALABLE_L2_STATE_KEY",
+        ADDRESS_SCALABLE_L2_ADDRESS_HASHED.to_fixed_bytes(),
+    );
     /// Taken from https://github.com/0xPolygonHermez/cdk-erigon/blob/61f0b6912055c73f6879ea7e9b5bac22ea5fc85c/zk/utils/global_exit_root.go#L17.
     pub const GLOBAL_EXIT_ROOT_STORAGE_POS: (&str, u64) = ("GLOBAL_EXIT_ROOT_STORAGE_POS", 0);
 
+    /// Taken from https://github.com/0xPolygonHermez/cdk-erigon/blob/dc3cbcc59a95769626056c7bc70aade501e7741d/core/state/intra_block_state_zkevm.go#L16.
+    pub const LAST_BLOCK_STORAGE_POS: (&str, u64) = ("LAST_BLOCK_STORAGE_POS", 0);
+    pub const STATE_ROOT_STORAGE_POS: (&str, u64) = ("STATE_ROOT_STORAGE_POS", 1);
+    pub const TIMESTAMP_STORAGE_POS: (&str, u64) = ("TIMESTAMP_STORAGE_POS", 2);
+    pub const BLOCK_INFO_ROOT_STORAGE_POS: (&str, u64) = ("BLOCK_INFO_ROOT_STORAGE_POS", 3);
+
     /// Taken from https://zkevm.polygonscan.com/address/0xa40D5f56745a118D0906a34E69aeC8C0Db1cB8fA#code.
     pub const GLOBAL_EXIT_ROOT_CONTRACT_CODE: [u8; 2112] = hex!("60806040526004361061004e5760003560e01c80633659cfe6146100655780634f1ef286146100855780635c60da1b146100985780638f283970146100c9578063f851a440146100e95761005d565b3661005d5761005b6100fe565b005b61005b6100fe565b34801561007157600080fd5b5061005b6100803660046106ca565b610118565b61005b6100933660046106e5565b61015f565b3480156100a457600080fd5b506100ad6101d0565b6040516001600160a01b03909116815260200160405180910390f35b3480156100d557600080fd5b5061005b6100e43660046106ca565b61020b565b3480156100f557600080fd5b506100ad610235565b610106610292565b610116610111610331565b61033b565b565b61012061035f565b6001600160a01b0316336001600160a01b031614156101575761015481604051806020016040528060008152506000610392565b50565b6101546100fe565b61016761035f565b6001600160a01b0316336001600160a01b031614156101c8576101c38383838080601f01602080910402602001604051908101604052809392919081815260200183838082843760009201919091525060019250610392915050565b505050565b6101c36100fe565b60006101da61035f565b6001600160a01b0316336001600160a01b03161415610200576101fb610331565b905090565b6102086100fe565b90565b61021361035f565b6001600160a01b0316336001600160a01b0316141561015757610154816103f1565b600061023f61035f565b6001600160a01b0316336001600160a01b03161415610200576101fb61035f565b606061028583836040518060600160405280602781526020016107e460279139610445565b9392505050565b3b151590565b61029a61035f565b6001600160a01b0316336001600160a01b031614156101165760405162461bcd60e51b815260206004820152604260248201527f5472616e73706172656e745570677261646561626c6550726f78793a2061646d60448201527f696e2063616e6e6f742066616c6c6261636b20746f2070726f78792074617267606482015261195d60f21b608482015260a4015b60405180910390fd5b60006101fb610519565b3660008037600080366000845af43d6000803e80801561035a573d6000f35b3d6000fd5b60007fb53127684a568b3173ae13b9f8a6016e243e63b6e8ee1178d6a717850b5d61035b546001600160a01b0316919050565b61039b83610541565b6040516001600160a01b038416907fbc7cd75a20ee27fd9adebab32041f755214dbc6bffa90cc0225b39da2e5c2d3b90600090a26000825111806103dc5750805b156101c3576103eb8383610260565b50505050565b7f7e644d79422f17c01e4894b5f4f588d331ebfa28653d42ae832dc59e38c9798f61041a61035f565b604080516001600160a01b03928316815291841660208301520160405180910390a1610154816105e9565b6060833b6104a45760405162461bcd60e51b815260206004820152602660248201527f416464726573733a2064656c65676174652063616c6c20746f206e6f6e2d636f6044820152651b9d1c9858dd60d21b6064820152608401610328565b600080856001600160a01b0316856040516104bf9190610794565b600060405180830381855af49150503d80600081146104fa576040519150601f19603f3d011682016040523d82523d6000602084013e6104ff565b606091505b509150915061050f828286610675565b9695505050505050565b60007f360894a13ba1a3210667c828492db98dca3e2076cc3735a920a3ca505d382bbc610383565b803b6105a55760405162461bcd60e51b815260206004820152602d60248201527f455243313936373a206e657720696d706c656d656e746174696f6e206973206e60448201526c1bdd08184818dbdb9d1c9858dd609a1b6064820152608401610328565b807f360894a13ba1a3210667c828492db98dca3e2076cc3735a920a3ca505d382bbc5b80546001600160a01b0319166001600160a01b039290921691909117905550565b6001600160a01b03811661064e5760405162461bcd60e51b815260206004820152602660248201527f455243313936373a206e65772061646d696e20697320746865207a65726f206160448201526564647265737360d01b6064820152608401610328565b807fb53127684a568b3173ae13b9f8a6016e243e63b6e8ee1178d6a717850b5d61036105c8565b60608315610684575081610285565b8251156106945782518084602001fd5b8160405162461bcd60e51b815260040161032891906107b0565b80356001600160a01b03811681146106c557600080fd5b919050565b6000602082840312156106dc57600080fd5b610285826106ae565b6000806000604084860312156106fa57600080fd5b610703846106ae565b9250602084013567ffffffffffffffff8082111561072057600080fd5b818601915086601f83011261073457600080fd5b81358181111561074357600080fd5b87602082850101111561075557600080fd5b6020830194508093505050509250925092565b60005b8381101561078357818101518382015260200161076b565b838111156103eb5750506000910152565b600082516107a6818460208701610768565b9190910192915050565b60208152600082518060208401526107cf816040850160208701610768565b601f01601f1916919091016040019291505056fe416464726573733a206c6f772d6c6576656c2064656c65676174652063616c6c206661696c6564a26469706673582212204675187caf3a43285d9a2c1844a981e977bd52a85ff073e7fc649f73847d70a464736f6c63430008090033");
     pub const GLOBAL_EXIT_ROOT_CONTRACT_CODE_HASH: [u8; 32] =
         hex!("6bec2bf64f7e824109f6ed55f77dd7665801d6195e461666ad6a5342a9f6daf5");
-    pub const GLOBAL_EXIT_ROOT_ADDRESS_HASHED: [u8; 32] =
-        hex!("1d5e9c22b4b1a781d0ef63e9c1293c2a45fee966809019aa9804b5e7148b0ca9");
 
     pub const GLOBAL_EXIT_ROOT_ACCOUNT: AccountRlp = AccountRlp {
         nonce: U256::zero(),
@@ -463,4 +510,16 @@ pub mod global_exit_root {
         )),
         code_hash: H256(GLOBAL_EXIT_ROOT_CONTRACT_CODE_HASH),
     };
+
+    #[test]
+    fn hashed() {
+        assert_eq!(
+            keccak_hash::keccak(GLOBAL_EXIT_ROOT_MANAGER_L2.1),
+            GLOBAL_EXIT_ROOT_ADDRESS_HASHED
+        );
+        assert_eq!(
+            keccak_hash::keccak(ADDRESS_SCALABLE_L2.1),
+            ADDRESS_SCALABLE_L2_ADDRESS_HASHED
+        );
+    }
 }