Skip to main content

reth_trie_sparse/
lower.rs

1use crate::parallel::SparseSubtrie;
2use alloc::boxed::Box;
3use reth_trie_common::Nibbles;
4
5/// Tracks the state of the lower subtries.
6///
7/// When a [`crate::ParallelSparseTrie`] is initialized/cleared then its `LowerSparseSubtrie`s are
8/// all blinded, meaning they have no nodes. A blinded `LowerSparseSubtrie` may hold onto a cleared
9/// [`SparseSubtrie`] in order to reuse allocations.
10#[derive(Clone, Debug, Eq, PartialEq)]
11pub(crate) enum LowerSparseSubtrie {
12    Blind(Option<Box<SparseSubtrie>>),
13    Revealed(Box<SparseSubtrie>),
14}
15
16impl Default for LowerSparseSubtrie {
17    /// Creates a new blinded subtrie with no allocated storage.
18    fn default() -> Self {
19        Self::Blind(None)
20    }
21}
22
23impl LowerSparseSubtrie {
24    /// Returns a reference to the underlying [`SparseSubtrie`] if this subtrie is revealed.
25    ///
26    /// Returns `None` if the subtrie is blinded (has no nodes).
27    pub(crate) fn as_revealed_ref(&self) -> Option<&SparseSubtrie> {
28        match self {
29            Self::Blind(_) => None,
30            Self::Revealed(subtrie) => Some(subtrie.as_ref()),
31        }
32    }
33
34    /// Returns a mutable reference to the underlying [`SparseSubtrie`] if this subtrie is revealed.
35    ///
36    /// Returns `None` if the subtrie is blinded (has no nodes).
37    pub(crate) fn as_revealed_mut(&mut self) -> Option<&mut SparseSubtrie> {
38        match self {
39            Self::Blind(_) => None,
40            Self::Revealed(subtrie) => Some(subtrie.as_mut()),
41        }
42    }
43
44    /// Reveals the lower [`SparseSubtrie`], transitioning it from the Blinded to the Revealed
45    /// variant, preserving allocations if possible.
46    ///
47    /// The given path is the path of a node which will be set into the [`SparseSubtrie`]'s `nodes`
48    /// map immediately upon being revealed. If the subtrie is blinded, or if its current root path
49    /// is longer than this one, than this one becomes the new root path of the subtrie.
50    pub(crate) fn reveal(&mut self, path: &Nibbles) {
51        match self {
52            Self::Blind(allocated) => {
53                debug_assert!(allocated.as_ref().is_none_or(|subtrie| subtrie.is_empty()));
54                *self = if let Some(mut subtrie) = allocated.take() {
55                    subtrie.path = *path;
56                    Self::Revealed(subtrie)
57                } else {
58                    Self::Revealed(Box::new(SparseSubtrie::new(*path)))
59                }
60            }
61            Self::Revealed(subtrie) => {
62                if path.len() < subtrie.path.len() {
63                    subtrie.path = *path;
64                }
65            }
66        };
67    }
68
69    /// Clears the subtrie and transitions it to the blinded state, preserving a cleared
70    /// [`SparseSubtrie`] if possible.
71    pub(crate) fn clear(&mut self) {
72        *self = match core::mem::take(self) {
73            Self::Blind(allocated) => {
74                debug_assert!(allocated.as_ref().is_none_or(|subtrie| subtrie.is_empty()));
75                Self::Blind(allocated)
76            }
77            Self::Revealed(mut subtrie) => {
78                subtrie.clear();
79                Self::Blind(Some(subtrie))
80            }
81        }
82    }
83
84    /// Takes ownership of the underlying [`SparseSubtrie`] if revealed, putting this
85    /// `LowerSparseSubtrie` will be put into the blinded state.
86    ///
87    /// Otherwise returns None.
88    #[cfg(feature = "std")]
89    pub(crate) fn take_revealed(&mut self) -> Option<Box<SparseSubtrie>> {
90        self.take_revealed_if(|_| true)
91    }
92
93    /// Takes ownership of the underlying [`SparseSubtrie`] if revealed and the predicate returns
94    /// true.
95    ///
96    /// If the subtrie is revealed, and the predicate function returns `true` when called with it,
97    /// then this method will take ownership of the subtrie and transition this `LowerSparseSubtrie`
98    /// to the blinded state. Otherwise, returns `None`.
99    pub(crate) fn take_revealed_if<P>(&mut self, predicate: P) -> Option<Box<SparseSubtrie>>
100    where
101        P: FnOnce(&SparseSubtrie) -> bool,
102    {
103        match self {
104            Self::Revealed(subtrie) if predicate(subtrie) => {
105                let Self::Revealed(subtrie) = core::mem::take(self) else { unreachable!() };
106                Some(subtrie)
107            }
108            Self::Revealed(_) | Self::Blind(_) => None,
109        }
110    }
111
112    /// Shrinks the capacity of the subtrie's node storage.
113    /// Works for both revealed and blind tries with allocated storage.
114    pub(crate) fn shrink_nodes_to(&mut self, size: usize) {
115        match self {
116            Self::Revealed(trie) | Self::Blind(Some(trie)) => {
117                trie.shrink_nodes_to(size);
118            }
119            Self::Blind(None) => {}
120        }
121    }
122
123    /// Shrinks the capacity of the subtrie's value storage.
124    /// Works for both revealed and blind tries with allocated storage.
125    pub(crate) fn shrink_values_to(&mut self, size: usize) {
126        match self {
127            Self::Revealed(trie) | Self::Blind(Some(trie)) => {
128                trie.shrink_values_to(size);
129            }
130            Self::Blind(None) => {}
131        }
132    }
133
134    /// Returns a heuristic for the in-memory size of this subtrie in bytes.
135    pub(crate) fn memory_size(&self) -> usize {
136        match self {
137            Self::Revealed(subtrie) | Self::Blind(Some(subtrie)) => subtrie.memory_size(),
138            Self::Blind(None) => 0,
139        }
140    }
141
142    pub(crate) fn wipe(&mut self) {
143        *self = Self::default();
144    }
145}