feat: Added detailed information on amount of pubdata per write (#248)

* moved storage log related stuff to separate file

* Added info about pubdata

* review

src/formatter.rs

@@ -197,7 +197,47 @@ pub fn print_call(call: &Call, padding: usize, show_calls: &ShowCalls, resolve_h
     }
 }
 
-pub fn print_logs(log_query: &StorageLogQuery) {
+/// Amount of pubdata that given write has cost.
+pub enum PubdataBytesInfo {
+    // This slot is free
+    FreeSlot,
+    // This slot costs this much.
+    Paid(u32),
+    // This happens when we already paid a litte for this slot in the past.
+    // This slots costs additional X, the total cost is Y.
+    AdditionalPayment(u32, u32),
+    // We already paid for this slot in this transaction.
+    PaidAlready,
+}
+
+impl std::fmt::Display for PubdataBytesInfo {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            PubdataBytesInfo::FreeSlot => write!(f, "free slot"),
+            PubdataBytesInfo::Paid(cost) => write!(f, "{:?} bytes", cost),
+            PubdataBytesInfo::AdditionalPayment(additional_cost, total_cost) => write!(
+                f,
+                "{:?} addditional bytes, {:?} total cost",
+                additional_cost, total_cost
+            ),
+            PubdataBytesInfo::PaidAlready => write!(f, "already paid"),
+        }
+    }
+}
+
+impl PubdataBytesInfo {
+    // Whether the slot incurs any cost
+    pub fn does_cost(&self) -> bool {
+        match self {
+            PubdataBytesInfo::FreeSlot => false,
+            PubdataBytesInfo::Paid(_) => true,
+            PubdataBytesInfo::AdditionalPayment(_, _) => true,
+            PubdataBytesInfo::PaidAlready => false,
+        }
+    }
+}
+
+pub fn print_logs(log_query: &StorageLogQuery, pubdata_bytes: Option<PubdataBytesInfo>) {
     let separator = "─".repeat(82);
     tracing::info!("{:<15} {:?}", "Type:", log_query.log_type);
     tracing::info!(
@@ -221,6 +261,9 @@ pub fn print_logs(log_query: &StorageLogQuery) {
             log_query.log_query.written_value
         );
     }
+    if let Some(pubdata_bytes) = pubdata_bytes {
+        tracing::info!("{:<15} {:}", "Pubdata bytes:", pubdata_bytes);
+    }
     tracing::info!("{}", separator);
 }
 

src/node/in_memory.rs

@@ -6,6 +6,7 @@ use crate::{
     filters::EthFilters,
     fork::{ForkDetails, ForkSource, ForkStorage},
     formatter,
+    node::storage_logs::print_storage_logs_details,
     observability::Observability,
     system_contracts::{self, SystemContracts},
     utils::{
@@ -166,6 +167,7 @@ pub enum ShowStorageLogs {
     None,
     Read,
     Write,
+    Paid,
     All,
 }
 
@@ -177,9 +179,10 @@ impl FromStr for ShowStorageLogs {
             "none" => Ok(ShowStorageLogs::None),
             "read" => Ok(ShowStorageLogs::Read),
             "write" => Ok(ShowStorageLogs::Write),
+            "paid" => Ok(ShowStorageLogs::Paid),
             "all" => Ok(ShowStorageLogs::All),
             _ => Err(format!(
-                "Unknown ShowStorageLogs value {} - expected one of none|read|write|all.",
+                "Unknown ShowStorageLogs value {} - expected one of none|read|write|paid|all.",
                 s
             )),
         }
@@ -1372,32 +1375,7 @@ impl<S: ForkSource + std::fmt::Debug + Clone> InMemoryNode<S> {
         }
 
         if inner.show_storage_logs != ShowStorageLogs::None {
-            tracing::info!("");
-            tracing::info!("┌──────────────────┐");
-            tracing::info!("│   STORAGE LOGS   │");
-            tracing::info!("└──────────────────┘");
-        }
-
-        for log_query in &tx_result.logs.storage_logs {
-            match inner.show_storage_logs {
-                ShowStorageLogs::Write => {
-                    if matches!(
-                        log_query.log_type,
-                        StorageLogQueryType::RepeatedWrite | StorageLogQueryType::InitialWrite
-                    ) {
-                        formatter::print_logs(log_query);
-                    }
-                }
-                ShowStorageLogs::Read => {
-                    if log_query.log_type == StorageLogQueryType::Read {
-                        formatter::print_logs(log_query);
-                    }
-                }
-                ShowStorageLogs::All => {
-                    formatter::print_logs(log_query);
-                }
-                _ => {}
-            }
+            print_storage_logs_details(&inner.show_storage_logs, &tx_result);
         }
 
         if inner.show_vm_details != ShowVMDetails::None {

src/node/mod.rs

@@ -8,6 +8,7 @@ mod hardhat;
 mod in_memory;
 mod in_memory_ext;
 mod net;
+mod storage_logs;
 mod web3;
 mod zks;
 

src/node/storage_logs.rs

@@ -0,0 +1,116 @@
+use std::collections::HashMap;
+
+use crate::formatter::{self, PubdataBytesInfo};
+
+use super::ShowStorageLogs;
+use multivm::vm_latest::VmExecutionResultAndLogs;
+use zksync_basic_types::AccountTreeId;
+use zksync_types::{
+    utils::storage_key_for_eth_balance,
+    writes::{
+        compression::compress_with_best_strategy, BYTES_PER_DERIVED_KEY,
+        BYTES_PER_ENUMERATION_INDEX,
+    },
+    StorageKey, StorageLogQuery, StorageLogQueryType, BOOTLOADER_ADDRESS, SYSTEM_CONTEXT_ADDRESS,
+};
+use zksync_utils::u256_to_h256;
+
+fn is_storage_key_free(key: &StorageKey) -> bool {
+    key.address() == &SYSTEM_CONTEXT_ADDRESS
+        || *key == storage_key_for_eth_balance(&BOOTLOADER_ADDRESS)
+}
+
+fn compute_and_update_pubdata_cost(
+    cost_paid: &mut HashMap<StorageKey, u32>,
+    log_query: &StorageLogQuery,
+) -> PubdataBytesInfo {
+    let storage_key = StorageKey::new(
+        AccountTreeId::new(log_query.log_query.address),
+        u256_to_h256(log_query.log_query.key),
+    );
+
+    if is_storage_key_free(&storage_key) {
+        PubdataBytesInfo::FreeSlot
+    } else {
+        // how many bytes it takes after compression.
+        let compressed_value_size = compress_with_best_strategy(
+            log_query.log_query.read_value,
+            log_query.log_query.written_value,
+        )
+        .len() as u32;
+
+        let final_pubdata_cost = if log_query.log_type == StorageLogQueryType::InitialWrite {
+            (BYTES_PER_DERIVED_KEY as u32) + compressed_value_size
+        } else {
+            (BYTES_PER_ENUMERATION_INDEX as u32) + compressed_value_size
+        };
+
+        let result = match cost_paid.get(&storage_key).copied() {
+            Some(already_paid) => {
+                let to_pay = final_pubdata_cost.saturating_sub(already_paid);
+                if to_pay > 0 {
+                    PubdataBytesInfo::AdditionalPayment(to_pay, final_pubdata_cost)
+                } else {
+                    PubdataBytesInfo::PaidAlready
+                }
+            }
+            None => PubdataBytesInfo::Paid(final_pubdata_cost),
+        };
+        cost_paid.insert(storage_key, final_pubdata_cost);
+        result
+    }
+}
+
+pub fn print_storage_logs_details(
+    show_storage_logs: &ShowStorageLogs,
+    result: &VmExecutionResultAndLogs,
+) {
+    tracing::info!("");
+    tracing::info!("┌──────────────────┐");
+    tracing::info!("│   STORAGE LOGS   │");
+    tracing::info!("└──────────────────┘");
+
+    let mut cost_paid = HashMap::<StorageKey, u32>::default();
+
+    for log_query in &result.logs.storage_logs {
+        let pubdata_bytes_info = if matches!(
+            log_query.log_type,
+            StorageLogQueryType::RepeatedWrite | StorageLogQueryType::InitialWrite
+        ) {
+            Some(compute_and_update_pubdata_cost(&mut cost_paid, log_query))
+        } else {
+            None
+        };
+
+        match show_storage_logs {
+            ShowStorageLogs::Write => {
+                if matches!(
+                    log_query.log_type,
+                    StorageLogQueryType::RepeatedWrite | StorageLogQueryType::InitialWrite
+                ) {
+                    formatter::print_logs(log_query, pubdata_bytes_info);
+                }
+            }
+            ShowStorageLogs::Paid => {
+                // Show only the logs that incur any cost.
+                if pubdata_bytes_info
+                    .as_ref()
+                    .map(|x| x.does_cost())
+                    .unwrap_or_default()
+                {
+                    formatter::print_logs(log_query, pubdata_bytes_info);
+                }
+            }
+            ShowStorageLogs::Read => {
+                if log_query.log_type == StorageLogQueryType::Read {
+                    formatter::print_logs(log_query, pubdata_bytes_info);
+                }
+            }
+            ShowStorageLogs::All => {
+                formatter::print_logs(log_query, pubdata_bytes_info);
+            }
+
+            _ => {}
+        }
+    }
+}