snailtracer-pyevm.py

import time
from typing import Final

import cProfile
import logging
import pathlib
import pstats

import eth.abc
import eth.consensus.pow
import eth.constants
import eth.chains.base
import eth.db.atomic
import eth.tools.builder.chain.builders
import eth.vm.forks.berlin
import eth_keys
import eth_typing
import eth_utils

CONTRACT_DATA_FILE_PATH: Final[pathlib.Path] = (
    pathlib.Path(__file__).absolute().parent / "snailtracer.evm"
)
ZERO_ADDRESS: Final[eth_typing.Address] = eth.constants.ZERO_ADDRESS

CALLER_PRIVATE_KEY: Final[eth_keys.keys.PrivateKey] = eth_keys.keys.PrivateKey(
    eth_utils.hexadecimal.decode_hex(
        "0x45a915e4d060149eb4365960e6a7a45f334393093061116b197e3240065ff2d8"
    )
)
CALLER_ADDRESS: Final[eth_typing.Address] = eth_typing.Address(
    CALLER_PRIVATE_KEY.public_key.to_canonical_address()
)


def _load_contract_data(data_file_path: pathlib.Path) -> bytes:
    with open(data_file_path, mode="r") as file:
        return bytes.fromhex(file.read())


def _construct_evm(
    contract_address: eth_typing.Address, contract_data: bytes
) -> eth.vm.forks.berlin.BerlinVM:
    evm_class = eth.chains.base.MiningChain.configure(
        __name__="TestChain",
        vm_configuration=(
            (eth.constants.GENESIS_BLOCK_NUMBER, eth.vm.forks.berlin.BerlinVM),
        ),
    )
    evm = evm_class.from_genesis(
        eth.db.atomic.AtomicDB(),
        genesis_params={
            "difficulty": 1,
            "gas_limit": 3141592,
        },
    ).get_vm()
    evm.state.set_code(contract_address, contract_data)

    return evm


def _benchmark(
    evm: eth.vm.forks.berlin.BerlinVM,
    caller_address: eth_typing.Address,
    caller_private_key: eth_keys.keys.PrivateKey,
    contract_address: eth_typing.Address,
    call_data: bytes,
    num_runs: int = 10,
    warmup_runs: int = 2,
) -> None:
    nonce = evm.state.get_nonce(caller_address)
    tx = evm.create_unsigned_transaction(
        nonce=nonce,
        gas_price=0,
        gas=1_000_000_000,
        to=contract_address,
        value=0,
        data=call_data,
    )
    signed_tx = tx.as_signed_transaction(caller_private_key)
    evm_message = evm.state.get_transaction_executor().build_evm_message(signed_tx)

    def bench() -> None:
        evm.state.get_transaction_executor().build_computation(evm_message, signed_tx)

    for _ in range(warmup_runs):
        bench()

    start = time.perf_counter_ns()
    for _ in range(num_runs):
        bench()
    end = time.perf_counter_ns()
    print(((end - start) * 1e-6) / num_runs)


def main() -> None:
    contract_data = _load_contract_data(CONTRACT_DATA_FILE_PATH)
    evm = _construct_evm(ZERO_ADDRESS, contract_data)

    _benchmark(
        evm,
        caller_address=CALLER_ADDRESS,
        caller_private_key=CALLER_PRIVATE_KEY,
        contract_address=ZERO_ADDRESS,
        call_data=bytes.fromhex("30627b7c"),
        num_runs=5,
        warmup_runs=1,
    )


if __name__ == "__main__":
    main()