Refactor DA setup

matter-labs · Jan 15, 2025 · 42b78b2 · 42b78b2
1 parent c23eac0
commit 42b78b2
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 103 deletions.
diff --git a/core/node/state_keeper/src/executor/tests/tester.rs b/core/node/state_keeper/src/executor/tests/tester.rs
@@ -1,7 +1,7 @@
 //! Testing harness for the batch executor.
 //! Contains helper functionality to initialize test context and perform tests without too much boilerplate.
 
-use std::{collections::HashMap, fmt::Debug, str::FromStr, sync::Arc};
+use std::{collections::HashMap, fmt::Debug, sync::Arc};
 
 use assert_matches::assert_matches;
 use tempfile::TempDir;
@@ -42,18 +42,13 @@ use zksync_vm_executor::batch::{MainBatchExecutorFactory, TraceCalls};
 
 use super::{read_storage_factory::RocksdbStorageFactory, StorageType};
 use crate::{
-    testonly::{self, BASE_SYSTEM_CONTRACTS},
+    testonly::{self, apply_genesis_logs, BASE_SYSTEM_CONTRACTS},
     tests::{default_l1_batch_env, default_system_env},
     AsyncRocksdbCache,
 };
 
 pub(super) const TRANSFER_VALUE: u64 = 123_456_789;
 
-// FIXME: what
-fn get_da_contract_address() -> Address {
-    Address::from_str("7726827caac94a7f9e1b160f7ea819f172f7b6f9").unwrap()
-}
-
 /// Representation of configuration parameters used by the state keeper.
 /// Has sensible defaults for most tests, each of which can be overridden.
 #[derive(Debug)]
@@ -297,7 +292,7 @@ impl Tester {
             .await
             .unwrap();
 
-            // Also setting up the da for tests
+            // Also setting up the DA for tests
             Self::setup_da(&mut storage).await;
         }
     }
@@ -337,37 +332,28 @@ impl Tester {
         }
     }
 
-    pub async fn setup_contract<'a>(
-        con: &mut Connection<'a, Core>,
-        address: Address,
-        code: Vec<u8>,
-    ) {
+    async fn setup_contract(conn: &mut Connection<'_, Core>, address: Address, code: Vec<u8>) {
         let hash: H256 = BytecodeHash::for_bytecode(&code).value();
         let known_code_key = get_known_code_key(&hash);
         let code_key = get_code_key(&address);
 
-        let logs = vec![
-            StorageLog::new_write_log(known_code_key, H256::from_low_u64_be(1u64)),
+        let logs = [
+            StorageLog::new_write_log(known_code_key, H256::from_low_u64_be(1)),
             StorageLog::new_write_log(code_key, hash),
         ];
+        apply_genesis_logs(conn, &logs).await;
 
-        for log in logs {
-            apply_genesis_log(con, log).await;
-        }
-
-        let mut factory_deps = HashMap::new();
-        factory_deps.insert(hash, code);
-
-        con.factory_deps_dal()
+        let factory_deps = HashMap::from([(hash, code)]);
+        conn.factory_deps_dal()
             .insert_factory_deps(L2BlockNumber(0), &factory_deps)
             .await
             .unwrap();
     }
 
-    async fn setup_da<'a>(con: &mut Connection<'a, Core>) {
+    async fn setup_da(conn: &mut Connection<'_, Core>) {
         Self::setup_contract(
-            con,
-            get_da_contract_address(),
+            conn,
+            Address::repeat_byte(0x23),
             l2_rollup_da_validator_bytecode(),
         )
         .await;
@@ -770,25 +756,3 @@ impl StorageSnapshot {
         snapshot
     }
 }
-
-async fn apply_genesis_log<'a>(storage: &mut Connection<'a, Core>, log: StorageLog) {
-    storage
-        .storage_logs_dal()
-        .append_storage_logs(L2BlockNumber(0), &[log])
-        .await
-        .unwrap();
-
-    if storage
-        .storage_logs_dedup_dal()
-        .filter_written_slots(&[log.key.hashed_key()])
-        .await
-        .unwrap()
-        .is_empty()
-    {
-        storage
-            .storage_logs_dedup_dal()
-            .insert_initial_writes(L1BatchNumber(0), &[log.key.hashed_key()])
-            .await
-            .unwrap();
-    }
-}
diff --git a/core/node/state_keeper/src/testonly/mod.rs b/core/node/state_keeper/src/testonly/mod.rs
@@ -1,7 +1,7 @@
 //! Test utilities that can be used for testing sequencer that may
 //! be useful outside of this crate.
 
-use std::collections::HashMap;
+use std::collections::HashSet;
 
 use async_trait::async_trait;
 use once_cell::sync::Lazy;
@@ -15,10 +15,10 @@ use zksync_multivm::interface::{
 };
 use zksync_state::OwnedStorage;
 use zksync_types::{
-    bytecode::BytecodeHash, commitment::PubdataParams, fee::Fee, get_code_key, get_known_code_key,
-    u256_to_h256, utils::storage_key_for_standard_token_balance, AccountTreeId, Address,
-    L1BatchNumber, L2BlockNumber, StorageLog, Transaction, H256, L2_BASE_TOKEN_ADDRESS,
-    SYSTEM_CONTEXT_MINIMAL_BASE_FEE, U256,
+    commitment::PubdataParams, fee::Fee, u256_to_h256,
+    utils::storage_key_for_standard_token_balance, AccountTreeId, Address, L1BatchNumber,
+    L2BlockNumber, StorageLog, Transaction, L2_BASE_TOKEN_ADDRESS, SYSTEM_CONTEXT_MINIMAL_BASE_FEE,
+    U256,
 };
 
 pub mod test_batch_executor;
@@ -76,25 +76,28 @@ impl BatchExecutor<OwnedStorage> for MockBatchExecutor {
     }
 }
 
-async fn apply_genesis_log<'a>(storage: &mut Connection<'a, Core>, log: StorageLog) {
+pub(crate) async fn apply_genesis_logs(storage: &mut Connection<'_, Core>, logs: &[StorageLog]) {
     storage
         .storage_logs_dal()
-        .append_storage_logs(L2BlockNumber(0), &[log])
+        .append_storage_logs(L2BlockNumber(0), logs)
         .await
         .unwrap();
-    if storage
+
+    let all_hashed_keys: Vec<_> = logs.iter().map(|log| log.key.hashed_key()).collect();
+    let repeated_writes = storage
         .storage_logs_dedup_dal()
-        .filter_written_slots(&[log.key.hashed_key()])
+        .filter_written_slots(&all_hashed_keys)
         .await
-        .unwrap()
-        .is_empty()
-    {
-        storage
-            .storage_logs_dedup_dal()
-            .insert_initial_writes(L1BatchNumber(0), &[log.key.hashed_key()])
-            .await
-            .unwrap();
-    }
+        .unwrap();
+    let initial_writes: Vec<_> = HashSet::from_iter(all_hashed_keys)
+        .difference(&repeated_writes)
+        .copied()
+        .collect();
+    storage
+        .storage_logs_dedup_dal()
+        .insert_initial_writes(L1BatchNumber(0), &initial_writes)
+        .await
+        .unwrap();
 }
 
 /// Adds funds for specified account list.
@@ -103,43 +106,18 @@ pub async fn fund(pool: &ConnectionPool<Core>, addresses: &[Address]) {
     let mut storage = pool.connection().await.unwrap();
 
     let eth_amount = U256::from(10u32).pow(U256::from(32)); //10^32 wei
-
-    for address in addresses {
-        let key = storage_key_for_standard_token_balance(
-            AccountTreeId::new(L2_BASE_TOKEN_ADDRESS),
-            address,
-        );
-        let value = u256_to_h256(eth_amount);
-        let storage_log = StorageLog::new_write_log(key, value);
-
-        apply_genesis_log(&mut storage, storage_log).await;
-    }
-}
-
-pub async fn setup_contract(pool: &ConnectionPool<Core>, address: Address, code: Vec<u8>) {
-    let mut storage = pool.connection().await.unwrap();
-
-    let hash: H256 = BytecodeHash::for_bytecode(&code).value();
-    let known_code_key = get_known_code_key(&hash);
-    let code_key = get_code_key(&address);
-
-    let logs = vec![
-        StorageLog::new_write_log(known_code_key, H256::from_low_u64_be(1u64)),
-        StorageLog::new_write_log(code_key, hash),
-    ];
-
-    for log in logs {
-        apply_genesis_log(&mut storage, log).await;
-    }
-
-    let mut factory_deps = HashMap::new();
-    factory_deps.insert(hash, code);
-
-    storage
-        .factory_deps_dal()
-        .insert_factory_deps(L2BlockNumber(0), &factory_deps)
-        .await
-        .unwrap();
+    let storage_logs: Vec<_> = addresses
+        .iter()
+        .map(|address| {
+            let key = storage_key_for_standard_token_balance(
+                AccountTreeId::new(L2_BASE_TOKEN_ADDRESS),
+                address,
+            );
+            StorageLog::new_write_log(key, u256_to_h256(eth_amount))
+        })
+        .collect();
+
+    apply_genesis_logs(&mut storage, &storage_logs).await;
 }
 
 pub(crate) const DEFAULT_GAS_PER_PUBDATA: u32 = 10000;