Add snailtracer benchmark

bench/snailtracer/snailtracer.py

@@ -0,0 +1,73 @@
+from typing import Final
+
+import cProfile
+import pathlib
+import pstats
+
+import pyrevm
+
+
+CONTRACT_DATA_FILE_PATH: Final[pathlib.Path] = (
+    pathlib.Path(__file__).absolute().parent / "./snailtracer.evm"
+)
+ZERO_ADDRESS: Final[str] = "0x0000000000000000000000000000000000000000"
+CALLER_ADDRESS: Final[str] = "0x1000000000000000000000000000000000000001"
+
+
+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: str, contract_data: bytes) -> pyrevm.EVM:
+    evm = pyrevm.EVM()
+    evm.insert_account_info(
+        contract_address,
+        pyrevm.AccountInfo(code=contract_data),
+    )
+    return evm
+
+
+def _benchmark(
+    evm: pyrevm.EVM,
+    caller_address: str,
+    contract_address: str,
+    call_data: list[int],
+    num_runs: int = 10,
+    warmup_runs: int = 2,
+) -> None:
+    def bench() -> None:
+        evm.call_raw(
+            caller=caller_address,
+            to=contract_address,
+            data=call_data,
+        )
+
+    for _ in range(warmup_runs):
+        bench()
+
+    with cProfile.Profile() as pr:
+        for _ in range(num_runs):
+            bench()
+
+        pr.disable()
+        p = pstats.Stats(pr)
+        p.sort_stats(pstats.SortKey.CUMULATIVE).print_stats(10)
+
+
+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,
+        contract_address=ZERO_ADDRESS,
+        call_data=list(bytes.fromhex("30627b7c")),
+        num_runs=10,
+        warmup_runs=2,
+    )
+
+
+if __name__ == "__main__":
+    main()

pyproject.toml

@@ -19,6 +19,8 @@ maturin = "^0.13.5"
 pytest = "^7.1.3"
 
 [tool.poetry.dev-dependencies]
+black = "^22.8.0"
+mypy = "^0.981"
 
 [build-system]
 requires = ["maturin>=0.13,<0.14"]

pyrevm.pyi

@@ -0,0 +1,91 @@
+from typing import Optional, Type
+
+class CfgEnv:
+    def __new__(cls: Type["CfgEnv"]) -> "CfgEnv": ...
+
+class BlockEnv:
+    def __new__(
+        cls: Type["BlockEnv"],
+        number: Optional[int] = None,
+        coinbase: Optional[str] = None,
+        timestamp: Optional[int] = None,
+        difficulty: Optional[int] = None,
+        basefee: Optional[int] = None,
+        gas_limit: Optional[int] = None,
+    ) -> "BlockEnv": ...
+
+class TxEnv:
+    def __new__(
+        cls: Type["TxEnv"],
+        caller: Optional[str] = None,
+        gas_limit: Optional[int] = None,
+        gas_price: Optional[int] = None,
+        gas_priority_fee: Optional[int] = None,
+        to: Optional[str] = None,
+        value: Optional[int] = None,
+        data: Optional[list[int]] = None,
+        chain_id: Optional[int] = None,
+        nonce: Optional[int] = None,
+    ) -> "TxEnv": ...
+
+class Env:
+    def __new__(
+        cls: Type["Env"],
+        cfg: Optional[CfgEnv] = None,
+        block: Optional[BlockEnv] = None,
+        tx: Optional[TxEnv] = None,
+    ) -> "Env": ...
+
+class AccountInfo:
+    @property
+    def balance(self: "AccountInfo") -> int: ...
+    @property
+    def nonce(self: "AccountInfo") -> int: ...
+    @property
+    def code(self: "AccountInfo") -> list[int]: ...
+    @property
+    def code_hash(self: "AccountInfo") -> list[int]: ...
+    def __new__(
+        cls: Type["AccountInfo"],
+        nonce: int = 0,
+        code_hash: Optional[bytes] = None,
+        code: Optional[bytes] = None,
+    ) -> "AccountInfo": ...
+
+class EvmOpts:
+    env: Env
+    fork_url: Optional[str]
+    fork_block_number: Optional[int]
+    gas_limit: int
+    tracing: bool
+    def __new__(
+        cls: Type["EvmOpts"],
+        env: Optional[Env],
+        fork_url: Optional[str],
+    ) -> "EvmOpts": ...
+
+class EVM:
+    def __new__(
+        cls: Type["EVM"],
+        env: Optional[Env] = None,
+        fork_url: Optional[str] = None,
+        fork_block_number: Optional[int] = None,
+        gas_limit: int = 18446744073709551615,
+        tracing: bool = False,
+    ) -> "EVM": ...
+    def basic(self: "EVM", address: str) -> Optional[AccountInfo]: ...
+    def insert_account_info(self: "EVM", address: str, info: AccountInfo) -> None: ...
+    def call_raw_committing(
+        self: "EVM",
+        caller: str,
+        to: str,
+        value: Optional[int] = None,
+        data: Optional[list[int]] = None,
+    ) -> None: ...
+    def call_raw(
+        self: "EVM",
+        caller: str,
+        to: str,
+        value: Optional[int] = None,
+        data: Optional[list[int]] = None,
+    ) -> None: ...