arb_evm/
state_overlay.rs

1//! Snapshots `(info, status)` before each helper cache mutation and at
2//! end of tx pushes one explicit `TransitionAccount` per real change
3//! through `State::apply_transition`.
4
5use std::{cell::RefCell, collections::HashMap};
6
7use alloy_primitives::Address;
8use revm::{database::State, Database};
9use revm_database::{AccountStatus as CacheAccountStatus, TransitionAccount};
10use revm_state::AccountInfo;
11
12#[derive(Clone, Debug)]
13struct Entry {
14    previous_info: Option<AccountInfo>,
15    previous_status: CacheAccountStatus,
16}
17
18thread_local! {
19    static OVERLAY: RefCell<HashMap<Address, Entry>> = RefCell::new(HashMap::new());
20}
21
22pub fn reset_tx() {
23    OVERLAY.with(|o| o.borrow_mut().clear());
24}
25
26/// Snapshot pre-mutation `(info, status)` for `addr`; idempotent per tx.
27pub fn record_pre_touch<DB: Database>(state: &mut State<DB>, addr: Address) {
28    let already = OVERLAY.with(|o| o.borrow().contains_key(&addr));
29    if already {
30        return;
31    }
32    let _ = state.load_cache_account(addr);
33    let cache_entry = state.cache.accounts.get(&addr);
34    let previous_info = cache_entry
35        .and_then(|c| c.account.as_ref())
36        .map(|a| a.info.clone());
37    let previous_status = cache_entry
38        .map(|c| c.status)
39        .unwrap_or(CacheAccountStatus::LoadedNotExisting);
40    OVERLAY.with(|o| {
41        o.borrow_mut().insert(
42            addr,
43            Entry {
44                previous_info,
45                previous_status,
46            },
47        )
48    });
49}
50
51/// Push one `TransitionAccount` per address whose info actually changed.
52pub fn drain_and_apply<DB: Database>(state: &mut State<DB>) {
53    let entries: Vec<(Address, Entry)> = OVERLAY.with(|o| o.borrow_mut().drain().collect());
54    if entries.is_empty() {
55        return;
56    }
57
58    let mut transitions = Vec::with_capacity(entries.len());
59    for (addr, entry) in entries {
60        let current_info = state
61            .cache
62            .accounts
63            .get(&addr)
64            .and_then(|c| c.account.as_ref())
65            .map(|a| a.info.clone());
66
67        if current_info == entry.previous_info {
68            continue;
69        }
70
71        let pre_empty = entry
72            .previous_info
73            .as_ref()
74            .map(|i| i.is_empty())
75            .unwrap_or(true);
76        let cur_empty = current_info.as_ref().map(|i| i.is_empty()).unwrap_or(true);
77
78        let new_status = match (pre_empty, cur_empty) {
79            (true, true) => continue,
80            (true, false) => match entry.previous_status {
81                CacheAccountStatus::Destroyed
82                | CacheAccountStatus::DestroyedAgain
83                | CacheAccountStatus::DestroyedChanged => CacheAccountStatus::DestroyedChanged,
84                _ => CacheAccountStatus::InMemoryChange,
85            },
86            (false, true) => match entry.previous_status {
87                CacheAccountStatus::LoadedNotExisting => continue,
88                CacheAccountStatus::DestroyedAgain | CacheAccountStatus::DestroyedChanged => {
89                    CacheAccountStatus::DestroyedAgain
90                }
91                _ => CacheAccountStatus::Destroyed,
92            },
93            (false, false) => match entry.previous_status {
94                CacheAccountStatus::Loaded => CacheAccountStatus::Changed,
95                CacheAccountStatus::LoadedNotExisting | CacheAccountStatus::LoadedEmptyEIP161 => {
96                    CacheAccountStatus::InMemoryChange
97                }
98                CacheAccountStatus::DestroyedAgain
99                | CacheAccountStatus::Destroyed
100                | CacheAccountStatus::DestroyedChanged => CacheAccountStatus::DestroyedChanged,
101                other => other,
102            },
103        };
104
105        let goes_destroyed = matches!(
106            new_status,
107            CacheAccountStatus::Destroyed | CacheAccountStatus::DestroyedAgain
108        );
109        let transition_info = if goes_destroyed { None } else { current_info };
110        let storage_was_destroyed = goes_destroyed && !pre_empty;
111
112        if let Some(cached) = state.cache.accounts.get_mut(&addr) {
113            cached.status = new_status;
114            if goes_destroyed {
115                cached.account = None;
116            }
117        }
118
119        transitions.push((
120            addr,
121            TransitionAccount {
122                info: transition_info,
123                status: new_status,
124                previous_info: entry.previous_info,
125                previous_status: entry.previous_status,
126                storage: Default::default(),
127                storage_was_destroyed,
128            },
129        ));
130    }
131
132    if !transitions.is_empty() {
133        state.apply_transition(transitions);
134    }
135}