reth_trie_db/
trie_cursor.rs

Help
1use alloy_primitives::B256;
2use reth_db_api::{
3    cursor::{DbCursorRO, DbCursorRW, DbDupCursorRO, DbDupCursorRW},
4    tables,
5    transaction::DbTx,
6    DatabaseError,
7};
8use reth_trie::{
9    trie_cursor::{TrieCursor, TrieCursorFactory},
10    updates::StorageTrieUpdates,
11    BranchNodeCompact, Nibbles, StorageTrieEntry, StoredNibbles, StoredNibblesSubKey,
12};
13
14/// Wrapper struct for database transaction implementing trie cursor factory trait.
15#[derive(Debug)]
16pub struct DatabaseTrieCursorFactory<'a, TX>(&'a TX);
17
18impl<TX> Clone for DatabaseTrieCursorFactory<'_, TX> {
19    fn clone(&self) -> Self {
20        Self(self.0)
21    }
22}
23
24impl<'a, TX> DatabaseTrieCursorFactory<'a, TX> {
25    /// Create new [`DatabaseTrieCursorFactory`].
26    pub const fn new(tx: &'a TX) -> Self {
27        Self(tx)
28    }
29}
30
31/// Implementation of the trie cursor factory for a database transaction.
32impl<TX: DbTx> TrieCursorFactory for DatabaseTrieCursorFactory<'_, TX> {
33    type AccountTrieCursor = DatabaseAccountTrieCursor<<TX as DbTx>::Cursor<tables::AccountsTrie>>;
34    type StorageTrieCursor =
35        DatabaseStorageTrieCursor<<TX as DbTx>::DupCursor<tables::StoragesTrie>>;
36
37    fn account_trie_cursor(&self) -> Result<Self::AccountTrieCursor, DatabaseError> {
38        Ok(DatabaseAccountTrieCursor::new(self.0.cursor_read::<tables::AccountsTrie>()?))
39    }
40
41    fn storage_trie_cursor(
42        &self,
43        hashed_address: B256,
44    ) -> Result<Self::StorageTrieCursor, DatabaseError> {
45        Ok(DatabaseStorageTrieCursor::new(
46            self.0.cursor_dup_read::<tables::StoragesTrie>()?,
47            hashed_address,
48        ))
49    }
50}
51
52/// A cursor over the account trie.
53#[derive(Debug)]
54pub struct DatabaseAccountTrieCursor<C>(pub(crate) C);
55
56impl<C> DatabaseAccountTrieCursor<C> {
57    /// Create a new account trie cursor.
58    pub const fn new(cursor: C) -> Self {
59        Self(cursor)
60    }
61}
62
63impl<C> TrieCursor for DatabaseAccountTrieCursor<C>
64where
65    C: DbCursorRO<tables::AccountsTrie> + Send + Sync,
66{
67    /// Seeks an exact match for the provided key in the account trie.
68    fn seek_exact(
69        &mut self,
70        key: Nibbles,
71    ) -> Result<Option<(Nibbles, BranchNodeCompact)>, DatabaseError> {
72        Ok(self.0.seek_exact(StoredNibbles(key))?.map(|value| (value.0 .0, value.1)))
73    }
74
75    /// Seeks a key in the account trie that matches or is greater than the provided key.
76    fn seek(
77        &mut self,
78        key: Nibbles,
79    ) -> Result<Option<(Nibbles, BranchNodeCompact)>, DatabaseError> {
80        Ok(self.0.seek(StoredNibbles(key))?.map(|value| (value.0 .0, value.1)))
81    }
82
83    /// Move the cursor to the next entry and return it.
84    fn next(&mut self) -> Result<Option<(Nibbles, BranchNodeCompact)>, DatabaseError> {
85        Ok(self.0.next()?.map(|value| (value.0 .0, value.1)))
86    }
87
88    /// Retrieves the current key in the cursor.
89    fn current(&mut self) -> Result<Option<Nibbles>, DatabaseError> {
90        Ok(self.0.current()?.map(|(k, _)| k.0))
91    }
92}
93
94/// A cursor over the storage tries stored in the database.
95#[derive(Debug)]
96pub struct DatabaseStorageTrieCursor<C> {
97    /// The underlying cursor.
98    pub cursor: C,
99    /// Hashed address used for cursor positioning.
100    hashed_address: B256,
101}
102
103impl<C> DatabaseStorageTrieCursor<C> {
104    /// Create a new storage trie cursor.
105    pub const fn new(cursor: C, hashed_address: B256) -> Self {
106        Self { cursor, hashed_address }
107    }
108}
109
110impl<C> DatabaseStorageTrieCursor<C>
111where
112    C: DbCursorRO<tables::StoragesTrie>
113        + DbCursorRW<tables::StoragesTrie>
114        + DbDupCursorRO<tables::StoragesTrie>
115        + DbDupCursorRW<tables::StoragesTrie>,
116{
117    /// Writes storage updates
118    pub fn write_storage_trie_updates(
119        &mut self,
120        updates: &StorageTrieUpdates,
121    ) -> Result<usize, DatabaseError> {
122        // The storage trie for this account has to be deleted.
123        if updates.is_deleted() && self.cursor.seek_exact(self.hashed_address)?.is_some() {
124            self.cursor.delete_current_duplicates()?;
125        }
126
127        // Merge updated and removed nodes. Updated nodes must take precedence.
128        let mut storage_updates = updates
129            .removed_nodes_ref()
130            .iter()
131            .filter_map(|n| (!updates.storage_nodes_ref().contains_key(n)).then_some((n, None)))
132            .collect::<Vec<_>>();
133        storage_updates.extend(
134            updates.storage_nodes_ref().iter().map(|(nibbles, node)| (nibbles, Some(node))),
135        );
136
137        // Sort trie node updates.
138        storage_updates.sort_unstable_by(|a, b| a.0.cmp(b.0));
139
140        let mut num_entries = 0;
141        for (nibbles, maybe_updated) in storage_updates.into_iter().filter(|(n, _)| !n.is_empty()) {
142            num_entries += 1;
143            let nibbles = StoredNibblesSubKey(nibbles.clone());
144            // Delete the old entry if it exists.
145            if self
146                .cursor
147                .seek_by_key_subkey(self.hashed_address, nibbles.clone())?
148                .filter(|e| e.nibbles == nibbles)
149                .is_some()
150            {
151                self.cursor.delete_current()?;
152            }
153
154            // There is an updated version of this node, insert new entry.
155            if let Some(node) = maybe_updated {
156                self.cursor.upsert(
157                    self.hashed_address,
158                    &StorageTrieEntry { nibbles, node: node.clone() },
159                )?;
160            }
161        }
162
163        Ok(num_entries)
164    }
165}
166
167impl<C> TrieCursor for DatabaseStorageTrieCursor<C>
168where
169    C: DbCursorRO<tables::StoragesTrie> + DbDupCursorRO<tables::StoragesTrie> + Send + Sync,
170{
171    /// Seeks an exact match for the given key in the storage trie.
172    fn seek_exact(
173        &mut self,
174        key: Nibbles,
175    ) -> Result<Option<(Nibbles, BranchNodeCompact)>, DatabaseError> {
176        Ok(self
177            .cursor
178            .seek_by_key_subkey(self.hashed_address, StoredNibblesSubKey(key.clone()))?
179            .filter(|e| e.nibbles == StoredNibblesSubKey(key))
180            .map(|value| (value.nibbles.0, value.node)))
181    }
182
183    /// Seeks the given key in the storage trie.
184    fn seek(
185        &mut self,
186        key: Nibbles,
187    ) -> Result<Option<(Nibbles, BranchNodeCompact)>, DatabaseError> {
188        Ok(self
189            .cursor
190            .seek_by_key_subkey(self.hashed_address, StoredNibblesSubKey(key))?
191            .map(|value| (value.nibbles.0, value.node)))
192    }
193
194    /// Move the cursor to the next entry and return it.
195    fn next(&mut self) -> Result<Option<(Nibbles, BranchNodeCompact)>, DatabaseError> {
196        Ok(self.cursor.next_dup()?.map(|(_, v)| (v.nibbles.0, v.node)))
197    }
198
199    /// Retrieves the current value in the storage trie cursor.
200    fn current(&mut self) -> Result<Option<Nibbles>, DatabaseError> {
201        Ok(self.cursor.current()?.map(|(_, v)| v.nibbles.0))
202    }
203}
204
205#[cfg(test)]
206mod tests {
207    use super::*;
208    use alloy_primitives::hex_literal::hex;
209    use reth_db_api::{cursor::DbCursorRW, transaction::DbTxMut};
210    use reth_provider::test_utils::create_test_provider_factory;
211
212    #[test]
213    fn test_account_trie_order() {
214        let factory = create_test_provider_factory();
215        let provider = factory.provider_rw().unwrap();
216        let mut cursor = provider.tx_ref().cursor_write::<tables::AccountsTrie>().unwrap();
217
218        let data = vec![
219            hex!("0303040e").to_vec(),
220            hex!("030305").to_vec(),
221            hex!("03030500").to_vec(),
222            hex!("0303050a").to_vec(),
223        ];
224
225        for key in data.clone() {
226            cursor
227                .upsert(
228                    key.into(),
229                    &BranchNodeCompact::new(
230                        0b0000_0010_0000_0001,
231                        0b0000_0010_0000_0001,
232                        0,
233                        Vec::default(),
234                        None,
235                    ),
236                )
237                .unwrap();
238        }
239
240        let db_data = cursor.walk_range(..).unwrap().collect::<Result<Vec<_>, _>>().unwrap();
241        assert_eq!(db_data[0].0 .0.to_vec(), data[0]);
242        assert_eq!(db_data[1].0 .0.to_vec(), data[1]);
243        assert_eq!(db_data[2].0 .0.to_vec(), data[2]);
244        assert_eq!(db_data[3].0 .0.to_vec(), data[3]);
245
246        assert_eq!(
247            cursor.seek(hex!("0303040f").to_vec().into()).unwrap().map(|(k, _)| k.0.to_vec()),
248            Some(data[1].clone())
249        );
250    }
251
252    // tests that upsert and seek match on the storage trie cursor
253    #[test]
254    fn test_storage_cursor_abstraction() {
255        let factory = create_test_provider_factory();
256        let provider = factory.provider_rw().unwrap();
257        let mut cursor = provider.tx_ref().cursor_dup_write::<tables::StoragesTrie>().unwrap();
258
259        let hashed_address = B256::random();
260        let key = StoredNibblesSubKey::from(vec![0x2, 0x3]);
261        let value = BranchNodeCompact::new(1, 1, 1, vec![B256::random()], None);
262
263        cursor
264            .upsert(hashed_address, &StorageTrieEntry { nibbles: key.clone(), node: value.clone() })
265            .unwrap();
266
267        let mut cursor = DatabaseStorageTrieCursor::new(cursor, hashed_address);
268        assert_eq!(cursor.seek(key.into()).unwrap().unwrap().1, value);
269    }
270}
reth_trie_db/trie_cursor.rs

reth_trie_db/
trie_cursor.rs
