From 8422d718f1b6e10f6571423648493cec7ccd4b79 Mon Sep 17 00:00:00 2001
From: Nisheeth Barthwal <nbaztec@gmail.com>
Date: Tue, 24 Sep 2024 14:41:34 +0200
Subject: [PATCH] feat: add additional gas usage info from bootloader (#590)

* add additional gas usage info from bootloader
---
 crates/zksync/core/src/vm/inspect.rs | 72 ++++++++++++++++++++++++++--
 1 file changed, 69 insertions(+), 3 deletions(-)

diff --git a/crates/zksync/core/src/vm/inspect.rs b/crates/zksync/core/src/vm/inspect.rs
index 967e43e11..d6e21089f 100644
--- a/crates/zksync/core/src/vm/inspect.rs
+++ b/crates/zksync/core/src/vm/inspect.rs
@@ -1,5 +1,6 @@
 use alloy_primitives::{hex, Log};
 use era_test_node::{
+    bootloader_debug::{BootloaderDebug, BootloaderDebugTracer},
     config::node::ShowCalls,
     formatter,
     system_contracts::{Options, SystemContracts},
@@ -188,8 +189,19 @@ where
     let storage_ptr =
         StorageView::new(&mut era_db, modified_storage_keys, tx.common_data.initiator_address)
             .into_rc_ptr();
-    let InnerZkVmResult { tx_result, bytecodes, modified_storage, call_traces, create_outcome } =
-        inspect_inner(tx, storage_ptr, chain_id, ccx, call_ctx);
+    let InnerZkVmResult {
+        tx_result,
+        bytecodes,
+        modified_storage,
+        call_traces,
+        create_outcome,
+        gas_usage,
+    } = inspect_inner(tx, storage_ptr, chain_id, ccx, call_ctx);
+
+    info!(
+        limit=?gas_usage.limit, execution=?gas_usage.execution, pubdata=?gas_usage.pubdata, refunded=?gas_usage.refunded,
+        "gas usage",
+    );
 
     if let Some(record) = &mut era_db.accesses {
         for k in modified_storage.keys() {
@@ -364,12 +376,28 @@ struct InnerCreateOutcome {
     bytecode: Vec<u8>,
 }
 
+#[allow(unused)]
+#[derive(Debug)]
+struct ZkVmGasUsage {
+    /// Gas limit set for the user excluding the reserved gas.
+    pub limit: U256,
+    /// Gas refunded after transaction execution by the operator.
+    pub refunded: U256,
+    /// Gas used for only on transaction execution (validation and execution).
+    pub execution: U256,
+    /// Gas used for publishing pubdata.
+    pub pubdata: U256,
+    /// Additional bootloader debug info for gas usage.
+    pub bootloader_debug: BootloaderDebug,
+}
+
 struct InnerZkVmResult {
     tx_result: VmExecutionResultAndLogs,
     bytecodes: HashMap<U256, Vec<U256>>,
     modified_storage: HashMap<StorageKey, H256>,
     call_traces: Vec<Call>,
     create_outcome: Option<InnerCreateOutcome>,
+    gas_usage: ZkVmGasUsage,
 }
 
 fn inspect_inner<S: ReadStorage>(
@@ -401,8 +429,11 @@ fn inspect_inner<S: ReadStorage>(
     }
     let is_static = call_ctx.is_static;
     let is_create = call_ctx.is_create;
+    let bootloader_debug_tracer_result = Arc::new(OnceCell::default());
     let tracers = vec![
         CallTracer::new(call_tracer_result.clone()).into_tracer_pointer(),
+        BootloaderDebugTracer { result: bootloader_debug_tracer_result.clone() }
+            .into_tracer_pointer(),
         CheatcodeTracer::new(
             ccx.mocked_calls.clone(),
             expected_calls,
@@ -434,6 +465,34 @@ fn inspect_inner<S: ReadStorage>(
         expected_calls.extend(cheatcode_result.expected_calls);
     }
 
+    // populate gas usage info
+    let bootloader_debug = Arc::try_unwrap(bootloader_debug_tracer_result)
+        .unwrap()
+        .take()
+        .and_then(|result| result.ok())
+        .expect("failed obtaining bootloader debug info");
+    trace!("{bootloader_debug:?}");
+
+    let total_gas_limit =
+        bootloader_debug.total_gas_limit_from_user.saturating_sub(bootloader_debug.reserved_gas);
+    let intrinsic_gas = total_gas_limit - bootloader_debug.gas_limit_after_intrinsic;
+    let gas_for_validation =
+        bootloader_debug.gas_limit_after_intrinsic - bootloader_debug.gas_after_validation;
+    let gas_used_tx_execution =
+        intrinsic_gas + gas_for_validation + bootloader_debug.gas_spent_on_execution;
+
+    let gas_used_pubdata = bootloader_debug
+        .gas_per_pubdata
+        .saturating_mul(tx_result.statistics.pubdata_published.into());
+
+    let gas_usage = ZkVmGasUsage {
+        limit: total_gas_limit,
+        execution: gas_used_tx_execution,
+        pubdata: gas_used_pubdata,
+        refunded: bootloader_debug.refund_by_operator,
+        bootloader_debug,
+    };
+
     formatter::print_vm_details(&tx_result);
 
     info!("=== Console Logs: ");
@@ -517,7 +576,14 @@ fn inspect_inner<S: ReadStorage>(
         None
     };
 
-    InnerZkVmResult { tx_result, bytecodes, modified_storage, call_traces, create_outcome }
+    InnerZkVmResult {
+        tx_result,
+        bytecodes,
+        modified_storage,
+        call_traces,
+        create_outcome,
+        gas_usage,
+    }
 }
 
 /// Patch CREATE traces with bytecode as the data is empty bytes.