feat(js/tracing): invoke js's step inside step_end

src/tracing/js/bindings.rs

@@ -38,6 +38,8 @@ macro_rules! js_value_getter {
     };
 }
 
+pub(crate) use js_value_getter;
+
 /// A macro that creates a native function that returns a captured JsValue
 macro_rules! js_value_capture_getter {
     ($value:ident, $ctx:ident) => {
@@ -235,12 +237,18 @@ impl StepLog {
 pub(crate) struct MemoryRef(GuardedNullableGc<SharedMemory>);
 
 impl MemoryRef {
-    /// Creates a new stack reference
+    /// Creates a new memory reference
     pub(crate) fn new(mem: &SharedMemory) -> (Self, GcGuard<'_, SharedMemory>) {
         let (inner, guard) = GuardedNullableGc::new_ref(mem);
         (Self(inner), guard)
     }
 
+    /// Creates a new owned memory
+    pub(crate) fn new_owned(mem: SharedMemory) -> (Self, GcGuard<'static, SharedMemory>) {
+        let (inner, guard) = GuardedNullableGc::new_owned(mem);
+        (Self(inner), guard)
+    }
+
     fn len(&self) -> usize {
         self.0.with_inner(|mem| mem.len()).unwrap_or_default()
     }
@@ -426,6 +434,12 @@ impl StackRef {
         (Self(inner), guard)
     }
 
+    /// Creates a new owned stack
+    pub(crate) fn new_owned(stack: Stack) -> (Self, GcGuard<'static, Stack>) {
+        let (inner, guard) = GuardedNullableGc::new_owned(stack);
+        (Self(inner), guard)
+    }
+
     fn peek(&self, idx: usize, ctx: &mut Context) -> JsResult<JsValue> {
         self.0
             .with_inner(|stack| {

src/tracing/js/mod.rs

@@ -3,7 +3,8 @@
 use crate::tracing::{
     js::{
         bindings::{
-            CallFrame, Contract, EvmDbRef, FrameResult, JsEvmContext, MemoryRef, StackRef, StepLog,
+            js_value_getter, CallFrame, Contract, EvmDbRef, FrameResult, GcGuard, JsEvmContext,
+            MemoryRef, StackRef, StepLog,
         },
         builtins::{register_builtins, to_serde_value, PrecompileList},
     },
@@ -12,11 +13,14 @@ use crate::tracing::{
 };
 use alloy_primitives::{Address, Bytes, Log, U256};
 pub use boa_engine::vm::RuntimeLimits;
-use boa_engine::{js_string, Context, JsError, JsObject, JsResult, JsValue, Source};
+use boa_engine::{
+    js_string, object::FunctionObjectBuilder, Context, JsError, JsObject, JsResult, JsValue,
+    NativeFunction, Source,
+};
 use revm::{
     interpreter::{
         return_revert, CallInputs, CallOutcome, CallScheme, CreateInputs, CreateOutcome, Gas,
-        InstructionResult, Interpreter, InterpreterResult,
+        InstructionResult, Interpreter, InterpreterResult, SharedMemory, Stack,
     },
     primitives::{Env, ExecutionResult, Output, ResultAndState, TransactTo},
     ContextPrecompiles, Database, DatabaseRef, EvmContext, Inspector,
@@ -72,6 +76,11 @@ pub struct JsInspector {
     call_stack: Vec<CallStackItem>,
     /// Marker to track whether the precompiles have been registered.
     precompiles_registered: bool,
+    /// Represents the current step log that is being processed, initialized in the `step`
+    /// function, and be updated and used in the `step_end` function.
+    step: Option<(JsObject, GcGuard<'static, Stack>, GcGuard<'static, SharedMemory>)>,
+    /// The gas remaining before the opcode execution.
+    gas_remaining: u64,
 }
 
 impl JsInspector {
@@ -177,6 +186,8 @@ impl JsInspector {
             step_fn,
             call_stack: Default::default(),
             precompiles_registered: false,
+            step: None,
+            gas_remaining: 0,
         })
     }
 
@@ -297,11 +308,16 @@ impl JsInspector {
         Ok(())
     }
 
-    fn try_step(&mut self, step: StepLog, db: EvmDbRef) -> JsResult<()> {
+    fn try_step(&mut self, step: JsObject, cost: u64, refund: u64, db: EvmDbRef) -> JsResult<()> {
         if let Some(step_fn) = &self.step_fn {
-            let step = step.into_js_object(&mut self.ctx)?;
-            let db = db.into_js_object(&mut self.ctx)?;
-            step_fn.call(&(self.obj.clone().into()), &[step.into(), db.into()], &mut self.ctx)?;
+            let ctx = &mut self.ctx;
+            let get_cost = js_value_getter!(cost, ctx);
+            let get_refund = js_value_getter!(refund, ctx);
+            step.set(js_string!("getCost"), get_cost, false, ctx)?;
+            step.set(js_string!("getRefund"), get_refund, false, ctx)?;
+
+            let db = db.into_js_object(ctx)?;
+            step_fn.call(&(self.obj.clone().into()), &[step.into(), db.into()], ctx)?;
         }
         Ok(())
     }
@@ -395,25 +411,32 @@ where
             return;
         }
 
-        let (db, _db_guard) = EvmDbRef::new(&context.journaled_state.state, &context.db);
+        // Update the gas remaining before the opcode execution
+        self.gas_remaining = interp.gas.remaining();
+
+        // Create and store a new step log. The `cost` and `refund` values cannot be calculated yet,
+        // as they depend on the opcode execution. These values will be updated in `step_end`
+        // after the opcode has been executed, and then the `step` function will be invoked.
+        // Initialize `cost` and `refund` as placeholders for later updates.
+        let (stack, stack_guard) = StackRef::new_owned(interp.stack.clone());
+        let (memory, memory_guard) = MemoryRef::new_owned(interp.shared_memory.clone());
 
-        let (stack, _stack_guard) = StackRef::new(&interp.stack);
-        let (memory, _memory_guard) = MemoryRef::new(&interp.shared_memory);
         let step = StepLog {
             stack,
             op: interp.current_opcode().into(),
             memory,
             pc: interp.program_counter() as u64,
             gas_remaining: interp.gas.remaining(),
-            cost: interp.gas.spent(),
+            cost: 0,
             depth: context.journaled_state.depth(),
-            refund: interp.gas.refunded() as u64,
+            refund: 0,
             error: None,
             contract: self.active_call().contract.clone(),
         };
 
-        if self.try_step(step, db).is_err() {
-            interp.instruction_result = InstructionResult::Revert;
+        match step.into_js_object(&mut self.ctx) {
+            Ok(step) => self.step = Some((step, stack_guard, memory_guard)),
+            Err(_) => interp.instruction_result = InstructionResult::Revert,
         }
     }
 
@@ -422,26 +445,37 @@ where
             return;
         }
 
-        if matches!(interp.instruction_result, return_revert!()) {
-            let (db, _db_guard) = EvmDbRef::new(&context.journaled_state.state, &context.db);
+        // Calculate the gas cost and refund after opcode execution
+        if let Some((step, _stack_guard, _memory_guard)) = self.step.take() {
+            let cost = self.gas_remaining.saturating_sub(interp.gas.remaining());
+            let refund = interp.gas.refunded() as u64;
 
-            let (stack, _stack_guard) = StackRef::new(&interp.stack);
-            let (memory, _memory_guard) = MemoryRef::new(&interp.shared_memory);
-            let step = StepLog {
-                stack,
-                op: interp.current_opcode().into(),
-                memory,
-                pc: interp.program_counter() as u64,
-                gas_remaining: interp.gas.remaining(),
-                cost: interp.gas.spent(),
-                depth: context.journaled_state.depth(),
-                refund: interp.gas.refunded() as u64,
-                error: Some(format!("{:?}", interp.instruction_result)),
-                contract: self.active_call().contract.clone(),
-            };
+            let (db, _db_guard) = EvmDbRef::new(&context.journaled_state.state, &context.db);
+            if self.try_step(step, cost, refund, db).is_err() {
+                interp.instruction_result = InstructionResult::Revert;
+            }
 
-            let _ = self.try_fault(step, db);
-        }
+            if matches!(interp.instruction_result, return_revert!()) {
+                let (db, _db_guard) = EvmDbRef::new(&context.journaled_state.state, &context.db);
+
+                let (stack, _stack_guard) = StackRef::new(&interp.stack);
+                let (memory, _memory_guard) = MemoryRef::new(&interp.shared_memory);
+                let step = StepLog {
+                    stack,
+                    op: interp.current_opcode().into(),
+                    memory,
+                    pc: interp.program_counter() as u64,
+                    gas_remaining: interp.gas.remaining(),
+                    cost,
+                    depth: context.journaled_state.depth(),
+                    refund,
+                    error: Some(format!("{:?}", interp.instruction_result)),
+                    contract: self.active_call().contract.clone(),
+                };
+
+                let _ = self.try_fault(step, db);
+            }
+        };
     }
 
     fn log(&mut self, _interp: &mut Interpreter, _context: &mut EvmContext<DB>, _log: &Log) {}

tests/it/geth_js.rs

@@ -70,3 +70,101 @@ fn test_geth_jstracer_revert() {
     // reverted operation
     assert!(result["error"].as_bool().unwrap());
 }
+
+#[test]
+fn test_geth_jstracer_op_gascost() {
+    /*
+    pragma solidity ^0.8.13;
+
+    contract Foo {
+        event Log(address indexed addr, uint256 value);
+
+        function foo() external {
+            emit Log(msg.sender, 0);
+        }
+
+        function bar() external {
+            emit Log(msg.sender, 0);
+            require(false, "barbarbar");
+        }
+    }
+    */
+
+    let code = hex!("608060405261023e806100115f395ff3fe608060405234801561000f575f80fd5b5060043610610034575f3560e01c8063c298557814610038578063febb0f7e14610042575b5f80fd5b61004061004c565b005b61004a61009c565b005b3373ffffffffffffffffffffffffffffffffffffffff167ff950957d2407bed19dc99b718b46b4ce6090c05589006dfb86fd22c34865b23e5f6040516100929190610177565b60405180910390a2565b3373ffffffffffffffffffffffffffffffffffffffff167ff950957d2407bed19dc99b718b46b4ce6090c05589006dfb86fd22c34865b23e5f6040516100e29190610177565b60405180910390a25f61012a576040517f08c379a0000000000000000000000000000000000000000000000000000000008152600401610121906101ea565b60405180910390fd5b565b5f819050919050565b5f819050919050565b5f819050919050565b5f61016161015c6101578461012c565b61013e565b610135565b9050919050565b61017181610147565b82525050565b5f60208201905061018a5f830184610168565b92915050565b5f82825260208201905092915050565b7f62617262617262617200000000000000000000000000000000000000000000005f82015250565b5f6101d4600983610190565b91506101df826101a0565b602082019050919050565b5f6020820190508181035f830152610201816101c8565b905091905056fea2646970667358221220e058dc2c4bd629d62405850cc8e08e6bfad0eea187260784445dfe8f3ee0bea564736f6c634300081a0033");
+
+    let (addr, mut evm) = deploy_contract(code.into(), Address::ZERO, SpecId::CANCUN);
+
+    let code = r#"
+{
+  data: [],
+  memoryInstructions: { "MSTORE": "W", "MSTORE8": "B", "MLOAD": "R" },
+  fault: function (_) {},
+  step: function (log) {
+    let op = log.op.toString();
+    let instructions = this.memoryInstructions;
+    if (Object.keys(instructions).includes(op)) {
+      this.data.push({
+        op: instructions[op],
+        depth: log.getDepth(),
+        offset: log.stack.peek(0),
+        gasCost: log.getCost(),
+        memorySize: log.memory.length(),
+      });
+    }
+  },
+  result: function (ctx, _) { return { error: !!ctx.error, data: this.data }; }
+}
+"#;
+
+    // test with normal operation
+    let env = evm.env_with_tx(TxEnv {
+        transact_to: TransactTo::Call(addr),
+        data: hex!("c2985578").into(), // call foo
+        ..Default::default()
+    });
+    let mut insp = JsInspector::new(code.to_string(), serde_json::Value::Null).unwrap();
+    let (res, _) = inspect(&mut evm.db, env.clone(), &mut insp).unwrap();
+    assert!(res.result.is_success());
+
+    let result = insp.json_result(res, &env, &evm.db).unwrap();
+
+    assert!(!result["error"].as_bool().unwrap());
+    assert_eq!(
+        result["data"],
+        serde_json::json!([
+            { "op": "W", "depth": 1, "offset": "64", "gasCost": 12, "memorySize": 0 },
+            { "op": "R", "depth": 1, "offset": "64", "gasCost": 3, "memorySize": 96 },
+            { "op": "W", "depth": 1, "offset": "128", "gasCost": 9, "memorySize": 96 },
+            { "op": "R", "depth": 1, "offset": "64", "gasCost": 3, "memorySize": 160 }
+        ])
+    );
+
+    // test with reverted operation
+    let env = evm.env_with_tx(TxEnv {
+        transact_to: TransactTo::Call(addr),
+        data: hex!("febb0f7e").into(), // call bar
+        ..Default::default()
+    });
+    let mut insp = JsInspector::new(code.to_string(), serde_json::Value::Null).unwrap();
+    let (res, _) = inspect(&mut evm.db, env.clone(), &mut insp).unwrap();
+    assert!(!res.result.is_success());
+
+    let result = insp.json_result(res, &env, &evm.db).unwrap();
+
+    assert!(result["error"].as_bool().unwrap());
+    assert_eq!(
+        result["data"],
+        serde_json::json!([
+            { "op": "W", "depth": 1, "offset": "64", "gasCost": 12, "memorySize": 0 },
+            { "op": "R", "depth": 1, "offset": "64", "gasCost": 3, "memorySize": 96 },
+            { "op": "W", "depth": 1, "offset": "128", "gasCost": 9, "memorySize": 96 },
+            { "op": "R", "depth": 1, "offset": "64", "gasCost": 3, "memorySize": 160 },
+            { "op": "R", "depth": 1, "offset": "64", "gasCost": 3, "memorySize": 160 },
+            { "op": "W", "depth": 1, "offset": "128", "gasCost": 3, "memorySize": 160 },
+            { "op": "W", "depth": 1, "offset": "132", "gasCost": 6, "memorySize": 160 },
+            { "op": "W", "depth": 1, "offset": "164", "gasCost": 6, "memorySize": 192 },
+            { "op": "W", "depth": 1, "offset": "196", "gasCost": 6, "memorySize": 224 },
+            { "op": "R", "depth": 1, "offset": "64", "gasCost": 3, "memorySize": 256 }
+        ])
+    );
+}