split persistent + treat raw import marker as a "poisoned state"

Cargo.toml

@@ -183,6 +183,7 @@ tycho-wu-tuner = { path = "./wu-tuner", version = "0.3.12" }
 
 [patch.crates-io]
 # patches here
+tycho-types = { git = "https://github.com/broxus/tycho-types.git", rev = "b529060a4f229bc17f9a30cd12f90ee44931e16a" }
 
 [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,74 @@ 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 {
+            return Err(Error::IntOverflow);
+        };
+
+        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,
+        )
+    }
+
+    if depth >= ShardIdent::MAX_SPLIT_DEPTH {
+        return Err(Error::IntOverflow);
+    }
+    let mut shards = Vec::with_capacity(1 << depth);
+
     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)
 }

block-util/src/state/mod.rs

@@ -1,3 +1,7 @@
+use anyhow::{Context, Result};
+use tycho_types::models::ShardIdent;
+use tycho_util::FastHashSet;
+
 pub use self::consensus_info::choose_genesis_info;
 pub use self::min_ref_mc_state::{MinRefMcStateTracker, RefMcStateHandle};
 pub use self::shard_state_stuff::ShardStateStuff;
@@ -7,3 +11,101 @@ mod consensus_info;
 mod min_ref_mc_state;
 mod shard_state_stuff;
 mod state_proof;
+
+/// Checks that provided prefixes is a subset of `shard_ident`
+/// split to `split_depth` (depth is relative to the shard).
+pub fn validate_shard_prefixes(
+    shard_ident: ShardIdent,
+    split_depth: u8,
+    prefixes: impl IntoIterator<IntoIter: ExactSizeIterator<Item = u64>>,
+) -> Result<()> {
+    let prefixes = prefixes.into_iter();
+
+    anyhow::ensure!(
+        !shard_ident.is_masterchain(),
+        "masterchain state cannot be split into parts"
+    );
+
+    let Some(max_prefixes) = 1usize.checked_shl(split_depth as u32) else {
+        anyhow::bail!("invalid split depth");
+    };
+
+    anyhow::ensure!(
+        prefixes.len() <= max_prefixes,
+        "too many prefixes: prefixes={}, max={max_prefixes}",
+        prefixes.len()
+    );
+
+    let base_depth = shard_ident.prefix_len();
+
+    let mut unique_prefixes = FastHashSet::default();
+    for prefix in prefixes {
+        let ident = ShardIdent::new(shard_ident.workchain(), prefix)
+            .with_context(|| format!("invalid shard prefix: {prefix:016x}"))?;
+
+        let prefix_len = ident.prefix_len();
+        let expected_len = base_depth + split_depth as u16;
+        anyhow::ensure!(
+            prefix_len == expected_len,
+            "invalid shard prefix: {prefix:016x} \
+            (prefix_len={prefix_len}, expected_len={expected_len})"
+        );
+        anyhow::ensure!(
+            shard_ident.is_ancestor_of(&ident),
+            "unrelated shard prefix: {prefix:016x}"
+        );
+
+        anyhow::ensure!(
+            unique_prefixes.insert(prefix),
+            "duplicate shard prefix: {prefix:016x}"
+        );
+    }
+
+    Ok(())
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn validate_shard_prefixes_works() -> Result<()> {
+        validate_shard_prefixes(ShardIdent::BASECHAIN, 0, Vec::<u64>::new()).unwrap();
+        validate_shard_prefixes(ShardIdent::BASECHAIN, 2, vec![
+            0x2000000000000000,
+            0xa000000000000000,
+        ])
+        .unwrap();
+        let non_full_shard = ShardIdent::new(0, 0x4000000000000000).unwrap();
+        validate_shard_prefixes(non_full_shard, 1, vec![
+            0x2000000000000000,
+            0x6000000000000000,
+        ])
+        .unwrap();
+        // `10...` is not a child of `0...`
+        validate_shard_prefixes(non_full_shard, 1, vec![0xa000000000000000]).unwrap_err();
+
+        let deep_shard = ShardIdent::new(0, 0x1000000000000000).unwrap();
+        // `0001...`  cannot be split into smaller prefixes like `001...`
+        validate_shard_prefixes(deep_shard, 1, vec![0x2000000000000000]).unwrap_err();
+
+        // Too many parts.
+        validate_shard_prefixes(ShardIdent::BASECHAIN, 1, vec![
+            0x4000000000000000,
+            0xc000000000000000,
+            0x4000000000000000,
+        ])
+        .unwrap_err();
+
+        // Duplicate parts
+        validate_shard_prefixes(ShardIdent::BASECHAIN, 2, vec![
+            0x2000000000000000,
+            0x2000000000000000,
+        ])
+        .unwrap_err();
+
+        // Part prefix at the wrong depth.
+        validate_shard_prefixes(ShardIdent::BASECHAIN, 2, vec![0x4000000000000000]).unwrap_err();
+        Ok(())
+    }
+}

cli/src/cmd/debug/mempool.rs

@@ -377,15 +377,18 @@ async fn load_mc_zerostate(
 
     let zerostate_block_id = mc_zerostate_id.as_block_id();
     tracing::info!("loading zerostate {:?}", zerostate_block_id);
+    let raw_import = storage.shard_state_storage().create_raw_import_session();
     let root_hash = storage
         .shard_state_storage()
         .store_state_bytes(
             &zerostate_block_id,
             masterchain_zerostate,
             Some(&mc_zerostate_id.root_hash),
+            &raw_import,
         )
         .await?;
     assert_eq!(root_hash, mc_zerostate_id.root_hash);
+    raw_import.finish()?;
 
     let masterchain_zerostate = storage
         .shard_state_storage()

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};
@@ -237,7 +238,7 @@ async fn load_states_from_dump(storage: &CoreStorage, dump_path: &Path) -> Resul
             let tempfile = std::fs::File::open(&temp_path)?;
             storage
                 .shard_state_storage()
-                .store_state_file(&block_id, tempfile, None)
+                .store_state_file_without_session(&block_id, tempfile, None)
                 .await?;
 
             if let Some(handle) = storage.block_handle_storage().load_handle(&block_id) {

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 }