feat(l1): make the database path-based

cmd/ethrex/initializers.rs

@@ -418,6 +418,8 @@ pub async fn init_l1(
         },
     );
 
+    regenerate_head_state(&store, &blockchain).await?;
+
     let signer = get_signer(datadir);
 
     let local_p2p_node = get_local_p2p_node(&opts, &signer);
@@ -484,3 +486,43 @@ pub async fn init_l1(
         local_node_record,
     ))
 }
+
+async fn regenerate_head_state(store: &Store, blockchain: &Arc<Blockchain>) -> eyre::Result<()> {
+    let head_block_number = store.get_latest_block_number().await?;
+    let Some(last_header) = store.get_block_header(head_block_number)? else {
+        unreachable!("Database is empty, genesis block should be present");
+    };
+
+    let mut current_last_header = last_header;
+
+    while !store.has_state_root(current_last_header.state_root)? {
+        let parent_number = current_last_header.number - 1;
+        debug!("Need to regenerate state for block {parent_number}");
+        let Some(parent_header) = store.get_block_header(parent_number)? else {
+            return Err(eyre::eyre!(
+                "Parent header for block {parent_number} not found"
+            ));
+        };
+        current_last_header = parent_header;
+    }
+
+    let last_state_number = current_last_header.number;
+
+    if last_state_number == head_block_number {
+        debug!("State is already up to date");
+        return Ok(());
+    }
+    info!("Regenerating state from block {last_state_number} to {head_block_number}");
+
+    for i in (last_state_number + 1)..=head_block_number {
+        debug!("Re-applying block {i} to regenerate state");
+
+        let block = store
+            .get_block_by_number(i)
+            .await?
+            .ok_or_else(|| eyre::eyre!("Block {i} not found"))?;
+        blockchain.add_block(block).await?;
+    }
+    info!("Finished regenerating state");
+    Ok(())
+}

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,30 @@ 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)?;
+
                     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 +469,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)
     }
 }