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 anyhow::{Context as _, Result};
20use serde::{Deserialize, Serialize};
21use tracing::{debug, info, warn};
22
23use crate::StorageAccessList;
24use crate::cache::{EvmCache, versioned};
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<()> {
98        if let Some(parent) = path.parent() {
99            std::fs::create_dir_all(parent).with_context(|| {
100                format!("failed to create prefetch registry directory {parent:?}")
101            })?;
102        }
103        let data = versioned::encode(
104            PREFETCH_REGISTRY_MAGIC,
105            PREFETCH_REGISTRY_VERSION,
106            self,
107            "prefetch registry",
108        )?;
109        std::fs::write(path, data)
110            .with_context(|| format!("failed to persist prefetch registry to {path:?}"))?;
111
112        let total_slots: usize = self.phases.values().map(|al| al.slots.len()).sum::<usize>()
113            + self
114                .keyed_phases
115                .values()
116                .flat_map(|m| m.values())
117                .map(|al| al.slots.len())
118                .sum::<usize>();
119        debug!(total_slots, "Saved prefetch registry");
120        Ok(())
121    }
122
123    /// Record the aggregated access list for `phase`, **overwriting** any access
124    /// list previously recorded for that phase.
125    ///
126    /// Each call wholesale replaces the phase's slot set; it does not merge with
127    /// the prior list. To accumulate per-address lists instead, use
128    /// [`record_keyed`](Self::record_keyed).
129    ///
130    /// ```
131    /// use evm_fork_cache::prefetch_registry::PrefetchRegistry;
132    /// use evm_fork_cache::StorageAccessList;
133    /// use alloy_primitives::{Address, U256};
134    ///
135    /// let mut registry = PrefetchRegistry::default();
136    /// let addr = Address::repeat_byte(0x01);
137    ///
138    /// let mut al = StorageAccessList::default();
139    /// al.slots.insert((addr, U256::from(1)));
140    /// registry.record("pool_refresh", al);
141    /// assert!(registry.phase_slots("pool_refresh").contains(&(addr, U256::from(1))));
142    ///
143    /// // A second record replaces the slot set rather than merging.
144    /// let mut al2 = StorageAccessList::default();
145    /// al2.slots.insert((addr, U256::from(2)));
146    /// registry.record("pool_refresh", al2);
147    /// let slots = registry.phase_slots("pool_refresh");
148    /// assert_eq!(slots.len(), 1);
149    /// assert!(slots.contains(&(addr, U256::from(2))));
150    /// ```
151    pub fn record(&mut self, phase: &str, access_list: StorageAccessList) {
152        self.phases.insert(phase.to_string(), access_list);
153    }
154
155    /// Record the access list for a single `key` within a keyed `phase`.
156    ///
157    /// Unlike [`record`](Self::record), this **inserts into** the phase's per-key
158    /// nested map: other keys already recorded under `phase` are preserved, and
159    /// only the entry for `key` is replaced. Pairs with
160    /// [`prefetch_keyed`](Self::prefetch_keyed).
161    pub fn record_keyed(&mut self, phase: &str, key: Address, access_list: StorageAccessList) {
162        self.keyed_phases
163            .entry(phase.to_string())
164            .or_default()
165            .insert(key, access_list);
166    }
167
168    /// Prefetch all slots for an aggregated phase.
169    ///
170    /// Returns `(fetched, errors)`.
171    pub fn prefetch_phase(&self, phase: &str, cache: &mut EvmCache) -> (usize, usize) {
172        let Some(access_list) = self.phases.get(phase) else {
173            debug!(phase, "No prefetch data for phase");
174            return (0, 0);
175        };
176
177        if access_list.slots.is_empty() {
178            return (0, 0);
179        }
180
181        batch_prefetch(cache, access_list.slots.iter().copied(), phase)
182    }
183
184    /// Prefetch slots for specific keys within a keyed phase, excluding slots
185    /// already warm from a previous prefetch stage.
186    ///
187    /// Pairs with [`record_keyed`](Self::record_keyed).
188    ///
189    /// Returns `(fetched, errors)`.
190    pub fn prefetch_keyed(
191        &self,
192        phase: &str,
193        keys: &[Address],
194        cache: &mut EvmCache,
195        exclude: &HashSet<(Address, U256)>,
196    ) -> (usize, usize) {
197        let Some(keyed_map) = self.keyed_phases.get(phase) else {
198            debug!(phase, "No keyed prefetch data for phase");
199            return (0, 0);
200        };
201
202        let slots: HashSet<(Address, U256)> = keys
203            .iter()
204            .filter_map(|addr| keyed_map.get(addr))
205            .flat_map(|al| al.slots.iter().copied())
206            .filter(|slot| !exclude.contains(slot))
207            .collect();
208
209        if slots.is_empty() {
210            debug!(
211                phase,
212                keys = keys.len(),
213                excluded = exclude.len(),
214                "All keyed slots excluded or empty"
215            );
216            return (0, 0);
217        }
218
219        batch_prefetch(cache, slots.into_iter(), phase)
220    }
221
222    /// Returns the set of `(address, slot)` pairs recorded for an aggregated
223    /// `phase`, or an empty set if the phase was never [`record`](Self::record)ed.
224    ///
225    /// Typically used to build the `exclude` set passed to
226    /// [`prefetch_keyed`](Self::prefetch_keyed) so a later stage skips slots a
227    /// prior aggregated prefetch already warmed.
228    pub fn phase_slots(&self, phase: &str) -> HashSet<(Address, U256)> {
229        self.phases
230            .get(phase)
231            .map(|al| al.slots.clone())
232            .unwrap_or_default()
233    }
234}
235
236/// Batch-fetch `slots` into `cache` via its `storage_batch_fetcher` and inject
237/// the results, returning `(fetched, errors)`.
238///
239/// Deduplicating, exclusion, and phase lookup are the caller's responsibility
240/// ([`PrefetchRegistry::prefetch_phase`] / [`PrefetchRegistry::prefetch_keyed`]).
241/// If `slots` is empty, or the cache has no batch fetcher configured, returns
242/// `(0, 0)` without fetching. Otherwise each slot that the fetcher resolves
243/// successfully is injected into the cache and counted in `fetched`; per-slot
244/// fetch errors are counted in `errors` and skipped.
245fn batch_prefetch(
246    cache: &mut EvmCache,
247    slots: impl Iterator<Item = (Address, U256)>,
248    phase: &str,
249) -> (usize, usize) {
250    let requests: Vec<(Address, U256)> = slots.collect();
251    if requests.is_empty() {
252        return (0, 0);
253    }
254
255    let fetcher = match cache.storage_batch_fetcher().cloned() {
256        Some(f) => f,
257        None => {
258            debug!(
259                "No batch fetcher available, skipping prefetch for {}",
260                phase
261            );
262            return (0, 0);
263        }
264    };
265
266    let start = std::time::Instant::now();
267    let total_requested = requests.len();
268    // `None`: fetch at the cache's currently-pinned block (synchronous, no repin race).
269    let results = fetcher(requests, None);
270
271    let mut successes: Vec<(Address, U256, U256)> = Vec::with_capacity(results.len());
272    let mut errors = 0usize;
273    for (addr, slot, result) in results {
274        match result {
275            Ok(value) => successes.push((addr, slot, value)),
276            Err(_) => errors += 1,
277        }
278    }
279
280    let fetched = successes.len();
281    cache.inject_storage_batch(&successes);
282
283    let ms = start.elapsed().as_millis() as u64;
284    info!(
285        phase,
286        slots_requested = total_requested,
287        slots_fetched = fetched,
288        errors,
289        prefetch_ms = ms,
290        "Prefetch complete"
291    );
292
293    (fetched, errors)
294}
295
296#[cfg(test)]
297mod tests {
298    use super::*;
299
300    #[test]
301    fn test_registry_record_and_phase_slots() {
302        let mut registry = PrefetchRegistry::default();
303        let mut al = StorageAccessList::default();
304        let addr = Address::repeat_byte(0x01);
305        al.slots.insert((addr, U256::from(1)));
306        al.slots.insert((addr, U256::from(2)));
307
308        registry.record("pool_refresh", al);
309
310        let slots = registry.phase_slots("pool_refresh");
311        assert_eq!(slots.len(), 2);
312        assert!(slots.contains(&(addr, U256::from(1))));
313        assert!(slots.contains(&(addr, U256::from(2))));
314
315        // Non-existent phase returns empty
316        assert!(registry.phase_slots("nonexistent").is_empty());
317    }
318
319    #[test]
320    fn test_registry_record_keyed() {
321        let mut registry = PrefetchRegistry::default();
322        let key_a = Address::repeat_byte(0x01);
323        let key_b = Address::repeat_byte(0x02);
324
325        let mut al1 = StorageAccessList::default();
326        al1.slots.insert((key_a, U256::from(10)));
327
328        let mut al2 = StorageAccessList::default();
329        al2.slots.insert((key_b, U256::from(20)));
330
331        registry.record_keyed("per_target", key_a, al1);
332        registry.record_keyed("per_target", key_b, al2);
333
334        // Verify keyed_phases has both
335        let map = registry.keyed_phases.get("per_target").unwrap();
336        assert_eq!(map.len(), 2);
337        assert!(
338            map.get(&key_a)
339                .unwrap()
340                .slots
341                .contains(&(key_a, U256::from(10)))
342        );
343    }
344
345    #[test]
346    fn test_save_load_round_trip() {
347        let dir = std::env::temp_dir().join("evm_fork_cache_test_prefetch_registry");
348        let path = dir.join("test_registry.bin");
349        let _ = std::fs::remove_file(&path);
350
351        let mut registry = PrefetchRegistry::default();
352        let addr = Address::repeat_byte(0xAA);
353        let mut al = StorageAccessList::default();
354        al.slots.insert((addr, U256::from(42)));
355        al.accounts.insert(addr);
356        registry.record("test_phase", al);
357
358        let key = Address::repeat_byte(0xBB);
359        let mut sal = StorageAccessList::default();
360        sal.slots.insert((key, U256::from(99)));
361        registry.record_keyed("per_target", key, sal);
362
363        registry.save(&path).expect("save registry");
364        let data = std::fs::read(&path).expect("read saved registry");
365        assert!(
366            data.starts_with(b"EFC-PFRG"),
367            "prefetch registry files must carry a magic/version header"
368        );
369
370        let loaded = PrefetchRegistry::load(&path);
371        assert_eq!(loaded.phases.len(), 1);
372        assert!(
373            loaded.phases["test_phase"]
374                .slots
375                .contains(&(addr, U256::from(42)))
376        );
377        assert_eq!(loaded.keyed_phases.len(), 1);
378        assert!(
379            loaded.keyed_phases["per_target"]
380                .get(&key)
381                .unwrap()
382                .slots
383                .contains(&(key, U256::from(99)))
384        );
385
386        let _ = std::fs::remove_file(&path);
387        let _ = std::fs::remove_dir(&dir);
388    }
389
390    #[test]
391    fn save_reports_write_failures() {
392        let dir = std::env::temp_dir().join("evm_fork_cache_test_prefetch_registry_write_error");
393        let _ = std::fs::remove_dir_all(&dir);
394        let _ = std::fs::remove_file(&dir);
395        std::fs::write(&dir, b"not a directory").expect("create file path conflict");
396
397        let registry = PrefetchRegistry::default();
398        let path = dir.join("registry.bin");
399        let err = registry
400            .save(&path)
401            .expect_err("save must report write failure");
402        assert!(
403            err.to_string().contains("directory") || err.to_string().contains("Not a directory"),
404            "unexpected error: {err:#}"
405        );
406
407        let _ = std::fs::remove_file(&dir);
408    }
409
410    #[test]
411    fn legacy_raw_bincode_loads_as_default() {
412        let dir = std::env::temp_dir().join("evm_fork_cache_test_prefetch_registry_legacy");
413        let path = dir.join("legacy_registry.bin");
414        let _ = std::fs::remove_dir_all(&dir);
415        std::fs::create_dir_all(&dir).expect("create temp dir");
416
417        let mut registry = PrefetchRegistry::default();
418        let addr = Address::repeat_byte(0xAA);
419        let mut al = StorageAccessList::default();
420        al.slots.insert((addr, U256::from(42)));
421        registry.record("legacy_phase", al);
422        let legacy = bincode::serialize(&registry).expect("serialize legacy registry");
423        std::fs::write(&path, legacy).expect("write legacy registry");
424
425        let loaded = PrefetchRegistry::load(&path);
426        assert!(
427            loaded.phases.is_empty() && loaded.keyed_phases.is_empty(),
428            "legacy raw bincode must be treated as a cache miss"
429        );
430
431        let _ = std::fs::remove_dir_all(&dir);
432    }
433
434    #[test]
435    fn test_load_missing_file_returns_default() {
436        let path = std::path::Path::new("/tmp/nonexistent_prefetch_registry.bin");
437        let registry = PrefetchRegistry::load(path);
438        assert!(registry.phases.is_empty());
439        assert!(registry.keyed_phases.is_empty());
440    }
441
442    #[test]
443    fn test_record_replaces_existing() {
444        let mut registry = PrefetchRegistry::default();
445        let addr = Address::repeat_byte(0x01);
446
447        let mut al1 = StorageAccessList::default();
448        al1.slots.insert((addr, U256::from(1)));
449        registry.record("phase", al1);
450
451        let mut al2 = StorageAccessList::default();
452        al2.slots.insert((addr, U256::from(2)));
453        registry.record("phase", al2);
454
455        let slots = registry.phase_slots("phase");
456        assert_eq!(slots.len(), 1);
457        assert!(slots.contains(&(addr, U256::from(2))));
458    }
459}
evm_fork_cache/prefetch_registry.rs

evm_fork_cache/
prefetch_registry.rs
