ethrex_trie/
rlp.rs

1#[cfg(not(feature = "std"))]
2use alloc::{string::ToString, vec::Vec};
3use core::array;
4
5// Contains RLP encoding and decoding implementations for Trie Nodes
6// This encoding is only used to store the nodes in the DB, it is not the encoding used for hash computation
7use ethrex_rlp::{
8    constants::RLP_NULL,
9    decode::{RLPDecode, decode_bytes},
10    encode::{RLPEncode, encode_length},
11    error::RLPDecodeError,
12    structs::{Decoder, Encoder},
13};
14
15use ethrex_crypto::NativeCrypto;
16
17use super::node::{BranchNode, ExtensionNode, LeafNode, Node};
18use crate::{Nibbles, NodeHash};
19
20// SAFETY: `NativeCrypto` is used here instead of a `&dyn Crypto` parameter because
21// `RLPEncode` is a fixed trait signature that cannot accept extra parameters.
22// This is safe in the `commit()` path: `NodeRef::commit()` recursively populates
23// child `OnceLock` hashes before calling `encode()`, so `compute_hash_ref` returns
24// cached values without invoking keccak. If `encode()` were called on uncommitted
25// nodes (e.g. from `put_batch_no_alloc`), `NativeCrypto` would be used and the
26// result stored in the `OnceLock` — but this only happens in native storage paths
27// where `NativeCrypto` is the correct provider.
28impl RLPEncode for BranchNode {
29    fn encode(&self, buf: &mut dyn bytes::BufMut) {
30        // Resolve each child's hash once: the length pass and the encode pass
31        // both needed it, so a 16-choice branch was paying 32 resolutions.
32        let hashes: [&NodeHash; 16] =
33            array::from_fn(|i| self.choices[i].compute_hash_ref(&NativeCrypto));
34
35        let value_len = <[u8] as RLPEncode>::length(&self.value);
36        let payload_len = hashes
37            .iter()
38            .fold(value_len, |acc, hash| acc + RLPEncode::length(*hash));
39
40        encode_length(payload_len, buf);
41        for hash in hashes {
42            match hash {
43                NodeHash::Hashed(hash) => hash.0.encode(buf),
44                NodeHash::Inline((_, 0)) => buf.put_u8(RLP_NULL),
45                NodeHash::Inline((encoded, len)) => buf.put_slice(&encoded[..*len as usize]),
46            }
47        }
48        <[u8] as RLPEncode>::encode(&self.value, buf);
49    }
50
51    // Duplicated to prealloc the buffer and avoid calculating the payload length twice
52    fn encode_to_vec(&self) -> Vec<u8> {
53        let value_len = <[u8] as RLPEncode>::length(&self.value);
54        let choices_len = self.choices.iter().fold(0, |acc, child| {
55            acc + RLPEncode::length(child.compute_hash_ref(&NativeCrypto))
56        });
57        let payload_len = choices_len + value_len;
58
59        let mut buf: Vec<u8> = Vec::with_capacity(payload_len + 3); // 3 byte prefix headroom
60
61        encode_length(payload_len, &mut buf);
62        for child in self.choices.iter() {
63            match child.compute_hash_ref(&NativeCrypto) {
64                NodeHash::Hashed(hash) => hash.0.encode(&mut buf),
65                NodeHash::Inline((_, 0)) => buf.push(RLP_NULL),
66                NodeHash::Inline((encoded, len)) => {
67                    buf.extend_from_slice(&encoded[..*len as usize])
68                }
69            }
70        }
71        <[u8] as RLPEncode>::encode(&self.value, &mut buf);
72
73        buf
74    }
75}
76
77impl BranchNode {
78    /// Concrete-typed sibling of `<Self as RLPEncode>::encode`, appending into a
79    /// `Vec<u8>` rather than a `&mut dyn BufMut`.
80    ///
81    /// `RLPEncode::encode` must take a trait object, so hashing a branch paid a
82    /// vtable dispatch for each of its ~41 `put_u8`/`put_slice` calls. Hashing is
83    /// the hot consumer (every `memoize_hashes` walk re-encodes each dirty
84    /// branch), so it gets a monomorphic path; `RLPEncode::encode` is untouched
85    /// for every other caller.
86    pub fn encode_into_vec(&self, buf: &mut Vec<u8>) {
87        let hashes: [&NodeHash; 16] =
88            array::from_fn(|i| self.choices[i].compute_hash_ref(&NativeCrypto));
89
90        let value_len = <[u8] as RLPEncode>::length(&self.value);
91        let payload_len = hashes
92            .iter()
93            .fold(value_len, |acc, hash| acc + RLPEncode::length(*hash));
94
95        encode_length(payload_len, buf);
96        for hash in hashes {
97            match hash {
98                NodeHash::Hashed(hash) => hash.0.encode(&mut *buf),
99                NodeHash::Inline((_, 0)) => buf.push(RLP_NULL),
100                NodeHash::Inline((encoded, len)) => {
101                    buf.extend_from_slice(&encoded[..*len as usize])
102                }
103            }
104        }
105        <[u8] as RLPEncode>::encode(&self.value, buf);
106    }
107}
108
109impl RLPEncode for ExtensionNode {
110    fn encode(&self, buf: &mut dyn bytes::BufMut) {
111        let mut encoder = Encoder::new(buf).encode_bytes(&self.prefix.encode_compact());
112        encoder = self.child.compute_hash(&NativeCrypto).encode(encoder);
113        encoder.finish();
114    }
115}
116
117impl RLPEncode for LeafNode {
118    fn encode(&self, buf: &mut dyn bytes::BufMut) {
119        Encoder::new(buf)
120            .encode_bytes(&self.partial.encode_compact())
121            .encode_bytes(&self.value)
122            .finish()
123    }
124}
125
126impl RLPEncode for Node {
127    fn encode(&self, buf: &mut dyn bytes::BufMut) {
128        match self {
129            Node::Branch(n) => n.encode(buf),
130            Node::Extension(n) => n.encode(buf),
131            Node::Leaf(n) => n.encode(buf),
132        }
133    }
134}
135
136impl RLPDecode for Node {
137    fn decode_unfinished(rlp: &[u8]) -> Result<(Self, &[u8]), RLPDecodeError> {
138        let mut rlp_items_len = 0;
139        let mut rlp_items: [Option<&[u8]>; 17] = Default::default();
140        let mut decoder = Decoder::new(rlp)?;
141        let mut item;
142        // Get encoded fields
143
144        // Check if we reached the end or if we decoded more items than the ones we need
145        while !decoder.is_done() && rlp_items_len < 17 {
146            (item, decoder) = decoder.get_encoded_item_ref()?;
147            rlp_items[rlp_items_len] = Some(item);
148            rlp_items_len += 1;
149        }
150        if !decoder.is_done() {
151            return Err(RLPDecodeError::Custom(
152                "Invalid arg count for Node, expected 2 or 17, got more than 17".to_string(),
153            ));
154        }
155        // Deserialize into node depending on the available fields
156        let node = match rlp_items_len {
157            // Leaf or Extension Node
158            2 => {
159                let (path, _) = decode_bytes(rlp_items[0].expect("we already checked the length"))?;
160                let path = Nibbles::decode_compact(path);
161                if path.is_leaf() {
162                    // Decode as Leaf
163                    let (value, _) =
164                        decode_bytes(rlp_items[1].expect("we already checked the length"))?;
165                    LeafNode {
166                        partial: path,
167                        value: value.to_vec(),
168                    }
169                    .into()
170                } else {
171                    // Decode as Extension
172                    ExtensionNode {
173                        prefix: path,
174                        child: decode_child(rlp_items[1].expect("we already checked the length"))
175                            .into(),
176                    }
177                    .into()
178                }
179            }
180            // Branch Node
181            17 => {
182                let choices = array::from_fn(|i| {
183                    decode_child(rlp_items[i].expect("we already checked the length")).into()
184                });
185                let (value, _) =
186                    decode_bytes(rlp_items[16].expect("we already checked the length"))?;
187                BranchNode {
188                    choices,
189                    value: value.to_vec(),
190                }
191                .into()
192            }
193            n => {
194                return Err(RLPDecodeError::Custom(format!(
195                    "Invalid arg count for Node, expected 2 or 17, got {n}"
196                )));
197            }
198        };
199        Ok((node, decoder.finish()?))
200    }
201}
202
203fn decode_child(rlp: &[u8]) -> NodeHash {
204    match decode_bytes(rlp) {
205        Ok((hash, &[])) if hash.len() == 32 => NodeHash::from_slice(hash),
206        Ok((&[], &[])) => NodeHash::default(),
207        _ => NodeHash::from_slice(rlp),
208    }
209}
ethrex_trie/rlp.rs

ethrex_trie/
rlp.rs
