reth_trie/trie_cursor/
in_memory.rs

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 const 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 exact && 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 ||
206            (exact && in_memory.as_ref().is_some_and(|entry| entry.0 == key))
207        {
208            return Ok(in_memory.filter(|(nibbles, _)| !exact || nibbles == &key))
209        }
210
211        // Reposition the cursor to the first greater or equal node that wasn't removed.
212        let mut db_entry = self.cursor.seek(key.clone())?;
213        while db_entry
214            .as_ref()
215            .is_some_and(|entry| self.removed_nodes.as_ref().is_some_and(|r| r.contains(&entry.0)))
216        {
217            db_entry = self.cursor.next()?;
218        }
219
220        // Compare two entries and return the lowest.
221        // If seek is exact, filter the entry for exact key match.
222        Ok(compare_trie_node_entries(in_memory, db_entry)
223            .filter(|(nibbles, _)| !exact || nibbles == &key))
224    }
225
226    fn next_inner(
227        &mut self,
228        last: Nibbles,
229    ) -> Result<Option<(Nibbles, BranchNodeCompact)>, DatabaseError> {
230        let in_memory = self.in_memory_cursor.as_mut().and_then(|c| c.first_after(&last));
231        if self.storage_trie_cleared {
232            return Ok(in_memory)
233        }
234
235        // Reposition the cursor to the first greater or equal node that wasn't removed.
236        let mut db_entry = self.cursor.seek(last.clone())?;
237        while db_entry.as_ref().is_some_and(|entry| {
238            entry.0 < last || self.removed_nodes.as_ref().is_some_and(|r| r.contains(&entry.0))
239        }) {
240            db_entry = self.cursor.next()?;
241        }
242
243        // Compare two entries and return the lowest.
244        Ok(compare_trie_node_entries(in_memory, db_entry))
245    }
246}
247
248impl<C: TrieCursor> TrieCursor for InMemoryStorageTrieCursor<'_, C> {
249    fn seek_exact(
250        &mut self,
251        key: Nibbles,
252    ) -> Result<Option<(Nibbles, BranchNodeCompact)>, DatabaseError> {
253        let entry = self.seek_inner(key, true)?;
254        self.last_key = entry.as_ref().map(|(nibbles, _)| nibbles.clone());
255        Ok(entry)
256    }
257
258    fn seek(
259        &mut self,
260        key: Nibbles,
261    ) -> Result<Option<(Nibbles, BranchNodeCompact)>, DatabaseError> {
262        let entry = self.seek_inner(key, false)?;
263        self.last_key = entry.as_ref().map(|(nibbles, _)| nibbles.clone());
264        Ok(entry)
265    }
266
267    fn next(&mut self) -> Result<Option<(Nibbles, BranchNodeCompact)>, DatabaseError> {
268        let next = match &self.last_key {
269            Some(last) => {
270                let entry = self.next_inner(last.clone())?;
271                self.last_key = entry.as_ref().map(|entry| entry.0.clone());
272                entry
273            }
274            // no previous entry was found
275            None => None,
276        };
277        Ok(next)
278    }
279
280    fn current(&mut self) -> Result<Option<Nibbles>, DatabaseError> {
281        match &self.last_key {
282            Some(key) => Ok(Some(key.clone())),
283            None => self.cursor.current(),
284        }
285    }
286}
287
288/// Return the node with the lowest nibbles.
289///
290/// Given the next in-memory and database entries, return the smallest of the two.
291/// If the node keys are the same, the in-memory entry is given precedence.
292fn compare_trie_node_entries(
293    mut in_memory_item: Option<(Nibbles, BranchNodeCompact)>,
294    mut db_item: Option<(Nibbles, BranchNodeCompact)>,
295) -> Option<(Nibbles, BranchNodeCompact)> {
296    if let Some((in_memory_entry, db_entry)) = in_memory_item.as_ref().zip(db_item.as_ref()) {
297        // If both are not empty, return the smallest of the two
298        // In-memory is given precedence if keys are equal
299        if in_memory_entry.0 <= db_entry.0 {
300            in_memory_item.take()
301        } else {
302            db_item.take()
303        }
304    } else {
305        // Return either non-empty entry
306        db_item.or(in_memory_item)
307    }
308}
reth_trie/trie_cursor/in_memory.rs

reth_trie/trie_cursor/
in_memory.rs
