Split persistent state on parts by shards

Cargo.toml

@@ -178,6 +178,7 @@ tycho-wu-tuner = { path = "./wu-tuner", version = "0.3.10" }
 
 [patch.crates-io]
 # patches here
+tycho-types = { git = "https://github.com/broxus/tycho-types.git", rev = "d6e3355d470bb26009006b2c86968758932eaa31" }
 
 [workspace.lints.rust]
 future_incompatible = "warn"

block-util/src/dict.rs

@@ -211,6 +211,15 @@ where
     K: DictKey,
     A: Default,
 {
+    let (dict_root, _) = dict.into_parts();
+    split_dict_raw(dict_root.into_root(), K::BITS, depth)
+}
+
+pub fn split_dict_raw(
+    dict: Option<Cell>,
+    key_bit_len: u16,
+    depth: u8,
+) -> Result<FastHashMap<HashBytes, Cell>, Error> {
     fn split_dict_impl(
         dict: Option<Cell>,
         key_bit_len: u16,
@@ -241,8 +250,72 @@ where
     let mut shards =
         FastHashMap::with_capacity_and_hasher(2usize.pow(depth as _), Default::default());
 
+    split_dict_impl(dict, key_bit_len, depth, &mut shards)?;
+
+    Ok(shards)
+}
+
+/// Splits aug dict by shards, preserving empty shards.
+/// E.g. if `depth == 1` and all entries are in the left shard,
+/// then will return `None` cell for the right shard.
+pub fn split_aug_dict_raw_by_shards<K, A, V>(
+    workchain: i32,
+    dict: AugDict<K, A, V>,
+    depth: u8,
+) -> Result<Vec<(ShardIdent, Option<Cell>)>, Error>
+where
+    K: DictKey,
+    A: Default,
+{
+    fn split_dict_impl(
+        shard: &ShardIdent,
+        dict: Option<Cell>,
+        key_bit_len: u16,
+        depth: u8,
+        shards: &mut Vec<(ShardIdent, Option<Cell>)>,
+    ) -> Result<(), Error> {
+        if depth == 0 {
+            shards.push((*shard, dict));
+            return Ok(());
+        }
+
+        let Some((left_shard, right_shard)) = shard.split() else {
+            shards.push((*shard, dict));
+            return Ok(());
+        };
+
+        let PartialSplitDict {
+            remaining_bit_len,
+            left_branch,
+            right_branch,
+        } = dict_split_raw(dict.as_ref(), key_bit_len, Cell::empty_context())?;
+
+        split_dict_impl(
+            &left_shard,
+            left_branch,
+            remaining_bit_len,
+            depth - 1,
+            shards,
+        )?;
+        split_dict_impl(
+            &right_shard,
+            right_branch,
+            remaining_bit_len,
+            depth - 1,
+            shards,
+        )
+    }
+
+    let mut shards = Vec::with_capacity(2usize.pow(depth as _));
+
     let (dict_root, _) = dict.into_parts();
-    split_dict_impl(dict_root.into_root(), K::BITS, depth, &mut shards)?;
+    split_dict_impl(
+        &ShardIdent::new_full(workchain),
+        dict_root.into_root(),
+        K::BITS,
+        depth,
+        &mut shards,
+    )?;
 
     Ok(shards)
 }

cli/src/cmd/tools/dump_state.rs

@@ -250,7 +250,7 @@ impl Dumper {
         let dir = Dir::new(self.output_dir.path().join("persistents"))?;
         self.storage
             .shard_state_storage()
-            .write_persistent_shard_state(dir, *block_id, *state.root_cell().repr_hash(), None)
+            .write_persistent_shard_state(dir, *block_id, *state.root_cell().repr_hash(), 0, None)
             .await
             .context(format!("Failed to write state for {}", block_id))?;
         println!(" - Persistent state saved");

cli/src/cmd/tools/hardfork.rs

@@ -65,7 +65,12 @@ impl Cmd {
             })
             .await?;
 
-            let storage = CoreStorage::open(ctx, CoreStorageConfig::default().without_gc()).await?;
+            let storage_config = CoreStorageConfig::default().without_gc();
+            anyhow::ensure!(
+                storage_config.persistent_state_split_depth == 0,
+                "hardfork creation supports only persistent_state_split_depth = 0"
+            );
+            let storage = CoreStorage::open(ctx, storage_config).await?;
 
             let Some(mc_block_id) = storage
                 .shard_state_storage()

collator/src/test_utils.rs

@@ -19,6 +19,7 @@ use tycho_storage::StorageContext;
 use tycho_types::boc::{Boc, BocRepr};
 use tycho_types::cell::CellBuilder;
 use tycho_types::models::{Block, BlockId, ShardStateUnsplit};
+#[cfg(any(test, feature = "test"))]
 use tycho_util::compression::zstd_decompress_simple;
 
 use crate::internal_queue::queue::{QueueConfig, QueueFactory, QueueFactoryStdImpl};

core/Cargo.toml

@@ -45,6 +45,7 @@ castaway = { workspace = true }
 clap = { workspace = true, optional = true }
 crc32c = { workspace = true }
 croaring = { workspace = true }
+crossbeam-queue = { workspace = true }
 dashmap = { workspace = true }
 futures-util = { workspace = true }
 governor = { workspace = true, optional = true }