refactor(l1)!: unify trie node encoding

crates/blockchain/blockchain.rs

@@ -347,7 +347,7 @@ impl Blockchain {
         if used_trie_nodes.is_empty()
             && let Some(root) = root_node
         {
-            used_trie_nodes.push(root.encode_raw());
+            used_trie_nodes.push(root.encode_to_vec());
         }
 
         let mut needed_block_numbers = block_hashes.keys().collect::<Vec<_>>();

crates/common/trie/logger.rs

@@ -3,7 +3,7 @@ use std::{
     sync::{Arc, Mutex},
 };
 
-use ethrex_rlp::decode::RLPDecode;
+use ethrex_rlp::{decode::RLPDecode, encode::RLPEncode};
 
 use crate::{Nibbles, Node, Trie, TrieDB, TrieError};
 
@@ -39,7 +39,7 @@ impl TrieDB for TrieLogger {
             && let Ok(decoded) = Node::decode(result)
         {
             let mut lock = self.witness.lock().map_err(|_| TrieError::LockError)?;
-            lock.insert(decoded.encode_raw());
+            lock.insert(decoded.encode_to_vec());
         }
         Ok(result)
     }

crates/common/trie/node.rs

@@ -2,18 +2,10 @@ mod branch;
 mod extension;
 mod leaf;
 
-use std::{
-    array,
-    sync::{Arc, OnceLock},
-};
+use std::sync::{Arc, OnceLock};
 
 pub use branch::BranchNode;
-use ethrex_rlp::{
-    decode::{RLPDecode, decode_bytes},
-    encode::RLPEncode,
-    error::RLPDecodeError,
-    structs::Decoder,
-};
+use ethrex_rlp::{decode::RLPDecode, encode::RLPEncode};
 pub use extension::ExtensionNode;
 pub use leaf::LeafNode;
 
@@ -35,7 +27,7 @@ impl NodeRef {
         match *self {
             NodeRef::Node(ref node, _) => Ok(Some(node.as_ref().clone())),
             NodeRef::Hash(NodeHash::Inline((data, len))) => {
-                Ok(Some(Node::decode_raw(&data[..len as usize])?))
+                Ok(Some(Node::decode(&data[..len as usize])?))
             }
             NodeRef::Hash(hash) => db
                 .get(path)?
@@ -69,11 +61,13 @@ impl NodeRef {
                     }
                     Node::Leaf(_) => {}
                 }
-                let hash = *hash.get_or_init(|| node.compute_hash());
+                let mut buf = Vec::new();
+                node.encode(&mut buf);
+                let hash = *hash.get_or_init(|| NodeHash::from_encoded(&buf));
+                acc.push((path.clone(), buf));
                 if let Node::Leaf(leaf) = node.as_ref() {
                     acc.push((path.concat(&leaf.partial), leaf.value.clone()));
                 }
-                acc.push((path.clone(), node.encode_to_vec()));
 
                 *self = hash.into();
 
@@ -218,70 +212,6 @@ impl Node {
         }
     }
 
-    /// Encodes the node
-    pub fn encode_raw(&self) -> Vec<u8> {
-        match self {
-            Node::Branch(n) => n.encode_raw(),
-            Node::Extension(n) => n.encode_raw(),
-            Node::Leaf(n) => n.encode_raw(),
-        }
-    }
-
-    /// Decodes the node
-    pub fn decode_raw(rlp: &[u8]) -> Result<Self, RLPDecodeError> {
-        let mut rlp_items = vec![];
-        let mut decoder = Decoder::new(rlp)?;
-        let mut item;
-        // Get encoded fields
-        loop {
-            (item, decoder) = decoder.get_encoded_item()?;
-            rlp_items.push(item);
-            // Check if we reached the end or if we decoded more items than the ones we need
-            if decoder.is_done() || rlp_items.len() > 17 {
-                break;
-            }
-        }
-        // Deserialize into node depending on the available fields
-        Ok(match rlp_items.len() {
-            // Leaf or Extension Node
-            2 => {
-                let (path, _) = decode_bytes(&rlp_items[0])?;
-                let path = Nibbles::decode_compact(path);
-                if path.is_leaf() {
-                    // Decode as Leaf
-                    let (value, _) = decode_bytes(&rlp_items[1])?;
-                    LeafNode {
-                        partial: path,
-                        value: value.to_vec(),
-                    }
-                    .into()
-                } else {
-                    // Decode as Extension
-                    ExtensionNode {
-                        prefix: path,
-                        child: decode_child(&rlp_items[1]).into(),
-                    }
-                    .into()
-                }
-            }
-            // Branch Node
-            17 => {
-                let choices = array::from_fn(|i| decode_child(&rlp_items[i]).into());
-                let (value, _) = decode_bytes(&rlp_items[16])?;
-                BranchNode {
-                    choices,
-                    value: value.to_vec(),
-                }
-                .into()
-            }
-            n => {
-                return Err(RLPDecodeError::Custom(format!(
-                    "Invalid arg count for Node, expected 2 or 17, got {n}"
-                )));
-            }
-        })
-    }
-
     /// Computes the node's hash
     pub fn compute_hash(&self) -> NodeHash {
         match self {
@@ -291,11 +221,3 @@ impl Node {
         }
     }
 }
-
-fn decode_child(rlp: &[u8]) -> NodeHash {
-    match decode_bytes(rlp) {
-        Ok((hash, &[])) if hash.len() == 32 => NodeHash::from_slice(hash),
-        Ok((&[], &[])) => NodeHash::default(),
-        _ => NodeHash::from_slice(rlp),
-    }
-}

crates/common/trie/node/branch.rs

@@ -1,4 +1,4 @@
-use ethrex_rlp::structs::Encoder;
+use ethrex_rlp::encode::RLPEncode;
 
 use crate::{
     InconsistentTreeError, TrieDB, ValueRLP, error::TrieError, nibbles::Nibbles,
@@ -248,25 +248,7 @@ impl BranchNode {
 
     /// Computes the node's hash
     pub fn compute_hash(&self) -> NodeHash {
-        NodeHash::from_encoded_raw(&self.encode_raw())
-    }
-
-    /// Encodes the node
-    pub fn encode_raw(&self) -> Vec<u8> {
-        let mut buf = vec![];
-        let mut encoder = Encoder::new(&mut buf);
-        for child in self.choices.iter() {
-            match child.compute_hash() {
-                NodeHash::Hashed(hash) => encoder = encoder.encode_bytes(&hash.0),
-                child @ NodeHash::Inline(raw) if raw.1 != 0 => {
-                    encoder = encoder.encode_raw(child.as_ref())
-                }
-                _ => encoder = encoder.encode_bytes(&[]),
-            }
-        }
-        encoder = encoder.encode_bytes(&self.value);
-        encoder.finish();
-        buf
+        NodeHash::from_encoded(&self.encode_to_vec())
     }
 
     /// Traverses own subtrie until reaching the node containing `path`
@@ -279,7 +261,7 @@ impl BranchNode {
         node_path: &mut Vec<Vec<u8>>,
     ) -> Result<(), TrieError> {
         // Add self to node_path (if not inlined in parent)
-        let encoded = self.encode_raw();
+        let encoded = self.encode_to_vec();
         if encoded.len() >= 32 {
             node_path.push(encoded);
         };
@@ -307,6 +289,7 @@ impl BranchNode {
 #[cfg(test)]
 mod test {
     use ethereum_types::H256;
+    use ethrex_rlp::{decode::RLPDecode, encode::RLPEncode};
 
     use super::*;
 
@@ -654,7 +637,7 @@ mod test {
             }
         }
         .into();
-        assert_eq!(Node::decode_raw(&node.encode_raw()).unwrap(), node)
+        assert_eq!(Node::decode(&node.encode_to_vec()).unwrap(), node)
     }
 
     #[test]
@@ -670,7 +653,7 @@ mod test {
             }
         }
         .into();
-        assert_eq!(Node::decode_raw(&node.encode_raw()).unwrap(), node)
+        assert_eq!(Node::decode(&node.encode_to_vec()).unwrap(), node)
     }
 
     #[test]
@@ -696,6 +679,6 @@ mod test {
             } with_leaf { &[0x1] => vec![0x1] }
         }
         .into();
-        assert_eq!(Node::decode_raw(&node.encode_raw()).unwrap(), node)
+        assert_eq!(Node::decode(&node.encode_to_vec()).unwrap(), node)
     }
 }

crates/common/trie/node/extension.rs

@@ -1,4 +1,4 @@
-use ethrex_rlp::structs::Encoder;
+use ethrex_rlp::encode::RLPEncode;
 
 use crate::ValueRLP;
 use crate::nibbles::Nibbles;
@@ -190,16 +190,7 @@ impl ExtensionNode {
 
     /// Computes the node's hash
     pub fn compute_hash(&self) -> NodeHash {
-        NodeHash::from_encoded_raw(&self.encode_raw())
-    }
-
-    /// Encodes the node
-    pub fn encode_raw(&self) -> Vec<u8> {
-        let mut buf = vec![];
-        let mut encoder = Encoder::new(&mut buf).encode_bytes(&self.prefix.encode_compact());
-        encoder = self.child.compute_hash().encode(encoder);
-        encoder.finish();
-        buf
+        NodeHash::from_encoded(&self.encode_to_vec())
     }
 
     /// Traverses own subtrie until reaching the node containing `path`
@@ -212,7 +203,7 @@ impl ExtensionNode {
         node_path: &mut Vec<Vec<u8>>,
     ) -> Result<(), TrieError> {
         // Add self to node_path (if not inlined in parent)
-        let encoded = self.encode_raw();
+        let encoded = self.encode_to_vec();
         if encoded.len() >= 32 {
             node_path.push(encoded);
         };
@@ -236,6 +227,8 @@ impl ExtensionNode {
 
 #[cfg(test)]
 mod test {
+    use ethrex_rlp::{decode::RLPDecode, encode::RLPEncode};
+
     use super::*;
     use crate::{Trie, node::LeafNode, pmt_node};
 
@@ -544,7 +537,7 @@ mod test {
             } }
         }
         .into();
-        assert_eq!(Node::decode_raw(&node.encode_raw()).unwrap(), node)
+        assert_eq!(Node::decode(&node.encode_to_vec()).unwrap(), node)
     }
 
     #[test]
@@ -560,7 +553,7 @@ mod test {
         }
         .into();
 
-        assert_eq!(Node::decode_raw(&node.encode_raw()).unwrap(), node)
+        assert_eq!(Node::decode(&node.encode_to_vec()).unwrap(), node)
     }
 
     #[test]
@@ -585,6 +578,6 @@ mod test {
             } }
         }
         .into();
-        assert_eq!(Node::decode_raw(&node.encode_raw()).unwrap(), node)
+        assert_eq!(Node::decode(&node.encode_to_vec()).unwrap(), node)
     }
 }

crates/common/trie/node/leaf.rs

@@ -1,4 +1,4 @@
-use ethrex_rlp::structs::Encoder;
+use ethrex_rlp::encode::RLPEncode;
 
 use crate::{ValueRLP, error::TrieError, nibbles::Nibbles, node::BranchNode, node_hash::NodeHash};
 
@@ -115,22 +115,12 @@ impl LeafNode {
 
     /// Computes the node's hash
     pub fn compute_hash(&self) -> NodeHash {
-        NodeHash::from_encoded_raw(&self.encode_raw())
-    }
-
-    /// Encodes the node
-    pub fn encode_raw(&self) -> Vec<u8> {
-        let mut buf = vec![];
-        Encoder::new(&mut buf)
-            .encode_bytes(&self.partial.encode_compact())
-            .encode_bytes(&self.value)
-            .finish();
-        buf
+        NodeHash::from_encoded(&self.encode_to_vec())
     }
 
     /// Encodes the node and appends it to `node_path` if the encoded node is 32 or more bytes long
     pub fn get_path(&self, node_path: &mut Vec<Vec<u8>>) -> Result<(), TrieError> {
-        let encoded = self.encode_raw();
+        let encoded = self.encode_to_vec();
         if encoded.len() >= 32 {
             node_path.push(encoded);
         }
@@ -140,6 +130,8 @@ impl LeafNode {
 
 #[cfg(test)]
 mod test {
+    use ethrex_rlp::{decode::RLPDecode, encode::RLPEncode};
+
     use super::*;
     use crate::{Trie, pmt_node};
 
@@ -321,7 +313,7 @@ mod test {
             b"a comparatively long value".to_vec(),
         )
         .into();
-        assert_eq!(Node::decode_raw(&node.encode_raw()).unwrap(), node)
+        assert_eq!(Node::decode(&node.encode_to_vec()).unwrap(), node)
     }
 
     #[test]
@@ -331,7 +323,7 @@ mod test {
             vec![0x12, 0x34, 0x56, 0x78],
         )
         .into();
-        assert_eq!(Node::decode_raw(&node.encode_raw()).unwrap(), node)
+        assert_eq!(Node::decode(&node.encode_to_vec()).unwrap(), node)
     }
 
     #[test]
@@ -341,6 +333,6 @@ mod test {
             vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 20],
         )
         .into();
-        assert_eq!(Node::decode_raw(&node.encode_raw()).unwrap(), node)
+        assert_eq!(Node::decode(&node.encode_to_vec()).unwrap(), node)
     }
 }

crates/common/trie/node_hash.rs

@@ -24,7 +24,7 @@ impl AsRef<[u8]> for NodeHash {
 
 impl NodeHash {
     /// Returns the `NodeHash` of an encoded node (encoded using the NodeEncoder)
-    pub fn from_encoded_raw(encoded: &[u8]) -> NodeHash {
+    pub fn from_encoded(encoded: &[u8]) -> NodeHash {
         if encoded.len() >= 32 {
             let hash = Keccak256::new_with_prefix(encoded).finalize();
             NodeHash::Hashed(H256::from_slice(hash.as_slice()))
@@ -35,7 +35,7 @@ impl NodeHash {
 
     /// Converts a slice of an already hashed data (in case it's not inlineable) to a NodeHash.
     /// Panics if the slice is over 32 bytes
-    /// If you need to hash it in case its len >= 32 see `from_encoded_raw`
+    /// If you need to hash it in case its len >= 32 see `from_encoded`
     pub(crate) fn from_slice(slice: &[u8]) -> NodeHash {
         match slice.len() {
             0..32 => {