TrieDB Reth Integration Changes

Cargo.toml

@@ -16,6 +16,7 @@ metrics-derive = "0.1.0"
 metrics = "0.24.1"
 zerocopy = { version = "0.8.24", features = ["derive"] }
 reth-trie-common = { git = "https://github.com/paradigmxyz/reth", tag = "v1.3.8" }
+reth-primitives-traits = { git = "https://github.com/paradigmxyz/reth", tag = "v1.3.8" }
 parking_lot = { version = "0.12.3", features = ["send_guard"] }
 fxhash = "0.2.1"
 static_assertions = "1.1.0"

benches/benchmark_common.rs

@@ -1,6 +1,7 @@
 use alloy_primitives::{Address, StorageKey, StorageValue, U256};
 use alloy_trie::{EMPTY_ROOT_HASH, KECCAK_EMPTY};
 use rand::prelude::*;
+use std::sync::Arc;
 use tempdir::TempDir;
 use triedb::{
     account::Account,
@@ -15,10 +16,10 @@ pub fn generate_random_address(rng: &mut StdRng) -> AddressPath {
     AddressPath::for_address(addr)
 }
 
-pub fn setup_database(size: usize) -> (TempDir, Database) {
+pub fn setup_database(size: usize) -> (TempDir, Arc<Database>) {
     let dir = TempDir::new("triedb_bench").unwrap();
     let db_path = dir.path().join("db");
-    let db = Database::create(db_path.to_str().unwrap()).unwrap();
+    let db = Arc::new(Database::create(db_path.to_str().unwrap()).unwrap());
 
     // Populate database with initial accounts
     let mut rng = StdRng::seed_from_u64(42);
@@ -38,10 +39,10 @@ pub fn setup_database(size: usize) -> (TempDir, Database) {
     (dir, db)
 }
 
-pub fn setup_database_with_storage(size: usize) -> (TempDir, Database) {
+pub fn setup_database_with_storage(size: usize) -> (TempDir, Arc<Database>) {
     let dir = TempDir::new("triedb_bench_storage").unwrap();
     let db_path = dir.path().join("db");
-    let db = Database::create(db_path.to_str().unwrap()).unwrap();
+    let db = Arc::new(Database::create(db_path.to_str().unwrap()).unwrap());
 
     // Populate database with initial accounts
     let mut rng = StdRng::seed_from_u64(42);

src/account.rs

@@ -2,6 +2,7 @@ use alloy_primitives::{B256, U256};
 use alloy_trie::{EMPTY_ROOT_HASH, KECCAK_EMPTY};
 use proptest::prelude::*;
 use proptest_derive::Arbitrary;
+use reth_primitives_traits::Account as RethAccount;
 
 #[derive(Debug, Clone, PartialEq, Eq, Arbitrary)]
 pub struct Account {
@@ -13,6 +14,17 @@ pub struct Account {
     pub code_hash: B256,
 }
 
+impl From<RethAccount> for Account {
+    fn from(reth_account: RethAccount) -> Self {
+        Account::new(
+            reth_account.nonce,
+            reth_account.balance,
+            EMPTY_ROOT_HASH,
+            reth_account.get_bytecode_hash(),
+        )
+    }
+}
+
 impl Account {
     pub fn new(nonce: u64, balance: U256, storage_root: B256, code_hash: B256) -> Self {
         Self { nonce, balance, storage_root, code_hash }

src/database.rs

@@ -9,7 +9,7 @@ use crate::{
 use alloy_primitives::B256;
 use parking_lot::RwLock;
 
-use std::{io, path::Path};
+use std::{io, path::Path, sync::Arc};
 
 #[derive(Debug)]
 pub struct Database {
@@ -23,17 +23,17 @@ pub(crate) struct Inner {
     pub(crate) transaction_manager: RwLock<TransactionManager>,
 }
 
-#[derive(Debug)]
+#[derive(Clone, Debug)]
 pub enum Error {
     PageError(PageError),
     EngineError(engine::Error),
 }
 
-#[derive(Debug)]
+#[derive(Clone, Debug)]
 pub enum OpenError {
     PageError(PageError),
     MetadataError(OpenMetadataError),
-    IO(io::Error),
+    IO(Arc<io::Error>),
 }
 
 impl Database {
@@ -50,7 +50,13 @@ impl Database {
             .open(db_file_path)
             .map_err(OpenError::PageError)?;
 
-        Ok(Self::new(StorageEngine::new(page_manager, meta_manager)))
+        let db = Self::new(StorageEngine::new(page_manager, meta_manager));
+        let db_arc = Arc::new(db);
+
+        let tx = db_arc.begin_rw().unwrap();
+        tx.commit().unwrap();
+
+        Ok(Arc::try_unwrap(db_arc).unwrap())
     }
 
     pub fn open(path: impl AsRef<Path>) -> Result<Self, OpenError> {
@@ -130,25 +136,26 @@ impl Database {
         Ok(())
     }
 
-    pub fn begin_rw(&self) -> Result<Transaction<'_, RW>, TransactionError> {
+    pub fn begin_rw(self: &Arc<Self>) -> Result<Transaction<RW>, TransactionError> {
         let mut transaction_manager = self.inner.transaction_manager.write();
         let mut storage_engine = self.inner.storage_engine.write();
         let metadata = storage_engine.metadata().dirty_slot();
         let min_snapshot_id = transaction_manager.begin_rw(metadata.snapshot_id())?;
         if min_snapshot_id > 0 {
             storage_engine.unlock(min_snapshot_id - 1);
         }
+
         let context = storage_engine.write_context();
-        Ok(Transaction::new(context, self))
+        Ok(Transaction::new(context, Arc::clone(self)))
     }
 
-    pub fn begin_ro(&self) -> Result<Transaction<'_, RO>, TransactionError> {
+    pub fn begin_ro(self: &Arc<Self>) -> Result<Transaction<RO>, TransactionError> {
         let mut transaction_manager = self.inner.transaction_manager.write();
         let storage_engine = self.inner.storage_engine.read();
         let metadata = storage_engine.metadata().active_slot();
         transaction_manager.begin_ro(metadata.snapshot_id())?;
         let context = storage_engine.read_context();
-        Ok(Transaction::new(context, self))
+        Ok(Transaction::new(context, Arc::clone(self)))
     }
 
     pub fn state_root(&self) -> B256 {
@@ -199,7 +206,7 @@ mod tests {
     fn test_set_get_account() {
         let tmp_dir = TempDir::new("test_db").unwrap();
         let file_path = tmp_dir.path().join("test.db");
-        let db = Database::create(file_path).unwrap();
+        let db = Arc::new(Database::create(file_path).unwrap());
 
         let address = address!("0xd8da6bf26964af9d7eed9e03e53415d37aa96045");
 
@@ -256,7 +263,7 @@ mod tests {
     fn test_data_persistence() {
         let tmp_dir = TempDir::new("test_db").unwrap();
         let file_path = tmp_dir.path().join("test.db");
-        let db = Database::create(&file_path).unwrap();
+        let db = Arc::new(Database::create(&file_path).unwrap());
 
         let address1 = address!("0xd8da6bf26964af9d7eed9e03e53415d37aa96045");
         let account1 = Account::new(1, U256::from(100), EMPTY_ROOT_HASH, KECCAK_EMPTY);
@@ -265,9 +272,9 @@ mod tests {
         tx.set_account(AddressPath::for_address(address1), Some(account1.clone())).unwrap();
 
         tx.commit().unwrap();
-        db.close().unwrap();
+        Arc::try_unwrap(db).unwrap().close().unwrap();
 
-        let db = Database::open(&file_path).unwrap();
+        let db = Arc::new(Database::open(&file_path).unwrap());
         let mut tx = db.begin_ro().unwrap();
         let account = tx.get_account(AddressPath::for_address(address1)).unwrap().unwrap();
         assert_eq!(account, account1);
@@ -280,9 +287,9 @@ mod tests {
         tx.set_account(AddressPath::for_address(address2), Some(account2.clone())).unwrap();
 
         tx.commit().unwrap();
-        db.close().unwrap();
+        Arc::try_unwrap(db).unwrap().close().unwrap();
 
-        let db = Database::open(&file_path).unwrap();
+        let db = Arc::new(Database::open(&file_path).unwrap());
         let mut tx = db.begin_ro().unwrap();
 
         let account = tx.get_account(AddressPath::for_address(address1)).unwrap().unwrap();

src/meta/error.rs

@@ -1,4 +1,4 @@
-use std::{error, fmt, io};
+use std::{error, fmt, io, sync::Arc};
 
 #[derive(Debug)]
 pub struct CorruptedMetadataError;
@@ -11,10 +11,10 @@ impl fmt::Display for CorruptedMetadataError {
 
 impl error::Error for CorruptedMetadataError {}
 
-#[derive(Debug)]
+#[derive(Clone, Debug)]
 pub enum OpenMetadataError {
     Corrupted,
-    IO(io::Error),
+    IO(Arc<io::Error>),
 }
 
 impl fmt::Display for OpenMetadataError {
@@ -36,6 +36,6 @@ impl From<CorruptedMetadataError> for OpenMetadataError {
 
 impl From<io::Error> for OpenMetadataError {
     fn from(e: io::Error) -> Self {
-        Self::IO(e)
+        Self::IO(e.into())
     }
 }

src/node.rs

@@ -53,7 +53,7 @@ pub enum Node {
     },
 }
 
-#[derive(Debug)]
+#[derive(Clone, Debug)]
 pub enum NodeError {
     ChildrenUnsupported,
     MaxPrefixLengthExceeded,

src/page/manager.rs

@@ -1,10 +1,17 @@
 use crate::page::PageId;
+use std::{io, sync::Arc};
 
 pub(super) mod mmap;
 pub(super) mod options;
 
+impl From<io::Error> for PageError {
+    fn from(error: io::Error) -> Self {
+        Self::IO(error.into())
+    }
+}
+
 /// Represents various errors that might arise from page operations.
-#[derive(Debug)]
+#[derive(Clone, Debug)]
 pub enum PageError {
     PageNotFound(PageId),
     PageLimitReached,
@@ -13,7 +20,7 @@ pub enum PageError {
     NoFreeCells,
     PageIsFull,
     PageSplitLimitReached,
-    IO(std::io::Error),
+    IO(Arc<io::Error>),
     InvalidValue,
     InvalidPageContents(PageId),
     // TODO: add more errors here for other cases.

src/page/manager/mmap.rs

@@ -3,7 +3,7 @@ use crate::{
     snapshot::SnapshotId,
 };
 use memmap2::{Advice, MmapMut, MmapOptions};
-use std::{fs::File, io, path::Path};
+use std::{fs::File, io, path::Path, sync::Arc};
 
 // Manages pages in a memory mapped file.
 #[derive(Debug)]
@@ -36,7 +36,7 @@ impl PageManager {
         opts: &PageManagerOptions,
         path: impl AsRef<Path>,
     ) -> Result<Self, PageError> {
-        let file = opts.open_options.open(path).map_err(PageError::IO)?;
+        let file = opts.open_options.open(path)?;
         Self::from_file_with_options(opts, file)
     }
 
@@ -73,11 +73,15 @@ impl PageManager {
 
         // SAFETY: we assume that we have full ownership of the file, even though in practice
         // there's no way to guarantee it
-        let mmap =
-            unsafe { MmapOptions::new().len(mmap_len).map_mut(&file).map_err(PageError::IO)? };
-        mmap.advise(Advice::Random).map_err(PageError::IO)?;
-
-        let file_len = file.metadata().map_err(PageError::IO)?.len();
+        let mmap = unsafe {
+            MmapOptions::new()
+                .len(mmap_len)
+                .map_mut(&file)
+                .map_err(|err| PageError::IO(Arc::new(err)))?
+        };
+        mmap.advise(Advice::Random)?;
+
+        let file_len = file.metadata()?.len();
         let min_file_len = (opts.page_count as u64) * (Page::SIZE as u64);
         assert!(
             file_len >= min_file_len,
@@ -142,7 +146,7 @@ impl PageManager {
 
         assert!(new_len > cur_len, "reached max capacity");
 
-        self.file.set_len(new_len).map_err(PageError::IO)?;
+        self.file.set_len(new_len)?;
         self.file_len = new_len;
         Ok(())
     }

src/page/manager/options.rs

@@ -76,7 +76,7 @@ impl PageManagerOptions {
     /// Opens a temporary file with the options specified by `self`.
     #[cfg(test)]
     pub fn open_temp_file(&self) -> Result<PageManager, PageError> {
-        let file = tempfile::tempfile().map_err(PageError::IO)?;
+        let file = tempfile::tempfile()?;
         self.wrap(file)
     }
 }

src/pointer.rs

@@ -67,6 +67,10 @@ impl Value for Pointer {
         } else if first_rlp_byte <= 0xa0 {
             let rlp_len = first_rlp_byte - 0x80;
 
+            RlpNode::from_raw(&arr[4..5 + rlp_len as usize]).unwrap()
+        } else if first_rlp_byte <= 0xf7 {
+            let rlp_len = first_rlp_byte - 0xc0;
+
             RlpNode::from_raw(&arr[4..5 + rlp_len as usize]).unwrap()
         } else {
             return Err(value::Error::InvalidEncoding);

src/storage/engine.rs

@@ -19,7 +19,7 @@ use crate::{
 };
 use alloy_primitives::StorageValue;
 use alloy_trie::{nodes::RlpNode, nybbles, Nibbles, EMPTY_ROOT_HASH};
-use std::{cmp::Ordering, fmt::Debug, io};
+use std::{cmp::Ordering, fmt::Debug, io, sync::Arc};
 
 /// The [StorageEngine] is responsible for managing the storage of data in the database.
 /// It handles reading and writing account and storage values, as well as managing the lifecycle of
@@ -1422,7 +1422,7 @@ impl StorageEngine {
                     if new_slotted_page.id() != slotted_page.id() && !print_whole_db {
                         let child_page_id = child.location().page_id().unwrap();
                         writeln!(buf, "{}Child on new page: {:?}", new_indent, child_page_id)?;
-                        return Ok(())
+                        return Ok(());
                     } else {
                         self.print_page_helper(
                             context,
@@ -1928,9 +1928,9 @@ impl StorageEngine {
     }
 }
 
-#[derive(Debug)]
+#[derive(Clone, Debug)]
 pub enum Error {
-    IO(io::Error),
+    IO(Arc<io::Error>),
     NodeError(NodeError),
     PageError(PageError),
     InvalidCommonPrefixIndex,
@@ -1953,7 +1953,7 @@ impl From<NodeError> for Error {
 
 impl From<io::Error> for Error {
     fn from(error: io::Error) -> Self {
-        Self::IO(error)
+        Self::IO(error.into())
     }
 }