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    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}