Skip to main content

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