arbiter_core/database/
mod.rs

1//! The [`ArbiterDB`] is a wrapper around a `CacheDB` that is used to provide
2//! access to the `Environment`'s database to multiple `Coprocessors`.
3//! It is also used to be able to write out the `Environment` database to a
4//! file.
5//!
6//! Further, it gives the ability to be generated from a [`fork::Fork`] so that
7//! you can preload an [`environment::Environment`] with a specific state.
8
9use std::{
10    fs,
11    io::{self, Read, Write},
12};
13
14use revm::{
15    primitives::{db::DatabaseRef, keccak256, Bytecode, B256},
16    DatabaseCommit,
17};
18use serde_json;
19
20use super::*;
21pub mod fork;
22pub mod inspector;
23
24/// A [`ArbiterDB`] is contains both a [`CacheDB`] that is used to provide
25/// state for the [`environment::Environment`]'s as well as for multiple
26/// [`coprocessor::Coprocessor`]s.
27/// The `logs` field is a [`HashMap`] to store [`ethers::types::Log`]s that can
28/// be queried from at any point.
29#[derive(Debug, Serialize, Deserialize)]
30pub struct ArbiterDB {
31    /// The state of the `ArbiterDB`. This is a `CacheDB` that is used to
32    /// provide a db for the `Environment` to use.
33    pub state: Arc<RwLock<CacheDB<EmptyDB>>>,
34
35    /// The logs of the `ArbiterDB`. This is a `HashMap` that is used to store
36    /// logs that can be queried from at any point.
37    pub logs: Arc<RwLock<HashMap<U256, Vec<eLog>>>>,
38}
39
40// Implement `Clone` by hand so we utilize the `Arc`'s `Clone` implementation.
41impl Clone for ArbiterDB {
42    fn clone(&self) -> Self {
43        Self {
44            state: self.state.clone(),
45            logs: self.logs.clone(),
46        }
47    }
48}
49
50impl ArbiterDB {
51    /// Create a new `ArbiterDB`.
52    pub fn new() -> Self {
53        Self {
54            state: Arc::new(RwLock::new(CacheDB::new(EmptyDB::new()))),
55            logs: Arc::new(RwLock::new(HashMap::new())),
56        }
57    }
58
59    /// Write the `ArbiterDB` to a file at the given path.``
60    pub fn write_to_file(&self, path: &str) -> io::Result<()> {
61        // Serialize the ArbiterDB
62        let serialized = serde_json::to_string(self)?;
63        // Write to file
64        let mut file = fs::File::create(path)?;
65        file.write_all(serialized.as_bytes())?;
66        Ok(())
67    }
68
69    /// Read the `ArbiterDB` from a file at the given path.
70    pub fn read_from_file(path: &str) -> io::Result<Self> {
71        // Read the file content
72        let mut file = fs::File::open(path)?;
73        let mut contents = String::new();
74        file.read_to_string(&mut contents)?;
75
76        // Deserialize the content into ArbiterDB
77        #[derive(Deserialize)]
78        struct TempDB {
79            state: Option<CacheDB<EmptyDB>>,
80            logs: Option<HashMap<U256, Vec<eLog>>>,
81        }
82        let temp_db: TempDB = serde_json::from_str(&contents)?;
83        Ok(Self {
84            state: Arc::new(RwLock::new(temp_db.state.unwrap_or_default())),
85            logs: Arc::new(RwLock::new(temp_db.logs.unwrap_or_default())),
86        })
87    }
88}
89
90impl Default for ArbiterDB {
91    fn default() -> Self {
92        Self::new()
93    }
94}
95
96// TODO: This is a BAD implementation of PartialEq, but it works for now as we
97// do not ever really need to compare DBs directly at the moment.
98// This is only used in the `Outcome` enum for `instruction.rs`.
99impl PartialEq for ArbiterDB {
100    fn eq(&self, _other: &Self) -> bool {
101        true
102    }
103}
104
105impl Database for ArbiterDB {
106    type Error = Infallible; // TODO: Not sure we want this, but it works for now.
107
108    fn basic(
109        &mut self,
110        address: revm::primitives::Address,
111    ) -> Result<Option<AccountInfo>, Self::Error> {
112        self.state.write().unwrap().basic(address)
113    }
114
115    fn code_by_hash(&mut self, code_hash: B256) -> Result<Bytecode, Self::Error> {
116        self.state.write().unwrap().code_by_hash(code_hash)
117    }
118
119    fn storage(
120        &mut self,
121        address: revm::primitives::Address,
122        index: U256,
123    ) -> Result<U256, Self::Error> {
124        self.state.write().unwrap().storage(address, index)
125    }
126
127    fn block_hash(&mut self, number: U256) -> Result<B256, Self::Error> {
128        self.state.write().unwrap().block_hash(number)
129    }
130}
131
132impl DatabaseRef for ArbiterDB {
133    type Error = Infallible; // TODO: Not sure we want this, but it works for now.
134
135    fn basic_ref(
136        &self,
137        address: revm::primitives::Address,
138    ) -> Result<Option<AccountInfo>, Self::Error> {
139        self.state.read().unwrap().basic_ref(address)
140    }
141
142    fn code_by_hash_ref(&self, code_hash: B256) -> Result<Bytecode, Self::Error> {
143        self.state.read().unwrap().code_by_hash_ref(code_hash)
144    }
145
146    fn storage_ref(
147        &self,
148        address: revm::primitives::Address,
149        index: U256,
150    ) -> Result<U256, Self::Error> {
151        self.state.read().unwrap().storage_ref(address, index)
152    }
153
154    fn block_hash_ref(&self, number: U256) -> Result<B256, Self::Error> {
155        self.state.read().unwrap().block_hash_ref(number)
156    }
157}
158
159impl DatabaseCommit for ArbiterDB {
160    fn commit(
161        &mut self,
162        changes: revm_primitives::HashMap<revm::primitives::Address, revm::primitives::Account>,
163    ) {
164        self.state.write().unwrap().commit(changes)
165    }
166}
167
168/// [AnvilDump] models the schema of an [anvil](https://github.com/foundry-rs/foundry) state dump.
169#[derive(Clone, Debug, Serialize, Deserialize)]
170pub struct AnvilDump {
171    /// Mapping of account addresses to [AccountRecord]s stored in the dump
172    /// file.
173    pub accounts: BTreeMap<Address, AccountRecord>,
174}
175
176/// [AccountRecord] describes metadata about an account within the state trie.
177#[derive(Clone, Debug, Serialize, Deserialize)]
178pub struct AccountRecord {
179    /// The nonce of the account.
180    pub nonce: u64,
181    /// The balance of the account.
182    pub balance: U256,
183    /// The bytecode of the account. If empty, the account is an EOA.
184    pub code: Bytes,
185    /// The storage mapping of the account.
186    pub storage: revm_primitives::HashMap<U256, U256>,
187}
188
189impl TryFrom<AnvilDump> for CacheDB<EmptyDB> {
190    type Error = <CacheDB<EmptyDB> as Database>::Error;
191
192    fn try_from(dump: AnvilDump) -> Result<Self, Self::Error> {
193        let mut db = CacheDB::default();
194
195        dump.accounts
196            .into_iter()
197            .try_for_each(|(address, account_record)| {
198                db.insert_account_info(
199                    address,
200                    AccountInfo {
201                        balance: account_record.balance,
202                        nonce: account_record.nonce,
203                        code_hash: keccak256(account_record.code.as_ref()),
204                        code: (!account_record.code.is_empty())
205                            .then(|| Bytecode::new_raw(account_record.code)),
206                    },
207                );
208                db.replace_account_storage(address, account_record.storage)
209            })?;
210
211        Ok(db)
212    }
213}
214
215#[cfg(test)]
216mod tests {
217    use revm_primitives::{address, bytes};
218
219    use super::*;
220
221    #[test]
222    fn read_write_to_file() {
223        let db = ArbiterDB::new();
224        db.write_to_file("test.json").unwrap();
225        let db = ArbiterDB::read_from_file("test.json").unwrap();
226        assert_eq!(db, ArbiterDB::new());
227        fs::remove_file("test.json").unwrap();
228    }
229
230    #[test]
231    fn load_anvil_dump_cachedb() {
232        const RAW_DUMP: &str = r#"
233        {
234            "accounts": {
235                "0x0000000000000000000000000000000000000000": {
236                    "nonce": 1234,
237                    "balance": "0xfacade",
238                    "code": "0x",
239                    "storage": {}
240                },
241                "0x0000000000000000000000000000000000000001": {
242                    "nonce": 555,
243                    "balance": "0xc0ffee",
244                    "code": "0xbadc0de0",
245                    "storage": {
246                        "0x0000000000000000000000000000000000000000000000000000000000000000": "0x000000000000000000000000000000000000000000000000000000000000deAD",
247                        "0x0000000000000000000000000000000000000000000000000000000000000001": "0x000000000000000000000000000000000000000000000000000000000000babe"
248                    }
249                }
250            }
251        }
252        "#;
253
254        let dump: AnvilDump = serde_json::from_str(RAW_DUMP).unwrap();
255        let mut db: CacheDB<EmptyDB> = dump.try_into().unwrap();
256
257        let account_a = db
258            .load_account(address!("0000000000000000000000000000000000000000"))
259            .unwrap();
260        assert_eq!(account_a.info.nonce, 1234);
261        assert_eq!(account_a.info.balance, U256::from(0xfacade));
262        assert_eq!(account_a.info.code, None);
263        assert_eq!(account_a.info.code_hash, keccak256([]));
264
265        let account_b = db
266            .load_account(address!("0000000000000000000000000000000000000001"))
267            .unwrap();
268        let b_bytecode = bytes!("badc0de0");
269        assert_eq!(account_b.info.nonce, 555);
270        assert_eq!(account_b.info.balance, U256::from(0xc0ffee));
271        assert_eq!(account_b.info.code_hash, keccak256(b_bytecode.as_ref()));
272        assert_eq!(account_b.info.code, Some(Bytecode::new_raw(b_bytecode)));
273        assert_eq!(
274            account_b.storage.get(&U256::ZERO),
275            Some(&U256::from(0xdead))
276        );
277        assert_eq!(
278            account_b.storage.get(&U256::from(1)),
279            Some(&U256::from(0xbabe))
280        );
281    }
282}
arbiter_core/database/mod.rs

arbiter_core/database/
mod.rs
