fix(fork-network): Re-introduce block hash

tools/fork-network/src/cli.rs

@@ -25,8 +25,7 @@ use near_primitives::types::{
 };
 use near_primitives::version::PROTOCOL_VERSION;
 use near_store::db::RocksDB;
-use near_store::flat::{store_helper, BlockInfo};
-use near_store::flat::{FlatStorageManager, FlatStorageStatus};
+use near_store::flat::{store_helper, BlockInfo, FlatStorageManager, FlatStorageStatus};
 use near_store::{
     checkpoint_hot_storage_and_cleanup_columns, DBCol, Store, TrieDBStorage, TrieStorage,
     FINAL_HEAD_KEY,
@@ -283,6 +282,7 @@ impl ForkNetworkCommand {
 
         let mut store_update = store.store_update();
         store_update.set_ser(DBCol::Misc, b"FORK_TOOL_EPOCH_ID", epoch_id)?;
+        store_update.set_ser(DBCol::Misc, b"FORK_TOOL_BLOCK_HASH", &desired_block_hash)?;
         store_update.set(DBCol::Misc, b"FORK_TOOL_BLOCK_HEIGHT", &block_height.to_le_bytes());
         for (shard_id, state_root) in state_roots.iter().enumerate() {
             store_update.set_ser(
@@ -307,8 +307,9 @@ impl ForkNetworkCommand {
         let storage = open_storage(&home_dir, near_config).unwrap();
         let store = storage.get_hot_store();
 
-        let (prev_state_roots, epoch_id, _block_height) =
+        let (prev_state_roots, prev_hash, epoch_id, _block_height) =
             self.get_state_roots_and_hash(store.clone())?;
+        tracing::info!(?prev_state_roots, ?epoch_id, ?prev_hash);
 
         let epoch_manager =
             EpochManager::new_arc_handle(store.clone(), &near_config.genesis.config);
@@ -321,7 +322,12 @@ impl ForkNetworkCommand {
 
         let make_storage_mutator: MakeSingleShardStorageMutatorFn =
             Arc::new(move |shard_id, prev_state_root| {
-                SingleShardStorageMutator::new(shard_id, &runtime.clone(), prev_state_root)
+                SingleShardStorageMutator::new(
+                    shard_id,
+                    &runtime.clone(),
+                    prev_hash,
+                    prev_state_root,
+                )
             });
 
         let new_state_roots = self.prepare_state(
@@ -350,7 +356,7 @@ impl ForkNetworkCommand {
         let storage = open_storage(&home_dir, near_config).unwrap();
         let store = storage.get_hot_store();
 
-        let (prev_state_roots, epoch_id, block_height) =
+        let (prev_state_roots, prev_hash, epoch_id, block_height) =
             self.get_state_roots_and_hash(store.clone())?;
 
         let epoch_manager =
@@ -366,6 +372,7 @@ impl ForkNetworkCommand {
             epoch_manager.clone(),
             &runtime,
             epoch_id.clone(),
+            prev_hash,
             prev_state_roots,
         )?;
         let (new_state_roots, new_validator_accounts) =
@@ -410,8 +417,9 @@ impl ForkNetworkCommand {
     fn get_state_roots_and_hash(
         &self,
         store: Store,
-    ) -> anyhow::Result<(Vec<StateRoot>, EpochId, BlockHeight)> {
+    ) -> anyhow::Result<(Vec<StateRoot>, CryptoHash, EpochId, BlockHeight)> {
         let epoch_id = EpochId(store.get_ser(DBCol::Misc, b"FORK_TOOL_EPOCH_ID")?.unwrap());
+        let block_hash = store.get_ser(DBCol::Misc, b"FORK_TOOL_BLOCK_HASH")?.unwrap();
         let block_height = store.get(DBCol::Misc, b"FORK_TOOL_BLOCK_HEIGHT")?.unwrap();
         let block_height = u64::from_le_bytes(block_height.as_slice().try_into().unwrap());
         let mut state_roots = vec![];
@@ -424,8 +432,8 @@ impl ForkNetworkCommand {
             assert_eq!(key, format!("FORK_TOOL_SHARD_ID:{shard_id}"));
             state_roots.push(state_root);
         }
-        tracing::info!(?state_roots, ?epoch_id, block_height);
-        Ok((state_roots, epoch_id, block_height))
+        tracing::info!(?state_roots, ?block_hash, ?epoch_id, block_height);
+        Ok((state_roots, block_hash, epoch_id, block_height))
     }
 
     /// Checks that `~/.near/data/fork-snapshot/data` exists.

tools/fork-network/src/single_shard_storage_mutator.rs

@@ -23,15 +23,10 @@ impl SingleShardStorageMutator {
     pub(crate) fn new(
         shard_id: ShardId,
         runtime: &NightshadeRuntime,
+        prev_block_hash: CryptoHash,
         state_root: StateRoot,
     ) -> anyhow::Result<Self> {
-        let block_hash = CryptoHash::new();
-        let trie = runtime.get_trie_for_shard(
-            shard_id,
-            &block_hash, // unused because of use_flat_storage=false
-            state_root,
-            false,
-        )?;
+        let trie = runtime.get_trie_for_shard(shard_id, &prev_block_hash, state_root, false)?;
         let trie_update = TrieUpdate::new(trie);
         Ok(Self { trie_update, shard_tries: runtime.get_tries(), num_changes: 0 })
     }

tools/fork-network/src/storage_mutator.rs

@@ -2,7 +2,8 @@ use crate::single_shard_storage_mutator::SingleShardStorageMutator;
 use near_crypto::PublicKey;
 use near_epoch_manager::EpochManagerAdapter;
 use near_primitives::account::{AccessKey, Account};
-use near_primitives::types::{AccountId, EpochId, StateRoot};
+use near_primitives::hash::CryptoHash;
+use near_primitives::types::{AccountId, EpochId, ShardId, StateRoot};
 use nearcore::NightshadeRuntime;
 use std::sync::Arc;
 
@@ -19,17 +20,19 @@ impl StorageMutator {
         epoch_manager: Arc<dyn EpochManagerAdapter>,
         runtime: &NightshadeRuntime,
         epoch_id: EpochId,
+        prev_block_hash: CryptoHash,
         state_roots: Vec<StateRoot>,
     ) -> anyhow::Result<Self> {
         let shard_layout = epoch_manager.get_shard_layout(&epoch_id)?;
-        let num_shards = shard_layout.num_shards();
+        assert_eq!(shard_layout.num_shards(), state_roots.len() as u64);
 
         let mut mutators = vec![];
-        for shard_id in 0..num_shards {
+        for (shard_id, state_root) in state_roots.iter().enumerate() {
             mutators.push(SingleShardStorageMutator::new(
-                shard_id,
+                shard_id as ShardId,
                 runtime,
-                state_roots[shard_id as usize],
+                prev_block_hash,
+                *state_root,
             )?);
         }
         Ok(Self { epoch_manager, epoch_id, mutators })