fix: block apply in the past

chain/chain/src/runtime/mod.rs

@@ -745,6 +745,14 @@ impl RuntimeAdapter for NightshadeRuntime {
             StorageDataSource::Db => {
                 self.tries.get_trie_for_shard(shard_uid, storage_config.state_root)
             }
+            StorageDataSource::DbTrieOnly => {
+                // If there is no flat storage on disk, use trie but simulate costs with enabled
+                // flat storage by not charging gas for trie nodes.
+                // WARNING: should never be used in production! Consider this option only for debugging or replaying blocks.
+                let mut trie = self.tries.get_trie_for_shard(shard_uid, storage_config.state_root);
+                trie.dont_charge_gas_for_trie_node_access();
+                trie
+            }
             StorageDataSource::Recorded(storage) => Trie::from_recorded_storage(
                 storage,
                 storage_config.state_root,
@@ -958,6 +966,19 @@ impl RuntimeAdapter for NightshadeRuntime {
                 storage_config.state_root,
                 storage_config.use_flat_storage,
             )?,
+            StorageDataSource::DbTrieOnly => {
+                // If there is no flat storage on disk, use trie but simulate costs with enabled
+                // flat storage by not charging gas for trie nodes.
+                // WARNING: should never be used in production! Consider this option only for debugging or replaying blocks.
+                let mut trie = self.get_trie_for_shard(
+                    shard_id,
+                    &block.prev_block_hash,
+                    storage_config.state_root,
+                    false,
+                )?;
+                trie.dont_charge_gas_for_trie_node_access();
+                trie
+            }
             StorageDataSource::Recorded(storage) => Trie::from_recorded_storage(
                 storage,
                 storage_config.state_root,

chain/chain/src/types.rs

@@ -261,6 +261,11 @@ impl ChainGenesis {
 pub enum StorageDataSource {
     /// Full state data is present in DB.
     Db,
+    /// Trie is present in DB and flat storage is not.
+    /// Used to reply past blocks and simulate gas costs as if flat storage
+    /// was present.
+    /// WARNING: do not use this variant in production!
+    DbTrieOnly,
     /// State data is supplied from state witness, there is no state data
     /// stored on disk.
     Recorded(PartialStorage),
@@ -282,6 +287,19 @@ impl RuntimeStorageConfig {
             state_patch: Default::default(),
         }
     }
+
+    /// Creates a [RuntimeStorageConfig] with [StorageDataSource::DbTrieOnly].
+    /// Flat storage is disabled because it is implied to be missing.
+    ///
+    /// This's meant to be used only to replay blocks.
+    pub fn with_db_trie_only(state_root: StateRoot) -> Self {
+        Self {
+            state_root,
+            use_flat_storage: false,
+            source: StorageDataSource::DbTrieOnly,
+            state_patch: Default::default(),
+        }
+    }
 }
 
 #[derive(Clone)]

core/store/src/trie/mod.rs

@@ -660,6 +660,11 @@ impl Trie {
         }
     }
 
+    /// Helper to simulate gas costs as if flat storage was present.
+    pub fn dont_charge_gas_for_trie_node_access(&mut self) {
+        self.charge_gas_for_trie_node_access = false;
+    }
+
     /// Makes a new trie that has everything the same except that access
     /// through that trie accumulates a state proof for all nodes accessed.
     pub fn recording_reads(&self) -> Self {

tools/state-viewer/src/apply_chain_range.rs

@@ -1,9 +1,8 @@
-use crate::cli::ApplyRangeMode;
+use crate::cli::{ApplyRangeMode, StorageSource};
 use near_chain::chain::collect_receipts_from_response;
 use near_chain::migrations::check_if_block_is_first_with_chunk_of_version;
 use near_chain::types::{
     ApplyChunkBlockContext, ApplyChunkResult, ApplyChunkShardContext, RuntimeAdapter,
-    RuntimeStorageConfig,
 };
 use near_chain::{ChainStore, ChainStoreAccess, ChainStoreUpdate};
 use near_chain_configs::Genesis;
@@ -125,7 +124,7 @@ fn apply_block_from_range(
     verbose_output: bool,
     csv_file_mutex: &Mutex<Option<&mut File>>,
     only_contracts: bool,
-    use_flat_storage: bool,
+    storage: StorageSource,
 ) {
     // normally save_trie_changes depends on whether the node is
     // archival, but here we don't care, and can just set it to false
@@ -232,9 +231,10 @@ fn apply_block_from_range(
                 return;
             }
         }
+
         runtime_adapter
             .apply_chunk(
-                RuntimeStorageConfig::new(*chunk_inner.prev_state_root(), use_flat_storage),
+                storage.create_runtime_storage(*chunk_inner.prev_state_root()),
                 ApplyChunkReason::UpdateTrackedShard,
                 ApplyChunkShardContext {
                     shard_id,
@@ -260,7 +260,7 @@ fn apply_block_from_range(
 
         runtime_adapter
             .apply_chunk(
-                RuntimeStorageConfig::new(*chunk_extra.state_root(), use_flat_storage),
+                storage.create_runtime_storage(*chunk_extra.state_root()),
                 ApplyChunkReason::UpdateTrackedShard,
                 ApplyChunkShardContext {
                     shard_id,
@@ -376,7 +376,7 @@ pub fn apply_chain_range(
     verbose_output: bool,
     csv_file: Option<&mut File>,
     only_contracts: bool,
-    use_flat_storage: bool,
+    storage: StorageSource,
 ) {
     let parent_span = tracing::debug_span!(
         target: "state_viewer",
@@ -386,14 +386,14 @@ pub fn apply_chain_range(
         ?end_height,
         %shard_id,
         only_contracts,
-        use_flat_storage)
+        ?storage)
     .entered();
     let chain_store = ChainStore::new(store.clone(), genesis.config.genesis_height, false);
     let (start_height, end_height) = match mode {
         ApplyRangeMode::Benchmarking => {
             // Benchmarking mode requires flat storage and retrieves start and
             // end heights from flat storage and chain.
-            assert!(use_flat_storage);
+            assert!(matches!(storage, StorageSource::FlatStorage));
             assert!(start_height.is_none());
             assert!(end_height.is_none());
 
@@ -457,7 +457,7 @@ pub fn apply_chain_range(
             verbose_output,
             &csv_file_mutex,
             only_contracts,
-            use_flat_storage,
+            storage,
         );
     };
 
@@ -542,7 +542,7 @@ mod test {
     use nearcore::NightshadeRuntime;
 
     use crate::apply_chain_range::apply_chain_range;
-    use crate::cli::ApplyRangeMode;
+    use crate::cli::{ApplyRangeMode, StorageSource};
 
     fn setup(epoch_length: NumBlocks) -> (Store, Genesis, TestEnv) {
         let mut genesis =
@@ -641,7 +641,7 @@ mod test {
             true,
             None,
             false,
-            false,
+            StorageSource::Trie,
         );
     }
 
@@ -685,7 +685,7 @@ mod test {
             true,
             Some(file.as_file_mut()),
             false,
-            false,
+            StorageSource::Trie,
         );
         let mut csv = String::new();
         file.as_file_mut().seek(SeekFrom::Start(0)).unwrap();