Add tracing

Cargo.toml

@@ -7,7 +7,7 @@ edition = "2021"
 crate-type = ["cdylib"]
 
 [dependencies]
-revm = { "version"="7.1.0", features=["ethersdb", "serde"] }
+revm = { "version"="7.1.0", features=["ethersdb", "serde", "serde-json", "std"] }
 revm-interpreter = "3.3.0"
 tracing = "0.1"
 pyo3 = { version = "0.19", features = ["extension-module"] }

pyrevm.pyi

@@ -198,3 +198,7 @@ class EVM:
     @property
     def env(self: "EVM") -> Env:
         """ Get the environment. """
+
+    @property
+    def tracing(self: "EVM") -> bool:
+        """ Whether tracing is enabled. """

pytest/test.py

@@ -1,3 +1,4 @@
+import json
 import os.path
 
 import pytest
@@ -8,8 +9,8 @@
 
 fork_url = "https://mainnet.infura.io/v3/c60b0bb42f8a4c6481ecd229eddaca27"
 
-KWARG_OPTS = [
-    {"fork_url": fork_url, "tracing": True},
+KWARG_CASES = [
+    {"fork_url": fork_url},
     {"fork_url": fork_url, "tracing": False, "fork_block_number": "latest"},
     {},
 ]
@@ -105,7 +106,7 @@ def test_balances_fork():
     assert evm.basic(address).balance == AMT
 
 
-@pytest.mark.parametrize("kwargs", KWARG_OPTS)
+@pytest.mark.parametrize("kwargs", KWARG_CASES)
 def test_call_raw(kwargs):
     evm = EVM(**kwargs)
     info = AccountInfo(code=load_contract_bin("full_math.bin"))
@@ -126,7 +127,7 @@ def test_call_raw(kwargs):
     assert changes[address2].nonce == 1
 
 
-@pytest.mark.parametrize("kwargs", KWARG_OPTS)
+@pytest.mark.parametrize("kwargs", KWARG_CASES)
 def test_call_committing(kwargs):
     evm = EVM(**kwargs)
     evm.insert_account_info(
@@ -145,7 +146,7 @@ def test_call_committing(kwargs):
     assert int.from_bytes(result, "big") == 171
 
 
-@pytest.mark.parametrize("kwargs", KWARG_OPTS)
+@pytest.mark.parametrize("kwargs", KWARG_CASES)
 def test_call_empty_result(kwargs):
     evm = EVM(**kwargs)
     evm.insert_account_info(address, AccountInfo(code=load_contract_bin("weth_9.bin")))
@@ -168,3 +169,24 @@ def test_call_empty_result(kwargs):
     )
 
     assert int.from_bytes(balance, "big") == 10000
+    assert not evm.tracing
+
+
+def test_tracing(capsys):
+    evm = EVM(tracing=True)
+    evm.insert_account_info(address, AccountInfo(code=load_contract_bin("weth_9.bin")))
+    evm.set_balance(address2, 10000)
+    evm.call_raw_committing(
+        caller=address2,
+        to=address,
+        value=10000,
+        calldata=bytes.fromhex("d0e30db0"),
+    )
+    assert evm.tracing
+    captured = capsys.readouterr()
+    traces = [json.loads(i) for i in captured.out.split("\n") if i]
+    assert {'gasUsed': '0xffffffffffff5011',
+            'output': '0x',
+            'pass': True,
+            'stateRoot': '0x0000000000000000000000000000000000000000000000000000000000000000'} == traces[-1]
+    assert len(traces) == 128

src/evm.rs

@@ -1,17 +1,21 @@
 use std::collections::HashMap;
 use std::fmt::Debug;
+use std::io::stdout;
 use pyo3::{PyErr, pymethods, PyResult, pyclass};
+use pyo3::ffi::PySys_WriteStdout;
 
-use revm::{Database, Evm, primitives::U256};
+use revm::{Database, Evm, inspector_handle_register, primitives::U256};
 use revm::DatabaseCommit;
 use revm::precompile::{Address, Bytes};
 use revm::precompile::B256;
 use revm::primitives::{AccountInfo as RevmAccountInfo, BlockEnv, CreateScheme, Env as RevmEnv, EnvWithHandlerCfg, HandlerCfg, Output, ResultAndState, SpecId, State, TransactTo, TxEnv};
 use revm::primitives::ExecutionResult::Success;
+use revm::inspectors::TracerEip3155;
 use tracing::{trace, warn};
 
 use crate::{types::{AccountInfo, Env}, utils::{addr, pydict, pyerr}};
 use crate::database::DB;
+use crate::pystdout::PySysStdout;
 
 
 #[derive(Clone, Debug)]
@@ -27,6 +31,9 @@ pub struct EVM {
     /// the passed in environment, as those limits are used by the EVM for certain opcodes like
     /// `gaslimit`.
     gas_limit: U256,
+
+    /// whether to trace the execution to stdout
+    tracing: bool,
 }
 
 #[pymethods]
@@ -46,7 +53,8 @@ impl EVM {
             db: fork_url.map(|url| DB::new_fork(url, fork_block_number)).unwrap_or(Ok(DB::new_memory()))?,
             env: env.unwrap_or_default().into(),
             gas_limit: U256::from(gas_limit),
-            handler_cfg: HandlerCfg::new(SpecId::from(spec_id))
+            handler_cfg: HandlerCfg::new(SpecId::from(spec_id)),
+            tracing,
         })
     }
 
@@ -118,7 +126,7 @@ impl EVM {
 
     #[pyo3(signature = (caller, to, calldata=None, value=None))]
     pub fn call_raw(
-        &self,
+        &mut self,
         caller: &str,
         to: &str,
         calldata: Option<Vec<u8>>,
@@ -152,8 +160,13 @@ impl EVM {
     }
 
     #[getter]
-    fn env(&self) -> PyResult<Env> {
-        Ok(self.env.clone().into())
+    fn env(&self) -> Env {
+        self.env.clone().into()
+    }
+
+    #[getter]
+    fn tracing(&self) -> bool {
+        self.tracing
     }
 }
 
@@ -197,7 +210,7 @@ impl EVM {
 
     /// Deploys a contract using the given `env` and commits the new state to the underlying
     /// database
-    fn deploy_with_env(&self, env: EnvWithHandlerCfg) -> PyResult<(Address, State)> {
+    fn deploy_with_env(&mut self, env: EnvWithHandlerCfg) -> PyResult<(Address, State)> {
         debug_assert!(
             matches!(env.tx.transact_to, TransactTo::Create(_)),
             "Expect create transaction"
@@ -220,7 +233,7 @@ impl EVM {
         }
     }
 
-    fn call_raw_with_env(&self, env: EnvWithHandlerCfg) -> PyResult<(Bytes, State)>
+    fn call_raw_with_env(&mut self, env: EnvWithHandlerCfg) -> PyResult<(Bytes, State)>
     {
         debug_assert!(
             matches!(env.tx.transact_to, TransactTo::Call(_)),
@@ -241,14 +254,26 @@ impl EVM {
         }
     }
 
-    fn run_env(&self, env: EnvWithHandlerCfg) -> Result<ResultAndState, PyErr>
+    fn run_env(&mut self, env: EnvWithHandlerCfg) -> Result<ResultAndState, PyErr>
     {
-        let evm = Evm::builder()
-            .with_db(self.db.clone())
-            .with_env_with_handler_cfg(env)
-            .build()
-            .transact()
-            .map_err(pyerr)?;
-        Ok(evm)
-    }
-}
+        let builder = Evm::builder()
+            .with_db(&mut self.db);
+
+        let result =
+            if self.tracing {
+                let tracer = TracerEip3155::new(Box::new(PySysStdout {}), true);
+                builder
+                    .with_external_context(tracer)
+                    .with_env_with_handler_cfg(env)
+                    .append_handler_register(inspector_handle_register)
+                    .build()
+                    .transact()
+            } else {
+                builder
+                    .with_env_with_handler_cfg(env)
+                    .build()
+                    .transact()
+            };
+        Ok(result.map_err(pyerr)?)
+    }
+}

src/lib.rs

@@ -16,6 +16,7 @@ pub use evm::EVM;
 mod utils;
 mod empty_db_wrapper;
 mod database;
+mod pystdout;
 
 #[pymodule]
 fn pyrevm(_py: Python<'_>, m: &PyModule) -> PyResult<()> {

src/pystdout.rs

@@ -0,0 +1,25 @@
+use std::io::Write;
+use pyo3::Python;
+use pyo3::types::PyDict;
+
+pub(crate) struct PySysStdout;
+
+// alloow us to capture stdout from python
+// based on https://github.com/PyO3/pyo3/discussions/1960#discussioncomment-8414724
+impl Write for PySysStdout {
+    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
+        let s = std::str::from_utf8(buf).unwrap();
+        Python::with_gil(|py| {
+            let locals = PyDict::new(py);
+            locals.set_item("s", s).unwrap();
+            py.run("print(s, end='')", None, Some(&locals)).unwrap();
+        });
+        Ok(buf.len())
+    }
+    fn flush(&mut self) -> std::io::Result<()> {
+        Python::with_gil(|py| {
+            py.run("import sys;sys.stdout.flush()", None, None).unwrap();
+        });
+        Ok(())
+    }
+}