plonky2/evm/tests/erc721.rs

use std::str::FromStr;
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::{Address, BigEndianHash, H160, H256, U256};
use hex_literal::hex;
use keccak_hash::keccak;
use plonky2::field::goldilocks_field::GoldilocksField;
use plonky2::plonk::config::KeccakGoldilocksConfig;
use plonky2::util::timing::TimingTree;
use plonky2_evm::all_stark::AllStark;
use plonky2_evm::config::StarkConfig;
use plonky2_evm::generation::mpt::{AccountRlp, LegacyReceiptRlp, LogRlp};
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;

type F = GoldilocksField;
const D: usize = 2;
type C = KeccakGoldilocksConfig;

/// Test a simple ERC721 token transfer.
/// Used the following Solidity code:
/// ```solidity
/// pragma solidity ^0.8.20;
///
/// import "@openzeppelin/contracts@5.0.1/token/ERC721/ERC721.sol";
/// import "@openzeppelin/contracts@5.0.1/access/Ownable.sol";
///
/// contract TestToken is ERC721, Ownable {
///     constructor(address initialOwner)
///         ERC721("TestToken", "TEST")
///         Ownable(initialOwner)
///     {}
///
///     function safeMint(address to, uint256 tokenId) public onlyOwner {
///         _safeMint(to, tokenId);
///     }
/// }
/// ```
///
/// The transaction calls the `safeTransferFrom` function to transfer token `1337` from address
/// `0x5B38Da6a701c568545dCfcB03FcB875f56beddC4` to address `0xAb8483F64d9C6d1EcF9b849Ae677dD3315835cb2`.
#[test]
fn test_erc721() -> anyhow::Result<()> {
    init_logger();

    let all_stark = AllStark::<F, D>::default();
    let config = StarkConfig::standard_fast_config();

    let beneficiary = hex!("deadbeefdeadbeefdeadbeefdeadbeefdeadbeef");
    let owner = hex!("5B38Da6a701c568545dCfcB03FcB875f56beddC4");
    let contract = hex!("f2B1114C644cBb3fF63Bf1dD284c8Cd716e95BE9");

    let owner_state_key = keccak(owner);
    let contract_state_key = keccak(contract);

    let owner_nibbles = Nibbles::from_bytes_be(owner_state_key.as_bytes()).unwrap();
    let contract_nibbles = Nibbles::from_bytes_be(contract_state_key.as_bytes()).unwrap();

    let mut state_trie_before = HashedPartialTrie::from(Node::Empty);
    state_trie_before.insert(owner_nibbles, rlp::encode(&owner_account()).to_vec());
    state_trie_before.insert(contract_nibbles, rlp::encode(&contract_account()).to_vec());

    let storage_tries = vec![(contract_state_key, contract_storage())];

    let tries_before = TrieInputs {
        state_trie: state_trie_before,
        transactions_trie: HashedPartialTrie::from(Node::Empty),
        receipts_trie: HashedPartialTrie::from(Node::Empty),
        storage_tries,
    };

    let txn = signed_tx();

    let gas_used = 58_418.into();

    let contract_code = [contract_bytecode(), vec![]]
        .map(|v| (keccak(v.clone()), v))
        .into();

    let expected_state_trie_after: HashedPartialTrie = {
        let mut state_trie_after = HashedPartialTrie::from(Node::Empty);
        let owner_account = owner_account();
        let owner_account_after = AccountRlp {
            nonce: owner_account.nonce + 1,
            balance: owner_account.balance - gas_used * 0xa,
            ..owner_account
        };
        state_trie_after.insert(owner_nibbles, rlp::encode(&owner_account_after).to_vec());
        let contract_account_after = AccountRlp {
            storage_root: contract_storage_after().hash(),
            ..contract_account()
        };
        state_trie_after.insert(
            contract_nibbles,
            rlp::encode(&contract_account_after).to_vec(),
        );

        state_trie_after
    };

    let logs = vec![LogRlp {
        address: H160::from_str("0xf2B1114C644cBb3fF63Bf1dD284c8Cd716e95BE9").unwrap(),
        topics: vec![
            H256::from_str("0xddf252ad1be2c89b69c2b068fc378daa952ba7f163c4a11628f55a4df523b3ef")
                .unwrap(),
            H256::from_str("0x0000000000000000000000005b38da6a701c568545dcfcb03fcb875f56beddc4")
                .unwrap(),
            H256::from_str("0x000000000000000000000000ab8483f64d9c6d1ecf9b849ae677dd3315835cb2")
                .unwrap(),
            H256::from_str("0x0000000000000000000000000000000000000000000000000000000000000539")
                .unwrap(),
        ],
        data: vec![].into(),
    }];

    let mut bloom_bytes = [0u8; 256];
    add_logs_to_bloom(&mut bloom_bytes, &logs);

    let receipt_0 = LegacyReceiptRlp {
        status: true,
        cum_gas_used: gas_used,
        bloom: bloom_bytes.to_vec().into(),
        logs,
    };
    let mut receipts_trie = HashedPartialTrie::from(Node::Empty);
    receipts_trie.insert(Nibbles::from_str("0x80").unwrap(), receipt_0.encode(0));
    let transactions_trie: HashedPartialTrie = Node::Leaf {
        nibbles: Nibbles::from_str("0x80").unwrap(),
        value: txn.to_vec(),
    }
    .into();

    let trie_roots_after = TrieRoots {
        state_root: expected_state_trie_after.hash(),
        transactions_root: transactions_trie.hash(),
        receipts_root: receipts_trie.hash(),
    };

    let bloom = bloom_bytes
        .chunks_exact(32)
        .map(U256::from_big_endian)
        .collect::<Vec<_>>();

    let block_metadata = BlockMetadata {
        block_beneficiary: Address::from(beneficiary),
        block_timestamp: 0x03e8.into(),
        block_number: 1.into(),
        block_difficulty: 0x020000.into(),
        block_random: H256::from_uint(&0x020000.into()),
        block_gaslimit: 0xff112233u32.into(),
        block_chain_id: 1.into(),
        block_base_fee: 0xa.into(),
        block_gas_used: gas_used,
        block_bloom: bloom.try_into().unwrap(),
    };

    let inputs = GenerationInputs {
        signed_txn: Some(txn.to_vec()),
        withdrawals: vec![],
        tries: tries_before,
        trie_roots_after,
        contract_code,
        checkpoint_state_trie_root: HashedPartialTrie::from(Node::Empty).hash(),
        block_metadata,
        txn_number_before: 0.into(),
        gas_used_before: 0.into(),
        gas_used_after: gas_used,
        block_hashes: BlockHashes {
            prev_hashes: vec![H256::default(); 256],
            cur_hash: H256::default(),
        },
    };

    let mut timing = TimingTree::new("prove", log::Level::Debug);
    let proof = prove::<F, C, D>(&all_stark, &config, inputs, &mut timing, None)?;
    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"));
}

fn contract_bytecode() -> Vec<u8> {
    hex!("608060405234801561000f575f80fd5b5060043610610109575f3560e01c8063715018a6116100a0578063a22cb4651161006f578063a22cb465146102a1578063b88d4fde146102bd578063c87b56dd146102d9578063e985e9c514610309578063f2fde38b1461033957610109565b8063715018a61461023f5780638da5cb5b1461024957806395d89b4114610267578063a14481941461028557610109565b806323b872dd116100dc57806323b872dd146101a757806342842e0e146101c35780636352211e146101df57806370a082311461020f57610109565b806301ffc9a71461010d57806306fdde031461013d578063081812fc1461015b578063095ea7b31461018b575b5f80fd5b61012760048036038101906101229190611855565b610355565b604051610134919061189a565b60405180910390f35b610145610436565b604051610152919061193d565b60405180910390f35b61017560048036038101906101709190611990565b6104c5565b60405161018291906119fa565b60405180910390f35b6101a560048036038101906101a09190611a3d565b6104e0565b005b6101c160048036038101906101bc9190611a7b565b6104f6565b005b6101dd60048036038101906101d89190611a7b565b6105f5565b005b6101f960048036038101906101f49190611990565b610614565b60405161020691906119fa565b60405180910390f35b61022960048036038101906102249190611acb565b610625565b6040516102369190611b05565b60405180910390f35b6102476106db565b005b6102516106ee565b60405161025e91906119fa565b60405180910390f35b61026f610716565b60405161027c919061193d565b60405180910390f35b61029f600480360381019061029a9190611a3d565b6107a6565b005b6102bb60048036038101906102b69190611b48565b6107bc565b005b6102d760048036038101906102d29190611cb2565b6107d2565b005b6102f360048036038101906102ee9190611990565b6107ef565b604051610300919061193d565b60405180910390f35b610323600480360381019061031e9190611d32565b610855565b604051610330919061189a565b60405180910390f35b610353600480360381019061034e9190611acb565b6108e3565b005b5f7f80ac58cd000000000000000000000000000000000000000000000000000000007bffffffffffffffffffffffffffffffffffffffffffffffffffffffff1916827bffffffffffffffffffffffffffffffffffffffffffffffffffffffff1916148061041f57507f5b5e139f000000000000000000000000000000000000000000000000000000007bffffffffffffffffffffffffffffffffffffffffffffffffffffffff1916827bffffffffffffffffffffffffffffffffffffffffffffffffffffffff1916145b8061042f575061042e82610967565b5b9050919050565b60605f805461044490611d9d565b80601f016020809104026020016040519081016040528092919081815260200182805461047090611d9d565b80156104bb5780601f10610492576101008083540402835291602001916104bb565b820191905f5260205f20905b81548152906001019060200180831161049e57829003601f168201915b5050505050905090565b5f6104cf826109d0565b506104d982610a56565b9050919050565b6104f282826104ed610a8f565b610a96565b5050565b5f73ffffffffffffffffffffffffffffffffffffffff168273ffffffffffffffffffffffffffffffffffffffff1603610566575f6040517f64a0ae9200000000000000000000000000000000000000000000000000000000815260040161055d91906119fa565b60405180910390fd5b5f6105798383610574610a8f565b610aa8565b90508373ffffffffffffffffffffffffffffffffffffffff168173ffffffffffffffffffffffffffffffffffffffff16146105ef578382826040517f64283d7b0000000000000000000000000000000000000000000000000000000081526004016105e693929190611dcd565b60405180910390fd5b50505050565b61060f83838360405180602001604052805f8152506107d2565b505050565b5f61061e826109d0565b9050919050565b5f8073ffffffffffffffffffffffffffffffffffffffff168273ffffffffffffffffffffffffffffffffffffffff1603610696575f6040517f89c62b6400000000000000000000000000000000000000000000000000000000815260040161068d91906119fa565b60405180910390fd5b60035f8373ffffffffffffffffffffffffffffffffffffffff1673ffffffffffffffffffffffffffffffffffffffff1681526020019081526020015f20549050919050565b6106e3610cb3565b6106ec5f610d3a565b565b5f60065f9054906101000a900473ffffffffffffffffffffffffffffffffffffffff16905090565b60606001805461072590611d9d565b80601f016020809104026020016040519081016040528092919081815260200182805461075190611d9d565b801561079c5780601f106107735761010080835404028352916020019161079c565b820191905f5260205f20905b81548152906001019060200180831161077f57829003601f168201915b5050505050905090565b6107ae610cb3565b6107b88282610dfd565b5050565b6107ce6107c7610a8f565b8383610e1a565b5050565b6107dd8484846104f6565b6107e984848484610f83565b50505050565b60606107fa826109d0565
}

fn insert_storage(trie: &mut HashedPartialTrie, slot: U256, value: U256) {
    let mut bytes = [0; 32];
    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);
    let r = r.freeze().to_vec();
    trie.insert(nibbles, r);
}

fn sd2u(s: &str) -> U256 {
    U256::from_dec_str(s).unwrap()
}

fn sh2u(s: &str) -> U256 {
    U256::from_str_radix(s, 16).unwrap()
}

fn contract_storage() -> HashedPartialTrie {
    let mut trie = HashedPartialTrie::from(Node::Empty);
    insert_storage(
        &mut trie,
        U256::zero(),
        sh2u("0x54657374546f6b656e0000000000000000000000000000000000000000000012"),
    );
    insert_storage(
        &mut trie,
        U256::one(),
        sh2u("0x5445535400000000000000000000000000000000000000000000000000000008"),
    );
    insert_storage(
        &mut trie,
        sd2u("6"),
        sh2u("0x5b38da6a701c568545dcfcb03fcb875f56beddc4"),
    );
    insert_storage(
        &mut trie,
        sh2u("0x343ff8127bd64f680be4e996254dc3528603c6ecd54364b4cf956ebdd28f0028"),
        sh2u("0x5b38da6a701c568545dcfcb03fcb875f56beddc4"),
    );
    insert_storage(
        &mut trie,
        sh2u("0x118c1ea466562cb796e30ef705e4db752f5c39d773d22c5efd8d46f67194e78a"),
        sd2u("1"),
    );
    trie
}

fn contract_storage_after() -> HashedPartialTrie {
    let mut trie = HashedPartialTrie::from(Node::Empty);
    insert_storage(
        &mut trie,
        U256::zero(),
        sh2u("0x54657374546f6b656e0000000000000000000000000000000000000000000012"),
    );
    insert_storage(
        &mut trie,
        U256::one(),
        sh2u("0x5445535400000000000000000000000000000000000000000000000000000008"),
    );
    insert_storage(
        &mut trie,
        sd2u("6"),
        sh2u("0x5b38da6a701c568545dcfcb03fcb875f56beddc4"),
    );
    insert_storage(
        &mut trie,
        sh2u("0x343ff8127bd64f680be4e996254dc3528603c6ecd54364b4cf956ebdd28f0028"),
        sh2u("0xab8483f64d9c6d1ecf9b849ae677dd3315835cb2"),
    );
    insert_storage(
        &mut trie,
        sh2u("0xf3aa6a8a9f7e3707e36cc99c499a27514922afe861ec3d80a1a314409cba92f9"),
        sd2u("1"),
    );
    trie
}

fn owner_account() -> AccountRlp {
    AccountRlp {
        nonce: 2.into(),
        balance: 0x1000000.into(),
        storage_root: HashedPartialTrie::from(Node::Empty).hash(),
        code_hash: keccak([]),
    }
}

fn contract_account() -> AccountRlp {
    AccountRlp {
        nonce: 0.into(),
        balance: 0.into(),
        storage_root: contract_storage().hash(),
        code_hash: keccak(contract_bytecode()),
    }
}

fn signed_tx() -> Vec<u8> {
    hex!("f8c5020a8307a12094f2b1114c644cbb3ff63bf1dd284c8cd716e95be980b86442842e0e0000000000000000000000005b38da6a701c568545dcfcb03fcb875f56beddc4000000000000000000000000ab8483f64d9c6d1ecf9b849ae677dd3315835cb2000000000000000000000000000000000000000000000000000000000000053925a0414867f13ac63d663e84099d52c8215615666ea37c969c69aa58a0fad26a3f6ea01a7160c6274969083b2316eb8ca6011b4bf6b00972159a78bf64d06fa40c1402").into()
}

fn add_logs_to_bloom(bloom: &mut [u8; 256], logs: &Vec<LogRlp>) {
    for log in logs {
        add_to_bloom(bloom, log.address.as_bytes());
        for topic in &log.topics {
            add_to_bloom(bloom, topic.as_bytes());
        }
    }
}

fn add_to_bloom(bloom: &mut [u8; 256], bloom_entry: &[u8]) {
    let bloom_hash = keccak(bloom_entry).to_fixed_bytes();

    for idx in 0..3 {
        let bit_pair = u16::from_be_bytes(bloom_hash[2 * idx..2 * (idx + 1)].try_into().unwrap());
        let bit_to_set = 0x07FF - (bit_pair & 0x07FF);
        let byte_index = bit_to_set / 8;
        let bit_value = 1 << (7 - bit_to_set % 8);
        bloom[byte_index as usize] |= bit_value;
    }
}