Use chunk sizes from the contract version

core/bin/data_restore/src/contract/default.rs

@@ -45,6 +45,7 @@ pub fn rollup_ops_blocks_from_bytes(input_data: Vec<u8>) -> Result<RollupOpsBloc
             fee_account,
             timestamp: None,
             previous_block_root_hash: H256::default(),
+            contract_version: None,
         };
         Ok(block)
     } else {

core/bin/data_restore/src/contract/v4.rs

@@ -91,6 +91,7 @@ pub fn rollup_ops_blocks_from_bytes(data: Vec<u8>) -> anyhow::Result<Vec<RollupO
                         fee_account: AccountId(fee_acc.as_u32()),
                         timestamp: Some(timestamp.as_u64()),
                         previous_block_root_hash,
+                        contract_version: None,
                     });
 
                     previous_block_root_hash = H256::from_slice(&root_hash);

core/bin/data_restore/src/contract/version.rs

@@ -45,11 +45,15 @@ impl ZkSyncContractVersion {
         data: Vec<u8>,
     ) -> anyhow::Result<Vec<RollupOpsBlock>> {
         use ZkSyncContractVersion::*;
-        let res = match self {
+        let mut blocks = match self {
             V0 | V1 | V2 | V3 => vec![contract::default::rollup_ops_blocks_from_bytes(data)?],
             V4 => contract::v4::rollup_ops_blocks_from_bytes(data)?,
         };
-        Ok(res)
+        // Set the contract version.
+        for block in blocks.iter_mut() {
+            block.contract_version = Some(*self);
+        }
+        Ok(blocks)
     }
 
     /// Increment the contract version by one.

core/bin/data_restore/src/data_restore_driver.rs

@@ -68,7 +68,6 @@ pub struct DataRestoreDriver<T: Transport, I> {
     /// Expected root hash to be observed after restoring process. Only
     /// available in finite mode, and intended for tests.
     pub final_hash: Option<Fr>,
-    pub available_block_chunk_sizes: Vec<usize>,
     phantom_data: PhantomData<I>,
 }
 
@@ -98,7 +97,6 @@ where
         finite_mode: bool,
         final_hash: Option<Fr>,
         zksync_contract: ZkSyncDeployedContract<T>,
-        available_block_chunk_sizes: Vec<usize>,
     ) -> Self {
         let governance_contract = {
             let abi = governance_contract();
@@ -122,7 +120,6 @@ where
             finite_mode,
             final_hash,
             phantom_data: Default::default(),
-            available_block_chunk_sizes,
         }
     }
 
@@ -330,9 +327,13 @@ where
         let mut updates = vec![];
         let mut count = 0;
         for op_block in new_ops_blocks {
+            let available_block_chunk_sizes = op_block
+                .contract_version
+                .expect("contract version must be set")
+                .available_block_chunk_sizes();
             let (block, acc_updates) = self
                 .tree_state
-                .update_tree_states_from_ops_block(&op_block, &self.available_block_chunk_sizes)
+                .update_tree_states_from_ops_block(&op_block, available_block_chunk_sizes)
                 .expect("Updating tree state: cant update tree from operations");
             blocks.push(block);
             updates.push(acc_updates);

core/bin/data_restore/src/main.rs

@@ -52,7 +52,6 @@ pub struct ContractsConfig {
     governance_addr: Address,
     genesis_tx_hash: H256,
     contract_addr: Address,
-    available_block_chunk_sizes: Vec<usize>,
 }
 
 impl ContractsConfig {
@@ -71,7 +70,6 @@ impl ContractsConfig {
             governance_addr: contracts_opts.governance_addr,
             genesis_tx_hash: contracts_opts.genesis_tx_hash,
             contract_addr: contracts_opts.contract_addr,
-            available_block_chunk_sizes: chain_opts.state_keeper.block_chunk_sizes,
         }
     }
 }
@@ -114,7 +112,6 @@ async fn main() {
         finite_mode,
         final_hash,
         contract,
-        config.available_block_chunk_sizes,
     );
 
     let mut interactor = DatabaseStorageInteractor::new(storage);

core/bin/data_restore/src/rollup_ops.rs

@@ -2,7 +2,7 @@ use web3::{Transport, Web3};
 
 use zksync_types::operations::ZkSyncOp;
 
-use crate::contract;
+use crate::contract::ZkSyncContractVersion;
 use crate::eth_tx_helpers::{get_ethereum_transaction, get_input_data_from_ethereum_transaction};
 use crate::events::BlockEvent;
 use zksync_types::{AccountId, BlockNumber, H256};
@@ -20,6 +20,10 @@ pub struct RollupOpsBlock {
     pub timestamp: Option<u64>,
     /// Previous block root hash.
     pub previous_block_root_hash: H256,
+    /// zkSync contract version for the given block.
+    /// Used to obtain block chunk sizes. Stored in the database
+    /// in the corresponding block event.
+    pub contract_version: Option<ZkSyncContractVersion>,
 }
 
 impl RollupOpsBlock {
@@ -37,13 +41,9 @@ impl RollupOpsBlock {
     ) -> anyhow::Result<Vec<Self>> {
         let transaction = get_ethereum_transaction(web3, &event_data.transaction_hash).await?;
         let input_data = get_input_data_from_ethereum_transaction(&transaction)?;
-        let blocks = if let Ok(block) =
-            contract::default::rollup_ops_blocks_from_bytes(input_data.clone())
-        {
-            vec![block]
-        } else {
-            contract::v4::rollup_ops_blocks_from_bytes(input_data)?
-        };
+        let blocks: Vec<RollupOpsBlock> = event_data
+            .contract_version
+            .rollup_ops_blocks_from_bytes(input_data)?;
         Ok(blocks)
     }
 }

core/bin/data_restore/src/storage_interactor.rs

@@ -156,5 +156,9 @@ pub fn stored_ops_block_into_ops_block(op_block: StoredRollupOpsBlock) -> Rollup
             .previous_block_root_hash
             .map(|h| H256::from_slice(&h))
             .unwrap_or_default(),
+        contract_version: Some(
+            ZkSyncContractVersion::try_from(op_block.contract_version as u32)
+                .expect("invalid contract version in the database"),
+        ),
     }
 }

core/bin/data_restore/src/tests/mod.rs

@@ -381,7 +381,6 @@ async fn test_run_state_update(mut storage: StorageProcessor<'_>) {
     ]);
 
     let eth = Eth::new(transport.clone());
-    let available_block_chunk_sizes = vec![10, 32, 72, 156, 322, 654];
     let mut driver = DataRestoreDriver::new(
         Web3::new(transport.clone()),
         [1u8; 20].into(),
@@ -390,7 +389,6 @@ async fn test_run_state_update(mut storage: StorageProcessor<'_>) {
         true,
         None,
         ZkSyncDeployedContract::version4(eth, [1u8; 20].into()),
-        available_block_chunk_sizes,
     );
 
     driver.run_state_update(&mut interactor).await;
@@ -416,7 +414,6 @@ async fn test_run_state_update(mut storage: StorageProcessor<'_>) {
     // Nullify the state of driver
     let eth = Eth::new(transport.clone());
 
-    let available_block_chunk_sizes = vec![10, 32, 72, 156, 322, 654];
     let mut driver = DataRestoreDriver::new(
         Web3::new(transport.clone()),
         [1u8; 20].into(),
@@ -425,7 +422,6 @@ async fn test_run_state_update(mut storage: StorageProcessor<'_>) {
         true,
         None,
         ZkSyncDeployedContract::version4(eth, [1u8; 20].into()),
-        available_block_chunk_sizes,
     );
 
     // Load state from db and check it
@@ -589,7 +585,6 @@ async fn test_with_inmemory_storage() {
     let web3 = Web3::new(transport.clone());
 
     let eth = Eth::new(transport.clone());
-    let available_block_chunk_sizes = vec![10, 32, 72, 156, 322, 654];
     let mut driver = DataRestoreDriver::new(
         web3.clone(),
         [1u8; 20].into(),
@@ -598,7 +593,6 @@ async fn test_with_inmemory_storage() {
         true,
         None,
         ZkSyncDeployedContract::version4(eth, [1u8; 20].into()),
-        available_block_chunk_sizes,
     );
 
     driver.run_state_update(&mut interactor).await;
@@ -617,7 +611,6 @@ async fn test_with_inmemory_storage() {
 
     // Nullify the state of driver
     let eth = Eth::new(transport.clone());
-    let available_block_chunk_sizes = vec![10, 32, 72, 156, 322, 654];
     let mut driver = DataRestoreDriver::new(
         web3.clone(),
         [1u8; 20].into(),
@@ -626,7 +619,6 @@ async fn test_with_inmemory_storage() {
         true,
         None,
         ZkSyncDeployedContract::version4(eth, [1u8; 20].into()),
-        available_block_chunk_sizes,
     );
 
     // Load state from db and check it

core/bin/data_restore/src/tree_state.rs

@@ -464,6 +464,7 @@ mod test {
             fee_account: AccountId(0),
             timestamp: None,
             previous_block_root_hash: Default::default(),
+            contract_version: None,
         };
 
         // Withdraw 20 with 1 fee from 7 to 10
@@ -490,6 +491,7 @@ mod test {
             fee_account: AccountId(0),
             timestamp: None,
             previous_block_root_hash: Default::default(),
+            contract_version: None,
         };
 
         // Transfer 40 with 1 fee from 7 to 8
@@ -517,6 +519,7 @@ mod test {
             fee_account: AccountId(0),
             timestamp: None,
             previous_block_root_hash: Default::default(),
+            contract_version: None,
         };
 
         // Transfer 19 with 1 fee from 8 to 7
@@ -544,6 +547,7 @@ mod test {
             fee_account: AccountId(0),
             timestamp: None,
             previous_block_root_hash: Default::default(),
+            contract_version: None,
         };
 
         let pub_key_hash_7 = PubKeyHash::from_hex("sync:8888888888888888888888888888888888888888")
@@ -571,6 +575,7 @@ mod test {
             fee_account: AccountId(0),
             timestamp: None,
             previous_block_root_hash: Default::default(),
+            contract_version: None,
         };
 
         // Full exit for 8
@@ -591,6 +596,7 @@ mod test {
             fee_account: AccountId(0),
             timestamp: None,
             previous_block_root_hash: Default::default(),
+            contract_version: None,
         };
 
         // Forced exit for 7
@@ -616,6 +622,7 @@ mod test {
             fee_account: AccountId(1),
             timestamp: None,
             previous_block_root_hash: Default::default(),
+            contract_version: None,
         };
         // This transaction have to be deleted, do not uncomment. Delete it after removing the corresponding code        // let tx6 = Close {
         //     account: Address::from_hex("sync:8888888888888888888888888888888888888888").unwrap(),
@@ -817,6 +824,7 @@ mod test {
             fee_account: AccountId(0),
             timestamp: None,
             previous_block_root_hash: Default::default(),
+            contract_version: None,
         };
 
         let mut tree = TreeState::new();

core/lib/storage/sqlx-data.json

@@ -1427,6 +1427,54 @@
       "nullable": []
     }
   },
+  "44946f944f3c496ba90d5843084acf36aa3db3dd9d90f72a9bbeee862af056e4": {
+    "query": "SELECT blocks.block_num AS block_num, ops, fee_account,\n            timestamp, previous_block_root_hash, contract_version\n            FROM data_restore_rollup_blocks AS blocks\n            JOIN (\n                SELECT block_num, array_agg(operation) as ops\n                FROM data_restore_rollup_block_ops\n                GROUP BY block_num\n            ) ops\n                ON blocks.block_num = ops.block_num\n            JOIN data_restore_events_state as events\n                ON blocks.block_num = events.block_num\n            ORDER BY blocks.block_num ASC",
+    "describe": {
+      "columns": [
+        {
+          "ordinal": 0,
+          "name": "block_num",
+          "type_info": "Int8"
+        },
+        {
+          "ordinal": 1,
+          "name": "ops",
+          "type_info": "JsonbArray"
+        },
+        {
+          "ordinal": 2,
+          "name": "fee_account",
+          "type_info": "Int8"
+        },
+        {
+          "ordinal": 3,
+          "name": "timestamp",
+          "type_info": "Int8"
+        },
+        {
+          "ordinal": 4,
+          "name": "previous_block_root_hash",
+          "type_info": "Bytea"
+        },
+        {
+          "ordinal": 5,
+          "name": "contract_version",
+          "type_info": "Int4"
+        }
+      ],
+      "parameters": {
+        "Left": []
+      },
+      "nullable": [
+        false,
+        null,
+        false,
+        true,
+        true,
+        false
+      ]
+    }
+  },
   "44b276fda62734e9c9d9853f493340265116ab7f13599674d27aafe3d3887391": {
     "query": "UPDATE eth_operations \n            SET last_used_gas_price = $1, last_deadline_block = $2\n            WHERE id = $3",
     "describe": {
@@ -3364,48 +3412,6 @@
       ]
     }
   },
-  "b6848163f7fc88b1fe751e723245e79555df16ce601a58e5ed53c78bfdb97fd5": {
-    "query": "SELECT blocks.block_num AS block_num,\n            ops, fee_account, timestamp, previous_block_root_hash\n            FROM data_restore_rollup_blocks AS blocks\n            JOIN (\n                SELECT block_num, array_agg(operation) as ops\n                FROM data_restore_rollup_block_ops\n                GROUP BY block_num\n            ) ops\n            ON blocks.block_num = ops.block_num\n            ORDER BY blocks.block_num ASC",
-    "describe": {
-      "columns": [
-        {
-          "ordinal": 0,
-          "name": "block_num",
-          "type_info": "Int8"
-        },
-        {
-          "ordinal": 1,
-          "name": "ops",
-          "type_info": "JsonbArray"
-        },
-        {
-          "ordinal": 2,
-          "name": "fee_account",
-          "type_info": "Int8"
-        },
-        {
-          "ordinal": 3,
-          "name": "timestamp",
-          "type_info": "Int8"
-        },
-        {
-          "ordinal": 4,
-          "name": "previous_block_root_hash",
-          "type_info": "Bytea"
-        }
-      ],
-      "parameters": {
-        "Left": []
-      },
-      "nullable": [
-        false,
-        null,
-        false,
-        true,
-        true
-      ]
-    }
-  },
   "b89088c6516e2db2e01bfdf0afa5a8fdd7e20fde80183884a9769eae9b635010": {
     "query": "DELETE FROM executed_priority_operations WHERE block_number > $1",
     "describe": {

core/lib/storage/src/data_restore/mod.rs

@@ -97,17 +97,20 @@ impl<'a, 'c> DataRestoreSchema<'a, 'c> {
         // For each block aggregate its operations from the
         // `data_restore_rollup_block_ops` table into array and
         // match it by the block number from `data_restore_rollup_blocks`.
+        // The contract version is obtained from block events.
         let stored_blocks = sqlx::query_as!(
             StoredRollupOpsBlock,
-            "SELECT blocks.block_num AS block_num,
-            ops, fee_account, timestamp, previous_block_root_hash
+            "SELECT blocks.block_num AS block_num, ops, fee_account,
+            timestamp, previous_block_root_hash, contract_version
             FROM data_restore_rollup_blocks AS blocks
             JOIN (
                 SELECT block_num, array_agg(operation) as ops
                 FROM data_restore_rollup_block_ops
                 GROUP BY block_num
             ) ops
-            ON blocks.block_num = ops.block_num
+                ON blocks.block_num = ops.block_num
+            JOIN data_restore_events_state as events
+                ON blocks.block_num = events.block_num
             ORDER BY blocks.block_num ASC"
         )
         .fetch_all(self.0.conn())