revm_context/journal/
warm_addresses.rs

1//! This module contains [`WarmAddresses`] struct that stores addresses that are warm loaded.
2//!
3//! It is used to optimize access to precompile addresses.
4
5use bitvec::{bitvec, vec::BitVec};
6use context_interface::journaled_state::JournalLoadError;
7use primitives::{short_address, Address, HashMap, HashSet, StorageKey, SHORT_ADDRESS_CAP};
8
9/// Stores addresses that are warm loaded. Contains precompiles and coinbase address.
10///
11/// It contains precompiles addresses that are not changed frequently and AccessList that
12/// is changed per transaction.
13///
14/// [WarmAddresses::precompiles] will always contain all precompile addresses.
15///
16/// As precompiles addresses are usually very small, precompile_short_addresses will
17/// contain bitset of shrunk precompile address.
18#[derive(Debug, Clone, PartialEq, Eq)]
19#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
20pub struct WarmAddresses {
21    /// Set of warm loaded precompile addresses.
22    precompile_set: HashSet<Address>,
23    /// Bit vector of precompile short addresses. If address is shorter than [`SHORT_ADDRESS_CAP`] it
24    /// will be stored in this bit vector for faster access.
25    precompile_short_addresses: BitVec,
26    /// `true` if all precompiles are short addresses.
27    precompile_all_short_addresses: bool,
28    /// Coinbase address.
29    coinbase: Option<Address>,
30    /// Access list
31    access_list: HashMap<Address, HashSet<StorageKey>>,
32}
33
34impl Default for WarmAddresses {
35    fn default() -> Self {
36        Self::new()
37    }
38}
39
40impl WarmAddresses {
41    /// Create a new warm addresses instance.
42    #[inline]
43    pub fn new() -> Self {
44        Self {
45            precompile_set: HashSet::default(),
46            precompile_short_addresses: bitvec![0; SHORT_ADDRESS_CAP],
47            precompile_all_short_addresses: true,
48            coinbase: None,
49            access_list: HashMap::default(),
50        }
51    }
52
53    /// Returns the precompile addresses.
54    #[inline]
55    pub fn precompiles(&self) -> &HashSet<Address> {
56        &self.precompile_set
57    }
58
59    /// Returns the coinbase address.
60    #[inline]
61    pub fn coinbase(&self) -> Option<Address> {
62        self.coinbase
63    }
64
65    /// Set the precompile addresses and short addresses.
66    #[inline]
67    pub fn set_precompile_addresses(&mut self, addresses: HashSet<Address>) {
68        self.precompile_short_addresses.fill(false);
69
70        let mut all_short_addresses = true;
71        for address in addresses.iter() {
72            if let Some(short_address) = short_address(address) {
73                self.precompile_short_addresses.set(short_address, true);
74            } else {
75                all_short_addresses = false;
76            }
77        }
78
79        self.precompile_all_short_addresses = all_short_addresses;
80        self.precompile_set = addresses;
81    }
82
83    /// Set the coinbase address.
84    #[inline]
85    pub fn set_coinbase(&mut self, address: Address) {
86        self.coinbase = Some(address);
87    }
88
89    /// Set the access list.
90    #[inline]
91    pub fn set_access_list(&mut self, access_list: HashMap<Address, HashSet<StorageKey>>) {
92        self.access_list = access_list;
93    }
94
95    /// Returns the access list.
96    #[inline]
97    pub fn access_list(&self) -> &HashMap<Address, HashSet<StorageKey>> {
98        &self.access_list
99    }
100
101    /// Clear the coinbase address.
102    #[inline]
103    pub fn clear_coinbase(&mut self) {
104        self.coinbase = None;
105    }
106
107    /// Clear the coinbase and access list.
108    #[inline]
109    pub fn clear_coinbase_and_access_list(&mut self) {
110        self.coinbase = None;
111        self.access_list.clear();
112    }
113
114    /// Returns true if the address is warm loaded.
115    #[inline]
116    pub fn is_warm(&self, address: &Address) -> bool {
117        // check if it is coinbase
118        if Some(*address) == self.coinbase {
119            return true;
120        }
121
122        // if it is part of access list.
123        if self.access_list.contains_key(address) {
124            return true;
125        }
126
127        // if there are no precompiles, it is cold loaded and bitvec is not set.
128        if self.precompile_set.is_empty() {
129            return false;
130        }
131
132        // check if it is short precompile address
133        if let Some(short_address) = short_address(address) {
134            return self.precompile_short_addresses[short_address];
135        }
136
137        if !self.precompile_all_short_addresses {
138            // in the end check if it is inside precompile set
139            return self.precompile_set.contains(address);
140        }
141
142        false
143    }
144
145    /// Returns true if the storage is warm loaded.
146    #[inline]
147    pub fn is_storage_warm(&self, address: &Address, key: &StorageKey) -> bool {
148        if let Some(access_list) = self.access_list.get(address) {
149            return access_list.contains(key);
150        }
151
152        false
153    }
154
155    /// Returns true if the address is cold loaded.
156    #[inline]
157    pub fn is_cold(&self, address: &Address) -> bool {
158        !self.is_warm(address)
159    }
160
161    /// Checks if the address is cold loaded and returns an error if it is and skip_cold_load is true.
162    #[inline(never)]
163    pub fn check_is_cold<E>(
164        &self,
165        address: &Address,
166        skip_cold_load: bool,
167    ) -> Result<bool, JournalLoadError<E>> {
168        let is_cold = self.is_cold(address);
169
170        if is_cold && skip_cold_load {
171            return Err(JournalLoadError::ColdLoadSkipped);
172        }
173
174        Ok(is_cold)
175    }
176}
177
178#[cfg(test)]
179mod tests {
180    use super::*;
181    use primitives::{address, Address};
182
183    #[test]
184    fn test_initialization() {
185        let warm_addresses = WarmAddresses::new();
186        assert!(warm_addresses.precompile_set.is_empty());
187        assert_eq!(
188            warm_addresses.precompile_short_addresses.len(),
189            SHORT_ADDRESS_CAP
190        );
191        assert!(!warm_addresses.precompile_short_addresses.any());
192        assert!(warm_addresses.coinbase.is_none());
193
194        // Test Default trait
195        let default_addresses = WarmAddresses::default();
196        assert_eq!(warm_addresses, default_addresses);
197    }
198
199    #[test]
200    fn test_coinbase_management() {
201        let mut warm_addresses = WarmAddresses::new();
202        let coinbase_addr = address!("1234567890123456789012345678901234567890");
203
204        // Test setting coinbase
205        warm_addresses.set_coinbase(coinbase_addr);
206        assert_eq!(warm_addresses.coinbase, Some(coinbase_addr));
207        assert!(warm_addresses.is_warm(&coinbase_addr));
208
209        // Test clearing coinbase
210        warm_addresses.clear_coinbase_and_access_list();
211        assert!(warm_addresses.coinbase.is_none());
212        assert!(!warm_addresses.is_warm(&coinbase_addr));
213    }
214
215    #[test]
216    fn test_short_address_precompiles() {
217        let mut warm_addresses = WarmAddresses::new();
218
219        // Create short addresses (18 leading zeros, last 2 bytes < 300)
220        let mut bytes1 = [0u8; 20];
221        bytes1[19] = 1u8;
222        let short_addr1 = Address::from(bytes1);
223
224        let mut bytes2 = [0u8; 20];
225        bytes2[19] = 5u8;
226        let short_addr2 = Address::from(bytes2);
227
228        let mut precompiles = HashSet::default();
229        precompiles.insert(short_addr1);
230        precompiles.insert(short_addr2);
231
232        warm_addresses.set_precompile_addresses(precompiles.clone());
233
234        // Verify storage
235        assert_eq!(warm_addresses.precompile_set, precompiles);
236        assert_eq!(
237            warm_addresses.precompile_short_addresses.len(),
238            SHORT_ADDRESS_CAP
239        );
240
241        // Verify bitvec optimization
242        assert!(warm_addresses.precompile_short_addresses[1]);
243        assert!(warm_addresses.precompile_short_addresses[5]);
244        assert!(!warm_addresses.precompile_short_addresses[0]);
245
246        // Verify warmth detection
247        assert!(warm_addresses.is_warm(&short_addr1));
248        assert!(warm_addresses.is_warm(&short_addr2));
249
250        // Test non-existent short address
251        let mut other_bytes = [0u8; 20];
252        other_bytes[19] = 20u8;
253        let other_short_addr = Address::from(other_bytes);
254        assert!(!warm_addresses.is_warm(&other_short_addr));
255    }
256
257    #[test]
258    fn test_regular_address_precompiles() {
259        let mut warm_addresses = WarmAddresses::new();
260
261        // Create non-short addresses
262        let regular_addr = address!("1234567890123456789012345678901234567890");
263        let mut bytes = [0u8; 20];
264        bytes[18] = 1u8;
265        bytes[19] = 44u8; // 300
266        let boundary_addr = Address::from(bytes);
267
268        let mut precompiles = HashSet::default();
269        precompiles.insert(regular_addr);
270        precompiles.insert(boundary_addr);
271
272        warm_addresses.set_precompile_addresses(precompiles.clone());
273
274        // Verify storage
275        assert_eq!(warm_addresses.precompile_set, precompiles);
276        assert!(!warm_addresses.precompile_short_addresses.any());
277
278        // Verify warmth detection
279        assert!(warm_addresses.is_warm(&regular_addr));
280        assert!(warm_addresses.is_warm(&boundary_addr));
281
282        // Test non-existent regular address
283        let other_addr = address!("0987654321098765432109876543210987654321");
284        assert!(!warm_addresses.is_warm(&other_addr));
285    }
286
287    #[test]
288    fn test_mixed_address_types() {
289        let mut warm_addresses = WarmAddresses::new();
290
291        let mut short_bytes = [0u8; 20];
292        short_bytes[19] = 7u8;
293        let short_addr = Address::from(short_bytes);
294        let regular_addr = address!("1234567890123456789012345678901234567890");
295
296        let mut precompiles = HashSet::default();
297        precompiles.insert(short_addr);
298        precompiles.insert(regular_addr);
299
300        warm_addresses.set_precompile_addresses(precompiles);
301
302        // Both types should be warm
303        assert!(warm_addresses.is_warm(&short_addr));
304        assert!(warm_addresses.is_warm(&regular_addr));
305
306        // Verify short address optimization is used
307        assert!(warm_addresses.precompile_short_addresses[7]);
308        assert!(!warm_addresses.precompile_short_addresses[8]);
309    }
310
311    #[test]
312    fn test_short_address_boundary() {
313        let mut warm_addresses = WarmAddresses::new();
314
315        // Address at boundary (SHORT_ADDRESS_CAP - 1)
316        let mut boundary_bytes = [0u8; 20];
317        let boundary_val = (SHORT_ADDRESS_CAP - 1) as u16;
318        boundary_bytes[18] = (boundary_val >> 8) as u8;
319        boundary_bytes[19] = boundary_val as u8;
320        let boundary_addr = Address::from(boundary_bytes);
321
322        let mut precompiles = HashSet::default();
323        precompiles.insert(boundary_addr);
324
325        warm_addresses.set_precompile_addresses(precompiles);
326
327        assert!(warm_addresses.is_warm(&boundary_addr));
328        assert!(warm_addresses.precompile_short_addresses[SHORT_ADDRESS_CAP - 1]);
329    }
330}