From 6ee2ebfaf8d327ba397098955f08cba779fa07ce Mon Sep 17 00:00:00 2001
From: Yash Atreya <44857776+yash-atreya@users.noreply.github.com>
Date: Thu, 12 Dec 2024 13:02:35 +0530
Subject: [PATCH] fix(`anvil`): set `best_number` correctly while loading state
 with fork activated

---
 crates/anvil/src/eth/backend/mem/mod.rs | 26 ++++++++++++++++---
 crates/anvil/tests/it/state.rs          | 33 ++++++++++++++++++++++++-
 2 files changed, 54 insertions(+), 5 deletions(-)

diff --git a/crates/anvil/src/eth/backend/mem/mod.rs b/crates/anvil/src/eth/backend/mem/mod.rs
index 26787cca60cf..dfd96c72391a 100644
--- a/crates/anvil/src/eth/backend/mem/mod.rs
+++ b/crates/anvil/src/eth/backend/mem/mod.rs
@@ -922,10 +922,28 @@ impl Backend {
             let fork_num_and_hash = self.get_fork().map(|f| (f.block_number(), f.block_hash()));
 
             if let Some((number, hash)) = fork_num_and_hash {
-                // If loading state file on a fork, set best number to the fork block number.
-                // Ref: https://github.com/foundry-rs/foundry/pull/9215#issue-2618681838
-                self.blockchain.storage.write().best_number = U64::from(number);
-                self.blockchain.storage.write().best_hash = hash;
+                let best_number = state.best_block_number.unwrap_or(block.number.to::<U64>());
+                trace!(target: "backend", state_block_number=?best_number, fork_block_number=?number);
+                // If the state.block_number is greater than the fork block number, set best number
+                // to the state block number.
+                // Ref: https://github.com/foundry-rs/foundry/issues/9539
+                if best_number.to::<u64>() > number {
+                    self.blockchain.storage.write().best_number = best_number;
+                    let best_hash =
+                        self.blockchain.storage.read().hash(best_number.into()).ok_or_else(
+                            || {
+                                BlockchainError::RpcError(RpcError::internal_error_with(format!(
+                                    "Best hash not found for best number {best_number}",
+                                )))
+                            },
+                        )?;
+                    self.blockchain.storage.write().best_hash = best_hash;
+                } else {
+                    // If loading state file on a fork, set best number to the fork block number.
+                    // Ref: https://github.com/foundry-rs/foundry/pull/9215#issue-2618681838
+                    self.blockchain.storage.write().best_number = U64::from(number);
+                    self.blockchain.storage.write().best_hash = hash;
+                }
             } else {
                 let best_number = state.best_block_number.unwrap_or(block.number.to::<U64>());
                 self.blockchain.storage.write().best_number = best_number;
diff --git a/crates/anvil/tests/it/state.rs b/crates/anvil/tests/it/state.rs
index f7736de2f403..5337b36b561f 100644
--- a/crates/anvil/tests/it/state.rs
+++ b/crates/anvil/tests/it/state.rs
@@ -2,7 +2,7 @@
 
 use crate::abi::Greeter;
 use alloy_network::{ReceiptResponse, TransactionBuilder};
-use alloy_primitives::{address, utils::Unit, Bytes, Uint, U256};
+use alloy_primitives::{address, utils::Unit, Bytes, Uint, U256, U64};
 use alloy_provider::Provider;
 use alloy_rpc_types::{BlockId, TransactionRequest};
 use alloy_serde::WithOtherFields;
@@ -245,3 +245,34 @@ async fn test_fork_load_state() {
 
     assert_eq!(balance_alice + value, latest_balance_alice);
 }
+
+// <https://github.com/foundry-rs/foundry/issues/9539>
+#[tokio::test(flavor = "multi_thread")]
+async fn test_fork_load_state_with_greater_state_block() {
+    let (api, _handle) = spawn(
+        NodeConfig::test()
+            .with_eth_rpc_url(Some(next_http_rpc_endpoint()))
+            .with_fork_block_number(Some(21070682u64)),
+    )
+    .await;
+
+    api.mine_one().await;
+
+    let block_number = api.block_number().unwrap();
+
+    let serialized_state = api.serialized_state(false).await.unwrap();
+
+    assert_eq!(serialized_state.best_block_number, Some(block_number.to::<U64>()));
+
+    let (api, _handle) = spawn(
+        NodeConfig::test()
+            .with_eth_rpc_url(Some(next_http_rpc_endpoint()))
+            .with_fork_block_number(Some(21070682u64)) // Forked chain has moved forward
+            .with_init_state(Some(serialized_state)),
+    )
+    .await;
+
+    let new_block_number = api.block_number().unwrap();
+
+    assert_eq!(new_block_number, block_number);
+}