cleanup test

cmd/ethrex/Cargo.toml

@@ -22,6 +22,7 @@ ethrex-l2.workspace = true
 ethrex-l2-common.workspace = true
 ethrex-sdk.workspace = true
 ethrex-l2-rpc.workspace = true
+ethrex-trie.workspace = true
 tikv-jemallocator = { version = "0.6.0", optional = true, features = ["stats", "unprefixed_malloc_on_supported_platforms"] }
 bytes.workspace = true
 hex.workspace = true

cmd/ethrex/initializers.rs

@@ -21,7 +21,9 @@ use ethrex_p2p::{
     types::{Node, NodeRecord},
     utils::public_key_from_signing_key,
 };
+use ethrex_rlp::decode::RLPDecode;
 use ethrex_storage::{EngineType, Store};
+use ethrex_trie::EMPTY_TRIE_HASH;
 use local_ip_address::{local_ip, local_ipv6};
 use rand::rngs::OsRng;
 use secp256k1::SecretKey;
@@ -385,6 +387,44 @@ async fn set_sync_block(store: &Store) {
     }
 }
 
+/// Rollback the chain to the latest block we have the state for.
+async fn reset_to_head(store: &Store) -> eyre::Result<()> {
+    let trie = store.open_direct_state_trie(*EMPTY_TRIE_HASH)?;
+    let Some(root) = trie.db().get(Default::default())? else {
+        return Ok(());
+    };
+    let root = ethrex_trie::Node::decode(&root)?;
+    let state_root = root.compute_hash().finalize();
+
+    // TODO: store latest state metadata in the DB to avoid this loop
+    for block_number in (0..=store.get_latest_block_number().await?).rev() {
+        if let Some(header) = store.get_block_header(block_number)?
+            && header.state_root == state_root
+        {
+            info!("Resetting head to {block_number}");
+            let last_kept_block = block_number;
+            let mut block_to_delete = last_kept_block + 1;
+            while store
+                .get_canonical_block_hash(block_to_delete)
+                .await?
+                .is_some()
+            {
+                debug!("Deleting block {block_to_delete}");
+                store.remove_block(block_to_delete).await?;
+                block_to_delete += 1;
+            }
+            let last_kept_header = store
+                .get_block_header(last_kept_block)?
+                .ok_or_else(|| eyre::eyre!("Block number {} not found", last_kept_block))?;
+            store
+                .forkchoice_update(None, last_kept_block, last_kept_header.hash(), None, None)
+                .await?;
+            break;
+        }
+    }
+    Ok(())
+}
+
 pub async fn init_l1(
     opts: Options,
     log_filter_handler: Option<reload::Handle<EnvFilter, Registry>>,
@@ -406,6 +446,8 @@ pub async fn init_l1(
 
     let store = init_store(datadir, genesis).await;
 
+    reset_to_head(&store).await?;
+
     #[cfg(feature = "sync-test")]
     set_sync_block(&store).await;
 

cmd/ethrex/l2/command.rs

@@ -20,7 +20,7 @@ use ethrex_l2_sdk::call_contract;
 use ethrex_rpc::{
     EthClient, clients::beacon::BeaconClient, types::block_identifier::BlockIdentifier,
 };
-use ethrex_storage::{EngineType, Store, UpdateBatch};
+use ethrex_storage::{EngineType, Store};
 use ethrex_storage_rollup::StoreRollup;
 use eyre::OptionExt;
 use itertools::Itertools;
@@ -405,11 +405,8 @@ impl Command {
 
                 // Get genesis
                 let genesis_header = store.get_block_header(0)?.expect("Genesis block not found");
-                let genesis_block_hash = genesis_header.hash();
 
-                let mut new_trie = store
-                    .state_trie(genesis_block_hash)?
-                    .expect("Genesis block not found");
+                let mut current_state_root = genesis_header.state_root;
 
                 let mut last_block_number = 0;
                 let mut new_canonical_blocks = vec![];
@@ -433,37 +430,32 @@ impl Command {
                     let state_diff = StateDiff::decode(&blob)?;
 
                     // Apply all account updates to trie
-                    let account_updates = state_diff.to_account_updates(&new_trie)?;
+                    let trie = store.open_direct_state_trie(current_state_root)?;
+
+                    let account_updates = state_diff.to_account_updates(&trie)?;
+
+                    dbg!(&account_updates);
+
                     let account_updates_list = store
-                        .apply_account_updates_from_trie_batch(new_trie, account_updates.values())
+                        .apply_account_updates_from_trie_batch(trie, account_updates.values())
                         .await
                         .map_err(|e| format!("Error applying account updates: {e}"))
                         .unwrap();
 
-                    let (new_state_root, state_updates, accounts_updates) = (
-                        account_updates_list.state_trie_hash,
-                        account_updates_list.state_updates,
-                        account_updates_list.storage_updates,
-                    );
+                    store
+                        .open_direct_state_trie(current_state_root)?
+                        .db()
+                        .put_batch(account_updates_list.state_updates)?;
 
-                    let pseudo_update_batch = UpdateBatch {
-                        account_updates: state_updates,
-                        storage_updates: accounts_updates,
-                        blocks: vec![],
-                        receipts: vec![],
-                        code_updates: vec![],
-                    };
+                    current_state_root = account_updates_list.state_trie_hash;
 
                     store
-                        .store_block_updates(pseudo_update_batch)
-                        .await
-                        .map_err(|e| format!("Error storing trie updates: {e}"))
-                        .unwrap();
+                        .write_storage_trie_nodes_batch(account_updates_list.storage_updates)
+                        .await?;
 
-                    new_trie = store
-                        .open_state_trie(new_state_root)
-                        .map_err(|e| format!("Error opening new state trie: {e}"))
-                        .unwrap();
+                    store
+                        .write_account_code_batch(account_updates_list.code_updates)
+                        .await?;
 
                     // Get withdrawal hashes
                     let message_hashes = state_diff
@@ -479,10 +471,7 @@ impl Command {
                     // Note that its state_root is the root of new_trie.
                     let new_block = BlockHeader {
                         coinbase,
-                        state_root: new_trie
-                            .hash()
-                            .map_err(|e| format!("Error committing state: {e}"))
-                            .unwrap(),
+                        state_root: account_updates_list.state_trie_hash,
                         ..state_diff.last_header
                     };
 

crates/blockchain/blockchain.rs

@@ -416,7 +416,6 @@ impl Blockchain {
         };
 
         self.storage
-            .clone()
             .store_block_updates(update_batch)
             .await
             .map_err(|e| e.into())

crates/blockchain/smoke_test.rs

@@ -310,7 +310,7 @@ mod blockchain_integration_test {
         };
 
         // Create blockchain
-        let blockchain = Blockchain::default_with_store(store.clone().clone());
+        let blockchain = Blockchain::default_with_store(store.clone());
 
         let block = create_payload(&args, store, Bytes::new()).unwrap();
         let result = blockchain.build_payload(block).await.unwrap();

crates/blockchain/tracing.rs

@@ -135,7 +135,7 @@ async fn get_missing_state_parents(
         let Some(parent_block) = store.get_block_by_hash(parent_hash).await? else {
             return Err(ChainError::Custom("Parent Block not Found".to_string()));
         };
-        if store.contains_state_node(parent_block.header.state_root)? {
+        if store.has_state_root(parent_block.header.state_root)? {
             break;
         }
         parent_hash = parent_block.header.parent_hash;

crates/common/trie/db.rs

@@ -1,42 +1,58 @@
 use ethereum_types::H256;
 use ethrex_rlp::encode::RLPEncode;
 
-use crate::{Node, NodeHash, NodeRLP, Trie, error::TrieError};
+use crate::{Nibbles, Node, NodeRLP, Trie, error::TrieError};
 use std::{
     collections::BTreeMap,
     sync::{Arc, Mutex},
 };
 
+// Nibbles -> encoded node
+pub type NodeMap = Arc<Mutex<BTreeMap<Vec<u8>, Vec<u8>>>>;
+
 pub trait TrieDB: Send + Sync {
-    fn get(&self, key: NodeHash) -> Result<Option<Vec<u8>>, TrieError>;
-    fn put_batch(&self, key_values: Vec<(NodeHash, Vec<u8>)>) -> Result<(), TrieError>;
+    fn get(&self, key: Nibbles) -> Result<Option<Vec<u8>>, TrieError>;
+    fn put_batch(&self, key_values: Vec<(Nibbles, Vec<u8>)>) -> Result<(), TrieError>;
     // TODO: replace putbatch with this function.
-    fn put_batch_no_alloc(&self, key_values: &[(NodeHash, Node)]) -> Result<(), TrieError> {
+    fn put_batch_no_alloc(&self, key_values: &[(Nibbles, Node)]) -> Result<(), TrieError> {
         self.put_batch(
             key_values
                 .iter()
-                .map(|node| (node.0, node.1.encode_to_vec()))
+                .map(|node| (node.0.clone(), node.1.encode_to_vec()))
                 .collect(),
         )
     }
-    fn put(&self, key: NodeHash, value: Vec<u8>) -> Result<(), TrieError> {
+    fn put(&self, key: Nibbles, value: Vec<u8>) -> Result<(), TrieError> {
         self.put_batch(vec![(key, value)])
     }
 }
 
 /// InMemory implementation for the TrieDB trait, with get and put operations.
 #[derive(Default)]
 pub struct InMemoryTrieDB {
-    pub inner: Arc<Mutex<BTreeMap<NodeHash, Vec<u8>>>>,
+    inner: NodeMap,
+    prefix: Option<Nibbles>,
 }
 
 impl InMemoryTrieDB {
-    pub const fn new(map: Arc<Mutex<BTreeMap<NodeHash, Vec<u8>>>>) -> Self {
-        Self { inner: map }
+    pub const fn new(map: NodeMap) -> Self {
+        Self {
+            inner: map,
+            prefix: None,
+        }
+    }
+
+    pub const fn new_with_prefix(map: NodeMap, prefix: Nibbles) -> Self {
+        Self {
+            inner: map,
+            prefix: Some(prefix),
+        }
     }
+
     pub fn new_empty() -> Self {
         Self {
             inner: Default::default(),
+            prefix: None,
         }
     }
 
@@ -45,33 +61,56 @@ impl InMemoryTrieDB {
         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<(NodeHash, Vec<u8>)> = vec![];
-        embedded_root.commit(&mut hashed_nodes);
+        let mut hashed_nodes = vec![];
+        embedded_root.commit(Nibbles::default(), &mut hashed_nodes);
 
-        let hashed_nodes = hashed_nodes.into_iter().collect();
+        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),
+            None => path,
+        }
+    }
 }
 
 impl TrieDB for InMemoryTrieDB {
-    fn get(&self, key: NodeHash) -> Result<Option<Vec<u8>>, TrieError> {
+    fn get(&self, key: Nibbles) -> Result<Option<Vec<u8>>, TrieError> {
         Ok(self
             .inner
             .lock()
             .map_err(|_| TrieError::LockError)?
-            .get(&key)
+            .get(self.apply_prefix(key).as_ref())
             .cloned())
     }
 
-    fn put_batch(&self, key_values: Vec<(NodeHash, Vec<u8>)>) -> Result<(), TrieError> {
+    fn put_batch(&self, key_values: Vec<(Nibbles, Vec<u8>)>) -> Result<(), TrieError> {
         let mut db = self.inner.lock().map_err(|_| TrieError::LockError)?;
 
         for (key, value) in key_values {
-            db.insert(key, value);
+            let prefixed_key = self.apply_prefix(key);
+            db.insert(prefixed_key.into_vec(), value);
         }
 
         Ok(())
     }
 }
+
+pub fn nibbles_to_fixed_size(nibbles: Nibbles) -> [u8; 33] {
+    let node_hash_ref = nibbles.to_bytes();
+    let original_len = node_hash_ref.len();
+
+    let mut buffer = [0u8; 33];
+
+    // Encode the node as [node_path..., original_len]
+    buffer[32] = nibbles.len() as u8;
+    buffer[..original_len].copy_from_slice(&node_hash_ref);
+    buffer
+}

crates/common/trie/logger.rs

@@ -5,7 +5,7 @@ use std::{
 
 use ethrex_rlp::decode::RLPDecode;
 
-use crate::{Node, NodeHash, Trie, TrieDB, TrieError};
+use crate::{Nibbles, Node, Trie, TrieDB, TrieError};
 
 pub type TrieWitness = Arc<Mutex<HashSet<Vec<u8>>>>;
 
@@ -33,7 +33,7 @@ impl TrieLogger {
 }
 
 impl TrieDB for TrieLogger {
-    fn get(&self, key: NodeHash) -> Result<Option<Vec<u8>>, TrieError> {
+    fn get(&self, key: Nibbles) -> Result<Option<Vec<u8>>, TrieError> {
         let result = self.inner_db.get(key)?;
         if let Some(result) = result.as_ref()
             && let Ok(decoded) = Node::decode(result)
@@ -44,11 +44,11 @@ impl TrieDB for TrieLogger {
         Ok(result)
     }
 
-    fn put(&self, key: NodeHash, value: Vec<u8>) -> Result<(), TrieError> {
+    fn put(&self, key: Nibbles, value: Vec<u8>) -> Result<(), TrieError> {
         self.inner_db.put(key, value)
     }
 
-    fn put_batch(&self, key_values: Vec<(NodeHash, Vec<u8>)>) -> Result<(), TrieError> {
+    fn put_batch(&self, key_values: Vec<(Nibbles, Vec<u8>)>) -> Result<(), TrieError> {
         self.inner_db.put_batch(key_values)
     }
 }