perf(state-keeper): Improve FilterWrittenSlots l1 batch seal stage (#…

…1854)

## What ❔

New versions of VM (1.4.2 and 1.5.0) return list of initially written
slots. State keeper doesn't need to query DB to filter slots anymore.

## Why ❔

Improve L1 batch seal performance.

## Checklist

<!-- Check your PR fulfills the following items. -->
<!-- For draft PRs check the boxes as you complete them. -->

- [ ] PR title corresponds to the body of PR (we generate changelog
entries from PRs).
- [ ] Tests for the changes have been added / updated.
- [ ] Documentation comments have been added / updated.
- [ ] Code has been formatted via `zk fmt` and `zk lint`.
- [ ] Spellcheck has been run via `zk spellcheck`.
- [ ] Linkcheck has been run via `zk linkcheck`.

core/lib/multivm/src/glue/types/vm/vm_block_result.rs

@@ -71,6 +71,7 @@ impl GlueFrom<crate::vm_m5::vm_instance::VmBlockResult> for crate::interface::Fi
             },
             final_bootloader_memory: None,
             pubdata_input: None,
+            initially_written_slots: None,
         }
     }
 }
@@ -130,6 +131,7 @@ impl GlueFrom<crate::vm_m6::vm_instance::VmBlockResult> for crate::interface::Fi
             },
             final_bootloader_memory: None,
             pubdata_input: None,
+            initially_written_slots: None,
         }
     }
 }
@@ -187,6 +189,7 @@ impl GlueFrom<crate::vm_1_3_2::vm_instance::VmBlockResult> for crate::interface:
             },
             final_bootloader_memory: None,
             pubdata_input: None,
+            initially_written_slots: None,
         }
     }
 }

core/lib/multivm/src/interface/traits/vm.rs

@@ -143,6 +143,7 @@ pub trait VmInterface<S, H: HistoryMode> {
             final_execution_state: execution_state,
             final_bootloader_memory: Some(bootloader_memory),
             pubdata_input: None,
+            initially_written_slots: None,
         }
     }
 }

core/lib/multivm/src/interface/types/outputs/finished_l1batch.rs

@@ -1,3 +1,5 @@
+use zksync_types::H256;
+
 use super::{BootloaderMemory, CurrentExecutionState, VmExecutionResultAndLogs};
 
 /// State of the VM after the batch execution.
@@ -7,7 +9,11 @@ pub struct FinishedL1Batch {
     pub block_tip_execution_result: VmExecutionResultAndLogs,
     /// State of the VM after the execution of the last transaction.
     pub final_execution_state: CurrentExecutionState,
-    /// Memory of the bootloader with all executed transactions. Could be optional for old versions of the VM.
+    /// Memory of the bootloader with all executed transactions. Could be none for old versions of the VM.
     pub final_bootloader_memory: Option<BootloaderMemory>,
+    /// Pubdata to be published on L1. Could be none for old versions of the VM.
     pub pubdata_input: Option<Vec<u8>>,
+    /// List of hashed keys of slots that were initially written in the batch.
+    /// Could be none for old versions of the VM.
+    pub initially_written_slots: Option<Vec<H256>>,
 }

core/lib/multivm/src/versions/vm_1_4_1/vm.rs

@@ -179,6 +179,7 @@ impl<S: WriteStorage, H: HistoryMode> VmInterface<S, H> for Vm<S, H> {
                     .clone()
                     .build_pubdata(false),
             ),
+            initially_written_slots: None,
         }
     }
 }

core/lib/multivm/src/versions/vm_1_4_2/vm.rs

@@ -3,7 +3,7 @@ use zksync_state::{StoragePtr, WriteStorage};
 use zksync_types::{
     event::extract_l2tol1logs_from_l1_messenger,
     l2_to_l1_log::{SystemL2ToL1Log, UserL2ToL1Log},
-    Transaction,
+    Transaction, H256,
 };
 use zksync_utils::bytecode::CompressedBytecodeInfo;
 
@@ -179,6 +179,18 @@ impl<S: WriteStorage, H: HistoryMode> VmInterface<S, H> for Vm<S, H> {
                     .clone()
                     .build_pubdata(false),
             ),
+            initially_written_slots: Some(
+                self.bootloader_state
+                    .get_pubdata_information()
+                    .state_diffs
+                    .iter()
+                    .filter_map(|record| {
+                        record
+                            .is_write_initial()
+                            .then_some(H256(record.derived_key))
+                    })
+                    .collect(),
+            ),
         }
     }
 }

core/lib/multivm/src/versions/vm_boojum_integration/vm.rs

@@ -179,6 +179,7 @@ impl<S: WriteStorage, H: HistoryMode> VmInterface<S, H> for Vm<S, H> {
                     .clone()
                     .build_pubdata(false),
             ),
+            initially_written_slots: None,
         }
     }
 }

core/lib/multivm/src/versions/vm_latest/vm.rs

@@ -3,7 +3,7 @@ use zksync_state::{StoragePtr, WriteStorage};
 use zksync_types::{
     event::extract_l2tol1logs_from_l1_messenger,
     l2_to_l1_log::{SystemL2ToL1Log, UserL2ToL1Log},
-    Transaction, VmVersion,
+    Transaction, VmVersion, H256,
 };
 use zksync_utils::bytecode::CompressedBytecodeInfo;
 
@@ -209,6 +209,18 @@ impl<S: WriteStorage, H: HistoryMode> VmInterface<S, H> for Vm<S, H> {
                     .clone()
                     .build_pubdata(false),
             ),
+            initially_written_slots: Some(
+                self.bootloader_state
+                    .get_pubdata_information()
+                    .state_diffs
+                    .iter()
+                    .filter_map(|record| {
+                        record
+                            .is_write_initial()
+                            .then_some(H256(record.derived_key))
+                    })
+                    .collect(),
+            ),
         }
     }
 }

core/lib/types/src/storage/writes/mod.rs

@@ -142,6 +142,10 @@ impl StateDiffRecord {
 
         comp_state_diff
     }
+
+    pub fn is_write_initial(&self) -> bool {
+        self.enumeration_index == 0
+    }
 }
 
 /// Compresses a vector of state diff records according to the following:

core/lib/zksync_core/src/state_keeper/io/persistence.rs

@@ -249,9 +249,10 @@ mod tests {
     use multivm::zk_evm_latest::ethereum_types::{H256, U256};
     use zksync_dal::CoreDal;
     use zksync_types::{
-        api::TransactionStatus, block::BlockGasCount, tx::ExecutionMetrics, L1BatchNumber,
-        L2BlockNumber,
+        api::TransactionStatus, block::BlockGasCount, tx::ExecutionMetrics, AccountTreeId,
+        L1BatchNumber, L2BlockNumber, StorageKey, StorageLogQueryType,
     };
+    use zksync_utils::u256_to_h256;
 
     use super::*;
     use crate::{
@@ -351,10 +352,22 @@ mod tests {
         let mut batch_result = default_vm_batch_result();
         batch_result
             .final_execution_state
-            .deduplicated_storage_log_queries = storage_logs
-            .into_iter()
-            .map(|query| query.log_query)
-            .collect();
+            .deduplicated_storage_log_queries =
+            storage_logs.iter().map(|query| query.log_query).collect();
+        batch_result.initially_written_slots = Some(
+            storage_logs
+                .into_iter()
+                .filter(|&log| log.log_type == StorageLogQueryType::InitialWrite)
+                .map(|log| {
+                    let key = StorageKey::new(
+                        AccountTreeId::new(log.log_query.address),
+                        u256_to_h256(log.log_query.key),
+                    );
+                    key.hashed_key()
+                })
+                .collect(),
+        );
+
         updates.finish_batch(batch_result);
         persistence.handle_l1_batch(&updates).await.unwrap();
 

core/lib/zksync_core/src/state_keeper/io/seal_logic.rs

@@ -172,35 +172,50 @@ impl UpdatesManager {
         }
 
         let progress = L1_BATCH_METRICS.start(L1BatchSealStage::FilterWrittenSlots);
-        let deduplicated_writes_hashed_keys: Vec<_> = deduplicated_writes
-            .iter()
-            .map(|log| {
-                H256(StorageKey::raw_hashed_key(
-                    &log.address,
-                    &u256_to_h256(log.key),
-                ))
-            })
-            .collect();
-        let non_initial_writes = transaction
-            .storage_logs_dedup_dal()
-            .filter_written_slots(&deduplicated_writes_hashed_keys)
-            .await?;
+        let written_storage_keys: Vec<_> =
+            if let Some(initially_written_slots) = &finished_batch.initially_written_slots {
+                deduplicated_writes
+                    .iter()
+                    .filter_map(|log| {
+                        let key =
+                            StorageKey::new(AccountTreeId::new(log.address), u256_to_h256(log.key));
+                        initially_written_slots
+                            .contains(&key.hashed_key())
+                            .then_some(key)
+                    })
+                    .collect()
+            } else {
+                let deduplicated_writes_hashed_keys: Vec<_> = deduplicated_writes
+                    .iter()
+                    .map(|log| {
+                        H256(StorageKey::raw_hashed_key(
+                            &log.address,
+                            &u256_to_h256(log.key),
+                        ))
+                    })
+                    .collect();
+                let non_initial_writes = transaction
+                    .storage_logs_dedup_dal()
+                    .filter_written_slots(&deduplicated_writes_hashed_keys)
+                    .await?;
+
+                deduplicated_writes
+                    .iter()
+                    .filter_map(|log| {
+                        let key =
+                            StorageKey::new(AccountTreeId::new(log.address), u256_to_h256(log.key));
+                        (!non_initial_writes.contains(&key.hashed_key())).then_some(key)
+                    })
+                    .collect()
+            };
         progress.observe(deduplicated_writes.len());
 
         let progress = L1_BATCH_METRICS.start(L1BatchSealStage::InsertInitialWrites);
-        let written_storage_keys: Vec<_> = deduplicated_writes
-            .iter()
-            .filter_map(|log| {
-                let key = StorageKey::new(AccountTreeId::new(log.address), u256_to_h256(log.key));
-                (!non_initial_writes.contains(&key.hashed_key())).then_some(key)
-            })
-            .collect();
-
         transaction
             .storage_logs_dedup_dal()
             .insert_initial_writes(self.l1_batch.number, &written_storage_keys)
             .await?;
-        progress.observe(deduplicated_writes.len());
+        progress.observe(written_storage_keys.len());
 
         let progress = L1_BATCH_METRICS.start(L1BatchSealStage::CommitL1Batch);
         transaction.commit().await?;

core/lib/zksync_core/src/state_keeper/tests/mod.rs

@@ -113,6 +113,7 @@ pub(super) fn default_vm_batch_result() -> FinishedL1Batch {
         },
         final_bootloader_memory: Some(vec![]),
         pubdata_input: Some(vec![]),
+        initially_written_slots: Some(vec![]),
     }
 }