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}