chore(l2): replace RLP decoding of tries with rkyv in stateless execution

crates/blockchain/blockchain.rs

@@ -7,6 +7,7 @@ mod smoke_test;
 pub mod tracing;
 pub mod vm;
 
+use ethrex_common::rkyv_utils::SizedNode;
 use ::tracing::{debug, info};
 use constants::{MAX_INITCODE_SIZE, MAX_TRANSACTION_DATA_SIZE, POST_OSAKA_GAS_LIMIT_CAP};
 use error::MempoolError;
@@ -397,7 +398,7 @@ impl Blockchain {
                     ChainError::WitnessGeneration("Failed to lock storage trie witness".to_string())
                 })?;
                 let witness = std::mem::take(&mut *witness);
-                let witness = witness.into_iter().collect::<Vec<_>>();
+                let witness = witness.into_values().collect::<Vec<_>>();
                 used_trie_nodes.extend_from_slice(&witness);
                 touched_account_storage_slots.entry(address).or_default();
             }
@@ -426,20 +427,22 @@ impl Blockchain {
                 })?
                 .iter()
             {
-                accumulated_state_trie_witness.insert(state_trie_witness.clone());
+                accumulated_state_trie_witness
+                    .insert(state_trie_witness.0.clone(), state_trie_witness.1.clone());
             }
 
             current_trie_witness = new_state_trie_witness;
         }
 
-        used_trie_nodes
-            .extend_from_slice(&Vec::from_iter(accumulated_state_trie_witness.into_iter()));
+        used_trie_nodes.extend_from_slice(&Vec::from_iter(
+            accumulated_state_trie_witness.into_values(),
+        ));
 
         // If the witness is empty at least try to store the root
         if used_trie_nodes.is_empty()
             && let Some(root) = root_node
         {
-            used_trie_nodes.push(root.encode_to_vec());
+            used_trie_nodes.push((*root).clone());
         }
 
         let mut needed_block_numbers = block_hashes.keys().collect::<Vec<_>>();
@@ -479,7 +482,7 @@ impl Blockchain {
 
         let chain_config = self.storage.get_chain_config();
 
-        let nodes = used_trie_nodes.into_iter().collect::<Vec<_>>();
+        let nodes = used_trie_nodes.into_iter().map(SizedNode::from).collect::<Vec<_>>();
 
         let mut keys = Vec::new();
 
@@ -497,6 +500,7 @@ impl Blockchain {
             chain_config,
             nodes,
             keys,
+            guest_program_state: None
         })
     }
 

crates/common/Cargo.toml

@@ -31,6 +31,7 @@ lazy_static.workspace = true
 rayon.workspace = true
 url.workspace = true
 rkyv.workspace = true
+serde_bytes = "0.11.9"
 
 [dev-dependencies]
 hex-literal.workspace = true

crates/common/rkyv_utils.rs

@@ -1,5 +1,9 @@
 use bytes::Bytes;
 use ethereum_types::{Bloom, H160, H256, U256};
+use ethrex_trie::{
+    Nibbles, Node, NodeHash, NodeRef, ValueRLP,
+    node::{BranchNode, ExtensionNode, LeafNode},
+};
 use rkyv::{
     Archive, Archived, Deserialize, Serialize,
     rancor::{Fallible, Source},
@@ -258,6 +262,79 @@ where
     }
 }
 
+/// Node representation that only references childs via hashes instead of pointers.
+#[derive(Archive, serde::Serialize, serde::Deserialize, Serialize, Deserialize, Clone)]
+pub enum SizedNode {
+    Branch {
+        choices: [NodeHash; 16],
+        value: ValueRLP, // 0
+    },
+    Extension {
+        child: NodeHash,
+        prefix: [u8; 32], // max 32 bytes
+    },
+    Leaf {
+        partial: [u8; 32], // max 32 bytes
+        #[serde(with = "serde_bytes")]
+        value: [u8; 108], // max 108 bytes for an account state
+    },
+}
+
+impl From<SizedNode> for Node {
+    fn from(value: SizedNode) -> Self {
+        match value {
+            SizedNode::Branch { choices, value } => {
+                let mut ref_choices = Vec::with_capacity(16);
+                for choice in choices {
+                    ref_choices.push(NodeRef::Hash(choice));
+                }
+                let choices = ref_choices.try_into().unwrap();
+                Node::Branch(Box::new(BranchNode { choices, value }))
+            }
+            SizedNode::Extension { child, prefix } => {
+                let child = NodeRef::Hash(child);
+                let prefix = Nibbles::from_bytes(&prefix);
+                Node::Extension(ExtensionNode { child, prefix })
+            }
+            SizedNode::Leaf { partial, value } =>  {
+                let partial = Nibbles::from_bytes(&partial);
+                let value = Vec::from(value);
+                Node::Leaf(LeafNode { partial, value })
+            }
+        }
+    }
+}
+
+impl From<Node> for SizedNode {
+    fn from(value: Node) -> Self {
+        match value {
+            Node::Branch(n) => {
+                let BranchNode { choices, value } = *n;
+                let mut ref_choices = Vec::with_capacity(16);
+                for choice in choices {
+                    ref_choices.push(choice.compute_hash());
+                }
+                let choices = ref_choices.try_into().unwrap();
+                SizedNode::Branch { choices, value }
+            }
+            Node::Extension(ExtensionNode { child, prefix }) => {
+                let child = child.compute_hash();
+                let mut prefix = prefix.into_vec();
+                prefix.resize(32, 0);
+                SizedNode::Extension { child, prefix: prefix.try_into().unwrap() }
+            }
+            Node::Leaf(LeafNode { partial, mut value }) => {
+                let mut partial = partial.into_vec();
+                value.resize(108, 0);
+                partial.resize(32, 0);
+                let partial = partial.try_into().unwrap();
+                let value = value.try_into().unwrap();
+                SizedNode::Leaf { partial, value }
+            },
+        }
+    }
+}
+
 #[cfg(test)]
 mod test {
     use ethereum_types::{H160, H256};

crates/common/trie/Cargo.toml

@@ -23,6 +23,7 @@ smallvec = { version = "1.10.0", features = ["const_generics", "union"] }
 digest = "0.10.6"
 lazy_static.workspace = true
 crossbeam.workspace = true
+rkyv.workspace = true
 
 [features]
 default = []

crates/common/trie/db.rs

@@ -59,23 +59,6 @@ impl InMemoryTrieDB {
         }
     }
 
-    pub fn from_nodes(
-        root_hash: H256,
-        state_nodes: &BTreeMap<H256, NodeRLP>,
-    ) -> Result<Self, TrieError> {
-        let mut embedded_root = Trie::get_embedded_root(state_nodes, root_hash)?;
-        let mut hashed_nodes = vec![];
-        embedded_root.commit(Nibbles::default(), &mut hashed_nodes);
-
-        let hashed_nodes = hashed_nodes
-            .into_iter()
-            .map(|(k, v)| (k.into_vec(), v))
-            .collect();
-
-        let in_memory_trie = Arc::new(Mutex::new(hashed_nodes));
-        Ok(Self::new(in_memory_trie))
-    }
-
     fn apply_prefix(&self, path: Nibbles) -> Nibbles {
         match &self.prefix {
             Some(prefix) => prefix.concat(&path),

crates/common/trie/logger.rs

@@ -1,29 +1,29 @@
 use std::{
-    collections::HashSet,
+    collections::HashMap,
     sync::{Arc, Mutex},
 };
 
-use ethrex_rlp::{decode::RLPDecode, encode::RLPEncode};
+use ethrex_rlp::decode::RLPDecode;
 
-use crate::{Nibbles, Node, Trie, TrieDB, TrieError};
+use crate::{Nibbles, Node, NodeHash, Trie, TrieDB, TrieError};
 
-pub type TrieWitness = Arc<Mutex<HashSet<Vec<u8>>>>;
+pub type TrieWitness = Arc<Mutex<HashMap<NodeHash, Node>>>;
 
 pub struct TrieLogger {
     inner_db: Box<dyn TrieDB>,
     witness: TrieWitness,
 }
 
 impl TrieLogger {
-    pub fn get_witness(&self) -> Result<HashSet<Vec<u8>>, TrieError> {
+    pub fn get_witness(&self) -> Result<HashMap<NodeHash, Node>, TrieError> {
         let lock = self.witness.lock().map_err(|_| TrieError::LockError)?;
         Ok(lock.clone())
     }
 
     pub fn open_trie(trie: Trie) -> (TrieWitness, Trie) {
         let root = trie.hash_no_commit();
         let db = trie.db;
-        let witness = Arc::new(Mutex::new(HashSet::new()));
+        let witness = Arc::new(Mutex::new(HashMap::new()));
         let logger = TrieLogger {
             inner_db: db,
             witness: witness.clone(),
@@ -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_to_vec());
+            lock.insert(decoded.compute_hash(), decoded);
         }
         Ok(result)
     }

crates/common/trie/nibbles.rs

@@ -1,16 +1,18 @@
 use std::{cmp, mem};
+use rkyv::{Archive, Deserialize as RDeserialize, Serialize as RSerialize};
 
 use ethrex_rlp::{
     decode::RLPDecode,
     encode::RLPEncode,
     error::RLPDecodeError,
     structs::{Decoder, Encoder},
 };
+use serde::{Deserialize, Serialize};
 
 // TODO: move path-tracking logic somewhere else
 // PERF: try using a stack-allocated array
 /// Struct representing a list of nibbles (half-bytes)
-#[derive(Debug, Clone, Default)]
+#[derive(Debug, Clone, Default, RDeserialize, RSerialize, Archive, Serialize, Deserialize)]
 pub struct Nibbles {
     data: Vec<u8>,
     /// Parts of the path that have already been consumed (used for tracking

crates/common/trie/node.rs

@@ -2,12 +2,17 @@ mod branch;
 mod extension;
 mod leaf;
 
-use std::sync::{Arc, OnceLock};
+use std::{
+    collections::BTreeMap,
+    sync::{Arc, OnceLock},
+};
 
 pub use branch::BranchNode;
+use ethereum_types::H256;
 use ethrex_rlp::{decode::RLPDecode, encode::RLPEncode};
 pub use extension::ExtensionNode;
 pub use leaf::LeafNode;
+use serde::{Deserialize, Deserializer, Serialize, Serializer};
 
 use crate::{TrieDB, error::TrieError, nibbles::Nibbles};
 
@@ -130,6 +135,42 @@ impl NodeRef {
             hash.take();
         }
     }
+
+    pub fn resolve_subtrie(
+        &mut self,
+        all_nodes: &mut BTreeMap<H256, Arc<Node>>,
+    ) -> Result<(), TrieError> {
+        // this will be a reference in the future, to avoid the memcpy here
+        let finalized_hash = self.compute_hash().finalize();
+
+        // remove and later re-insert instead of get_mut because we can't keep the mut borrow
+        // maybe there is a better solution, this remove might be expensive (memmove)
+        let Some(mut node_ref) = all_nodes.remove(&finalized_hash) else {
+            return Ok(());
+        };
+
+        match Arc::make_mut(&mut node_ref) {
+            Node::Branch(node) => {
+                for choice in &mut node.choices {
+                    if let NodeRef::Hash(hash) = choice
+                        && hash.is_valid()
+                    {
+                        choice.resolve_subtrie(all_nodes)?;
+                    }
+                }
+            }
+            Node::Extension(node) => {
+                if let NodeRef::Hash(_) = node.child {
+                    node.child.resolve_subtrie(all_nodes)?;
+                }
+            }
+            Node::Leaf(_) => {}
+        };
+
+        *self = node_ref.clone().into();
+        all_nodes.insert(finalized_hash, node_ref);
+        Ok(())
+    }
 }
 
 impl Default for NodeRef {
@@ -144,6 +185,12 @@ impl From<Node> for NodeRef {
     }
 }
 
+impl From<Arc<Node>> for NodeRef {
+    fn from(value: Arc<Node>) -> Self {
+        Self::Node(value, OnceLock::new())
+    }
+}
+
 impl From<NodeHash> for NodeRef {
     fn from(value: NodeHash) -> Self {
         Self::Hash(value)
@@ -173,8 +220,8 @@ impl From<NodeHash> for ValueOrHash {
     }
 }
 
-/// A Node in an Ethereum Compatible Patricia Merkle Trie
 #[derive(Debug, Clone, PartialEq)]
+/// A Node in an Ethereum Compatible Patricia Merkle Trie
 pub enum Node {
     Branch(Box<BranchNode>),
     Extension(ExtensionNode),
@@ -289,3 +336,23 @@ pub enum NodeRemoveResult {
     Mutated,
     New(Node),
 }
+
+impl Serialize for Node {
+    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+    where
+        S: Serializer,
+    {
+        let rlp = self.encode_to_vec();
+        serializer.serialize_bytes(&rlp)
+    }
+}
+impl<'de> Deserialize<'de> for Node {
+    fn deserialize<D>(d: D) -> Result<Node, D::Error>
+    where
+        D: Deserializer<'de>,
+    {
+        let rlp: Vec<u8> = Deserialize::deserialize(d)?;
+        Node::decode(&rlp)
+            .map_err(|e| serde::de::Error::custom(format!("could not decode rlp bytes: {e}")))
+    }
+}

crates/common/trie/node/leaf.rs

@@ -13,7 +13,7 @@ use crate::{
 use super::{ExtensionNode, Node, ValueOrHash};
 /// Leaf Node of an an Ethereum Compatible Patricia Merkle Trie
 /// Contains the node's hash, value & path
-#[derive(Debug, Clone, Default, PartialEq)]
+#[derive(Debug, Clone, Default, PartialEq, Eq)]
 pub struct LeafNode {
     pub partial: Nibbles,
     pub value: ValueRLP,