Skip to main content

phos_data_network_precompiles/storage/
hashmap.rs

1use std::collections::HashMap;
2
3use alloy_primitives::{Address, U256};
4
5use crate::{
6    error::{DataNetworkPrecompileError, Result},
7    storage::PrecompileStorageProvider,
8};
9
10/// In-memory [`PrecompileStorageProvider`] for tests and isolated execution.
11#[derive(Debug, Default)]
12pub struct HashMapStorageProvider {
13    internals: HashMap<(Address, U256), U256>,
14    fail_on_sload: Option<(Address, U256)>,
15    is_static: bool,
16    counter_sload: u64,
17    counter_sstore: u64,
18}
19
20impl HashMapStorageProvider {
21    /// Creates an empty provider.
22    pub fn new() -> Self {
23        Self::default()
24    }
25
26    /// Returns this provider with the static-call flag set.
27    pub fn with_static(mut self, is_static: bool) -> Self {
28        self.is_static = is_static;
29        self
30    }
31}
32
33impl PrecompileStorageProvider for HashMapStorageProvider {
34    fn sload(&mut self, address: Address, key: U256) -> Result<U256> {
35        if self.fail_on_sload == Some((address, key)) {
36            return Err(DataNetworkPrecompileError::Fatal(
37                "injected sload failure".into(),
38            ));
39        }
40
41        self.counter_sload += 1;
42        Ok(self
43            .internals
44            .get(&(address, key))
45            .copied()
46            .unwrap_or(U256::ZERO))
47    }
48
49    fn sstore(&mut self, address: Address, key: U256, value: U256) -> Result<()> {
50        self.counter_sstore += 1;
51        self.internals.insert((address, key), value);
52        Ok(())
53    }
54
55    fn is_static(&self) -> bool {
56        self.is_static
57    }
58}
59
60impl HashMapStorageProvider {
61    /// Makes the next load at `address` and `slot` fail.
62    pub fn fail_next_sload_at(&mut self, address: Address, slot: U256) {
63        self.fail_on_sload = Some((address, slot));
64    }
65
66    /// Returns the number of SLOAD operations.
67    pub fn counter_sload(&self) -> u64 {
68        self.counter_sload
69    }
70
71    /// Returns the number of SSTORE operations.
72    pub fn counter_sstore(&self) -> u64 {
73        self.counter_sstore
74    }
75
76    /// Resets the storage-operation counters.
77    pub fn reset_counters(&mut self) {
78        self.counter_sload = 0;
79        self.counter_sstore = 0;
80    }
81
82    /// Returns all storage entries as `(address, slot, value)` tuples.
83    pub fn into_storage(self) -> impl Iterator<Item = (Address, U256, U256)> {
84        self.internals
85            .into_iter()
86            .map(|((address, slot), value)| (address, slot, value))
87    }
88}