WIP: split merkle updates

Author: pashinov

Cargo.lock

@@ -135,7 +135,7 @@ tycho-storage = { path = "./storage", version = "0.2.7" }
 tycho-util = { path = "./util", version = "0.2.7" }
 
 [patch.crates-io]
-everscale-types = { git = "https://github.com/broxus/everscale-types.git", rev = "b55da379aba126d4496c88bc7b8cc7af553ab37a" }
+everscale-types = { git = "https://github.com/broxus/everscale-types.git", rev = "9b0694bf66cfbbd1dad6e664435a3b9abe4e7e5a" }
 
 [workspace.lints.rust]
 future_incompatible = "warn"

Cargo.toml

@@ -1,5 +1,7 @@
 pub use self::min_ref_mc_state::{MinRefMcStateTracker, RefMcStateHandle};
+pub use self::shard_state_data::{split_shard, ShardStateData};
 pub use self::shard_state_stuff::ShardStateStuff;
 
 mod min_ref_mc_state;
+mod shard_state_data;
 mod shard_state_stuff;

block-util/src/state/mod.rs

@@ -0,0 +1,96 @@
+use std::collections::BTreeMap;
+use std::sync::Arc;
+
+use anyhow::Result;
+use everscale_types::models::*;
+use everscale_types::prelude::*;
+
+/// Parsed shard state accounts.
+#[derive(Clone)]
+#[repr(transparent)]
+pub struct ShardStateData {
+    inner: Arc<Inner>,
+}
+
+impl ShardStateData {
+    pub fn from_root(root: Cell) -> Result<Self> {
+        let accounts = root.parse::<ShardAccounts>()?;
+        Self::from_accounts_and_root(accounts, root)
+    }
+
+    pub fn from_accounts(accounts: ShardAccounts) -> Result<Self> {
+        let root = CellBuilder::build_from(&accounts)?;
+        Self::from_accounts_and_root(accounts, root)
+    }
+
+    pub fn from_accounts_and_root(accounts: ShardAccounts, root: Cell) -> Result<Self> {
+        Ok(Self {
+            inner: Arc::new(Inner { accounts, root }),
+        })
+    }
+
+    pub fn root_cell(&self) -> &Cell {
+        &self.inner.root
+    }
+
+    pub fn accounts(&self) -> &ShardAccounts {
+        &self.inner.accounts
+    }
+}
+
+#[doc(hidden)]
+pub struct Inner {
+    accounts: ShardAccounts,
+    root: Cell,
+}
+
+pub fn split_shard(
+    shard: &ShardIdent,
+    accounts: &ShardAccounts,
+    depth: u8,
+    shards: &mut BTreeMap<u64, ShardAccounts>,
+) -> Result<()> {
+    fn split_shard_impl(
+        shard: &ShardIdent,
+        accounts: &ShardAccounts,
+        depth: u8,
+        shards: &mut BTreeMap<u64, ShardAccounts>,
+        builder: &mut CellBuilder,
+    ) -> Result<()> {
+        let (left_shard_ident, right_shard_ident) = 'split: {
+            if depth > 0 {
+                if let Some((left, right)) = shard.split() {
+                    break 'split (left, right);
+                }
+            }
+            shards.insert(shard.prefix(), accounts.clone());
+            return Ok(());
+        };
+
+        let (left_accounts, right_accounts) = {
+            builder.clear_bits();
+            let prefix_len = shard.prefix_len();
+            if prefix_len > 0 {
+                builder.store_uint(shard.prefix() >> (64 - prefix_len), prefix_len)?;
+            }
+            accounts.split_by_prefix(&builder.as_data_slice())?
+        };
+
+        split_shard_impl(
+            &left_shard_ident,
+            &left_accounts,
+            depth - 1,
+            shards,
+            builder,
+        )?;
+        split_shard_impl(
+            &right_shard_ident,
+            &right_accounts,
+            depth - 1,
+            shards,
+            builder,
+        )
+    }
+
+    split_shard_impl(shard, accounts, depth, shards, &mut CellBuilder::new())
+}

block-util/src/state/shard_state_data.rs

@@ -1,10 +1,12 @@
+use std::collections::BTreeMap;
 use std::sync::Arc;
 
 use anyhow::Result;
 use everscale_types::cell::Lazy;
 use everscale_types::models::*;
 use everscale_types::prelude::*;
 
+use crate::state::shard_state_data::ShardStateData;
 use crate::state::{MinRefMcStateTracker, RefMcStateHandle};
 
 /// Parsed shard state.
@@ -26,25 +28,40 @@ impl ShardStateStuff {
     pub fn from_root(
         block_id: &BlockId,
         root: Cell,
+        data_roots: BTreeMap<u64, Cell>,
         tracker: &MinRefMcStateTracker,
     ) -> Result<Self> {
         let shard_state = root.parse::<Box<ShardStateUnsplit>>()?;
-        Self::from_state_and_root(block_id, shard_state, root, tracker)
+
+        let shard_state_data = data_roots
+            .into_iter()
+            .map(|(k, cell)| ShardStateData::from_root(cell).map(|v| (k, v)))
+            .collect::<Result<BTreeMap<u64, ShardStateData>>>()?;
+
+        Self::from_state_and_root(block_id, root, shard_state, shard_state_data, tracker)
     }
 
-    pub fn from_state(
+    pub fn from_root_and_accounts(
         block_id: &BlockId,
-        shard_state: Box<ShardStateUnsplit>,
+        root: Cell,
+        state_data: BTreeMap<u64, ShardAccounts>,
         tracker: &MinRefMcStateTracker,
     ) -> Result<Self> {
-        let root = CellBuilder::build_from(&shard_state)?;
-        ShardStateStuff::from_state_and_root(block_id, shard_state, root, tracker)
+        let shard_state = root.parse::<Box<ShardStateUnsplit>>()?;
+
+        let shard_state_data = state_data
+            .into_iter()
+            .map(|(k, acc)| ShardStateData::from_accounts(acc).map(|v| (k, v)))
+            .collect::<Result<BTreeMap<_, _>>>()?;
+
+        Self::from_state_and_root(block_id, root, shard_state, shard_state_data, tracker)
     }
 
     pub fn from_state_and_root(
         block_id: &BlockId,
-        shard_state: Box<ShardStateUnsplit>,
         root: Cell,
+        shard_state: Box<ShardStateUnsplit>,
+        shard_state_data: BTreeMap<u64, ShardStateData>,
         tracker: &MinRefMcStateTracker,
     ) -> Result<Self> {
         anyhow::ensure!(
@@ -68,38 +85,13 @@ impl ShardStateStuff {
                 block_id: *block_id,
                 shard_state_extra: shard_state.load_custom()?,
                 shard_state,
+                shard_state_data,
                 root,
                 handle,
             }),
         })
     }
 
-    pub fn deserialize_zerostate(
-        zerostate_id: &BlockId,
-        bytes: &[u8],
-        tracker: &MinRefMcStateTracker,
-    ) -> Result<Self> {
-        anyhow::ensure!(zerostate_id.seqno == 0, "given id has a non-zero seqno");
-
-        let file_hash = Boc::file_hash_blake(bytes);
-        anyhow::ensure!(
-            zerostate_id.file_hash.as_slice() == file_hash.as_slice(),
-            "file_hash mismatch. Expected: {}, got: {}",
-            hex::encode(file_hash),
-            zerostate_id.file_hash,
-        );
-
-        let root = Boc::decode(bytes)?;
-        anyhow::ensure!(
-            &zerostate_id.root_hash == root.repr_hash(),
-            "root_hash mismatch for {zerostate_id}. Expected: {expected}, got: {got}",
-            expected = zerostate_id.root_hash,
-            got = root.repr_hash(),
-        );
-
-        Self::from_root(zerostate_id, root, tracker)
-    }
-
     pub fn block_id(&self) -> &BlockId {
         &self.inner.block_id
     }
@@ -123,6 +115,22 @@ impl ShardStateStuff {
         &self.inner.root
     }
 
+    pub fn data_root_cells(&self) -> BTreeMap<u64, Cell> {
+        self.inner
+            .shard_state_data
+            .iter()
+            .map(|(k, v)| (*k, v.root_cell().clone()))
+            .collect()
+    }
+
+    pub fn load_accounts(&self) -> BTreeMap<u64, ShardAccounts> {
+        self.inner
+            .shard_state_data
+            .iter()
+            .map(|(k, v)| (*k, v.accounts().clone()))
+            .collect()
+    }
+
     pub fn shards(&self) -> Result<&ShardHashes> {
         Ok(&self.state_extra()?.shards)
     }
@@ -167,6 +175,7 @@ unsafe impl arc_swap::RefCnt for ShardStateStuff {
 pub struct Inner {
     block_id: BlockId,
     shard_state: Box<ShardStateUnsplit>,
+    shard_state_data: BTreeMap<u64, ShardStateData>,
     shard_state_extra: Option<McStateExtra>,
     handle: RefMcStateHandle,
     root: Cell,

block-util/src/state/shard_state_stuff.rs

@@ -20,6 +20,7 @@ async-trait = { workspace = true }
 base64 = { workspace = true }
 bytes = { workspace = true }
 bytesize = { workspace = true }
+bumpalo = { workspace = true }
 clap = { workspace = true }
 clap-markdown = { workspace = true }
 dirs = { workspace = true }

cli/Cargo.toml

@@ -1,15 +1,18 @@
-use std::collections::hash_map;
+use std::collections::{hash_map, BTreeMap};
 use std::path::PathBuf;
 use std::sync::OnceLock;
 
 use anyhow::{Context, Result};
+use bumpalo::Bump;
 use everscale_crypto::ed25519;
+use everscale_types::boc;
 use everscale_types::cell::Lazy;
 use everscale_types::models::*;
 use everscale_types::num::Tokens;
 use everscale_types::prelude::*;
 use serde::{Deserialize, Serialize};
-use tycho_util::{FastHashMap, FastHashSet};
+use tycho_block_util::state::split_shard;
+use tycho_util::{FastHashMap, FastHashSet, FastHasherState};
 
 use crate::util::{compute_storage_used, print_json};
 
@@ -88,14 +91,37 @@ fn generate_zerostate(
         .add_required_accounts()
         .context("failed to add required accounts")?;
 
-    let state = config
+    let bump = Bump::new();
+
+    let mut boc = boc::ser::BocHeader::<FastHasherState>::with_capacity(33);
+
+    let (state, accounts) = config
         .build_masterchain_state(now)
         .context("failed to build masterchain zerostate")?;
 
-    let boc = CellBuilder::build_from(&state).context("failed to serialize zerostate")?;
+    let state_cell = CellBuilder::build_from(&state).context("failed to serialize zerostate")?;
+    let root_hash = *state_cell.repr_hash();
+
+    boc.add_root(state_cell.as_ref());
+
+    for (prefix, shard_accounts) in &accounts {
+        let prefix_cell = CellBuilder::build_from(prefix).context("failed to serialize prefix")?;
+
+        let cell = &*bump.alloc(prefix_cell);
+        boc.add_root(cell.as_ref());
+
+        let shard_accounts_cell =
+            CellBuilder::build_from(shard_accounts).context("failed to serialize accounts")?;
+
+        let cell = &*bump.alloc(shard_accounts_cell);
+        boc.add_root(cell.as_ref());
+    }
+
+    let boc_size = boc.compute_stats().total_size as usize;
+
+    let mut data = Vec::with_capacity(boc_size);
+    boc.encode(&mut data);
 
-    let root_hash = *boc.repr_hash();
-    let data = Boc::encode(&boc);
     let file_hash = Boc::file_hash_blake(&data);
 
     std::fs::write(output_path, data).context("failed to write masterchain zerostate")?;
@@ -322,16 +348,20 @@ impl ZerostateConfig {
         Ok(())
     }
 
-    fn build_masterchain_state(self, now: u32) -> Result<ShardStateUnsplit> {
+    fn build_masterchain_state(
+        self,
+        now: u32,
+    ) -> Result<(ShardStateUnsplit, BTreeMap<u64, ShardAccounts>)> {
         let mut state = make_shard_state(self.global_id, ShardIdent::MASTERCHAIN, now);
 
         let config = BlockchainConfig {
             address: self.params.get::<ConfigParam0>()?.unwrap(),
             params: self.params.clone(),
         };
 
+        let mut accounts = ShardAccounts::new();
+        let mut splitted_accounts = BTreeMap::new();
         {
-            let mut accounts = ShardAccounts::new();
             let mut libraries = FastHashMap::<HashBytes, (Cell, FastHashSet<HashBytes>)>::default();
             for (account, mut account_state) in self.accounts {
                 let balance = match account_state.as_mut() {
@@ -383,7 +413,21 @@ impl ZerostateConfig {
             update_config_account(&mut accounts, &config)?;
 
             assert_eq!(state.total_balance, accounts.root_extra().balance);
-            state.accounts = Lazy::new(&accounts)?;
+
+            // Split to 16 virtual shards
+            split_shard(
+                &ShardIdent::MASTERCHAIN,
+                &accounts,
+                4,
+                &mut splitted_accounts,
+            )?;
+
+            let accounts_roots = splitted_accounts
+                .iter()
+                .map(|(k, v)| (*k, *CellBuilder::build_from(v).unwrap().repr_hash()))
+                .collect::<BTreeMap<_, _>>();
+
+            state.accounts = Dict::try_from_btree(&accounts_roots)?;
 
             // Build lib dict
             let mut libs = Dict::new();
@@ -465,7 +509,7 @@ impl ZerostateConfig {
             global_balance: state.total_balance.clone(),
         })?);
 
-        Ok(state)
+        Ok((state, splitted_accounts))
     }
 }
 
@@ -485,9 +529,22 @@ impl Default for ZerostateConfig {
 }
 
 fn make_shard_state(global_id: i32, shard_ident: ShardIdent, now: u32) -> ShardStateUnsplit {
+    let accounts = ShardAccounts::new();
+    let mut splitted_accounts = BTreeMap::new();
+
+    split_shard(&shard_ident, &accounts, 4, &mut splitted_accounts).unwrap();
+
+    let accounts_roots = splitted_accounts
+        .into_iter()
+        .map(|(k, v)| (k, *CellBuilder::build_from(v).unwrap().repr_hash()))
+        .collect::<BTreeMap<_, _>>();
+
+    let accounts = Dict::try_from_btree(&accounts_roots).unwrap();
+
     ShardStateUnsplit {
         global_id,
         shard_ident,
+        accounts,
         gen_utime: now,
         min_ref_mc_seqno: u32::MAX,
         ..Default::default()