alloy_trie/nodes/
extension.rs

1use super::{super::Nibbles, encode_path_leaf, unpack_path_to_nibbles, RlpNode};
2use alloy_primitives::{hex, Bytes};
3use alloy_rlp::{length_of_length, BufMut, Decodable, Encodable, Header};
4use core::fmt;
5
6#[allow(unused_imports)]
7use alloc::vec::Vec;
8
9/// An extension node in an Ethereum Merkle Patricia Trie.
10///
11/// An intermediate node that exists solely to compress the trie's paths. It contains a path segment
12/// (a shared prefix of keys) and a single child pointer. Essentially, an extension node can be
13/// thought of as a shortcut within the trie to reduce its overall depth.
14///
15/// The purpose of an extension node is to optimize the trie structure by collapsing multiple nodes
16/// with a single child into one node. This simplification reduces the space and computational
17/// complexity when performing operations on the trie.
18#[derive(PartialEq, Eq, Clone)]
19#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
20pub struct ExtensionNode {
21    /// The key for this extension node.
22    pub key: Nibbles,
23    /// A pointer to the child node.
24    pub child: RlpNode,
25}
26
27impl fmt::Debug for ExtensionNode {
28    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
29        f.debug_struct("ExtensionNode")
30            .field("key", &self.key)
31            .field("child", &hex::encode(&self.child))
32            .finish()
33    }
34}
35
36impl Encodable for ExtensionNode {
37    #[inline]
38    fn encode(&self, out: &mut dyn BufMut) {
39        self.as_ref().encode(out)
40    }
41
42    #[inline]
43    fn length(&self) -> usize {
44        self.as_ref().length()
45    }
46}
47
48impl Decodable for ExtensionNode {
49    fn decode(buf: &mut &[u8]) -> alloy_rlp::Result<Self> {
50        let mut bytes = Header::decode_bytes(buf, true)?;
51        let encoded_key = Bytes::decode(&mut bytes)?;
52        if encoded_key.is_empty() {
53            return Err(alloy_rlp::Error::Custom("extension node key empty"));
54        }
55
56        // Retrieve first byte. If it's [Some], then the nibbles are odd.
57        let first = match encoded_key[0] & 0xf0 {
58            Self::ODD_FLAG => Some(encoded_key[0] & 0x0f),
59            Self::EVEN_FLAG => None,
60            _ => return Err(alloy_rlp::Error::Custom("node is not extension")),
61        };
62
63        let key = unpack_path_to_nibbles(first, &encoded_key[1..]);
64        let child = RlpNode::from_raw_rlp(bytes)?;
65        Ok(Self { key, child })
66    }
67}
68
69impl ExtensionNode {
70    /// The flag representing the even number of nibbles in the extension key.
71    pub const EVEN_FLAG: u8 = 0x00;
72
73    /// The flag representing the odd number of nibbles in the extension key.
74    pub const ODD_FLAG: u8 = 0x10;
75
76    /// Creates a new extension node with the given key and a pointer to the child.
77    pub const fn new(key: Nibbles, child: RlpNode) -> Self {
78        Self { key, child }
79    }
80
81    /// Return extension node as [ExtensionNodeRef].
82    pub fn as_ref(&self) -> ExtensionNodeRef<'_> {
83        ExtensionNodeRef { key: &self.key, child: &self.child }
84    }
85}
86
87/// Reference to the extension node. See [ExtensionNode] from more information.
88pub struct ExtensionNodeRef<'a> {
89    /// The key for this extension node.
90    pub key: &'a Nibbles,
91    /// A pointer to the child node.
92    pub child: &'a [u8],
93}
94
95impl fmt::Debug for ExtensionNodeRef<'_> {
96    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
97        f.debug_struct("ExtensionNodeRef")
98            .field("key", &self.key)
99            .field("node", &hex::encode(self.child))
100            .finish()
101    }
102}
103
104impl Encodable for ExtensionNodeRef<'_> {
105    #[inline]
106    fn encode(&self, out: &mut dyn BufMut) {
107        Header { list: true, payload_length: self.rlp_payload_length() }.encode(out);
108        encode_path_leaf(self.key, false).as_slice().encode(out);
109        // Pointer to the child is already RLP encoded.
110        out.put_slice(self.child);
111    }
112
113    #[inline]
114    fn length(&self) -> usize {
115        let payload_length = self.rlp_payload_length();
116        payload_length + length_of_length(payload_length)
117    }
118}
119
120impl<'a> ExtensionNodeRef<'a> {
121    /// Creates a new extension node with the given key and a pointer to the child.
122    #[inline]
123    pub const fn new(key: &'a Nibbles, child: &'a [u8]) -> Self {
124        Self { key, child }
125    }
126
127    /// RLP-encodes the node and returns either `rlp(node)` or `rlp(keccak(rlp(node)))`.
128    #[inline]
129    pub fn rlp(&self, rlp: &mut Vec<u8>) -> RlpNode {
130        self.encode(rlp);
131        RlpNode::from_rlp(rlp)
132    }
133
134    /// Returns the length of RLP encoded fields of extension node.
135    #[inline]
136    fn rlp_payload_length(&self) -> usize {
137        let mut encoded_key_len = self.key.len() / 2 + 1;
138        // For extension nodes the first byte cannot be greater than 0x80.
139        if encoded_key_len != 1 {
140            encoded_key_len += length_of_length(encoded_key_len);
141        }
142        encoded_key_len + self.child.len()
143    }
144}
145
146#[cfg(test)]
147mod tests {
148    use super::*;
149
150    #[test]
151    fn rlp_extension_node_roundtrip() {
152        let nibble = Nibbles::from_nibbles_unchecked(hex!("0604060f"));
153        let val = hex!("76657262");
154        let mut child = vec![];
155        val.to_vec().as_slice().encode(&mut child);
156        let extension = ExtensionNode::new(nibble, RlpNode::from_raw(&child).unwrap());
157        let rlp = extension.as_ref().rlp(&mut vec![]);
158        assert_eq!(rlp.as_ref(), hex!("c98300646f8476657262"));
159        assert_eq!(ExtensionNode::decode(&mut &rlp[..]).unwrap(), extension);
160    }
161}