From 75082c3f860fc3542b17f3dbe5a1c2ade494f1ae Mon Sep 17 00:00:00 2001
From: zombie-einstein <13398815+zombie-einstein@users.noreply.github.com>
Date: Sat, 10 Feb 2024 22:04:38 +0000
Subject: [PATCH] Track start block and time with request cache

---
 core/db/src/fork_db.rs   | 16 +++++++++++++++-
 core/db/src/types.rs     |  4 +++-
 rust/src/sim/fork_env.rs |  2 --
 rust/src/sim/snapshot.rs | 13 ++++---------
 4 files changed, 22 insertions(+), 13 deletions(-)

diff --git a/core/db/src/fork_db.rs b/core/db/src/fork_db.rs
index 363c33f..8699711 100644
--- a/core/db/src/fork_db.rs
+++ b/core/db/src/fork_db.rs
@@ -52,6 +52,15 @@ impl ForkDb {
         let mut contracts = HashMap::new();
         contracts.insert(KECCAK_EMPTY, Bytecode::new());
         contracts.insert(B256::ZERO, Bytecode::new());
+
+        // Track the original time and block for when we want
+        //  to run a sim using the same cache
+        let timestamp = U256::try_from(block.timestamp.as_u128()).unwrap();
+        let block_number = match block.number {
+            Some(n) => U256::try_from(n.as_u64()).unwrap(),
+            None => U256::ZERO,
+        };
+
         Self {
             accounts: HashMap::new(),
             contracts,
@@ -60,7 +69,12 @@ impl ForkDb {
             provider,
             block_id: Some(block.number.unwrap().into()),
             block,
-            requests: Requests::default(),
+            requests: Requests {
+                start_timestamp: timestamp,
+                start_block_number: block_number,
+                accounts: Vec::new(),
+                storage: Vec::new(),
+            },
         }
     }
 
diff --git a/core/db/src/types.rs b/core/db/src/types.rs
index 4067712..fab30f4 100644
--- a/core/db/src/types.rs
+++ b/core/db/src/types.rs
@@ -49,8 +49,10 @@ impl ToEthers for B256 {
     }
 }
 
-#[derive(Debug, Clone, Default)]
+#[derive(Debug, Clone)]
 pub struct Requests {
+    pub start_timestamp: AlloyU256,
+    pub start_block_number: AlloyU256,
     pub accounts: Vec<(Address, AccountInfo)>,
     pub storage: Vec<(Address, AlloyU256, AlloyU256)>,
 }
diff --git a/rust/src/sim/fork_env.rs b/rust/src/sim/fork_env.rs
index 4243823..e4ed3ea 100644
--- a/rust/src/sim/fork_env.rs
+++ b/rust/src/sim/fork_env.rs
@@ -92,8 +92,6 @@ impl ForkEnv {
     pub fn export_cache<'a>(&mut self, py: Python<'a>) -> PyResult<snapshot::PyRequests<'a>> {
         Ok(snapshot::create_py_request_history(
             py,
-            self.0.network.evm.env.block.timestamp,
-            self.0.network.evm.env.block.number,
             self.0.network.get_request_history(),
         ))
     }
diff --git a/rust/src/sim/snapshot.rs b/rust/src/sim/snapshot.rs
index bd05ff3..cf0c09f 100644
--- a/rust/src/sim/snapshot.rs
+++ b/rust/src/sim/snapshot.rs
@@ -87,14 +87,9 @@ pub type PyRequests<'a> = (
     Vec<(&'a PyBytes, &'a PyBytes, &'a PyBytes)>,
 );
 
-pub fn create_py_request_history<'a>(
-    py: Python<'a>,
-    time_stamp: U256,
-    block_number: U256,
-    requests: &Requests,
-) -> PyRequests<'a> {
-    let time_stamp: u128 = time_stamp.try_into().unwrap();
-    let block_number: u128 = block_number.try_into().unwrap();
+pub fn create_py_request_history<'a>(py: Python<'a>, requests: &Requests) -> PyRequests<'a> {
+    let timestamp: u128 = requests.start_timestamp.try_into().unwrap();
+    let block_number: u128 = requests.start_block_number.try_into().unwrap();
 
     let py_accounts: Vec<(&'a PyBytes, PyAccountInfo)> = requests
         .accounts
@@ -114,7 +109,7 @@ pub fn create_py_request_history<'a>(
         })
         .collect();
 
-    (time_stamp, block_number, py_accounts, py_storage)
+    (timestamp, block_number, py_accounts, py_storage)
 }
 
 pub fn create_py_snapshot<'a, D: DB>(py: Python<'a>, network: &mut Network<D>) -> PyDbState<'a> {