Add withdrawals (0xPolygonZero#1322)

* Withdrawals

* Remove AllRecursiveCircuits in withdrawals test

* Fix ERC20 test

evm/src/cpu/kernel/aggregator.rs

@@ -43,6 +43,7 @@ pub(crate) fn combined_kernel() -> Kernel {
         include_str!("asm/core/log.asm"),
         include_str!("asm/core/selfdestruct_list.asm"),
         include_str!("asm/core/touched_addresses.asm"),
+        include_str!("asm/core/withdrawals.asm"),
         include_str!("asm/core/precompiles/main.asm"),
         include_str!("asm/core/precompiles/ecrec.asm"),
         include_str!("asm/core/precompiles/sha256.asm"),

evm/src/cpu/kernel/asm/core/withdrawals.asm

@@ -0,0 +1,25 @@
+%macro withdrawals
+    // stack: (empty)
+    PUSH %%after
+    %jump(withdrawals)
+%%after:
+    // stack: (empty)
+%endmacro
+
+global withdrawals:
+    // stack: retdest
+    PROVER_INPUT(withdrawal)
+    // stack: address, retdest
+    PROVER_INPUT(withdrawal)
+    // stack: amount, address, retdest
+    DUP2 %eq_const(@U256_MAX) %jumpi(withdrawals_end)
+    SWAP1
+    // stack: address, amount, retdest
+    %add_eth
+    // stack: retdest
+    %jump(withdrawals)
+
+withdrawals_end:
+    // stack: amount, address, retdest
+    %pop2
+    JUMP

evm/src/cpu/kernel/asm/main.asm

@@ -32,7 +32,7 @@ global start_txns:
 txn_loop:
     // If the prover has no more txns for us to process, halt.
     PROVER_INPUT(end_of_txns)
-    %jumpi(hash_final_tries)
+    %jumpi(execute_withdrawals)
 
     // Call route_txn. When we return, continue the txn loop.
     PUSH txn_loop_after
@@ -48,6 +48,9 @@ global txn_loop_after:
     SWAP3 %increment SWAP3
     %jump(txn_loop)
 
+global execute_withdrawals:
+    // stack: cum_gas, txn_counter, num_nibbles, txn_nb
+    %withdrawals
 global hash_final_tries:
     // stack: cum_gas, txn_counter, num_nibbles, txn_nb
     // Check that we end up with the correct `cum_gas`, `txn_nb` and bloom filter.

evm/src/generation/mod.rs

@@ -47,6 +47,8 @@ pub struct GenerationInputs {
     pub block_bloom_after: [U256; 8],
 
     pub signed_txns: Vec<Vec<u8>>,
+    // Withdrawal pairs `(addr, amount)`. At the end of the txs, `amount` is added to `addr`'s balance. See EIP-4895.
+    pub withdrawals: Vec<(Address, U256)>,
     pub tries: TrieInputs,
     /// Expected trie roots after the transactions are executed.
     pub trie_roots_after: TrieRoots,

evm/src/generation/prover_input.rs

@@ -44,6 +44,7 @@ impl<F: Field> GenerationState<F> {
             "current_hash" => self.run_current_hash(),
             "account_code" => self.run_account_code(input_fn),
             "bignum_modmul" => self.run_bignum_modmul(),
+            "withdrawal" => self.run_withdrawal(),
             _ => Err(ProgramError::ProverInputError(InvalidFunction)),
         }
     }
@@ -219,6 +220,13 @@ impl<F: Field> GenerationState<F> {
 
         (biguint_to_mem_vec(rem), biguint_to_mem_vec(quo))
     }
+
+    /// Withdrawal data.
+    fn run_withdrawal(&mut self) -> Result<U256, ProgramError> {
+        self.withdrawal_prover_inputs
+            .pop()
+            .ok_or(ProgramError::ProverInputError(OutOfWithdrawalData))
+    }
 }
 
 enum EvmField {

evm/src/generation/state.rs

@@ -39,6 +39,8 @@ pub(crate) struct GenerationState<F: Field> {
     /// via `pop()`.
     pub(crate) rlp_prover_inputs: Vec<U256>,
 
+    pub(crate) withdrawal_prover_inputs: Vec<U256>,
+
     /// The state trie only stores state keys, which are hashes of addresses, but sometimes it is
     /// useful to see the actual addresses for debugging. Here we store the mapping for all known
     /// addresses.
@@ -63,6 +65,7 @@ impl<F: Field> GenerationState<F> {
         log::debug!("Input contract_code: {:?}", &inputs.contract_code);
         let mpt_prover_inputs = all_mpt_prover_inputs_reversed(&inputs.tries)?;
         let rlp_prover_inputs = all_rlp_prover_inputs_reversed(&inputs.signed_txns);
+        let withdrawal_prover_inputs = all_withdrawals_prover_inputs_reversed(&inputs.withdrawals);
         let bignum_modmul_result_limbs = Vec::new();
 
         Ok(Self {
@@ -73,6 +76,7 @@ impl<F: Field> GenerationState<F> {
             next_txn_index: 0,
             mpt_prover_inputs,
             rlp_prover_inputs,
+            withdrawal_prover_inputs,
             state_key_to_address: HashMap::new(),
             bignum_modmul_result_limbs,
         })
@@ -148,3 +152,16 @@ impl<F: Field> GenerationState<F> {
             .collect()
     }
 }
+
+/// Withdrawals prover input array is of the form `[addr0, amount0, ..., addrN, amountN, U256::MAX, U256::MAX]`.
+/// Returns the reversed array.
+fn all_withdrawals_prover_inputs_reversed(withdrawals: &[(Address, U256)]) -> Vec<U256> {
+    let mut withdrawal_prover_inputs = withdrawals
+        .iter()
+        .flat_map(|w| [U256::from((w.0).0.as_slice()), w.1])
+        .collect::<Vec<_>>();
+    withdrawal_prover_inputs.push(U256::MAX);
+    withdrawal_prover_inputs.push(U256::MAX);
+    withdrawal_prover_inputs.reverse();
+    withdrawal_prover_inputs
+}

evm/src/witness/errors.rs

@@ -31,6 +31,7 @@ pub enum MemoryError {
 pub enum ProverInputError {
     OutOfMptData,
     OutOfRlpData,
+    OutOfWithdrawalData,
     CodeHashNotFound,
     InvalidMptInput,
     InvalidInput,

evm/tests/add11_yml.rs

@@ -152,6 +152,7 @@ fn add11_yml() -> anyhow::Result<()> {
     };
     let inputs = GenerationInputs {
         signed_txns: vec![txn.to_vec()],
+        withdrawals: vec![],
         tries: tries_before,
         trie_roots_after,
         contract_code,

evm/tests/basic_smart_contract.rs

@@ -184,6 +184,7 @@ fn test_basic_smart_contract() -> anyhow::Result<()> {
     };
     let inputs = GenerationInputs {
         signed_txns: vec![txn.to_vec()],
+        withdrawals: vec![],
         tries: tries_before,
         trie_roots_after,
         contract_code,

evm/tests/empty_txn_list.rs

@@ -49,6 +49,7 @@ fn test_empty_txn_list() -> anyhow::Result<()> {
     };
     let inputs = GenerationInputs {
         signed_txns: vec![],
+        withdrawals: vec![],
         tries: TrieInputs {
             state_trie,
             transactions_trie,

evm/tests/erc20.rs

@@ -160,6 +160,7 @@ fn test_erc20() -> anyhow::Result<()> {
     };
     let inputs = GenerationInputs {
         signed_txns: vec![txn.to_vec()],
+        withdrawals: vec![],
         tries: tries_before,
         trie_roots_after,
         contract_code,

evm/tests/log_opcode.rs

@@ -230,6 +230,7 @@ fn test_log_opcodes() -> anyhow::Result<()> {
     ];
     let inputs = GenerationInputs {
         signed_txns: vec![txn.to_vec()],
+        withdrawals: vec![],
         tries: tries_before,
         trie_roots_after,
         contract_code,
@@ -437,6 +438,7 @@ fn test_log_with_aggreg() -> anyhow::Result<()> {
 
     let inputs_first = GenerationInputs {
         signed_txns: vec![txn.to_vec()],
+        withdrawals: vec![],
         tries: tries_before,
         trie_roots_after: tries_after,
         contract_code,
@@ -581,6 +583,7 @@ fn test_log_with_aggreg() -> anyhow::Result<()> {
     ];
     let inputs = GenerationInputs {
         signed_txns: vec![txn_2.to_vec()],
+        withdrawals: vec![],
         tries: tries_before,
         trie_roots_after,
         contract_code,
@@ -892,6 +895,7 @@ fn test_two_txn() -> anyhow::Result<()> {
     };
     let inputs = GenerationInputs {
         signed_txns: vec![txn_0.to_vec(), txn_1.to_vec()],
+        withdrawals: vec![],
         tries: tries_before,
         trie_roots_after,
         contract_code,

evm/tests/many_transactions.rs

@@ -212,6 +212,7 @@ fn test_four_transactions() -> anyhow::Result<()> {
     };
     let inputs = GenerationInputs {
         signed_txns: vec![txn1.to_vec(), txn2.to_vec(), txn3.to_vec(), txn4.to_vec()],
+        withdrawals: vec![],
         tries: tries_before,
         trie_roots_after,
         genesis_state_trie_root: HashedPartialTrie::from(Node::Empty).hash(),

evm/tests/self_balance_gas_cost.rs

@@ -171,6 +171,7 @@ fn self_balance_gas_cost() -> anyhow::Result<()> {
     };
     let inputs = GenerationInputs {
         signed_txns: vec![txn.to_vec()],
+        withdrawals: vec![],
         tries: tries_before,
         trie_roots_after,
         contract_code,

evm/tests/selfdestruct.rs

@@ -123,6 +123,7 @@ fn test_selfdestruct() -> anyhow::Result<()> {
     };
     let inputs = GenerationInputs {
         signed_txns: vec![txn.to_vec()],
+        withdrawals: vec![],
         tries: tries_before,
         trie_roots_after,
         contract_code,

evm/tests/simple_transfer.rs

@@ -139,6 +139,7 @@ fn test_simple_transfer() -> anyhow::Result<()> {
     };
     let inputs = GenerationInputs {
         signed_txns: vec![txn.to_vec()],
+        withdrawals: vec![],
         tries: tries_before,
         trie_roots_after,
         contract_code,

evm/tests/withdrawals.rs

@@ -0,0 +1,99 @@
+use std::collections::HashMap;
+use std::time::Duration;
+
+use env_logger::{try_init_from_env, Env, DEFAULT_FILTER_ENV};
+use eth_trie_utils::nibbles::Nibbles;
+use eth_trie_utils::partial_trie::{HashedPartialTrie, PartialTrie};
+use ethereum_types::{H160, H256, U256};
+use keccak_hash::keccak;
+use plonky2::field::goldilocks_field::GoldilocksField;
+use plonky2::plonk::config::PoseidonGoldilocksConfig;
+use plonky2::util::timing::TimingTree;
+use plonky2_evm::all_stark::AllStark;
+use plonky2_evm::config::StarkConfig;
+use plonky2_evm::generation::mpt::AccountRlp;
+use plonky2_evm::generation::{GenerationInputs, TrieInputs};
+use plonky2_evm::proof::{BlockHashes, BlockMetadata, TrieRoots};
+use plonky2_evm::prover::prove;
+use plonky2_evm::verifier::verify_proof;
+use plonky2_evm::Node;
+use rand::random;
+
+type F = GoldilocksField;
+const D: usize = 2;
+type C = PoseidonGoldilocksConfig;
+
+/// Execute 0 txns and 1 withdrawal.
+#[test]
+fn test_withdrawals() -> anyhow::Result<()> {
+    init_logger();
+
+    let all_stark = AllStark::<F, D>::default();
+    let config = StarkConfig::standard_fast_config();
+
+    let block_metadata = BlockMetadata::default();
+
+    let state_trie_before = HashedPartialTrie::from(Node::Empty);
+    let transactions_trie = HashedPartialTrie::from(Node::Empty);
+    let receipts_trie = HashedPartialTrie::from(Node::Empty);
+    let storage_tries = vec![];
+
+    let mut contract_code = HashMap::new();
+    contract_code.insert(keccak(vec![]), vec![]);
+
+    // Just one withdrawal.
+    let withdrawals = vec![(H160(random()), U256(random()))];
+
+    let state_trie_after = {
+        let mut trie = HashedPartialTrie::from(Node::Empty);
+        let addr_state_key = keccak(withdrawals[0].0);
+        let addr_nibbles = Nibbles::from_bytes_be(addr_state_key.as_bytes()).unwrap();
+        let account = AccountRlp {
+            balance: withdrawals[0].1,
+            ..AccountRlp::default()
+        };
+        trie.insert(addr_nibbles, rlp::encode(&account).to_vec());
+        trie
+    };
+
+    let trie_roots_after = TrieRoots {
+        state_root: state_trie_after.hash(),
+        transactions_root: transactions_trie.hash(),
+        receipts_root: receipts_trie.hash(),
+    };
+
+    let inputs = GenerationInputs {
+        signed_txns: vec![],
+        withdrawals,
+        tries: TrieInputs {
+            state_trie: state_trie_before,
+            transactions_trie,
+            receipts_trie,
+            storage_tries,
+        },
+        trie_roots_after,
+        contract_code,
+        genesis_state_trie_root: HashedPartialTrie::from(Node::Empty).hash(),
+        block_metadata,
+        txn_number_before: 0.into(),
+        gas_used_before: 0.into(),
+        gas_used_after: 0.into(),
+        block_bloom_before: [0.into(); 8],
+        block_bloom_after: [0.into(); 8],
+        block_hashes: BlockHashes {
+            prev_hashes: vec![H256::default(); 256],
+            cur_hash: H256::default(),
+        },
+        addresses: vec![],
+    };
+
+    let mut timing = TimingTree::new("prove", log::Level::Debug);
+    let proof = prove::<F, C, D>(&all_stark, &config, inputs, &mut timing)?;
+    timing.filter(Duration::from_millis(100)).print();
+
+    verify_proof(&all_stark, proof, &config)
+}
+
+fn init_logger() {
+    let _ = try_init_from_env(Env::default().filter_or(DEFAULT_FILTER_ENV, "info"));
+}