reth_trie_db/
trie_cursor.rs

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

reth_trie_db/
trie_cursor.rs
