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 pub(crate) fn take_revealed(&mut self) -> Option<Box<SparseSubtrie>> {
89 self.take_revealed_if(|_| true)
90 }
91
92 /// Takes ownership of the underlying [`SparseSubtrie`] if revealed and the predicate returns
93 /// true.
94 ///
95 /// If the subtrie is revealed, and the predicate function returns `true` when called with it,
96 /// then this method will take ownership of the subtrie and transition this `LowerSparseSubtrie`
97 /// to the blinded state. Otherwise, returns `None`.
98 pub(crate) fn take_revealed_if<P>(&mut self, predicate: P) -> Option<Box<SparseSubtrie>>
99 where
100 P: FnOnce(&SparseSubtrie) -> bool,
101 {
102 match self {
103 Self::Revealed(subtrie) if predicate(subtrie) => {
104 let Self::Revealed(subtrie) = core::mem::take(self) else { unreachable!() };
105 Some(subtrie)
106 }
107 Self::Revealed(_) | Self::Blind(_) => None,
108 }
109 }
110
111 /// Shrinks the capacity of the subtrie's node storage.
112 /// Works for both revealed and blind tries with allocated storage.
113 pub(crate) fn shrink_nodes_to(&mut self, size: usize) {
114 match self {
115 Self::Revealed(trie) | Self::Blind(Some(trie)) => {
116 trie.shrink_nodes_to(size);
117 }
118 Self::Blind(None) => {}
119 }
120 }
121
122 /// Shrinks the capacity of the subtrie's value storage.
123 /// Works for both revealed and blind tries with allocated storage.
124 pub(crate) fn shrink_values_to(&mut self, size: usize) {
125 match self {
126 Self::Revealed(trie) | Self::Blind(Some(trie)) => {
127 trie.shrink_values_to(size);
128 }
129 Self::Blind(None) => {}
130 }
131 }
132
133 /// Returns a heuristic for the in-memory size of this subtrie in bytes.
134 pub(crate) fn memory_size(&self) -> usize {
135 match self {
136 Self::Revealed(subtrie) | Self::Blind(Some(subtrie)) => subtrie.memory_size(),
137 Self::Blind(None) => 0,
138 }
139 }
140
141 pub(crate) fn wipe(&mut self) {
142 *self = Self::default();
143 }
144}