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 fn temp_path(tag: &str) -> std::path::PathBuf {
200 let dir = std::env::temp_dir().join(format!(
201 "evm_fork_cache_binary_state_{tag}_{}",
202 std::process::id()
203 ));
204 let _ = std::fs::remove_dir_all(&dir);
205 std::fs::create_dir_all(&dir).expect("create temp dir");
206 dir.join("state.bin")
207 }
208
209 #[test]
210 fn test_save_load_round_trip() {
211 let path = temp_path("round_trip");
212
213 let meta = BlockchainDbMeta::default();
215 let db = BlockchainDb::new(meta, None);
216
217 let addr1 = Address::repeat_byte(0x01);
218 let addr2 = Address::repeat_byte(0x02);
219
220 {
222 let mut accounts = db.accounts().write();
223 accounts.insert(
224 addr1,
225 AccountInfo {
226 balance: U256::from(1000),
227 nonce: 5,
228 code_hash: B256::repeat_byte(0xAA),
229 code: None,
230 account_id: None,
231 },
232 );
233 accounts.insert(
234 addr2,
235 AccountInfo {
236 balance: U256::from(2000),
237 nonce: 10,
238 code_hash: B256::repeat_byte(0xBB),
239 code: None,
240 account_id: None,
241 },
242 );
243 }
244
245 {
247 let mut storage = db.storage().write();
248 let mut slots1 = HashMap::default();
249 slots1.insert(U256::from(0), U256::from(42));
250 slots1.insert(U256::from(1), U256::from(99));
251 storage.insert(addr1, slots1);
252
253 let mut slots2 = HashMap::default();
254 slots2.insert(U256::from(4), U256::from(777));
255 storage.insert(addr2, slots2);
256 }
257
258 save_binary_state(&db, &path).expect("save binary state");
260 assert!(path.exists());
261 let bytes = std::fs::read(&path).expect("read saved state");
262 assert!(
263 bytes.starts_with(b"EFCSTAT\0"),
264 "binary state cache must carry a magic header"
265 );
266 assert_eq!(
267 &bytes[8..12],
268 &1u32.to_le_bytes(),
269 "binary state cache must carry an explicit version"
270 );
271
272 let meta2 = BlockchainDbMeta::default();
274 let db2 = BlockchainDb::new(meta2, None);
275 assert!(load_binary_state(&db2, &path));
276
277 {
279 let accounts = db2.accounts().read();
280 assert_eq!(accounts.len(), 2);
281 let info1 = accounts.get(&addr1).unwrap();
282 assert_eq!(info1.balance, U256::from(1000));
283 assert_eq!(info1.nonce, 5);
284 assert!(info1.code.is_none()); }
286
287 {
289 let storage = db2.storage().read();
290 assert_eq!(storage.len(), 2);
291 assert_eq!(
292 *storage.get(&addr1).unwrap().get(&U256::from(0)).unwrap(),
293 U256::from(42)
294 );
295 assert_eq!(
296 *storage.get(&addr2).unwrap().get(&U256::from(4)).unwrap(),
297 U256::from(777)
298 );
299 }
300 }
301
302 #[test]
303 fn save_binary_state_reports_write_failures() {
304 let dir = std::env::temp_dir().join(format!(
307 "evm_fork_cache_binary_state_write_error_{}",
308 std::process::id()
309 ));
310 let _ = std::fs::remove_dir_all(&dir);
311 let _ = std::fs::remove_file(&dir);
312 std::fs::write(&dir, b"not a directory").expect("create file path conflict");
313
314 let db = BlockchainDb::new(BlockchainDbMeta::default(), None);
315 let path = dir.join("state.bin");
316 let err = save_binary_state(&db, &path).expect_err("save must report write failure");
317 assert!(
318 err.to_string().contains("directory") || err.to_string().contains("Not a directory"),
319 "unexpected error: {err:#}"
320 );
321
322 let _ = std::fs::remove_file(&dir);
323 }
324
325 #[test]
326 fn test_load_missing_file_returns_false() {
327 let meta = BlockchainDbMeta::default();
328 let db = BlockchainDb::new(meta, None);
329 assert!(!load_binary_state(
330 &db,
331 std::path::Path::new("/tmp/nonexistent_binary_state.bin")
332 ));
333 }
334
335 #[test]
336 fn test_load_corrupt_file_returns_false() {
337 let path = temp_path("corrupt");
338 std::fs::write(&path, b"not valid bincode").unwrap();
339
340 let meta = BlockchainDbMeta::default();
341 let db = BlockchainDb::new(meta, None);
342 assert!(!load_binary_state(&db, &path));
343 }
344
345 #[test]
346 fn load_legacy_raw_bincode_returns_false() {
347 let path = temp_path("legacy");
348 let legacy = BinaryEvmState {
349 accounts: Vec::new(),
350 storage: Vec::new(),
351 };
352 std::fs::write(&path, bincode::serialize(&legacy).unwrap()).unwrap();
353
354 let db = BlockchainDb::new(BlockchainDbMeta::default(), None);
355 assert!(
356 !load_binary_state(&db, &path),
357 "unversioned legacy bincode must be treated as a cache miss"
358 );
359 }
360}