1use std::path::Path;
12use std::time::Instant;
13
14use alloy_primitives::map::HashMap;
15use alloy_primitives::{Address, B256, U256};
16use foundry_fork_db::BlockchainDb;
17use revm::state::AccountInfo;
18use serde::{Deserialize, Serialize};
19use tracing::{debug, warn};
20
21use super::versioned;
22use crate::errors::PersistenceError;
23
24const BINARY_STATE_MAGIC: &[u8; 8] = b"EFCSTAT\0";
25const BINARY_STATE_VERSION: u32 = 1;
26
27#[derive(Serialize, Deserialize)]
30struct BinaryEvmState {
31 accounts: Vec<(Address, BinaryAccountInfo)>,
32 storage: Vec<(Address, Vec<(U256, U256)>)>,
33}
34
35#[derive(Serialize, Deserialize)]
37struct BinaryAccountInfo {
38 balance: U256,
39 nonce: u64,
40 code_hash: B256,
41}
42
43pub fn save_binary_state(
58 blockchain_db: &BlockchainDb,
59 path: &Path,
60) -> Result<(), PersistenceError> {
61 let start = Instant::now();
62
63 let accounts: Vec<(Address, BinaryAccountInfo)> = blockchain_db
64 .accounts()
65 .read()
66 .iter()
67 .map(|(addr, info)| {
68 (
69 *addr,
70 BinaryAccountInfo {
71 balance: info.balance,
72 nonce: info.nonce,
73 code_hash: info.code_hash,
74 },
75 )
76 })
77 .collect();
78
79 let storage: Vec<(Address, Vec<(U256, U256)>)> = blockchain_db
80 .storage()
81 .read()
82 .iter()
83 .map(|(addr, slots)| (*addr, slots.iter().map(|(k, v)| (*k, *v)).collect()))
84 .collect();
85
86 let state = BinaryEvmState { accounts, storage };
87
88 let data = versioned::encode(
89 BINARY_STATE_MAGIC,
90 BINARY_STATE_VERSION,
91 &state,
92 "binary EVM state",
93 )?;
94 if let Some(parent) = path.parent() {
95 std::fs::create_dir_all(parent).map_err(|err| PersistenceError::create_dir(parent, err))?;
96 }
97 std::fs::write(path, &data).map_err(|err| PersistenceError::write(path, err))?;
98
99 let ms = start.elapsed().as_millis();
100 debug!(
101 accounts = state.accounts.len(),
102 storage_contracts = state.storage.len(),
103 bytes = data.len(),
104 save_ms = ms,
105 "Saved binary EVM state"
106 );
107 Ok(())
108}
109
110pub fn load_binary_state(blockchain_db: &BlockchainDb, path: &Path) -> bool {
122 let start = Instant::now();
123
124 let data = match std::fs::read(path) {
125 Ok(d) => d,
126 Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
127 debug!("No binary EVM state file found, starting fresh");
128 return false;
129 }
130 Err(e) => {
131 warn!(error = %e, "Failed to read binary EVM state, starting fresh");
132 return false;
133 }
134 };
135
136 let Some(state) = versioned::decode::<BinaryEvmState>(
137 &data,
138 BINARY_STATE_MAGIC,
139 BINARY_STATE_VERSION,
140 "binary EVM state",
141 ) else {
142 warn!("Failed to decode binary EVM state, starting fresh");
143 return false;
144 };
145
146 let account_count = state.accounts.len();
147 let storage_contract_count = state.storage.len();
148 let mut total_slots = 0usize;
149
150 {
152 let mut accounts = blockchain_db.accounts().write();
153 for (addr, info) in state.accounts {
154 accounts.insert(
155 addr,
156 AccountInfo {
157 balance: info.balance,
158 nonce: info.nonce,
159 code_hash: info.code_hash,
160 code: None,
161 account_id: None,
162 },
163 );
164 }
165 }
166
167 {
169 let mut storage = blockchain_db.storage().write();
170 for (addr, slots) in state.storage {
171 total_slots += slots.len();
172 let map: HashMap<U256, U256> = slots.into_iter().collect();
173 storage.insert(addr, map);
174 }
175 }
176
177 let ms = start.elapsed().as_millis();
178 debug!(
179 accounts = account_count,
180 storage_contracts = storage_contract_count,
181 total_slots,
182 bytes = data.len(),
183 load_ms = ms,
184 "Loaded binary EVM state"
185 );
186
187 true
188}
189
190#[cfg(test)]
191mod tests {
192 use super::*;
193 use foundry_fork_db::cache::BlockchainDbMeta;
194
195 #[test]
196 fn test_save_load_round_trip() {
197 let dir = std::env::temp_dir().join("evm_fork_cache_test_binary_state");
198 let path = dir.join("test_state.bin");
199 let _ = std::fs::remove_file(&path);
200
201 let meta = BlockchainDbMeta::default();
203 let db = BlockchainDb::new(meta, None);
204
205 let addr1 = Address::repeat_byte(0x01);
206 let addr2 = Address::repeat_byte(0x02);
207
208 {
210 let mut accounts = db.accounts().write();
211 accounts.insert(
212 addr1,
213 AccountInfo {
214 balance: U256::from(1000),
215 nonce: 5,
216 code_hash: B256::repeat_byte(0xAA),
217 code: None,
218 account_id: None,
219 },
220 );
221 accounts.insert(
222 addr2,
223 AccountInfo {
224 balance: U256::from(2000),
225 nonce: 10,
226 code_hash: B256::repeat_byte(0xBB),
227 code: None,
228 account_id: None,
229 },
230 );
231 }
232
233 {
235 let mut storage = db.storage().write();
236 let mut slots1 = HashMap::default();
237 slots1.insert(U256::from(0), U256::from(42));
238 slots1.insert(U256::from(1), U256::from(99));
239 storage.insert(addr1, slots1);
240
241 let mut slots2 = HashMap::default();
242 slots2.insert(U256::from(4), U256::from(777));
243 storage.insert(addr2, slots2);
244 }
245
246 save_binary_state(&db, &path).expect("save binary state");
248 assert!(path.exists());
249 let bytes = std::fs::read(&path).expect("read saved state");
250 assert!(
251 bytes.starts_with(b"EFCSTAT\0"),
252 "binary state cache must carry a magic header"
253 );
254 assert_eq!(
255 &bytes[8..12],
256 &1u32.to_le_bytes(),
257 "binary state cache must carry an explicit version"
258 );
259
260 let meta2 = BlockchainDbMeta::default();
262 let db2 = BlockchainDb::new(meta2, None);
263 assert!(load_binary_state(&db2, &path));
264
265 {
267 let accounts = db2.accounts().read();
268 assert_eq!(accounts.len(), 2);
269 let info1 = accounts.get(&addr1).unwrap();
270 assert_eq!(info1.balance, U256::from(1000));
271 assert_eq!(info1.nonce, 5);
272 assert!(info1.code.is_none()); }
274
275 {
277 let storage = db2.storage().read();
278 assert_eq!(storage.len(), 2);
279 assert_eq!(
280 *storage.get(&addr1).unwrap().get(&U256::from(0)).unwrap(),
281 U256::from(42)
282 );
283 assert_eq!(
284 *storage.get(&addr2).unwrap().get(&U256::from(4)).unwrap(),
285 U256::from(777)
286 );
287 }
288
289 let _ = std::fs::remove_file(&path);
290 let _ = std::fs::remove_dir(&dir);
291 }
292
293 #[test]
294 fn save_binary_state_reports_write_failures() {
295 let dir = std::env::temp_dir().join("evm_fork_cache_test_binary_state_write_error");
296 let _ = std::fs::remove_dir_all(&dir);
297 let _ = std::fs::remove_file(&dir);
298 std::fs::write(&dir, b"not a directory").expect("create file path conflict");
299
300 let db = BlockchainDb::new(BlockchainDbMeta::default(), None);
301 let path = dir.join("state.bin");
302 let err = save_binary_state(&db, &path).expect_err("save must report write failure");
303 assert!(
304 err.to_string().contains("directory") || err.to_string().contains("Not a directory"),
305 "unexpected error: {err:#}"
306 );
307
308 let _ = std::fs::remove_file(&dir);
309 }
310
311 #[test]
312 fn test_load_missing_file_returns_false() {
313 let meta = BlockchainDbMeta::default();
314 let db = BlockchainDb::new(meta, None);
315 assert!(!load_binary_state(
316 &db,
317 std::path::Path::new("/tmp/nonexistent_binary_state.bin")
318 ));
319 }
320
321 #[test]
322 fn test_load_corrupt_file_returns_false() {
323 let dir = std::env::temp_dir().join("evm_fork_cache_test_binary_state_corrupt");
324 let path = dir.join("corrupt.bin");
325 let _ = std::fs::create_dir_all(&dir);
326 std::fs::write(&path, b"not valid bincode").unwrap();
327
328 let meta = BlockchainDbMeta::default();
329 let db = BlockchainDb::new(meta, None);
330 assert!(!load_binary_state(&db, &path));
331
332 let _ = std::fs::remove_file(&path);
333 let _ = std::fs::remove_dir(&dir);
334 }
335
336 #[test]
337 fn load_legacy_raw_bincode_returns_false() {
338 let dir = std::env::temp_dir().join("evm_fork_cache_test_binary_state_legacy");
339 let path = dir.join("legacy.bin");
340 let _ = std::fs::create_dir_all(&dir);
341 let legacy = BinaryEvmState {
342 accounts: Vec::new(),
343 storage: Vec::new(),
344 };
345 std::fs::write(&path, bincode::serialize(&legacy).unwrap()).unwrap();
346
347 let db = BlockchainDb::new(BlockchainDbMeta::default(), None);
348 assert!(
349 !load_binary_state(&db, &path),
350 "unversioned legacy bincode must be treated as a cache miss"
351 );
352
353 let _ = std::fs::remove_file(&path);
354 let _ = std::fs::remove_dir(&dir);
355 }
356}