reth_trie_common/
trie_node_v2.rs

1//! Version 2 types related to representing nodes in an MPT.
2
3use crate::BranchNodeMasks;
4use alloc::vec::Vec;
5use alloy_primitives::hex;
6use alloy_rlp::{bytes, Decodable, Encodable, EMPTY_STRING_CODE};
7use alloy_trie::{
8    nodes::{BranchNodeRef, ExtensionNode, ExtensionNodeRef, LeafNode, RlpNode, TrieNode},
9    Nibbles, TrieMask,
10};
11use core::fmt;
12
13/// Carries all information needed by a sparse trie to reveal a particular node.
14#[derive(Debug, Clone, PartialEq, Eq)]
15pub struct ProofTrieNodeV2 {
16    /// Path of the node.
17    pub path: Nibbles,
18    /// The node itself.
19    pub node: TrieNodeV2,
20    /// Tree and hash masks for the node, if known.
21    /// Both masks are always set together (from database branch nodes).
22    pub masks: Option<BranchNodeMasks>,
23}
24
25impl ProofTrieNodeV2 {
26    /// Creates an empty `ProofTrieNodeV2` with an empty root node. Useful as a placeholder when
27    /// taking a node out of a slice via [`core::mem::replace`].
28    pub fn empty() -> Self {
29        Self { path: Nibbles::default(), node: TrieNodeV2::EmptyRoot, masks: None }
30    }
31
32    /// Converts an iterator of `(path, TrieNode, masks)` tuples into `Vec<ProofTrieNodeV2>`,
33    /// merging extension nodes into their child branch nodes.
34    ///
35    /// The input **must** be sorted in depth-first order (children before parents) for extension
36    /// merging to work correctly.
37    pub fn from_sorted_trie_nodes(
38        iter: impl IntoIterator<Item = (Nibbles, TrieNode, Option<BranchNodeMasks>)>,
39    ) -> Vec<Self> {
40        let iter = iter.into_iter();
41        let mut result = Vec::with_capacity(iter.size_hint().0);
42
43        for (path, node, masks) in iter {
44            match node {
45                TrieNode::EmptyRoot => {
46                    result.push(Self { path, node: TrieNodeV2::EmptyRoot, masks });
47                }
48                TrieNode::Leaf(leaf) => {
49                    result.push(Self { path, node: TrieNodeV2::Leaf(leaf), masks });
50                }
51                TrieNode::Branch(branch) => {
52                    result.push(Self {
53                        path,
54                        node: TrieNodeV2::Branch(BranchNodeV2 {
55                            key: Nibbles::new(),
56                            branch_rlp_node: None,
57                            stack: branch.stack,
58                            state_mask: branch.state_mask,
59                        }),
60                        masks,
61                    });
62                }
63                TrieNode::Extension(ext) => {
64                    // In depth-first order, the child branch comes BEFORE the parent
65                    // extension. The child branch should be the last item we added to
66                    // result, at path extension.path + extension.key.
67                    let expected_branch_path = path.join(&ext.key);
68
69                    // Check if the last item in result is the child branch
70                    if let Some(last) = result.last_mut() &&
71                        last.path == expected_branch_path &&
72                        let TrieNodeV2::Branch(branch_v2) = &mut last.node
73                    {
74                        debug_assert!(
75                            branch_v2.key.is_empty(),
76                            "Branch at {:?} already has extension key {:?}",
77                            last.path,
78                            branch_v2.key
79                        );
80                        branch_v2.key = ext.key;
81                        branch_v2.branch_rlp_node = Some(ext.child);
82                        last.path = path;
83                    }
84
85                    // If we reach here, the extension's child is not a branch in the
86                    // result. This happens when the child branch is hashed (not revealed
87                    // in the proof). In V2 format, extension nodes are always combined
88                    // with their child branch, so we skip extension nodes whose child
89                    // isn't revealed.
90                }
91            }
92        }
93
94        result
95    }
96}
97
98/// Enum representing an MPT trie node.
99///
100/// This is a V2 representiation, differing from [`TrieNode`] in that branch and extension nodes are
101/// compressed into a single node.
102#[derive(PartialEq, Eq, Clone, Debug)]
103#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
104pub enum TrieNodeV2 {
105    /// Variant representing empty root node.
106    EmptyRoot,
107    /// Variant representing a [`BranchNodeV2`].
108    Branch(BranchNodeV2),
109    /// Variant representing a [`LeafNode`].
110    Leaf(LeafNode),
111    /// Variant representing an [`ExtensionNode`].
112    ///
113    /// This will only be used for extension nodes for which child is not inlined. This variant
114    /// will never be produced by proof workers that will always reveal a full path to a requested
115    /// leaf.
116    Extension(ExtensionNode),
117}
118
119impl Encodable for TrieNodeV2 {
120    fn encode(&self, out: &mut dyn bytes::BufMut) {
121        match self {
122            Self::EmptyRoot => {
123                out.put_u8(EMPTY_STRING_CODE);
124            }
125            Self::Leaf(leaf) => {
126                leaf.as_ref().encode(out);
127            }
128            Self::Branch(branch) => branch.encode(out),
129            Self::Extension(ext) => {
130                ext.encode(out);
131            }
132        }
133    }
134}
135
136impl Decodable for TrieNodeV2 {
137    fn decode(buf: &mut &[u8]) -> Result<Self, alloy_rlp::Error> {
138        match TrieNode::decode(buf)? {
139            TrieNode::EmptyRoot => Ok(Self::EmptyRoot),
140            TrieNode::Leaf(leaf) => Ok(Self::Leaf(leaf)),
141            TrieNode::Branch(branch) => Ok(Self::Branch(BranchNodeV2::new(
142                Default::default(),
143                branch.stack,
144                branch.state_mask,
145                None,
146            ))),
147            TrieNode::Extension(ext) => {
148                if ext.child.is_hash() {
149                    Ok(Self::Extension(ext))
150                } else {
151                    let Self::Branch(mut branch) = Self::decode(&mut ext.child.as_ref())? else {
152                        return Err(alloy_rlp::Error::Custom(
153                            "extension node child is not a branch",
154                        ));
155                    };
156
157                    branch.key = ext.key;
158
159                    Ok(Self::Branch(branch))
160                }
161            }
162        }
163    }
164}
165
166/// A branch node in an Ethereum Merkle Patricia Trie.
167///
168/// Branch node is a 17-element array consisting of 16 slots that correspond to each hexadecimal
169/// character and an additional slot for a value. We do exclude the node value since all paths have
170/// a fixed size.
171///
172/// This node also encompasses the possible parent extension node of a branch via the `key` field.
173#[derive(PartialEq, Eq, Clone, Default)]
174#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
175pub struct BranchNodeV2 {
176    /// The key for the branch's parent extension. if key is empty then the branch does not have a
177    /// parent extension.
178    pub key: Nibbles,
179    /// The collection of RLP encoded children.
180    pub stack: Vec<RlpNode>,
181    /// The bitmask indicating the presence of children at the respective nibble positions
182    pub state_mask: TrieMask,
183    /// [`RlpNode`] encoding of the branch node. Always provided when `key` is not empty (i.e this
184    /// is an extension node).
185    pub branch_rlp_node: Option<RlpNode>,
186}
187
188impl fmt::Debug for BranchNodeV2 {
189    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
190        f.debug_struct("BranchNode")
191            .field("key", &self.key)
192            .field("stack", &self.stack.iter().map(hex::encode).collect::<Vec<_>>())
193            .field("state_mask", &self.state_mask)
194            .field("branch_rlp_node", &self.branch_rlp_node)
195            .finish()
196    }
197}
198
199impl BranchNodeV2 {
200    /// Creates a new branch node with the given short key, stack, and state mask.
201    pub const fn new(
202        key: Nibbles,
203        stack: Vec<RlpNode>,
204        state_mask: TrieMask,
205        branch_rlp_node: Option<RlpNode>,
206    ) -> Self {
207        Self { key, stack, state_mask, branch_rlp_node }
208    }
209}
210
211impl Encodable for BranchNodeV2 {
212    fn encode(&self, out: &mut dyn bytes::BufMut) {
213        if self.key.is_empty() {
214            BranchNodeRef::new(&self.stack, self.state_mask).encode(out);
215            return;
216        }
217
218        let branch_rlp_node = self
219            .branch_rlp_node
220            .as_ref()
221            .expect("branch_rlp_node must always be present for extension nodes");
222
223        ExtensionNodeRef::new(&self.key, branch_rlp_node.as_slice()).encode(out);
224    }
225
226    fn length(&self) -> usize {
227        if self.key.is_empty() {
228            return BranchNodeRef::new(&self.stack, self.state_mask).length()
229        }
230
231        let branch_rlp_node = self
232            .branch_rlp_node
233            .as_ref()
234            .expect("branch_rlp_node must always be present for extension nodes");
235
236        ExtensionNodeRef::new(&self.key, branch_rlp_node.as_slice()).length()
237    }
238}
reth_trie_common/trie_node_v2.rs

reth_trie_common/
trie_node_v2.rs
