Create tool to migrate caches

Cargo.toml

@@ -8,6 +8,7 @@ members = [
     "core/bin/parse_pub_data",
     "core/bin/block_revert",
     "core/bin/remove_proofs",
+    "core/bin/tree_cache_updater",
 
     # Server micro-services
     "core/bin/zksync_api",

core/bin/tree_cache_updater/Cargo.toml

@@ -0,0 +1,23 @@
+[package]
+name = "tree_cache_updater"
+version = "1.0.0"
+edition = "2018"
+authors = ["The Matter Labs Team <[email protected]>"]
+homepage = "https://zksync.io/"
+repository = "https://github.com/matter-labs/zksync"
+license = "Apache-2.0"
+keywords = ["blockchain", "zksync"]
+categories = ["cryptography"]
+publish = false # We don't want to publish our binaries.
+
+[dependencies]
+zksync_types = { path = "../../lib/types", version = "1.0" }
+zksync_storage = { path = "../../lib/storage", version = "1.0" }
+zksync_crypto = { path = "../../lib/crypto", version = "1.0" }
+zksync_config = { path = "../../lib/config", version = "1.0" }
+
+tokio = { version = "1", features = ["full"] }
+anyhow = "1.0"
+structopt = "0.3.20"
+serde_json = "1"
+

core/bin/tree_cache_updater/src/main.rs

@@ -0,0 +1,67 @@
+use structopt::StructOpt;
+use zksync_crypto::merkle_tree::parallel_smt::SparseMerkleTreeSerializableCacheBN256;
+use zksync_storage::StorageProcessor;
+use zksync_types::BlockNumber;
+
+#[derive(Debug, StructOpt)]
+#[structopt(
+    name = "tree_cache_updater",
+    about = "Tool to migrate server tree caches to the binary format."
+)]
+struct Opt {
+    /// Maximum amount of blocks to convert.
+    #[structopt(long)]
+    max_blocks: usize,
+    /// Whether to remove all the old JSON caches or not.
+    #[structopt(long)]
+    clear_old: bool,
+}
+
+#[tokio::main]
+async fn main() -> anyhow::Result<()> {
+    let opt = Opt::from_args();
+
+    let mut storage = StorageProcessor::establish_connection().await?;
+    let mut transaction = storage.start_transaction().await?;
+
+    let max_block = transaction
+        .chain()
+        .block_schema()
+        .get_last_saved_block()
+        .await?;
+
+    let min_block = std::cmp::max(max_block.0.saturating_sub(opt.max_blocks as u32), 1); // We can't go below the 1st block.
+
+    // Go through the suggested blocks range. For each block in this range, if the cachce exists, we will load it, convert to the bincode cache,
+    // and store to the binary schema.
+    for block in min_block..(max_block.0) {
+        if let Some(cache) = transaction
+            .chain()
+            .tree_cache_schema_json()
+            .get_account_tree_cache_block(BlockNumber(block))
+            .await?
+        {
+            let cache: SparseMerkleTreeSerializableCacheBN256 = serde_json::from_value(cache)?;
+            let binary_cache = cache.encode_bincode();
+            transaction
+                .chain()
+                .tree_cache_schema_bincode()
+                .store_account_tree_cache(BlockNumber(block), binary_cache)
+                .await?;
+        }
+    }
+
+    // We've processed all the blocks. Now, if user requested, we'll remove all the old caches.
+    if opt.clear_old {
+        // BlockNumber(0) because range is not inclusive. Everything starting from block 1 will be removed.
+        transaction
+            .chain()
+            .tree_cache_schema_json()
+            .remove_new_account_tree_cache(BlockNumber(0))
+            .await?;
+    }
+
+    transaction.commit().await?;
+
+    Ok(())
+}