From 32b3697f6c4a8ed83672acdad781ddc757013689 Mon Sep 17 00:00:00 2001
From: Dan Cline <6798349+Rjected@users.noreply.github.com>
Date: Mon, 30 Sep 2024 10:02:44 -0400
Subject: [PATCH] fix(providers): collect BlockState before constructing DB
 provider

---
 crates/chain-state/src/in_memory.rs               | 15 +++++++++++++++
 .../provider/src/providers/blockchain_provider.rs |  2 +-
 2 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/crates/chain-state/src/in_memory.rs b/crates/chain-state/src/in_memory.rs
index fc142dd03a71..1e5c7ed1c134 100644
--- a/crates/chain-state/src/in_memory.rs
+++ b/crates/chain-state/src/in_memory.rs
@@ -494,6 +494,21 @@ impl CanonicalInMemoryState {
         self.inner.canon_state_notification_sender.send(event).ok();
     }
 
+    /// Return state provider with reference to in-memory blocks that overlay database state.
+    ///
+    /// This merges the state of all blocks that are part of the chain that the requested block is
+    /// the head of. This includes all blocks that connect back to the canonical block on disk.
+    pub fn state_provider_from_state(
+        &self,
+        state: &BlockState,
+        historical: StateProviderBox,
+    ) -> MemoryOverlayStateProvider {
+        let in_memory: Vec<_> =
+            state.chain().into_iter().map(|block_state| block_state.block()).collect();
+
+        MemoryOverlayStateProvider::new(historical, in_memory)
+    }
+
     /// Return state provider with reference to in-memory blocks that overlay database state.
     ///
     /// This merges the state of all blocks that are part of the chain that the requested block is
diff --git a/crates/storage/provider/src/providers/blockchain_provider.rs b/crates/storage/provider/src/providers/blockchain_provider.rs
index cb90cc08a61c..0f7c3a0b66f1 100644
--- a/crates/storage/provider/src/providers/blockchain_provider.rs
+++ b/crates/storage/provider/src/providers/blockchain_provider.rs
@@ -177,7 +177,7 @@ impl<N: ProviderNodeTypes> BlockchainProvider2<N> {
         let state = state.as_ref();
         let anchor_hash = state.anchor().hash;
         let latest_historical = self.database.history_by_block_hash(anchor_hash)?;
-        Ok(self.canonical_in_memory_state.state_provider(state.hash(), latest_historical))
+        Ok(self.canonical_in_memory_state.state_provider_from_state(state, latest_historical))
     }
 
     /// Returns: