feat(core): add feature "test-utils"

core/Cargo.toml

@@ -70,3 +70,4 @@ ethers-core = { workspace = true }
 sp1 = ["dep:sp1-driver", "sp1-driver/enable"]
 risc0 = ["dep:risc0-driver", "risc0-driver/enable"]
 sgx = ["dep:sgx-prover", "sgx-prover/enable"]
+test-utils = []

core/src/provider/mod.rs

@@ -11,6 +11,7 @@ use crate::{
 };
 
 pub mod db;
+pub mod persistent_map;
 pub mod rpc;
 
 #[allow(async_fn_in_trait)]

core/src/provider/persistent_map.rs

@@ -0,0 +1,156 @@
+use alloy_rpc_types::EIP1186AccountProofResponse;
+use alloy_rpc_types::EIP1186StorageProof;
+use reth_primitives::{Address, U256};
+use reth_revm::primitives::AccountInfo;
+use serde::{Deserialize, Serialize};
+use std::{collections::HashMap, fs::File, hash::Hash, io::Write, path::PathBuf};
+
+#[derive(Hash, Eq, PartialEq)]
+pub struct StorageSlotKey {
+    address: Address,
+    slot: U256,
+}
+
+impl Serialize for StorageSlotKey {
+    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+    where
+        S: serde::Serializer,
+    {
+        format!("{}:{}", self.address, self.slot).serialize(serializer)
+    }
+}
+
+impl<'de> Deserialize<'de> for StorageSlotKey {
+    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
+    where
+        D: serde::Deserializer<'de>,
+    {
+        let s = String::deserialize(deserializer)?;
+        let mut parts = s.split(':');
+        let address = parts
+            .next()
+            .ok_or_else(|| serde::de::Error::custom("Missing address"))?
+            .parse()
+            .map_err(serde::de::Error::custom)?;
+        let slot = parts
+            .next()
+            .ok_or_else(|| serde::de::Error::custom("Missing slot"))?
+            .parse()
+            .map_err(serde::de::Error::custom)?;
+        Ok(Self { address, slot })
+    }
+}
+
+impl From<(Address, U256)> for StorageSlotKey {
+    fn from((address, slot): (Address, U256)) -> Self {
+        Self { address, slot }
+    }
+}
+
+impl From<StorageSlotKey> for (Address, U256) {
+    fn from(key: StorageSlotKey) -> Self {
+        (key.address, key.slot)
+    }
+}
+
+pub struct PersistentBlockData {
+    base: PathBuf,
+}
+
+impl PersistentBlockData {
+    pub fn new(base: impl Into<PathBuf>) -> Self {
+        Self { base: base.into() }
+    }
+
+    pub fn accounts(&self, block_number: u64) -> PersistentMap<Address, AccountInfo> {
+        PersistentMap::new(format!(
+            "{}/{}-accounts.json",
+            self.base.display(),
+            block_number
+        ))
+    }
+
+    pub fn storage_values(&self, block_number: u64) -> PersistentMap<StorageSlotKey, U256> {
+        PersistentMap::new(format!(
+            "{}/{}-storage_values.json",
+            self.base.display(),
+            block_number
+        ))
+    }
+
+    pub fn account_proofs(
+        &self,
+        block_number: u64,
+    ) -> PersistentMap<Address, EIP1186AccountProofResponse> {
+        PersistentMap::new(format!(
+            "{}/{}-account_proofs.json",
+            self.base.display(),
+            block_number
+        ))
+    }
+
+    pub fn account_storage_proofs(
+        &self,
+        block_number: u64,
+    ) -> PersistentMap<StorageSlotKey, EIP1186StorageProof> {
+        PersistentMap::new(format!(
+            "{}/{}-account_storage_proofs.json",
+            self.base.display(),
+            block_number
+        ))
+    }
+}
+
+/// A simple cache that implements most of the methods of `HashMap` and ensures that the data is persisted to a file.
+pub struct PersistentMap<
+    K: Hash + Eq + Serialize + for<'de> Deserialize<'de>,
+    V: Serialize + for<'de> Deserialize<'de>,
+> {
+    file_path: PathBuf,
+    map: HashMap<K, V>,
+}
+
+impl<
+        K: Hash + Eq + Serialize + for<'de> Deserialize<'de>,
+        V: Serialize + for<'de> Deserialize<'de>,
+    > PersistentMap<K, V>
+{
+    pub fn new(file_path: impl Into<PathBuf>) -> Self {
+        let file_path = file_path.into();
+        if let Some(parent) = file_path.parent() {
+            std::fs::create_dir_all(parent).unwrap_or_else(|e| {
+                tracing::warn!("Failed to create directory {}: {}", parent.display(), e);
+            });
+        }
+
+        // Load the storage from the file
+        let map = File::open(&file_path)
+            .and_then(|file| {
+                serde_json::from_reader::<_, HashMap<K, V>>(file)
+                    .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))
+            })
+            .unwrap_or(HashMap::new());
+
+        Self { file_path, map }
+    }
+
+    pub fn contains_key(&self, key: &K) -> bool {
+        self.map.contains_key(key)
+    }
+
+    pub fn get(&self, key: &K) -> Option<&V> {
+        self.map.get(key)
+    }
+
+    pub fn insert(&mut self, key: K, value: V) -> Option<V> {
+        self.map.insert(key, value)
+    }
+
+    pub fn save(&self) -> std::io::Result<()> {
+        let mut file = File::create(&self.file_path)?;
+        serde_json::to_writer(&mut file, &self.map)
+            .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))?;
+        file.flush()?;
+        Ok(())
+    }
+}