Skip to main content

anvil/eth/backend/mem/
storage.rs

1//! In-memory blockchain storage
2use crate::eth::{
3    backend::{
4        db::{
5            MaybeFullDatabase, SerializableBlock, SerializableHistoricalStates,
6            SerializableTransaction, StateDb,
7        },
8        mem::cache::DiskStateCache,
9    },
10    pool::transactions::PoolTransaction,
11};
12use alloy_consensus::{BlockHeader, Header, constants::EMPTY_WITHDRAWALS};
13use alloy_eips::eip7685::EMPTY_REQUESTS_HASH;
14use alloy_evm::EvmEnv;
15use alloy_network::Network;
16use alloy_primitives::{
17    B256, Bytes, U256,
18    map::{B256HashMap, HashMap},
19};
20use alloy_rpc_types::{
21    BlockId, BlockNumberOrTag, TransactionInfo as RethTransactionInfo,
22    trace::{
23        otterscan::{InternalOperation, OperationType},
24        parity::LocalizedTransactionTrace,
25    },
26};
27use anvil_core::eth::{
28    block::{Block, create_block},
29    transaction::{MaybeImpersonatedTransaction, TransactionInfo},
30};
31use foundry_evm::{
32    backend::MemDb,
33    traces::{CallKind, ParityTraceBuilder, TracingInspectorConfig},
34};
35#[cfg(test)]
36use foundry_primitives::FoundryNetwork;
37use foundry_primitives::{FoundryReceiptEnvelope, FoundryTxEnvelope};
38use parking_lot::RwLock;
39use revm::{context::Block as RevmBlock, primitives::hardfork::SpecId};
40use std::{collections::VecDeque, fmt, path::PathBuf, sync::Arc, time::Duration};
41// use yansi::Paint;
42
43// === various limits in number of blocks ===
44
45pub const DEFAULT_HISTORY_LIMIT: usize = 500;
46const MIN_HISTORY_LIMIT: usize = 10;
47// 1hr of up-time at lowest 1s interval
48const MAX_ON_DISK_HISTORY_LIMIT: usize = 3_600;
49
50/// Represents the complete state of single block
51pub struct InMemoryBlockStates {
52    /// The states at a certain block
53    states: B256HashMap<StateDb>,
54    /// states which data is moved to disk
55    on_disk_states: B256HashMap<StateDb>,
56    /// How many states to store at most
57    in_memory_limit: usize,
58    /// minimum amount of states we keep in memory
59    min_in_memory_limit: usize,
60    /// maximum amount of states we keep on disk
61    ///
62    /// Limiting the states will prevent disk blow up, especially in interval mining mode
63    max_on_disk_limit: usize,
64    /// the oldest states written to disk
65    oldest_on_disk: VecDeque<B256>,
66    /// all states present, used to enforce `in_memory_limit`
67    present: VecDeque<B256>,
68    /// Stores old states on disk
69    disk_cache: DiskStateCache,
70}
71
72impl InMemoryBlockStates {
73    /// Creates a new instance with limited slots
74    pub fn new(in_memory_limit: usize, on_disk_limit: usize) -> Self {
75        let in_memory_limit = in_memory_limit.max(1);
76        Self {
77            states: Default::default(),
78            on_disk_states: Default::default(),
79            in_memory_limit,
80            min_in_memory_limit: in_memory_limit.min(MIN_HISTORY_LIMIT),
81            max_on_disk_limit: on_disk_limit,
82            oldest_on_disk: Default::default(),
83            present: Default::default(),
84            disk_cache: Default::default(),
85        }
86    }
87
88    /// Configures no disk caching
89    pub const fn memory_only(mut self) -> Self {
90        self.max_on_disk_limit = 0;
91        self
92    }
93
94    /// Configures the path on disk where the states will cached.
95    pub fn disk_path(mut self, path: PathBuf) -> Self {
96        self.disk_cache = self.disk_cache.with_path(path);
97        self
98    }
99
100    /// This modifies the `limit` what to keep stored in memory.
101    ///
102    /// This will ensure the new limit adjusts based on the block time.
103    /// The lowest blocktime is 1s which should increase the limit slightly
104    pub fn update_interval_mine_block_time(&mut self, block_time: Duration) {
105        let block_time = block_time.as_secs();
106        // for block times lower than 2s we increase the mem limit since we're mining _small_ blocks
107        // very fast
108        // this will gradually be decreased once the max limit was reached
109        if block_time <= 2 {
110            self.in_memory_limit = DEFAULT_HISTORY_LIMIT * 3;
111            self.enforce_limits();
112        }
113    }
114
115    /// Returns true if only memory caching is supported.
116    const fn is_memory_only(&self) -> bool {
117        self.max_on_disk_limit == 0
118    }
119
120    /// Inserts a new (hash -> state) pair
121    ///
122    /// When the configured limit for the number of states that can be stored in memory is reached,
123    /// the oldest state is removed.
124    ///
125    /// Since we keep a snapshot of the entire state as history, the size of the state will increase
126    /// with the transactions processed. To counter this, we gradually decrease the cache limit with
127    /// the number of states/blocks until we reached the `min_limit`.
128    ///
129    /// When a state that was previously written to disk is requested, it is simply read from disk.
130    pub fn insert(&mut self, hash: B256, state: StateDb) {
131        if !self.is_memory_only() && self.present.len() >= self.in_memory_limit {
132            // once we hit the max limit we gradually decrease it
133            self.in_memory_limit =
134                self.in_memory_limit.saturating_sub(1).max(self.min_in_memory_limit);
135        }
136
137        self.enforce_limits();
138
139        self.states.insert(hash, state);
140        self.present.push_back(hash);
141    }
142
143    /// Enforces configured limits
144    fn enforce_limits(&mut self) {
145        // enforce memory limits
146        while self.present.len() >= self.in_memory_limit {
147            // evict the oldest block
148            if let Some((hash, mut state)) = self
149                .present
150                .pop_front()
151                .and_then(|hash| self.states.remove(&hash).map(|state| (hash, state)))
152            {
153                // only write to disk if supported
154                if !self.is_memory_only() {
155                    let state_snapshot = state.0.clear_into_state_snapshot();
156                    if self.disk_cache.write(hash, &state_snapshot) {
157                        // Write succeeded, move state to on-disk tracking
158                        self.on_disk_states.insert(hash, state);
159                        self.oldest_on_disk.push_back(hash);
160                    } else {
161                        // Write failed, restore state to memory to avoid data loss
162                        state.init_from_state_snapshot(state_snapshot);
163                        self.states.insert(hash, state);
164                        self.present.push_front(hash);
165                        // Increase limit temporarily to prevent infinite retry loop
166                        self.in_memory_limit = self.in_memory_limit.saturating_add(1);
167                        break;
168                    }
169                }
170            }
171        }
172
173        // enforce on disk limit and purge the oldest state cached on disk
174        while !self.is_memory_only() && self.oldest_on_disk.len() >= self.max_on_disk_limit {
175            // evict the oldest block
176            if let Some(hash) = self.oldest_on_disk.pop_front() {
177                self.on_disk_states.remove(&hash);
178                self.disk_cache.remove(hash);
179            }
180        }
181    }
182
183    /// Returns the in-memory state for the given `hash` if present
184    pub fn get_state(&self, hash: &B256) -> Option<&StateDb> {
185        self.states.get(hash)
186    }
187
188    /// Returns on-disk state for the given `hash` if present
189    pub fn get_on_disk_state(&mut self, hash: &B256) -> Option<&StateDb> {
190        if let Some(state) = self.on_disk_states.get_mut(hash)
191            && let Some(cached) = self.disk_cache.read(*hash)
192        {
193            state.init_from_state_snapshot(cached);
194            return Some(state);
195        }
196
197        None
198    }
199
200    /// Sets the maximum number of stats we keep in memory
201    pub const fn set_cache_limit(&mut self, limit: usize) {
202        let limit = if limit == 0 { 1 } else { limit };
203        self.in_memory_limit = limit;
204        self.min_in_memory_limit =
205            if limit < MIN_HISTORY_LIMIT { limit } else { MIN_HISTORY_LIMIT };
206    }
207
208    /// Clears all entries
209    pub fn clear(&mut self) {
210        self.states.clear();
211        self.on_disk_states.clear();
212        self.present.clear();
213        for on_disk in std::mem::take(&mut self.oldest_on_disk) {
214            self.disk_cache.remove(on_disk)
215        }
216    }
217
218    /// Removes states for the given block hashes.
219    ///
220    /// This is used during chain rollback to clean up states for blocks that are no longer part
221    /// of the canonical chain.
222    pub fn remove_block_states(&mut self, hashes: &[B256]) {
223        for hash in hashes {
224            self.states.remove(hash);
225            self.on_disk_states.remove(hash);
226            self.disk_cache.remove(*hash);
227        }
228        self.present.retain(|h| !hashes.contains(h));
229        self.oldest_on_disk.retain(|h| !hashes.contains(h));
230    }
231
232    /// Serialize all states to a list of serializable historical states
233    pub fn serialized_states(&mut self) -> SerializableHistoricalStates {
234        // Get in-memory states
235        let mut states = self
236            .states
237            .iter_mut()
238            .map(|(hash, state)| (*hash, state.serialize_state()))
239            .collect::<Vec<_>>();
240
241        // Get on-disk state snapshots
242        for hash in self.on_disk_states.keys() {
243            if let Some(state_snapshot) = self.disk_cache.read(*hash) {
244                states.push((*hash, state_snapshot));
245            }
246        }
247
248        SerializableHistoricalStates::new(states)
249    }
250
251    /// Load states from serialized data
252    pub fn load_states(&mut self, states: SerializableHistoricalStates) {
253        for (hash, state_snapshot) in states {
254            let mut state_db = StateDb::new(MemDb::default());
255            state_db.init_from_state_snapshot(state_snapshot);
256            self.insert(hash, state_db);
257        }
258    }
259}
260
261impl fmt::Debug for InMemoryBlockStates {
262    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
263        f.debug_struct("InMemoryBlockStates")
264            .field("in_memory_limit", &self.in_memory_limit)
265            .field("min_in_memory_limit", &self.min_in_memory_limit)
266            .field("max_on_disk_limit", &self.max_on_disk_limit)
267            .field("oldest_on_disk", &self.oldest_on_disk)
268            .field("present", &self.present)
269            .finish_non_exhaustive()
270    }
271}
272
273impl Default for InMemoryBlockStates {
274    fn default() -> Self {
275        // enough in memory to store `DEFAULT_HISTORY_LIMIT` blocks in memory
276        Self::new(DEFAULT_HISTORY_LIMIT, MAX_ON_DISK_HISTORY_LIMIT)
277    }
278}
279
280/// Stores the blockchain data (blocks, transactions)
281#[derive(Clone, Debug)]
282pub struct BlockchainStorage<N: Network> {
283    /// all stored blocks (block hash -> block)
284    pub blocks: B256HashMap<Block>,
285    /// mapping from block number -> block hash
286    pub hashes: HashMap<u64, B256>,
287    /// The current best hash
288    pub best_hash: B256,
289    /// The current best block number
290    pub best_number: u64,
291    /// genesis hash of the chain
292    pub genesis_hash: B256,
293    /// Mapping from the transaction hash to a tuple containing the transaction as well as the
294    /// transaction receipt
295    pub transactions: B256HashMap<MinedTransaction<N>>,
296    /// The total difficulty of the chain until this block
297    pub total_difficulty: U256,
298}
299
300impl<N: Network> BlockchainStorage<N> {
301    /// Creates a new storage with a genesis block
302    pub fn new(
303        evm_env: &EvmEnv,
304        base_fee: Option<u64>,
305        timestamp: u64,
306        genesis_number: u64,
307    ) -> Self {
308        let is_shanghai = *evm_env.spec_id() >= SpecId::SHANGHAI;
309        let is_cancun = *evm_env.spec_id() >= SpecId::CANCUN;
310        let is_prague = *evm_env.spec_id() >= SpecId::PRAGUE;
311
312        // create a dummy genesis block
313        let header = Header {
314            timestamp,
315            base_fee_per_gas: base_fee,
316            gas_limit: evm_env.block_env.gas_limit,
317            beneficiary: evm_env.block_env.beneficiary,
318            difficulty: evm_env.block_env.difficulty,
319            blob_gas_used: evm_env.block_env.blob_excess_gas_and_price.as_ref().map(|_| 0),
320            excess_blob_gas: evm_env.block_env.blob_excess_gas(),
321            number: genesis_number,
322            parent_beacon_block_root: is_cancun.then_some(Default::default()),
323            withdrawals_root: is_shanghai.then_some(EMPTY_WITHDRAWALS),
324            requests_hash: is_prague.then_some(EMPTY_REQUESTS_HASH),
325            ..Default::default()
326        };
327        let block =
328            create_block(header, Vec::<MaybeImpersonatedTransaction<FoundryTxEnvelope>>::new());
329        let genesis_hash = block.header.hash_slow();
330        let best_hash = genesis_hash;
331        let best_number = genesis_number;
332
333        let mut blocks = B256HashMap::default();
334        blocks.insert(genesis_hash, block);
335
336        let mut hashes = HashMap::default();
337        hashes.insert(best_number, genesis_hash);
338        Self {
339            blocks,
340            hashes,
341            best_hash,
342            best_number,
343            genesis_hash,
344            transactions: Default::default(),
345            total_difficulty: Default::default(),
346        }
347    }
348
349    pub fn forked(block_number: u64, block_hash: B256, total_difficulty: U256) -> Self {
350        let mut hashes = HashMap::default();
351        hashes.insert(block_number, block_hash);
352
353        Self {
354            blocks: B256HashMap::default(),
355            hashes,
356            best_hash: block_hash,
357            best_number: block_number,
358            genesis_hash: Default::default(),
359            transactions: Default::default(),
360            total_difficulty,
361        }
362    }
363
364    /// Unwind the chain state back to the given block in storage.
365    ///
366    /// The block identified by `block_number` and `block_hash` is __non-inclusive__, i.e. it will
367    /// remain in the state.
368    pub fn unwind_to(&mut self, block_number: u64, block_hash: B256) -> Vec<Block> {
369        let mut removed = vec![];
370        let best_num: u64 = self.best_number;
371        for i in (block_number + 1)..=best_num {
372            if let Some(hash) = self.hashes.get(&i).copied() {
373                // First remove the block's transactions while the mappings still exist
374                self.remove_block_transactions_by_number(i);
375
376                // Now remove the block from storage (may already be empty of txs) and drop mapping
377                if let Some(block) = self.blocks.remove(&hash) {
378                    removed.push(block);
379                }
380                self.hashes.remove(&i);
381            }
382        }
383        self.best_hash = block_hash;
384        self.best_number = block_number;
385        removed
386    }
387
388    pub fn empty() -> Self {
389        Self {
390            blocks: Default::default(),
391            hashes: Default::default(),
392            best_hash: Default::default(),
393            best_number: Default::default(),
394            genesis_hash: Default::default(),
395            transactions: Default::default(),
396            total_difficulty: Default::default(),
397        }
398    }
399
400    /// Removes all stored transactions for the given block number
401    pub fn remove_block_transactions_by_number(&mut self, num: u64) {
402        if let Some(hash) = self.hashes.get(&num).copied() {
403            self.remove_block_transactions(hash);
404        }
405    }
406
407    /// Removes all stored transactions for the given block hash
408    pub fn remove_block_transactions(&mut self, block_hash: B256) {
409        if let Some(block) = self.blocks.get_mut(&block_hash) {
410            for tx in &block.body.transactions {
411                self.transactions.remove(&tx.hash());
412            }
413            block.body.transactions.clear();
414        }
415    }
416
417    /// Serialize all blocks in storage
418    pub fn serialized_blocks(&self) -> Vec<SerializableBlock> {
419        self.blocks.values().map(|block| block.clone().into()).collect()
420    }
421
422    /// Deserialize and add all blocks data to the backend storage
423    pub fn load_blocks(&mut self, serializable_blocks: Vec<SerializableBlock>) {
424        for serializable_block in &serializable_blocks {
425            let block: Block = serializable_block.clone().into();
426            let block_hash = block.header.hash_slow();
427            let block_number = block.header.number();
428            self.blocks.insert(block_hash, block);
429            self.hashes.insert(block_number, block_hash);
430
431            // Update genesis_hash if we are loading block 0, so that Finalized/Safe/Earliest
432            // block tag lookups return the correct hash.
433            // See: https://github.com/foundry-rs/foundry/issues/12645
434            if block_number == 0 {
435                self.genesis_hash = block_hash;
436            }
437        }
438    }
439
440    /// Returns the hash for [BlockNumberOrTag]
441    pub fn hash(&self, number: BlockNumberOrTag, slots_in_an_epoch: u64) -> Option<B256> {
442        match number {
443            BlockNumberOrTag::Latest => Some(self.best_hash),
444            BlockNumberOrTag::Earliest => Some(self.genesis_hash),
445            BlockNumberOrTag::Pending => None,
446            BlockNumberOrTag::Number(num) => self.hashes.get(&num).copied(),
447            BlockNumberOrTag::Safe => {
448                if self.best_number > slots_in_an_epoch {
449                    self.hashes.get(&(self.best_number - slots_in_an_epoch)).copied()
450                } else {
451                    Some(self.genesis_hash)
452                }
453            }
454            BlockNumberOrTag::Finalized => {
455                if self.best_number > slots_in_an_epoch * 2 {
456                    self.hashes.get(&(self.best_number - slots_in_an_epoch * 2)).copied()
457                } else {
458                    Some(self.genesis_hash)
459                }
460            }
461        }
462    }
463}
464
465impl<N: Network<ReceiptEnvelope = FoundryReceiptEnvelope>> BlockchainStorage<N> {
466    pub fn serialized_transactions(&self) -> Vec<SerializableTransaction> {
467        self.transactions.values().map(|tx: &MinedTransaction<N>| tx.clone().into()).collect()
468    }
469
470    /// Deserialize and add all transactions data to the backend storage
471    pub fn load_transactions(&mut self, serializable_transactions: Vec<SerializableTransaction>) {
472        for serializable_transaction in &serializable_transactions {
473            let transaction: MinedTransaction<N> = serializable_transaction.clone().into();
474            self.transactions.insert(transaction.info.transaction_hash, transaction);
475        }
476    }
477}
478
479/// A simple in-memory blockchain
480#[derive(Clone, Debug)]
481pub struct Blockchain<N: Network> {
482    /// underlying storage that supports concurrent reads
483    pub storage: Arc<RwLock<BlockchainStorage<N>>>,
484}
485
486impl<N: Network> Blockchain<N> {
487    /// Creates a new storage with a genesis block
488    pub fn new(
489        evm_env: &EvmEnv,
490        base_fee: Option<u64>,
491        timestamp: u64,
492        genesis_number: u64,
493    ) -> Self {
494        Self {
495            storage: Arc::new(RwLock::new(BlockchainStorage::new(
496                evm_env,
497                base_fee,
498                timestamp,
499                genesis_number,
500            ))),
501        }
502    }
503
504    pub fn forked(block_number: u64, block_hash: B256, total_difficulty: U256) -> Self {
505        Self {
506            storage: Arc::new(RwLock::new(BlockchainStorage::forked(
507                block_number,
508                block_hash,
509                total_difficulty,
510            ))),
511        }
512    }
513
514    /// returns the header hash of given block
515    pub fn hash(&self, id: BlockId, slots_in_an_epoch: u64) -> Option<B256> {
516        match id {
517            BlockId::Hash(h) => Some(h.block_hash),
518            BlockId::Number(num) => self.storage.read().hash(num, slots_in_an_epoch),
519        }
520    }
521
522    pub fn get_block_by_hash(&self, hash: &B256) -> Option<Block> {
523        self.storage.read().blocks.get(hash).cloned()
524    }
525
526    pub fn get_transaction_by_hash(&self, hash: &B256) -> Option<MinedTransaction<N>> {
527        self.storage.read().transactions.get(hash).cloned()
528    }
529
530    /// Returns the total number of blocks
531    pub fn blocks_count(&self) -> usize {
532        self.storage.read().blocks.len()
533    }
534}
535
536/// Represents the outcome of mining a new block
537pub struct MinedBlockOutcome<T> {
538    /// The block that was mined
539    pub block_number: u64,
540    /// All transactions included in the block
541    pub included: Vec<Arc<PoolTransaction<T>>>,
542    /// All transactions that were attempted to be included but were invalid at the time of
543    /// execution
544    pub invalid: Vec<Arc<PoolTransaction<T>>>,
545    /// Transactions skipped because they're not yet valid (e.g., valid_after in the future).
546    /// These remain in the pool and should be retried later.
547    pub not_yet_valid: Vec<Arc<PoolTransaction<T>>>,
548}
549
550impl<T> Clone for MinedBlockOutcome<T> {
551    fn clone(&self) -> Self {
552        Self {
553            block_number: self.block_number,
554            included: self.included.clone(),
555            invalid: self.invalid.clone(),
556            not_yet_valid: self.not_yet_valid.clone(),
557        }
558    }
559}
560
561impl<T> fmt::Debug for MinedBlockOutcome<T> {
562    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
563        f.debug_struct("MinedBlockOutcome")
564            .field("block_number", &self.block_number)
565            .field("included", &self.included.len())
566            .field("invalid", &self.invalid.len())
567            .field("not_yet_valid", &self.not_yet_valid.len())
568            .finish()
569    }
570}
571
572/// Container type for a mined transaction
573#[derive(Clone, Debug)]
574pub struct MinedTransaction<N: Network> {
575    pub info: TransactionInfo,
576    pub receipt: N::ReceiptEnvelope,
577    pub block_hash: B256,
578    pub block_number: u64,
579}
580
581impl<N: Network> MinedTransaction<N> {
582    /// Returns the traces of the transaction for `trace_transaction`
583    pub fn parity_traces(&self) -> Vec<LocalizedTransactionTrace> {
584        ParityTraceBuilder::new(
585            self.info.traces.clone(),
586            None,
587            TracingInspectorConfig::default_parity(),
588        )
589        .into_localized_transaction_traces(RethTransactionInfo {
590            hash: Some(self.info.transaction_hash),
591            index: Some(self.info.transaction_index),
592            block_hash: Some(self.block_hash),
593            block_number: Some(self.block_number),
594            base_fee: None,
595            block_timestamp: None,
596        })
597    }
598
599    pub fn ots_internal_operations(&self) -> Vec<InternalOperation> {
600        self.info
601            .traces
602            .iter()
603            .filter_map(|node| {
604                let r#type = match node.trace.kind {
605                    _ if node.is_selfdestruct() => OperationType::OpSelfDestruct,
606                    CallKind::Call if !node.trace.value.is_zero() => OperationType::OpTransfer,
607                    CallKind::Create => OperationType::OpCreate,
608                    CallKind::Create2 => OperationType::OpCreate2,
609                    _ => return None,
610                };
611                let (from, to, value) = if node.is_selfdestruct() {
612                    (
613                        node.trace.address,
614                        node.trace.selfdestruct_refund_target.unwrap_or_default(),
615                        node.trace.selfdestruct_transferred_value.unwrap_or_default(),
616                    )
617                } else {
618                    (node.trace.caller, node.trace.address, node.trace.value)
619                };
620                Some(InternalOperation { r#type, from, to, value })
621            })
622            .collect()
623    }
624}
625
626/// Intermediary Anvil representation of a receipt
627#[derive(Clone, Debug)]
628pub struct MinedTransactionReceipt<N: Network> {
629    /// The actual json rpc receipt object
630    pub inner: N::ReceiptResponse,
631    /// Output data for the transaction
632    pub out: Option<Bytes>,
633}
634
635#[cfg(test)]
636mod tests {
637    use super::*;
638    use crate::eth::backend::db::Db;
639    use alloy_primitives::{Address, hex};
640    use alloy_rlp::Decodable;
641    use revm::{database::DatabaseRef, state::AccountInfo};
642
643    #[test]
644    fn test_interval_update() {
645        let mut storage = InMemoryBlockStates::default();
646        storage.update_interval_mine_block_time(Duration::from_secs(1));
647        assert_eq!(storage.in_memory_limit, DEFAULT_HISTORY_LIMIT * 3);
648    }
649
650    #[test]
651    fn test_init_state_limits() {
652        let mut storage = InMemoryBlockStates::default();
653        assert_eq!(storage.in_memory_limit, DEFAULT_HISTORY_LIMIT);
654        assert_eq!(storage.min_in_memory_limit, MIN_HISTORY_LIMIT);
655        assert_eq!(storage.max_on_disk_limit, MAX_ON_DISK_HISTORY_LIMIT);
656
657        storage = storage.memory_only();
658        assert!(storage.is_memory_only());
659
660        storage = InMemoryBlockStates::new(1, 0);
661        assert!(storage.is_memory_only());
662        assert_eq!(storage.in_memory_limit, 1);
663        assert_eq!(storage.min_in_memory_limit, 1);
664        assert_eq!(storage.max_on_disk_limit, 0);
665
666        storage = InMemoryBlockStates::new(1, 2);
667        assert!(!storage.is_memory_only());
668        assert_eq!(storage.in_memory_limit, 1);
669        assert_eq!(storage.min_in_memory_limit, 1);
670        assert_eq!(storage.max_on_disk_limit, 2);
671
672        storage = InMemoryBlockStates::new(0, 0);
673        assert!(storage.is_memory_only());
674        assert_eq!(storage.in_memory_limit, 1);
675        assert_eq!(storage.min_in_memory_limit, 1);
676        assert_eq!(storage.max_on_disk_limit, 0);
677
678        storage.set_cache_limit(0);
679        assert_eq!(storage.in_memory_limit, 1);
680        assert_eq!(storage.min_in_memory_limit, 1);
681    }
682
683    #[tokio::test(flavor = "multi_thread")]
684    async fn can_read_write_cached_state() {
685        let mut storage = InMemoryBlockStates::new(1, MAX_ON_DISK_HISTORY_LIMIT);
686        let one = B256::from(U256::from(1));
687        let two = B256::from(U256::from(2));
688
689        let mut state = MemDb::default();
690        let addr = Address::random();
691        let info = AccountInfo::from_balance(U256::from(1337));
692        state.insert_account(addr, info);
693        storage.insert(one, StateDb::new(state));
694        storage.insert(two, StateDb::new(MemDb::default()));
695
696        // wait for files to be flushed
697        tokio::time::sleep(std::time::Duration::from_secs(1)).await;
698
699        assert_eq!(storage.on_disk_states.len(), 1);
700        assert!(storage.on_disk_states.contains_key(&one));
701
702        let loaded = storage.get_on_disk_state(&one).unwrap();
703
704        let acc = loaded.basic_ref(addr).unwrap().unwrap();
705        assert_eq!(acc.balance, U256::from(1337u64));
706    }
707
708    #[tokio::test(flavor = "multi_thread")]
709    async fn can_decrease_state_cache_size() {
710        let limit = 15;
711        let mut storage = InMemoryBlockStates::new(limit, MAX_ON_DISK_HISTORY_LIMIT);
712
713        let num_states = 30;
714        for idx in 0..num_states {
715            let mut state = MemDb::default();
716            let hash = B256::from(U256::from(idx));
717            let addr = Address::from_word(hash);
718            let balance = (idx * 2) as u64;
719            let info = AccountInfo::from_balance(U256::from(balance));
720            state.insert_account(addr, info);
721            storage.insert(hash, StateDb::new(state));
722        }
723
724        // wait for files to be flushed
725        tokio::time::sleep(std::time::Duration::from_secs(1)).await;
726
727        let on_disk_states_len = num_states - storage.min_in_memory_limit;
728
729        assert_eq!(storage.on_disk_states.len(), on_disk_states_len);
730        assert_eq!(storage.present.len(), storage.min_in_memory_limit);
731
732        for idx in 0..num_states {
733            let hash = B256::from(U256::from(idx));
734            let addr = Address::from_word(hash);
735
736            let loaded = if idx < on_disk_states_len {
737                storage.get_on_disk_state(&hash).unwrap()
738            } else {
739                storage.get_state(&hash).unwrap()
740            };
741
742            let acc = loaded.basic_ref(addr).unwrap().unwrap();
743            let balance = (idx * 2) as u64;
744            assert_eq!(acc.balance, U256::from(balance));
745        }
746    }
747
748    #[test]
749    fn test_remove_block_states_on_rollback() {
750        let mut storage = InMemoryBlockStates::new(10, MAX_ON_DISK_HISTORY_LIMIT);
751
752        // Insert 5 states
753        let hashes: Vec<B256> = (0..5)
754            .map(|i| {
755                let hash = B256::from(U256::from(i));
756                let mut state = MemDb::default();
757                let addr = Address::from_word(hash);
758                state.insert_account(addr, AccountInfo::from_balance(U256::from(i * 100)));
759                storage.insert(hash, StateDb::new(state));
760                hash
761            })
762            .collect();
763
764        assert_eq!(storage.present.len(), 5);
765
766        // Simulate rollback: remove the last 3 blocks
767        let removed_hashes = &hashes[2..];
768        storage.remove_block_states(removed_hashes);
769
770        // Only the first 2 states should remain
771        assert_eq!(storage.present.len(), 2);
772        assert!(storage.get_state(&hashes[0]).is_some());
773        assert!(storage.get_state(&hashes[1]).is_some());
774        for h in removed_hashes {
775            assert!(storage.get_state(h).is_none());
776            assert!(!storage.present.contains(h));
777        }
778    }
779
780    #[tokio::test(flavor = "multi_thread")]
781    async fn test_remove_block_states_cleans_disk_cache() {
782        // Use limit=1 to force states to disk
783        let mut storage = InMemoryBlockStates::new(1, MAX_ON_DISK_HISTORY_LIMIT);
784
785        let hash_a = B256::from(U256::from(1));
786        let hash_b = B256::from(U256::from(2));
787
788        storage.insert(hash_a, StateDb::new(MemDb::default()));
789        storage.insert(hash_b, StateDb::new(MemDb::default()));
790
791        // Wait for disk flush
792        tokio::time::sleep(std::time::Duration::from_secs(1)).await;
793
794        assert!(storage.on_disk_states.contains_key(&hash_a));
795
796        // Remove hash_a (on disk)
797        storage.remove_block_states(&[hash_a]);
798
799        assert!(!storage.on_disk_states.contains_key(&hash_a));
800        assert!(!storage.oldest_on_disk.contains(&hash_a));
801        assert!(storage.get_on_disk_state(&hash_a).is_none());
802    }
803
804    // verifies that blocks and transactions in BlockchainStorage remain the same when dumped and
805    // reloaded
806    #[test]
807    fn test_storage_dump_reload_cycle() {
808        let mut dump_storage = BlockchainStorage::<FoundryNetwork>::empty();
809
810        let header = Header { gas_limit: 123456, ..Default::default() };
811        let bytes_first = &mut &hex::decode("f86b02843b9aca00830186a094d3e8763675e4c425df46cc3b5c0f6cbdac39604687038d7ea4c68000802ba00eb96ca19e8a77102767a41fc85a36afd5c61ccb09911cec5d3e86e193d9c5aea03a456401896b1b6055311536bf00a718568c744d8c1f9df59879e8350220ca18").unwrap()[..];
812        let tx: MaybeImpersonatedTransaction<FoundryTxEnvelope> =
813            FoundryTxEnvelope::decode(&mut &bytes_first[..]).unwrap().into();
814        let block = create_block(header.clone(), vec![tx.clone()]);
815        let block_hash = block.header.hash_slow();
816        dump_storage.blocks.insert(block_hash, block);
817
818        let serialized_blocks = dump_storage.serialized_blocks();
819        let serialized_transactions = dump_storage.serialized_transactions();
820
821        let mut load_storage = BlockchainStorage::<FoundryNetwork>::empty();
822
823        load_storage.load_blocks(serialized_blocks);
824        load_storage.load_transactions(serialized_transactions);
825
826        let loaded_block = load_storage.blocks.get(&block_hash).unwrap();
827        assert_eq!(loaded_block.header.gas_limit(), header.gas_limit());
828        let loaded_tx = loaded_block.body.transactions.first().unwrap();
829        assert_eq!(loaded_tx, &tx);
830    }
831}