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    /// A unique, freshly-created temp dir keyed by pid so concurrent `cargo
344    /// test` processes never share (and never `remove_dir_all`) each other's
345    /// directory; each test passes a distinct `tag`. Returns the file path to
346    /// write within it.
347    fn temp_path(tag: &str) -> std::path::PathBuf {
348        let dir = std::env::temp_dir().join(format!(
349            "evm_fork_cache_prefetch_registry_{tag}_{}",
350            std::process::id()
351        ));
352        let _ = std::fs::remove_dir_all(&dir);
353        std::fs::create_dir_all(&dir).expect("create temp dir");
354        dir.join("registry.bin")
355    }
356
357    #[test]
358    fn test_save_load_round_trip() {
359        let path = temp_path("round_trip");
360
361        let mut registry = PrefetchRegistry::default();
362        let addr = Address::repeat_byte(0xAA);
363        let mut al = StorageAccessList::default();
364        al.slots.insert((addr, U256::from(42)));
365        al.accounts.insert(addr);
366        registry.record("test_phase", al);
367
368        let key = Address::repeat_byte(0xBB);
369        let mut sal = StorageAccessList::default();
370        sal.slots.insert((key, U256::from(99)));
371        registry.record_keyed("per_target", key, sal);
372
373        registry.save(&path).expect("save registry");
374        let data = std::fs::read(&path).expect("read saved registry");
375        assert!(
376            data.starts_with(b"EFC-PFRG"),
377            "prefetch registry files must carry a magic/version header"
378        );
379
380        let loaded = PrefetchRegistry::load(&path);
381        assert_eq!(loaded.phases.len(), 1);
382        assert!(
383            loaded.phases["test_phase"]
384                .slots
385                .contains(&(addr, U256::from(42)))
386        );
387        assert_eq!(loaded.keyed_phases.len(), 1);
388        assert!(
389            loaded.keyed_phases["per_target"]
390                .get(&key)
391                .unwrap()
392                .slots
393                .contains(&(key, U256::from(99)))
394        );
395    }
396
397    #[test]
398    fn save_reports_write_failures() {
399        // Deliberately a *file* (not a dir) at a pid-unique path, so saving into
400        // a child path fails with "not a directory".
401        let dir = std::env::temp_dir().join(format!(
402            "evm_fork_cache_prefetch_registry_write_error_{}",
403            std::process::id()
404        ));
405        let _ = std::fs::remove_dir_all(&dir);
406        let _ = std::fs::remove_file(&dir);
407        std::fs::write(&dir, b"not a directory").expect("create file path conflict");
408
409        let registry = PrefetchRegistry::default();
410        let path = dir.join("registry.bin");
411        let err = registry
412            .save(&path)
413            .expect_err("save must report write failure");
414        assert!(
415            err.to_string().contains("directory") || err.to_string().contains("Not a directory"),
416            "unexpected error: {err:#}"
417        );
418
419        let _ = std::fs::remove_file(&dir);
420    }
421
422    #[test]
423    fn legacy_raw_bincode_loads_as_default() {
424        let path = temp_path("legacy");
425
426        let mut registry = PrefetchRegistry::default();
427        let addr = Address::repeat_byte(0xAA);
428        let mut al = StorageAccessList::default();
429        al.slots.insert((addr, U256::from(42)));
430        registry.record("legacy_phase", al);
431        let legacy = bincode::serialize(&registry).expect("serialize legacy registry");
432        std::fs::write(&path, legacy).expect("write legacy registry");
433
434        let loaded = PrefetchRegistry::load(&path);
435        assert!(
436            loaded.phases.is_empty() && loaded.keyed_phases.is_empty(),
437            "legacy raw bincode must be treated as a cache miss"
438        );
439    }
440
441    #[test]
442    fn test_load_missing_file_returns_default() {
443        let path = std::path::Path::new("/tmp/nonexistent_prefetch_registry.bin");
444        let registry = PrefetchRegistry::load(path);
445        assert!(registry.phases.is_empty());
446        assert!(registry.keyed_phases.is_empty());
447    }
448
449    #[test]
450    fn test_record_replaces_existing() {
451        let mut registry = PrefetchRegistry::default();
452        let addr = Address::repeat_byte(0x01);
453
454        let mut al1 = StorageAccessList::default();
455        al1.slots.insert((addr, U256::from(1)));
456        registry.record("phase", al1);
457
458        let mut al2 = StorageAccessList::default();
459        al2.slots.insert((addr, U256::from(2)));
460        registry.record("phase", al2);
461
462        let slots = registry.phase_slots("phase");
463        assert_eq!(slots.len(), 1);
464        assert!(slots.contains(&(addr, U256::from(2))));
465    }
466}
evm_fork_cache/prefetch_registry.rs

evm_fork_cache/
prefetch_registry.rs
