reth_trie_db/
trie_cursor.rs

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

reth_trie_db/
trie_cursor.rs
