Use ZkSyncState in init params directly

core/bin/zksync_core/src/state_keeper/init_params.rs

@@ -1,9 +1,10 @@
 use std::collections::{HashMap, VecDeque};
 // External uses
 // Workspace uses
+use zksync_state::state::ZkSyncState;
 use zksync_types::{
     block::{IncompleteBlock, PendingBlock as SendablePendingBlock},
-    Account, AccountId, AccountTree, Address, BlockNumber, TokenId, NFT,
+    AccountId, AccountTree, Address, BlockNumber, TokenId, NFT,
 };
 
 use super::{
@@ -13,14 +14,25 @@ use super::{
 
 #[derive(Debug, Clone)]
 pub struct ZkSyncStateInitParams {
-    pub tree: AccountTree,
-    pub acc_id_by_addr: HashMap<Address, AccountId>,
-    pub nfts: HashMap<TokenId, NFT>,
+    /// Restored zkSync state.
+    /// Corresponds to the latest **completed** block.
+    /// This state is used for two purposes: to initialize State Keeper (where we will
+    /// update it to match the latest **incomplete** block), and to initialize Root Hash
+    /// Calculator (where we'll be updating it from this point to the most relevant state,
+    /// yielding completed blocks).
+    pub state: ZkSyncState,
+    /// Block number to which `state` corresponds to`.
     pub last_block_number: BlockNumber,
+    /// ID of the next priority operation.
+    /// Corresponds to the latest observable state, including incomplete blocks.
     pub unprocessed_priority_op: u64,
-
+    /// Partially created block we should start with. May not exist if there were no
+    /// new transactions since the last block was sealed and before the restart of the server.
     pub pending_block: Option<SendablePendingBlock>,
+    /// Data on the incomplete blocks that were created by the state keeper, but not yet processed
+    /// by the root hash calculator.
     pub root_hash_jobs: Vec<BlockRootHashJob>,
+    /// Reverted blocks that we should process first (normally empty).
     pub reverted_blocks: VecDeque<IncompleteBlock>,
 }
 
@@ -32,10 +44,12 @@ impl Default for ZkSyncStateInitParams {
 
 impl ZkSyncStateInitParams {
     pub fn new() -> Self {
+        let tree = AccountTree::new(zksync_crypto::params::account_tree_depth());
+        let acc_id_by_addr = HashMap::new();
+        let nfts = HashMap::new();
+
         Self {
-            tree: AccountTree::new(zksync_crypto::params::account_tree_depth()),
-            acc_id_by_addr: HashMap::new(),
-            nfts: HashMap::new(),
+            state: ZkSyncState::new(tree, acc_id_by_addr, nfts),
             last_block_number: BlockNumber(0),
             unprocessed_priority_op: 0,
 
@@ -66,9 +80,7 @@ impl ZkSyncStateInitParams {
             Self::load_reverted_blocks(storage, fee_account_id, available_chunk_sizes).await;
 
         let init_params = Self {
-            tree,
-            acc_id_by_addr,
-            nfts,
+            state: ZkSyncState::new(tree, acc_id_by_addr, nfts),
             last_block_number,
             unprocessed_priority_op,
             pending_block,
@@ -110,6 +122,7 @@ impl ZkSyncStateInitParams {
             .await
             .unwrap()
     }
+
     async fn load_pending_block(
         storage: &mut zksync_storage::StorageProcessor<'_>,
         last_block_number: BlockNumber,
@@ -172,11 +185,6 @@ impl ZkSyncStateInitParams {
         }
     }
 
-    pub fn insert_account(&mut self, id: AccountId, acc: Account) {
-        self.acc_id_by_addr.insert(acc.address, id);
-        self.tree.insert(*id, acc);
-    }
-
     async fn load_nft_tokens(
         storage: &mut zksync_storage::StorageProcessor<'_>,
         block_number: BlockNumber,

core/bin/zksync_core/src/state_keeper/mod.rs

@@ -93,12 +93,8 @@ impl ZkSyncStateKeeper {
         //    separately below.
         // 2. For root hash calculator (`rhc_state`). It will require the state at *last finished block*, so it can keep
         //    working on calculating root hashes for incomplete blocks that we had before the restart.
-        let mut sk_state = ZkSyncState::new(
-            initial_state.tree,
-            initial_state.acc_id_by_addr,
-            initial_state.nfts,
-        );
-        let rhc_state = sk_state.clone();
+        let mut sk_state = initial_state.state.clone();
+        let rhc_state = initial_state.state.clone();
 
         // Update the state keeper copy of state.
         let mut last_block = initial_state.last_block_number;
@@ -880,9 +876,7 @@ impl ZkSyncStateKeeper {
 
     pub fn get_current_state(&self) -> ZkSyncStateInitParams {
         ZkSyncStateInitParams {
-            tree: self.state.get_balance_tree(),
-            acc_id_by_addr: self.state.get_account_addresses(),
-            nfts: self.state.nfts.clone(),
+            state: self.state.clone(),
             last_block_number: self.pending_block.number - 1,
             unprocessed_priority_op: self.pending_block.unprocessed_priority_op_current,
 

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

@@ -24,7 +24,9 @@ fn test_create_incorrect_state_keeper() {
     let fee_collector = Account::default_with_address(&H160::random());
 
     let mut init_params = ZkSyncStateInitParams::default();
-    init_params.insert_account(AccountId(0), fee_collector.clone());
+    init_params
+        .state
+        .insert_account(AccountId(0), fee_collector.clone());
 
     // should panic
     ZkSyncStateKeeper::new(

core/bin/zksync_core/src/state_keeper/tests/utils.rs

@@ -31,7 +31,9 @@ impl StateKeeperTester {
         let fee_collector = Account::default_with_address(&H160::random());
 
         let mut init_params = ZkSyncStateInitParams::default();
-        init_params.insert_account(AccountId(0), fee_collector.clone());
+        init_params
+            .state
+            .insert_account(AccountId(0), fee_collector.clone());
 
         let (state_keeper, _root_hash_calculator) = ZkSyncStateKeeper::new(
             init_params,

core/tests/testkit/src/bin/block_sizes_test.rs

@@ -101,7 +101,7 @@ async fn main() {
     let (sk_thread_handle, stop_state_keeper_sender, sk_channels) =
         spawn_state_keeper(&fee_account.address, genesis_state(&fee_account.address));
 
-    let genesis_root = genesis_state(&fee_account.address).tree.root_hash();
+    let genesis_root = genesis_state(&fee_account.address).state.root_hash();
 
     let contracts = deploy_contracts(true, genesis_root);
 

core/tests/testkit/src/bin/exodus_test.rs

@@ -413,7 +413,7 @@ async fn exit_test() {
     let (sk_thread_handle, stop_state_keeper_sender, sk_channels) =
         spawn_state_keeper(&fee_account.address, genesis_state(&fee_account.address));
 
-    let initial_root_hash = genesis_state(&fee_account.address).tree.root_hash();
+    let initial_root_hash = genesis_state(&fee_account.address).state.root_hash();
 
     let deploy_timer = Instant::now();
     info!("deploying contracts");

core/tests/testkit/src/bin/gas_price_test.rs

@@ -192,7 +192,7 @@ async fn gas_price_test() {
     let (sk_thread_handle, stop_state_keeper_sender, sk_channels) =
         spawn_state_keeper(&fee_account.address, genesis_state(&fee_account.address));
 
-    let genesis_root = genesis_state(&fee_account.address).tree.root_hash();
+    let genesis_root = genesis_state(&fee_account.address).state.root_hash();
 
     let contracts = deploy_contracts(false, genesis_root);
 

core/tests/testkit/src/bin/migration_test.rs

@@ -14,7 +14,7 @@ async fn migration_test() {
     let fee_account = ZkSyncAccount::rand();
     let (sk_thread_handle, stop_state_keeper_sender, sk_channels) =
         spawn_state_keeper(&fee_account.address, genesis_state(&fee_account.address));
-    let genesis_root = genesis_state(&fee_account.address).tree.root_hash();
+    let genesis_root = genesis_state(&fee_account.address).state.root_hash();
 
     let deploy_timer = Instant::now();
     println!("deploying contracts");

core/tests/testkit/src/bin/revert_blocks_test.rs

@@ -1,7 +1,7 @@
 use web3::transports::Http;
 
 use zksync_core::state_keeper::ZkSyncStateInitParams;
-use zksync_types::{block::Block, AccountId, AccountMap, AccountTree};
+use zksync_types::block::Block;
 
 use zksync_testkit::zksync_account::ZkSyncETHAccountData;
 use zksync_testkit::*;
@@ -144,28 +144,20 @@ async fn execute_blocks(
     (reverted_state, test_setup_accounts, executed_block)
 }
 
-fn balance_tree_to_account_map(balance_tree: &AccountTree) -> AccountMap {
-    let mut account_map = AccountMap::default();
-    for (id, account) in balance_tree.items.iter() {
-        account_map.insert(AccountId(*id as u32), account.clone());
-    }
-    account_map
-}
-
 async fn revert_blocks_test() {
     let fee_account = ZkSyncAccount::rand();
     let test_config = TestkitConfig::from_env();
 
     let state = genesis_state(&fee_account.address);
 
     println!("deploying contracts");
-    let contracts = deploy_contracts(false, state.tree.root_hash());
+    let contracts = deploy_contracts(false, state.state.root_hash());
     println!("contracts deployed");
 
     let (commit_account, account_set) =
         create_test_setup_state(&test_config, &contracts, &fee_account);
 
-    let hash = state.tree.root_hash();
+    let hash = state.state.root_hash();
     let (handler, sender, channels) = spawn_state_keeper(&fee_account.address, state);
     let mut test_setup = TestSetup::new(
         channels,
@@ -190,7 +182,7 @@ async fn revert_blocks_test() {
 
     sender.send(()).expect("sk stop send");
     handler.join().expect("sk thread join");
-    let hash = state.tree.root_hash();
+    let hash = state.state.root_hash();
     let start_block_number = state.last_block_number;
 
     let (handler, sender, channels) = spawn_state_keeper(&fee_account.address, state);
@@ -217,7 +209,7 @@ async fn revert_blocks_test() {
     sender.send(()).expect("sk stop send");
     handler.join().expect("sk thread join");
 
-    let hash = state.tree.root_hash();
+    let hash = state.state.root_hash();
     let start_block_number = state.last_block_number;
 
     let (handler, sender, channels) = spawn_state_keeper(&fee_account.address, state);
@@ -247,7 +239,7 @@ async fn revert_blocks_test() {
         &test_config,
         &contracts,
         fee_account.address,
-        balance_tree_to_account_map(&state.tree),
+        state.state.get_accounts(),
         vec![TokenId(0)],
         test_setup.current_state_root.unwrap(),
     )

core/tests/testkit/src/data_restore.rs

@@ -5,7 +5,7 @@ use zksync_data_restore::{
     data_restore_driver::DataRestoreDriver, inmemory_storage_interactor::InMemoryStorageInteractor,
     ETH_BLOCKS_STEP,
 };
-use zksync_types::{Account, AccountMap, TokenId};
+use zksync_types::{Account, AccountId, TokenId};
 
 use crate::{external_commands::Contracts, TestkitConfig};
 
@@ -16,7 +16,7 @@ pub async fn verify_restore(
     testkit_config: &TestkitConfig,
     contracts: &Contracts,
     fee_account_address: Address,
-    acc_state_from_test_setup: AccountMap,
+    acc_state_from_test_setup: Vec<(u32, Account)>,
     tokens: Vec<TokenId>,
     root_hash: Fr,
 ) {
@@ -53,8 +53,11 @@ pub async fn verify_restore(
     };
 
     for (id, account) in acc_state_from_test_setup {
-        let driver_acc = driver.tree_state.get_account(id).expect("Should exist");
-        let inter_acc = db.get_account(&id).expect("Should exist");
+        let driver_acc = driver
+            .tree_state
+            .get_account(AccountId(id))
+            .expect("Should exist");
+        let inter_acc = db.get_account(&AccountId(id)).expect("Should exist");
         for id in &tokens {
             assert_eq!(driver_acc.address, inter_acc.address);
             assert_eq!(account.address, inter_acc.address);

core/tests/testkit/src/lib.rs

@@ -25,9 +25,11 @@ pub mod types;
 pub fn genesis_state(fee_account_address: &Address) -> ZkSyncStateInitParams {
     let operator_account = Account::default_with_address(fee_account_address);
     let mut params = ZkSyncStateInitParams::new();
-    params.insert_account(AccountId(0), operator_account);
+    params.state.insert_account(AccountId(0), operator_account);
     let mut nft_storage = Account::default_with_address(&NFT_STORAGE_ACCOUNT_ADDRESS);
     nft_storage.set_balance(NFT_TOKEN_ID, BigUint::from(MIN_NFT_TOKEN_ID));
-    params.insert_account(NFT_STORAGE_ACCOUNT_ID, nft_storage);
+    params
+        .state
+        .insert_account(NFT_STORAGE_ACCOUNT_ID, nft_storage);
     params
 }

core/tests/testkit/src/scenarios.rs

@@ -31,7 +31,7 @@ pub async fn perform_basic_tests() {
     let (sk_thread_handle, stop_state_keeper_sender, sk_channels) =
         spawn_state_keeper(&fee_account_address, genesis_state(&fee_account_address));
 
-    let initial_root = genesis_state(&fee_account.address).tree.root_hash();
+    let initial_root = genesis_state(&fee_account.address).state.root_hash();
 
     let deploy_timer = Instant::now();
     println!("deploying contracts");
@@ -122,11 +122,17 @@ pub async fn perform_basic_tests() {
     assert!(operations_num > 0);
     assert_eq!(operations_num, expected_operations_num);
 
+    let acc_state_from_test_setup = test_setup
+        .get_accounts_state()
+        .await
+        .into_iter()
+        .map(|(id, acc)| (id.0, acc))
+        .collect();
     verify_restore(
         &testkit_config,
         &contracts,
         fee_account_address,
-        test_setup.get_accounts_state().await,
+        acc_state_from_test_setup,
         tokens,
         test_setup.last_committed_block.new_root_hash,
     )