evm_fork_cache/
prefetch_registry.rs

1//! Generalized storage prefetch registry for EVM cache pre-warming.
2//!
3//! Captures access lists from EVM interactions (multicall batches, simulations)
4//! and persists them across cycles. On the next cycle, batch-fetches the recorded
5//! slots into BlockchainDb before the EVM touches them, converting N individual
6//! `eth_getStorageAt` RPC calls into a small number of batched HTTP requests
7//! (the batch size is governed by the cache's speed mode).
8//!
9//! Supports two storage shapes:
10//! - **Aggregated phases** (e.g., a `pool_refresh` phase): one access list per
11//!   phase.
12//! - **Keyed phases**: per-address access lists, enabling selective prefetch
13//!   for only the addresses that will be simulated.
14
15use std::collections::{HashMap, HashSet};
16use std::path::Path;
17
18use alloy_primitives::{Address, U256};
19use serde::{Deserialize, Serialize};
20use tracing::{debug, info, warn};
21
22use crate::StorageAccessList;
23use crate::cache::{EvmCache, versioned};
24use crate::errors::PersistenceError;
25
26const PREFETCH_REGISTRY_MAGIC: &[u8; 8] = b"EFC-PFRG";
27const PREFETCH_REGISTRY_VERSION: u32 = 1;
28
29/// Registry of access lists keyed by phase, persisted across cycles.
30///
31/// On disk, the registry is stored as a crate-specific magic/version envelope
32/// followed by a bincode payload. Unknown or legacy unversioned files are
33/// treated as cache misses and loaded as an empty registry.
34#[derive(Debug, Default, Clone, Serialize, Deserialize)]
35pub struct PrefetchRegistry {
36    /// Phases with a single aggregated access list (e.g., a `pool_refresh` phase).
37    phases: HashMap<String, StorageAccessList>,
38    /// Phases with per-address access lists.
39    /// Stored by address so callers can selectively prefetch only ready targets.
40    keyed_phases: HashMap<String, HashMap<Address, StorageAccessList>>,
41}
42
43impl PrefetchRegistry {
44    /// Load a registry from `path` (versioned binary format).
45    ///
46    /// Returns [`Default`] (an empty registry) on any error — a missing file, an
47    /// unreadable file, an unrecognized magic/version header, or corrupt
48    /// contents. These cases are not distinguished by the return value: an
49    /// incompatible registry is indistinguishable from a fresh start, so it is
50    /// treated as a cache miss (logged at `warn`).
51    pub fn load(path: &Path) -> Self {
52        match std::fs::read(path) {
53            Ok(data) => {
54                if let Some(registry) = versioned::decode::<PrefetchRegistry>(
55                    &data,
56                    PREFETCH_REGISTRY_MAGIC,
57                    PREFETCH_REGISTRY_VERSION,
58                    "prefetch registry",
59                ) {
60                    let phase_count = registry.phases.len();
61                    let keyed_phase_count = registry.keyed_phases.len();
62                    let total_slots: usize = registry
63                        .phases
64                        .values()
65                        .map(|al| al.slots.len())
66                        .sum::<usize>()
67                        + registry
68                            .keyed_phases
69                            .values()
70                            .flat_map(|m| m.values())
71                            .map(|al| al.slots.len())
72                            .sum::<usize>();
73                    info!(
74                        phases = phase_count,
75                        keyed_phases = keyed_phase_count,
76                        total_slots,
77                        "Loaded prefetch registry"
78                    );
79                    registry
80                } else {
81                    warn!("Prefetch registry cache miss, starting fresh");
82                    Self::default()
83                }
84            }
85            Err(_) => {
86                debug!("No prefetch registry file found, starting fresh");
87                Self::default()
88            }
89        }
90    }
91
92    /// Persist the registry to `path` in versioned binary format, creating
93    /// parent directories as needed.
94    ///
95    /// Returns an error if the parent directory cannot be created, serialization
96    /// fails, or the write fails.
97    pub fn save(&self, path: &Path) -> Result<(), PersistenceError> {
98        if let Some(parent) = path.parent() {
99            std::fs::create_dir_all(parent)
100                .map_err(|err| PersistenceError::create_dir(parent, err))?;
101        }
102        let data = versioned::encode(
103            PREFETCH_REGISTRY_MAGIC,
104            PREFETCH_REGISTRY_VERSION,
105            self,
106            "prefetch registry",
107        )?;
108        std::fs::write(path, data).map_err(|err| PersistenceError::write(path, err))?;
109
110        let total_slots: usize = self.phases.values().map(|al| al.slots.len()).sum::<usize>()
111            + self
112                .keyed_phases
113                .values()
114                .flat_map(|m| m.values())
115                .map(|al| al.slots.len())
116                .sum::<usize>();
117        debug!(total_slots, "Saved prefetch registry");
118        Ok(())
119    }
120
121    /// Record the aggregated access list for `phase`, **overwriting** any access
122    /// list previously recorded for that phase.
123    ///
124    /// Each call wholesale replaces the phase's slot set; it does not merge with
125    /// the prior list. To accumulate per-address lists instead, use
126    /// [`record_keyed`](Self::record_keyed).
127    ///
128    /// ```
129    /// use evm_fork_cache::prefetch_registry::PrefetchRegistry;
130    /// use evm_fork_cache::StorageAccessList;
131    /// use alloy_primitives::{Address, U256};
132    ///
133    /// let mut registry = PrefetchRegistry::default();
134    /// let addr = Address::repeat_byte(0x01);
135    ///
136    /// let mut al = StorageAccessList::default();
137    /// al.slots.insert((addr, U256::from(1)));
138    /// registry.record("pool_refresh", al);
139    /// assert!(registry.phase_slots("pool_refresh").contains(&(addr, U256::from(1))));
140    ///
141    /// // A second record replaces the slot set rather than merging.
142    /// let mut al2 = StorageAccessList::default();
143    /// al2.slots.insert((addr, U256::from(2)));
144    /// registry.record("pool_refresh", al2);
145    /// let slots = registry.phase_slots("pool_refresh");
146    /// assert_eq!(slots.len(), 1);
147    /// assert!(slots.contains(&(addr, U256::from(2))));
148    /// ```
149    pub fn record(&mut self, phase: &str, access_list: StorageAccessList) {
150        self.phases.insert(phase.to_string(), access_list);
151    }
152
153    /// Record the access list for a single `key` within a keyed `phase`.
154    ///
155    /// Unlike [`record`](Self::record), this **inserts into** the phase's per-key
156    /// nested map: other keys already recorded under `phase` are preserved, and
157    /// only the entry for `key` is replaced. Pairs with
158    /// [`prefetch_keyed`](Self::prefetch_keyed).
159    pub fn record_keyed(&mut self, phase: &str, key: Address, access_list: StorageAccessList) {
160        self.keyed_phases
161            .entry(phase.to_string())
162            .or_default()
163            .insert(key, access_list);
164    }
165
166    /// Prefetch all slots for an aggregated phase.
167    ///
168    /// Returns `(fetched, errors)`.
169    pub fn prefetch_phase(&self, phase: &str, cache: &mut EvmCache) -> (usize, usize) {
170        let Some(access_list) = self.phases.get(phase) else {
171            debug!(phase, "No prefetch data for phase");
172            return (0, 0);
173        };
174
175        if access_list.slots.is_empty() {
176            return (0, 0);
177        }
178
179        batch_prefetch(cache, access_list.slots.iter().copied(), phase)
180    }
181
182    /// Prefetch slots for specific keys within a keyed phase, excluding slots
183    /// already warm from a previous prefetch stage.
184    ///
185    /// Pairs with [`record_keyed`](Self::record_keyed).
186    ///
187    /// Returns `(fetched, errors)`.
188    pub fn prefetch_keyed(
189        &self,
190        phase: &str,
191        keys: &[Address],
192        cache: &mut EvmCache,
193        exclude: &HashSet<(Address, U256)>,
194    ) -> (usize, usize) {
195        let Some(keyed_map) = self.keyed_phases.get(phase) else {
196            debug!(phase, "No keyed prefetch data for phase");
197            return (0, 0);
198        };
199
200        let slots: HashSet<(Address, U256)> = keys
201            .iter()
202            .filter_map(|addr| keyed_map.get(addr))
203            .flat_map(|al| al.slots.iter().copied())
204            .filter(|slot| !exclude.contains(slot))
205            .collect();
206
207        if slots.is_empty() {
208            debug!(
209                phase,
210                keys = keys.len(),
211                excluded = exclude.len(),
212                "All keyed slots excluded or empty"
213            );
214            return (0, 0);
215        }
216
217        batch_prefetch(cache, slots.into_iter(), phase)
218    }
219
220    /// Returns the set of `(address, slot)` pairs recorded for an aggregated
221    /// `phase`, or an empty set if the phase was never [`record`](Self::record)ed.
222    ///
223    /// Typically used to build the `exclude` set passed to
224    /// [`prefetch_keyed`](Self::prefetch_keyed) so a later stage skips slots a
225    /// prior aggregated prefetch already warmed.
226    pub fn phase_slots(&self, phase: &str) -> HashSet<(Address, U256)> {
227        self.phases
228            .get(phase)
229            .map(|al| al.slots.clone())
230            .unwrap_or_default()
231    }
232}
233
234/// Batch-fetch `slots` into `cache` via its `storage_batch_fetcher` and inject
235/// the results, returning `(fetched, errors)`.
236///
237/// Deduplicating, exclusion, and phase lookup are the caller's responsibility
238/// ([`PrefetchRegistry::prefetch_phase`] / [`PrefetchRegistry::prefetch_keyed`]).
239/// If `slots` is empty, or the cache has no batch fetcher configured, returns
240/// `(0, 0)` without fetching. Otherwise each slot that the fetcher resolves
241/// successfully is injected into the cache and counted in `fetched`; per-slot
242/// fetch errors are counted in `errors` and skipped.
243fn batch_prefetch(
244    cache: &mut EvmCache,
245    slots: impl Iterator<Item = (Address, U256)>,
246    phase: &str,
247) -> (usize, usize) {
248    let requests: Vec<(Address, U256)> = slots.collect();
249    if requests.is_empty() {
250        return (0, 0);
251    }
252
253    let fetcher = match cache.storage_batch_fetcher().cloned() {
254        Some(f) => f,
255        None => {
256            debug!(
257                "No batch fetcher available, skipping prefetch for {}",
258                phase
259            );
260            return (0, 0);
261        }
262    };
263
264    let start = std::time::Instant::now();
265    let total_requested = requests.len();
266    // Fetch at the cache's currently-pinned block.
267    let results = fetcher(requests, cache.block());
268
269    let mut successes: Vec<(Address, U256, U256)> = Vec::with_capacity(results.len());
270    let mut errors = 0usize;
271    for (addr, slot, result) in results {
272        match result {
273            Ok(value) => successes.push((addr, slot, value)),
274            Err(_) => errors += 1,
275        }
276    }
277
278    let fetched = successes.len();
279    cache.inject_storage_batch(&successes);
280
281    let ms = start.elapsed().as_millis() as u64;
282    info!(
283        phase,
284        slots_requested = total_requested,
285        slots_fetched = fetched,
286        errors,
287        prefetch_ms = ms,
288        "Prefetch complete"
289    );
290
291    (fetched, errors)
292}
293
294#[cfg(test)]
295mod tests {
296    use super::*;
297
298    #[test]
299    fn test_registry_record_and_phase_slots() {
300        let mut registry = PrefetchRegistry::default();
301        let mut al = StorageAccessList::default();
302        let addr = Address::repeat_byte(0x01);
303        al.slots.insert((addr, U256::from(1)));
304        al.slots.insert((addr, U256::from(2)));
305
306        registry.record("pool_refresh", al);
307
308        let slots = registry.phase_slots("pool_refresh");
309        assert_eq!(slots.len(), 2);
310        assert!(slots.contains(&(addr, U256::from(1))));
311        assert!(slots.contains(&(addr, U256::from(2))));
312
313        // Non-existent phase returns empty
314        assert!(registry.phase_slots("nonexistent").is_empty());
315    }
316
317    #[test]
318    fn test_registry_record_keyed() {
319        let mut registry = PrefetchRegistry::default();
320        let key_a = Address::repeat_byte(0x01);
321        let key_b = Address::repeat_byte(0x02);
322
323        let mut al1 = StorageAccessList::default();
324        al1.slots.insert((key_a, U256::from(10)));
325
326        let mut al2 = StorageAccessList::default();
327        al2.slots.insert((key_b, U256::from(20)));
328
329        registry.record_keyed("per_target", key_a, al1);
330        registry.record_keyed("per_target", key_b, al2);
331
332        // Verify keyed_phases has both
333        let map = registry.keyed_phases.get("per_target").unwrap();
334        assert_eq!(map.len(), 2);
335        assert!(
336            map.get(&key_a)
337                .unwrap()
338                .slots
339                .contains(&(key_a, U256::from(10)))
340        );
341    }
342
343    #[test]
344    fn test_save_load_round_trip() {
345        let dir = std::env::temp_dir().join("evm_fork_cache_test_prefetch_registry");
346        let path = dir.join("test_registry.bin");
347        let _ = std::fs::remove_file(&path);
348
349        let mut registry = PrefetchRegistry::default();
350        let addr = Address::repeat_byte(0xAA);
351        let mut al = StorageAccessList::default();
352        al.slots.insert((addr, U256::from(42)));
353        al.accounts.insert(addr);
354        registry.record("test_phase", al);
355
356        let key = Address::repeat_byte(0xBB);
357        let mut sal = StorageAccessList::default();
358        sal.slots.insert((key, U256::from(99)));
359        registry.record_keyed("per_target", key, sal);
360
361        registry.save(&path).expect("save registry");
362        let data = std::fs::read(&path).expect("read saved registry");
363        assert!(
364            data.starts_with(b"EFC-PFRG"),
365            "prefetch registry files must carry a magic/version header"
366        );
367
368        let loaded = PrefetchRegistry::load(&path);
369        assert_eq!(loaded.phases.len(), 1);
370        assert!(
371            loaded.phases["test_phase"]
372                .slots
373                .contains(&(addr, U256::from(42)))
374        );
375        assert_eq!(loaded.keyed_phases.len(), 1);
376        assert!(
377            loaded.keyed_phases["per_target"]
378                .get(&key)
379                .unwrap()
380                .slots
381                .contains(&(key, U256::from(99)))
382        );
383
384        let _ = std::fs::remove_file(&path);
385        let _ = std::fs::remove_dir(&dir);
386    }
387
388    #[test]
389    fn save_reports_write_failures() {
390        let dir = std::env::temp_dir().join("evm_fork_cache_test_prefetch_registry_write_error");
391        let _ = std::fs::remove_dir_all(&dir);
392        let _ = std::fs::remove_file(&dir);
393        std::fs::write(&dir, b"not a directory").expect("create file path conflict");
394
395        let registry = PrefetchRegistry::default();
396        let path = dir.join("registry.bin");
397        let err = registry
398            .save(&path)
399            .expect_err("save must report write failure");
400        assert!(
401            err.to_string().contains("directory") || err.to_string().contains("Not a directory"),
402            "unexpected error: {err:#}"
403        );
404
405        let _ = std::fs::remove_file(&dir);
406    }
407
408    #[test]
409    fn legacy_raw_bincode_loads_as_default() {
410        let dir = std::env::temp_dir().join("evm_fork_cache_test_prefetch_registry_legacy");
411        let path = dir.join("legacy_registry.bin");
412        let _ = std::fs::remove_dir_all(&dir);
413        std::fs::create_dir_all(&dir).expect("create temp dir");
414
415        let mut registry = PrefetchRegistry::default();
416        let addr = Address::repeat_byte(0xAA);
417        let mut al = StorageAccessList::default();
418        al.slots.insert((addr, U256::from(42)));
419        registry.record("legacy_phase", al);
420        let legacy = bincode::serialize(&registry).expect("serialize legacy registry");
421        std::fs::write(&path, legacy).expect("write legacy registry");
422
423        let loaded = PrefetchRegistry::load(&path);
424        assert!(
425            loaded.phases.is_empty() && loaded.keyed_phases.is_empty(),
426            "legacy raw bincode must be treated as a cache miss"
427        );
428
429        let _ = std::fs::remove_dir_all(&dir);
430    }
431
432    #[test]
433    fn test_load_missing_file_returns_default() {
434        let path = std::path::Path::new("/tmp/nonexistent_prefetch_registry.bin");
435        let registry = PrefetchRegistry::load(path);
436        assert!(registry.phases.is_empty());
437        assert!(registry.keyed_phases.is_empty());
438    }
439
440    #[test]
441    fn test_record_replaces_existing() {
442        let mut registry = PrefetchRegistry::default();
443        let addr = Address::repeat_byte(0x01);
444
445        let mut al1 = StorageAccessList::default();
446        al1.slots.insert((addr, U256::from(1)));
447        registry.record("phase", al1);
448
449        let mut al2 = StorageAccessList::default();
450        al2.slots.insert((addr, U256::from(2)));
451        registry.record("phase", al2);
452
453        let slots = registry.phase_slots("phase");
454        assert_eq!(slots.len(), 1);
455        assert!(slots.contains(&(addr, U256::from(2))));
456    }
457}
evm_fork_cache/prefetch_registry.rs

evm_fork_cache/
prefetch_registry.rs
