reth_trie/trie_cursor/
in_memory.rs

Help
1use super::{TrieCursor, TrieCursorFactory};
2use crate::{
3    forward_cursor::ForwardInMemoryCursor,
4    updates::{StorageTrieUpdatesSorted, TrieUpdatesSorted},
5};
6use alloy_primitives::{map::HashSet, B256};
7use reth_storage_errors::db::DatabaseError;
8use reth_trie_common::{BranchNodeCompact, Nibbles};
9
10/// The trie cursor factory for the trie updates.
11#[derive(Debug, Clone)]
12pub struct InMemoryTrieCursorFactory<'a, CF> {
13    /// Underlying trie cursor factory.
14    cursor_factory: CF,
15    /// Reference to sorted trie updates.
16    trie_updates: &'a TrieUpdatesSorted,
17}
18
19impl<'a, CF> InMemoryTrieCursorFactory<'a, CF> {
20    /// Create a new trie cursor factory.
21    pub const fn new(cursor_factory: CF, trie_updates: &'a TrieUpdatesSorted) -> Self {
22        Self { cursor_factory, trie_updates }
23    }
24}
25
26impl<'a, CF: TrieCursorFactory> TrieCursorFactory for InMemoryTrieCursorFactory<'a, CF> {
27    type AccountTrieCursor = InMemoryAccountTrieCursor<'a, CF::AccountTrieCursor>;
28    type StorageTrieCursor = InMemoryStorageTrieCursor<'a, CF::StorageTrieCursor>;
29
30    fn account_trie_cursor(&self) -> Result<Self::AccountTrieCursor, DatabaseError> {
31        let cursor = self.cursor_factory.account_trie_cursor()?;
32        Ok(InMemoryAccountTrieCursor::new(cursor, self.trie_updates))
33    }
34
35    fn storage_trie_cursor(
36        &self,
37        hashed_address: B256,
38    ) -> Result<Self::StorageTrieCursor, DatabaseError> {
39        let cursor = self.cursor_factory.storage_trie_cursor(hashed_address)?;
40        Ok(InMemoryStorageTrieCursor::new(
41            hashed_address,
42            cursor,
43            self.trie_updates.storage_tries.get(&hashed_address),
44        ))
45    }
46}
47
48/// The cursor to iterate over account trie updates and corresponding database entries.
49/// It will always give precedence to the data from the trie updates.
50#[derive(Debug)]
51pub struct InMemoryAccountTrieCursor<'a, C> {
52    /// The underlying cursor.
53    cursor: C,
54    /// Forward-only in-memory cursor over storage trie nodes.
55    in_memory_cursor: ForwardInMemoryCursor<'a, Nibbles, BranchNodeCompact>,
56    /// Collection of removed trie nodes.
57    removed_nodes: &'a HashSet<Nibbles>,
58    /// Last key returned by the cursor.
59    last_key: Option<Nibbles>,
60}
61
62impl<'a, C: TrieCursor> InMemoryAccountTrieCursor<'a, C> {
63    /// Create new account trie cursor from underlying cursor and reference to
64    /// [`TrieUpdatesSorted`].
65    pub fn new(cursor: C, trie_updates: &'a TrieUpdatesSorted) -> Self {
66        let in_memory_cursor = ForwardInMemoryCursor::new(&trie_updates.account_nodes);
67        Self {
68            cursor,
69            in_memory_cursor,
70            removed_nodes: &trie_updates.removed_nodes,
71            last_key: None,
72        }
73    }
74
75    fn seek_inner(
76        &mut self,
77        key: Nibbles,
78        exact: bool,
79    ) -> Result<Option<(Nibbles, BranchNodeCompact)>, DatabaseError> {
80        let in_memory = self.in_memory_cursor.seek(&key);
81        if in_memory.as_ref().is_some_and(|entry| entry.0 == key) {
82            return Ok(in_memory)
83        }
84
85        // Reposition the cursor to the first greater or equal node that wasn't removed.
86        let mut db_entry = self.cursor.seek(key.clone())?;
87        while db_entry.as_ref().is_some_and(|entry| self.removed_nodes.contains(&entry.0)) {
88            db_entry = self.cursor.next()?;
89        }
90
91        // Compare two entries and return the lowest.
92        // If seek is exact, filter the entry for exact key match.
93        Ok(compare_trie_node_entries(in_memory, db_entry)
94            .filter(|(nibbles, _)| !exact || nibbles == &key))
95    }
96
97    fn next_inner(
98        &mut self,
99        last: Nibbles,
100    ) -> Result<Option<(Nibbles, BranchNodeCompact)>, DatabaseError> {
101        let in_memory = self.in_memory_cursor.first_after(&last);
102
103        // Reposition the cursor to the first greater or equal node that wasn't removed.
104        let mut db_entry = self.cursor.seek(last.clone())?;
105        while db_entry
106            .as_ref()
107            .is_some_and(|entry| entry.0 < last || self.removed_nodes.contains(&entry.0))
108        {
109            db_entry = self.cursor.next()?;
110        }
111
112        // Compare two entries and return the lowest.
113        Ok(compare_trie_node_entries(in_memory, db_entry))
114    }
115}
116
117impl<C: TrieCursor> TrieCursor for InMemoryAccountTrieCursor<'_, C> {
118    fn seek_exact(
119        &mut self,
120        key: Nibbles,
121    ) -> Result<Option<(Nibbles, BranchNodeCompact)>, DatabaseError> {
122        let entry = self.seek_inner(key, true)?;
123        self.last_key = entry.as_ref().map(|(nibbles, _)| nibbles.clone());
124        Ok(entry)
125    }
126
127    fn seek(
128        &mut self,
129        key: Nibbles,
130    ) -> Result<Option<(Nibbles, BranchNodeCompact)>, DatabaseError> {
131        let entry = self.seek_inner(key, false)?;
132        self.last_key = entry.as_ref().map(|(nibbles, _)| nibbles.clone());
133        Ok(entry)
134    }
135
136    fn next(&mut self) -> Result<Option<(Nibbles, BranchNodeCompact)>, DatabaseError> {
137        let next = match &self.last_key {
138            Some(last) => {
139                let entry = self.next_inner(last.clone())?;
140                self.last_key = entry.as_ref().map(|entry| entry.0.clone());
141                entry
142            }
143            // no previous entry was found
144            None => None,
145        };
146        Ok(next)
147    }
148
149    fn current(&mut self) -> Result<Option<Nibbles>, DatabaseError> {
150        match &self.last_key {
151            Some(key) => Ok(Some(key.clone())),
152            None => self.cursor.current(),
153        }
154    }
155}
156
157/// The cursor to iterate over storage trie updates and corresponding database entries.
158/// It will always give precedence to the data from the trie updates.
159#[derive(Debug)]
160#[allow(dead_code)]
161pub struct InMemoryStorageTrieCursor<'a, C> {
162    /// The hashed address of the account that trie belongs to.
163    hashed_address: B256,
164    /// The underlying cursor.
165    cursor: C,
166    /// Forward-only in-memory cursor over storage trie nodes.
167    in_memory_cursor: Option<ForwardInMemoryCursor<'a, Nibbles, BranchNodeCompact>>,
168    /// Reference to the set of removed storage node keys.
169    removed_nodes: Option<&'a HashSet<Nibbles>>,
170    /// The flag indicating whether the storage trie was cleared.
171    storage_trie_cleared: bool,
172    /// Last key returned by the cursor.
173    last_key: Option<Nibbles>,
174}
175
176impl<'a, C> InMemoryStorageTrieCursor<'a, C> {
177    /// Create new storage trie cursor from underlying cursor and reference to
178    /// [`StorageTrieUpdatesSorted`].
179    pub fn new(
180        hashed_address: B256,
181        cursor: C,
182        updates: Option<&'a StorageTrieUpdatesSorted>,
183    ) -> Self {
184        let in_memory_cursor = updates.map(|u| ForwardInMemoryCursor::new(&u.storage_nodes));
185        let removed_nodes = updates.map(|u| &u.removed_nodes);
186        let storage_trie_cleared = updates.is_some_and(|u| u.is_deleted);
187        Self {
188            hashed_address,
189            cursor,
190            in_memory_cursor,
191            removed_nodes,
192            storage_trie_cleared,
193            last_key: None,
194        }
195    }
196}
197
198impl<C: TrieCursor> InMemoryStorageTrieCursor<'_, C> {
199    fn seek_inner(
200        &mut self,
201        key: Nibbles,
202        exact: bool,
203    ) -> Result<Option<(Nibbles, BranchNodeCompact)>, DatabaseError> {
204        let in_memory = self.in_memory_cursor.as_mut().and_then(|c| c.seek(&key));
205        if self.storage_trie_cleared || in_memory.as_ref().is_some_and(|entry| entry.0 == key) {
206            return Ok(in_memory.filter(|(nibbles, _)| !exact || nibbles == &key))
207        }
208
209        // Reposition the cursor to the first greater or equal node that wasn't removed.
210        let mut db_entry = self.cursor.seek(key.clone())?;
211        while db_entry
212            .as_ref()
213            .is_some_and(|entry| self.removed_nodes.as_ref().is_some_and(|r| r.contains(&entry.0)))
214        {
215            db_entry = self.cursor.next()?;
216        }
217
218        // Compare two entries and return the lowest.
219        // If seek is exact, filter the entry for exact key match.
220        Ok(compare_trie_node_entries(in_memory, db_entry)
221            .filter(|(nibbles, _)| !exact || nibbles == &key))
222    }
223
224    fn next_inner(
225        &mut self,
226        last: Nibbles,
227    ) -> Result<Option<(Nibbles, BranchNodeCompact)>, DatabaseError> {
228        let in_memory = self.in_memory_cursor.as_mut().and_then(|c| c.first_after(&last));
229        if self.storage_trie_cleared {
230            return Ok(in_memory)
231        }
232
233        // Reposition the cursor to the first greater or equal node that wasn't removed.
234        let mut db_entry = self.cursor.seek(last.clone())?;
235        while db_entry.as_ref().is_some_and(|entry| {
236            entry.0 < last || self.removed_nodes.as_ref().is_some_and(|r| r.contains(&entry.0))
237        }) {
238            db_entry = self.cursor.next()?;
239        }
240
241        // Compare two entries and return the lowest.
242        Ok(compare_trie_node_entries(in_memory, db_entry))
243    }
244}
245
246impl<C: TrieCursor> TrieCursor for InMemoryStorageTrieCursor<'_, C> {
247    fn seek_exact(
248        &mut self,
249        key: Nibbles,
250    ) -> Result<Option<(Nibbles, BranchNodeCompact)>, DatabaseError> {
251        let entry = self.seek_inner(key, true)?;
252        self.last_key = entry.as_ref().map(|(nibbles, _)| nibbles.clone());
253        Ok(entry)
254    }
255
256    fn seek(
257        &mut self,
258        key: Nibbles,
259    ) -> Result<Option<(Nibbles, BranchNodeCompact)>, DatabaseError> {
260        let entry = self.seek_inner(key, false)?;
261        self.last_key = entry.as_ref().map(|(nibbles, _)| nibbles.clone());
262        Ok(entry)
263    }
264
265    fn next(&mut self) -> Result<Option<(Nibbles, BranchNodeCompact)>, DatabaseError> {
266        let next = match &self.last_key {
267            Some(last) => {
268                let entry = self.next_inner(last.clone())?;
269                self.last_key = entry.as_ref().map(|entry| entry.0.clone());
270                entry
271            }
272            // no previous entry was found
273            None => None,
274        };
275        Ok(next)
276    }
277
278    fn current(&mut self) -> Result<Option<Nibbles>, DatabaseError> {
279        match &self.last_key {
280            Some(key) => Ok(Some(key.clone())),
281            None => self.cursor.current(),
282        }
283    }
284}
285
286/// Return the node with the lowest nibbles.
287///
288/// Given the next in-memory and database entries, return the smallest of the two.
289/// If the node keys are the same, the in-memory entry is given precedence.
290fn compare_trie_node_entries(
291    mut in_memory_item: Option<(Nibbles, BranchNodeCompact)>,
292    mut db_item: Option<(Nibbles, BranchNodeCompact)>,
293) -> Option<(Nibbles, BranchNodeCompact)> {
294    if let Some((in_memory_entry, db_entry)) = in_memory_item.as_ref().zip(db_item.as_ref()) {
295        // If both are not empty, return the smallest of the two
296        // In-memory is given precedence if keys are equal
297        if in_memory_entry.0 <= db_entry.0 {
298            in_memory_item.take()
299        } else {
300            db_item.take()
301        }
302    } else {
303        // Return either non-empty entry
304        db_item.or(in_memory_item)
305    }
306}
reth_trie/trie_cursor/in_memory.rs

reth_trie/trie_cursor/
in_memory.rs
