From f03f1200ee2fc633ddf4169eba09ed13bc55d8ef Mon Sep 17 00:00:00 2001 From: James Date: Tue, 13 Jan 2026 16:21:23 -0500 Subject: [PATCH 01/16] feat: add storage api --- Cargo.toml | 2 + crates/storage/Cargo.toml | 17 +++ crates/storage/README.md | 17 +++ crates/storage/src/cold/mod.rs | 1 + crates/storage/src/hot/db.rs | 87 ++++++++++++++ crates/storage/src/hot/db_traits.rs | 167 +++++++++++++++++++++++++++ crates/storage/src/hot/error.rs | 28 +++++ crates/storage/src/hot/mod.rs | 13 +++ crates/storage/src/hot/reth_impl.rs | 32 +++++ crates/storage/src/hot/traits.rs | 127 ++++++++++++++++++++ crates/storage/src/lib.rs | 24 ++++ crates/storage/src/ser/error.rs | 46 ++++++++ crates/storage/src/ser/impls.rs | 163 ++++++++++++++++++++++++++ crates/storage/src/ser/mod.rs | 9 ++ crates/storage/src/ser/reth_impls.rs | 105 +++++++++++++++++ crates/storage/src/ser/traits.rs | 99 ++++++++++++++++ crates/storage/src/tables/cold.rs | 0 crates/storage/src/tables/hot.rs | 86 ++++++++++++++ crates/storage/src/tables/macros.rs | 21 ++++ crates/storage/src/tables/mod.rs | 21 ++++ 20 files changed, 1065 insertions(+) create mode 100644 crates/storage/Cargo.toml create mode 100644 crates/storage/README.md create mode 100644 crates/storage/src/cold/mod.rs create mode 100644 crates/storage/src/hot/db.rs create mode 100644 crates/storage/src/hot/db_traits.rs create mode 100644 crates/storage/src/hot/error.rs create mode 100644 crates/storage/src/hot/mod.rs create mode 100644 crates/storage/src/hot/reth_impl.rs create mode 100644 crates/storage/src/hot/traits.rs create mode 100644 crates/storage/src/lib.rs create mode 100644 crates/storage/src/ser/error.rs create mode 100644 crates/storage/src/ser/impls.rs create mode 100644 crates/storage/src/ser/mod.rs create mode 100644 crates/storage/src/ser/reth_impls.rs create mode 100644 crates/storage/src/ser/traits.rs create mode 100644 crates/storage/src/tables/cold.rs create mode 100644 crates/storage/src/tables/hot.rs create mode 100644 crates/storage/src/tables/macros.rs create mode 100644 crates/storage/src/tables/mod.rs diff --git a/Cargo.toml b/Cargo.toml index cc89e52..856dba7 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -77,7 +77,9 @@ alloy-contract = { version = "1.4.0", features = ["pubsub"] } reth = { git = "https://github.com/paradigmxyz/reth", tag = "v1.9.1" } reth-ethereum = { git = "https://github.com/paradigmxyz/reth", tag = "v1.9.1" } reth-chainspec = { git = "https://github.com/paradigmxyz/reth", tag = "v1.9.1" } +reth-codecs = { git = "https://github.com/paradigmxyz/reth", tag = "v1.9.1" } reth-db = { git = "https://github.com/paradigmxyz/reth", tag = "v1.9.1" } +reth-db-api = { git = "https://github.com/paradigmxyz/reth", tag = "v1.9.1" } reth-db-common = { git = "https://github.com/paradigmxyz/reth", tag = "v1.9.1" } reth-eth-wire-types = { git = "https://github.com/paradigmxyz/reth", tag = "v1.9.1" } reth-evm-ethereum = { git = "https://github.com/paradigmxyz/reth", tag = "v1.9.1" } diff --git a/crates/storage/Cargo.toml b/crates/storage/Cargo.toml new file mode 100644 index 0000000..0384e1d --- /dev/null +++ b/crates/storage/Cargo.toml @@ -0,0 +1,17 @@ +[package] +name = "signet-storage" +version.workspace = true +edition.workspace = true +rust-version.workspace = true +authors.workspace = true +license.workspace = true +homepage.workspace = true +repository.workspace = true + +[dependencies] +alloy.workspace = true +bytes = "1.11.0" +reth.workspace = true +reth-db.workspace = true +reth-db-api.workspace = true +thiserror.workspace = true diff --git a/crates/storage/README.md b/crates/storage/README.md new file mode 100644 index 0000000..a7d6488 --- /dev/null +++ b/crates/storage/README.md @@ -0,0 +1,17 @@ +# Signet Storage + +High-level API for Signet's storage layer + +This library contains the following: + +- Traits for serializing and deserializing Signet data structures as DB keys/ + value. +- Traits for hot and cold storage operations. +- Relevant KV table definitions. + +## Significant Traits + +- `HotKv` - Encapsulates logic for reading and writing to hot storage. +- `ColdKv` - Encapsulates logic for reading and writing to cold storage. +- `KeySer` - Provides methods for serializing a type as a DB key. +- `ValueSer` - Provides methods for serializing a type as a DB value. diff --git a/crates/storage/src/cold/mod.rs b/crates/storage/src/cold/mod.rs new file mode 100644 index 0000000..9b9d9be --- /dev/null +++ b/crates/storage/src/cold/mod.rs @@ -0,0 +1 @@ +//! Placeholder module for cold storage implementation. diff --git a/crates/storage/src/hot/db.rs b/crates/storage/src/hot/db.rs new file mode 100644 index 0000000..a4911da --- /dev/null +++ b/crates/storage/src/hot/db.rs @@ -0,0 +1,87 @@ +use crate::hot::{HotKv, HotKvError, HotKvRead, HotKvWrite}; +use std::{ + borrow::Cow, + sync::atomic::{AtomicBool, Ordering}, +}; + +/// Hot database wrapper around a key-value storage. +#[derive(Debug)] +pub struct HotDb { + inner: Inner, + + write_locked: AtomicBool, +} + +impl HotDb { + /// Create a new HotDb wrapping the given inner KV storage. + pub const fn new(inner: Inner) -> Self { + Self { inner, write_locked: AtomicBool::new(false) } + } + + /// Get a read-only handle. + pub fn reader(&self) -> Result + where + Inner: HotKv, + { + self.inner.reader() + } + + /// Get a write handle, if available. If not available, returns + /// [`HotKvError::WriteLocked`]. + pub fn writer(&self) -> Result, HotKvError> + where + Inner: HotKv, + { + if self.write_locked.swap(true, Ordering::AcqRel) { + return Err(HotKvError::WriteLocked); + } + self.inner.writer().map(Some).map(|tx| WriteGuard { tx, db: self }) + } +} + +/// Write guard for a write transaction. +#[derive(Debug)] +pub struct WriteGuard<'a, Inner> +where + Inner: HotKv, +{ + tx: Option, + db: &'a HotDb, +} + +impl Drop for WriteGuard<'_, Inner> +where + Inner: HotKv, +{ + fn drop(&mut self) { + self.db.write_locked.store(false, Ordering::Release); + } +} + +impl HotKvRead for WriteGuard<'_, Inner> +where + Inner: HotKv, +{ + type Error = <::RwTx as HotKvRead>::Error; + + fn get_raw<'a>( + &'a self, + table: &str, + key: &[u8], + ) -> Result>, Self::Error> { + self.tx.as_ref().expect("present until drop").get_raw(table, key) + } +} + +impl HotKvWrite for WriteGuard<'_, Inner> +where + Inner: HotKv, +{ + fn queue_raw_put(&mut self, table: &str, key: &[u8], value: &[u8]) -> Result<(), Self::Error> { + self.tx.as_mut().expect("present until drop").queue_raw_put(table, key, value) + } + + fn raw_commit(mut self) -> Result<(), Self::Error> { + self.tx.take().expect("present until drop").raw_commit() + } +} diff --git a/crates/storage/src/hot/db_traits.rs b/crates/storage/src/hot/db_traits.rs new file mode 100644 index 0000000..b71b637 --- /dev/null +++ b/crates/storage/src/hot/db_traits.rs @@ -0,0 +1,167 @@ +use crate::{ + hot::{HotKvRead, HotKvWrite}, + tables::hot::{self as tables, AccountStorageKey}, +}; +use alloy::primitives::{Address, B256, U256}; +use reth::primitives::{Account, Bytecode, Header, SealedHeader, StorageEntry}; +use std::borrow::Cow; + +/// Trait for database read operations. +pub trait HotDbReader: sealed::Sealed { + /// The error type for read operations + type Error: std::error::Error + Send + Sync + 'static + From; + + /// Read a block header by its number. + fn get_header(&self, number: u64) -> Result, Self::Error>; + + /// Read a block number by its hash. + fn get_header_number(&self, hash: &B256) -> Result, Self::Error>; + + /// Read the canonical hash by block number. + fn get_canonical_hash(&self, number: u64) -> Result, Self::Error>; + + /// Read contract Bytecode by its hash. + fn get_bytecode(&self, code_hash: &B256) -> Result, Self::Error>; + + /// Read an account by its address. + fn get_account(&self, address: &Address) -> Result, Self::Error>; + + /// Read a storage slot by its address and key. + fn get_storage(&self, address: &Address, key: &B256) -> Result, Self::Error>; + + /// Read a [`StorageEntry`] by its address and key. + fn get_storage_entry( + &self, + address: &Address, + key: &B256, + ) -> Result, Self::Error> { + let opt = self.get_storage(address, key)?; + Ok(opt.map(|value| StorageEntry { key: *key, value })) + } +} + +impl HotDbReader for T +where + T: HotKvRead, +{ + type Error = ::Error; + + fn get_header(&self, number: u64) -> Result, Self::Error> { + self.get::(&number) + } + + fn get_header_number(&self, hash: &B256) -> Result, Self::Error> { + self.get::(hash) + } + + fn get_canonical_hash(&self, number: u64) -> Result, Self::Error> { + self.get::(&number) + } + + fn get_bytecode(&self, code_hash: &B256) -> Result, Self::Error> { + self.get::(code_hash) + } + + fn get_account(&self, address: &Address) -> Result, Self::Error> { + self.get::(address) + } + + fn get_storage(&self, address: &Address, key: &B256) -> Result, Self::Error> { + let storage_key = AccountStorageKey { + address: std::borrow::Cow::Borrowed(address), + key: std::borrow::Cow::Borrowed(key), + }; + let key = storage_key.encode_key(); + self.get::(&key) + } +} + +/// Trait for database write operations. +pub trait HotDbWriter: sealed::Sealed { + /// The error type for write operations + type Error: std::error::Error + Send + Sync + 'static + From; + + /// Read the latest block header. + fn put_header(&mut self, header: &Header) -> Result<(), Self::Error>; + + /// Write a block number by its hash. + fn put_header_number(&mut self, hash: &B256, number: u64) -> Result<(), Self::Error>; + + /// Write the canonical hash by block number. + fn put_canonical_hash(&mut self, number: u64, hash: &B256) -> Result<(), Self::Error>; + + /// Write contract Bytecode by its hash. + fn put_bytecode(&mut self, code_hash: &B256, bytecode: &Bytecode) -> Result<(), Self::Error>; + + /// Write an account by its address. + fn put_account(&mut self, address: &Address, account: &Account) -> Result<(), Self::Error>; + + /// Write a storage entry by its address and key. + fn put_storage( + &mut self, + address: &Address, + key: &B256, + entry: &U256, + ) -> Result<(), Self::Error>; + + /// Commit the write transaction. + fn commit(self) -> Result<(), Self::Error>; + + /// Write a canonical header (header, number mapping, and canonical hash). + fn put_canonical(&mut self, header: &SealedHeader) -> Result<(), Self::Error> { + self.put_header(header)?; + self.put_header_number(&header.hash(), header.number)?; + self.put_canonical_hash(header.number, &header.hash()) + } +} + +impl HotDbWriter for T +where + T: HotKvWrite, +{ + type Error = ::Error; + + fn put_header(&mut self, header: &Header) -> Result<(), Self::Error> { + self.queue_put::(&header.number, header) + } + + fn put_header_number(&mut self, hash: &B256, number: u64) -> Result<(), Self::Error> { + self.queue_put::(hash, &number) + } + + fn put_canonical_hash(&mut self, number: u64, hash: &B256) -> Result<(), Self::Error> { + self.queue_put::(&number, hash) + } + + fn put_bytecode(&mut self, code_hash: &B256, bytecode: &Bytecode) -> Result<(), Self::Error> { + self.queue_put::(code_hash, bytecode) + } + + fn put_account(&mut self, address: &Address, account: &Account) -> Result<(), Self::Error> { + self.queue_put::(address, account) + } + + fn put_storage( + &mut self, + address: &Address, + key: &B256, + entry: &U256, + ) -> Result<(), Self::Error> { + let storage_key = + AccountStorageKey { address: Cow::Borrowed(address), key: Cow::Borrowed(key) }; + self.queue_put::(&storage_key.encode_key(), entry) + } + + fn commit(self) -> Result<(), Self::Error> { + HotKvWrite::raw_commit(self) + } +} + +mod sealed { + use crate::hot::HotKvRead; + + /// Sealed trait to prevent external implementations of HotDbReader and HotDbWriter. + #[allow(dead_code, unreachable_pub)] + pub trait Sealed {} + impl Sealed for T where T: HotKvRead {} +} diff --git a/crates/storage/src/hot/error.rs b/crates/storage/src/hot/error.rs new file mode 100644 index 0000000..a5331ce --- /dev/null +++ b/crates/storage/src/hot/error.rs @@ -0,0 +1,28 @@ +/// Trait for hot storage read/write errors. +#[derive(thiserror::Error, Debug)] +pub enum HotKvError { + /// Boxed error. Indicates an issue with the DB backend. + #[error(transparent)] + Inner(#[from] Box), + + /// Deserialization error. Indicates an issue deserializing a key or value. + #[error("Deserialization error: {0}")] + DeserError(#[from] crate::ser::DeserError), + + /// Indicates that a write transaction is already in progress. + #[error("A write transaction is already in progress")] + WriteLocked, +} + +impl HotKvError { + /// Internal helper to create a `HotKvError::Inner` from any error. + pub fn from_err(err: E) -> Self + where + E: std::error::Error + Send + Sync + 'static, + { + HotKvError::Inner(Box::new(err)) + } +} + +/// Result type for hot storage operations. +pub type HotKvResult = Result; diff --git a/crates/storage/src/hot/mod.rs b/crates/storage/src/hot/mod.rs new file mode 100644 index 0000000..2055543 --- /dev/null +++ b/crates/storage/src/hot/mod.rs @@ -0,0 +1,13 @@ +mod db; +pub use db::{HotDb, WriteGuard}; + +mod db_traits; +pub use db_traits::{HotDbReader, HotDbWriter}; + +mod error; +pub use error::{HotKvError, HotKvResult}; + +mod reth_impl; + +mod traits; +pub use traits::{HotKv, HotKvRead, HotKvWrite}; diff --git a/crates/storage/src/hot/reth_impl.rs b/crates/storage/src/hot/reth_impl.rs new file mode 100644 index 0000000..706fad6 --- /dev/null +++ b/crates/storage/src/hot/reth_impl.rs @@ -0,0 +1,32 @@ +use std::borrow::Cow; + +use crate::{hot::HotKvRead, ser::DeserError}; +use reth_db::mdbx::{TransactionKind, tx::Tx}; +use reth_db_api::DatabaseError; + +impl From for DatabaseError { + fn from(value: DeserError) -> Self { + DatabaseError::Other(value.to_string()) + } +} + +impl HotKvRead for Tx +where + K: TransactionKind, +{ + type Error = DatabaseError; + + fn get_raw<'a>( + &'a self, + table: &str, + key: &[u8], + ) -> Result>, Self::Error> { + let dbi = self + .inner + .open_db(Some(table)) + .map(|db| db.dbi()) + .map_err(|e| DatabaseError::Open(e.into()))?; + + self.inner.get(dbi, key.as_ref()).map_err(|err| DatabaseError::Read(err.into())) + } +} diff --git a/crates/storage/src/hot/traits.rs b/crates/storage/src/hot/traits.rs new file mode 100644 index 0000000..11f03c8 --- /dev/null +++ b/crates/storage/src/hot/traits.rs @@ -0,0 +1,127 @@ +use std::borrow::Cow; + +use crate::{ + hot::HotKvError, + ser::{KeySer, MAX_KEY_SIZE, ValSer}, + tables::Table, +}; + +/// Trait for hot storage. This is a KV store with read/write transactions. +pub trait HotKv { + /// The read-only transaction type. + type RoTx: HotKvRead; + /// The read-write transaction type. + type RwTx: HotKvWrite; + + /// Create a read-only transaction. + fn reader(&self) -> Result; + + /// Create a read-write transaction. + /// + /// # Returns + /// + /// - `Ok(Some(tx))` if the write transaction was created successfully. + /// - [`Err(HotKvError::WriteLocked)`] if there is already a write + /// transaction in progress. + /// - [`Err(HotKvError::Inner)`] if there was an error creating the + /// transaction. + /// + /// [`Err(HotKvError::Inner)`]: HotKvError::Inner + /// [`Err(HotKvError::WriteLocked)`]: HotKvError::WriteLocked + fn writer(&self) -> Result; +} + +/// Trait for hot storage read transactions. +pub trait HotKvRead { + /// Error type for read operations. + type Error: std::error::Error + From + Send + Sync + 'static; + + /// Get a raw value from a specific table. + fn get_raw<'a>(&'a self, table: &str, key: &[u8]) + -> Result>, Self::Error>; + + /// Get a value from a specific table. + fn get(&self, key: &T::Key) -> Result, Self::Error> { + let mut key_buf = [0u8; MAX_KEY_SIZE]; + let key_bytes = key.encode_key(&mut key_buf); + debug_assert!( + key_bytes.len() == T::Key::SIZE, + "Encoded key length does not match expected size" + ); + + let Some(value_bytes) = self.get_raw(T::NAME, key_bytes)? else { + return Ok(None); + }; + let data = &value_bytes[..]; + T::Value::decode_value(data).map(Some).map_err(Into::into) + } + + /// Get many values from a specific table. + /// + /// # Arguments + /// + /// * `keys` - An iterator over keys to retrieve. + /// + /// # Returns + /// + /// A vector of `Option`, where each element corresponds to the + /// value for the respective key in the input iterator. If a key does not + /// exist in the table, the corresponding element will be `None`. + /// + /// If any error occurs during retrieval or deserialization, the entire + /// operation will return an error. + fn get_many<'a, T, I>(&self, keys: I) -> Result>, Self::Error> + where + T::Key: 'a, + T: Table, + I: IntoIterator, + { + let mut key_buf = [0u8; MAX_KEY_SIZE]; + + keys.into_iter() + .map(|key| self.get_raw(T::NAME, key.encode_key(&mut key_buf))) + .map(|maybe_val| { + maybe_val + .and_then(|val| ValSer::maybe_decode_value(val.as_deref()).map_err(Into::into)) + }) + .collect() + } +} + +/// Trait for hot storage write transactions. +pub trait HotKvWrite: HotKvRead { + /// Queue a raw put operation. + fn queue_raw_put(&mut self, table: &str, key: &[u8], value: &[u8]) -> Result<(), Self::Error>; + + /// Queue a put operation for a specific table. + fn queue_put(&mut self, key: &T::Key, value: &T::Value) -> Result<(), Self::Error> { + let mut key_buf = [0u8; MAX_KEY_SIZE]; + let key_bytes = key.encode_key(&mut key_buf); + let value_bytes = value.encoded(); + + self.queue_raw_put(T::NAME, key_bytes, &value_bytes) + } + + /// Queue many put operations for a specific table. + fn queue_put_many<'a, 'b, T, I>(&mut self, entries: I) -> Result<(), Self::Error> + where + T: Table, + T::Key: 'a, + T::Value: 'b, + I: IntoIterator, + { + let mut key_buf = [0u8; MAX_KEY_SIZE]; + + for (key, value) in entries { + let key_bytes = key.encode_key(&mut key_buf); + let value_bytes = value.encoded(); + + self.queue_raw_put(T::NAME, key_bytes, &value_bytes)?; + } + + Ok(()) + } + + /// Commit the queued operations. + fn raw_commit(self) -> Result<(), Self::Error>; +} diff --git a/crates/storage/src/lib.rs b/crates/storage/src/lib.rs new file mode 100644 index 0000000..381014c --- /dev/null +++ b/crates/storage/src/lib.rs @@ -0,0 +1,24 @@ +#![doc = include_str!("../README.md")] +#![warn( + missing_copy_implementations, + missing_debug_implementations, + missing_docs, + unreachable_pub, + clippy::missing_const_for_fn, + rustdoc::all +)] +#![cfg_attr(not(test), warn(unused_crate_dependencies))] +#![deny(unused_must_use, rust_2018_idioms)] +#![cfg_attr(docsrs, feature(doc_cfg, doc_auto_cfg))] + +/// Cold storage module. +pub mod cold; + +/// Hot storage module. +pub mod hot; + +/// Serialization module. +pub mod ser; + +/// Predefined tables module. +pub mod tables; diff --git a/crates/storage/src/ser/error.rs b/crates/storage/src/ser/error.rs new file mode 100644 index 0000000..de83f75 --- /dev/null +++ b/crates/storage/src/ser/error.rs @@ -0,0 +1,46 @@ +/// Error type for deserialization errors. +/// +/// Erases the underlying error type to a boxed trait object or a string +/// message, for convenience. +#[derive(thiserror::Error, Debug)] +pub enum DeserError { + /// Boxed error. + #[error(transparent)] + Boxed(Box), + + /// String error message. + #[error("{0}")] + String(String), + + /// Deserialization ended with extra bytes remaining. + #[error("inexact deserialization: {extra_bytes} extra bytes remaining")] + InexactDeser { + /// Number of extra bytes remaining after deserialization. + extra_bytes: usize, + }, + + /// Not enough data to complete deserialization. + #[error("insufficient data: needed {needed} bytes, but only {available} available")] + InsufficientData { + /// Number of bytes needed. + needed: usize, + /// Number of bytes available. + available: usize, + }, +} + +impl From<&str> for DeserError { + fn from(err: &str) -> Self { + DeserError::String(err.to_string()) + } +} + +impl DeserError { + /// Box an error into a `DeserError`. + pub fn from(err: E) -> Self + where + E: std::error::Error + Send + Sync + 'static, + { + DeserError::Boxed(Box::new(err)) + } +} diff --git a/crates/storage/src/ser/impls.rs b/crates/storage/src/ser/impls.rs new file mode 100644 index 0000000..a5d31ec --- /dev/null +++ b/crates/storage/src/ser/impls.rs @@ -0,0 +1,163 @@ +use crate::ser::{DeserError, KeySer, MAX_KEY_SIZE, ValSer}; +use bytes::BufMut; + +macro_rules! delegate_val_to_key { + ($ty:ty) => { + impl ValSer for $ty { + fn encode_value_to(&self, buf: &mut B) + where + B: BufMut + AsMut<[u8]>, + { + let mut key_buf = [0u8; MAX_KEY_SIZE]; + let key_bytes = KeySer::encode_key(self, &mut key_buf); + buf.put_slice(key_bytes); + } + + fn decode_value(data: &[u8]) -> Result + where + Self: Sized, + { + KeySer::decode_key(&data) + } + } + }; +} + +macro_rules! ser_alloy_fixed { + ($size:expr) => { + impl KeySer for alloy::primitives::FixedBytes<$size> { + const SIZE: usize = $size; + + fn encode_key<'a: 'c, 'b: 'c, 'c>(&'a self, _buf: &'b mut [u8; MAX_KEY_SIZE]) -> &'c [u8] { + self.as_ref() + } + + fn decode_key(data: &[u8]) -> Result + where + Self: Sized, + { + if data.len() < $size { + return Err(DeserError::InsufficientData { + needed: $size, + available: data.len(), + }); + } + let mut this = Self::default(); + this.as_mut_slice().copy_from_slice(&data[..$size]); + Ok(this) + } + } + + delegate_val_to_key!(alloy::primitives::FixedBytes<$size>); + }; + + ($($size:expr),* $(,)?) => { + $( + ser_alloy_fixed!($size); + )+ + }; +} + +macro_rules! ser_be_num { + ($ty:ty, $size:expr) => { + impl KeySer for $ty { + const SIZE: usize = $size; + + fn encode_key<'a: 'c, 'b: 'c, 'c>(&'a self, buf: &'b mut [u8; MAX_KEY_SIZE]) -> &'c [u8] { + let be_bytes: [u8; $size] = self.to_be_bytes(); + buf[..$size].copy_from_slice(&be_bytes); + &buf[..$size] + } + + fn decode_key(data: &[u8]) -> Result + where + Self: Sized, + { + if data.len() < $size { + return Err(DeserError::InsufficientData { + needed: $size, + available: data.len(), + }); + } + let bytes: [u8; $size] = data[..$size].try_into().map_err(DeserError::from)?; + Ok(<$ty>::from_be_bytes(bytes)) + } + } + + delegate_val_to_key!($ty); + }; + ($($ty:ty, $size:expr);* $(;)?) => { + $( + ser_be_num!($ty, $size); + )+ + }; +} + +ser_be_num!( + u8, 1; + i8, 1; + u16, 2; + u32, 4; + u64, 8; + u128, 16; + i16, 2; + i32, 4; + i64, 8; + i128, 16; + usize, std::mem::size_of::(); + isize, std::mem::size_of::(); + alloy::primitives::U160, 20; + alloy::primitives::U256, 32; +); + +// NB: 52 is for AccountStorageKey which is (20 + 32) +ser_alloy_fixed!(16, 20, 32, 52); + +impl KeySer for alloy::primitives::Address { + const SIZE: usize = 20; + + fn encode_key<'a: 'c, 'b: 'c, 'c>(&'a self, _buf: &'b mut [u8; MAX_KEY_SIZE]) -> &'c [u8] { + self.as_ref() + } + + fn decode_key(data: &[u8]) -> Result { + if data.len() < Self::SIZE { + return Err(DeserError::InsufficientData { needed: Self::SIZE, available: data.len() }); + } + let mut addr = Self::default(); + addr.copy_from_slice(&data[..Self::SIZE]); + Ok(addr) + } +} + +impl ValSer for bytes::Bytes { + fn encode_value_to(&self, buf: &mut B) + where + B: BufMut + AsMut<[u8]>, + { + buf.put_slice(self); + } + + fn decode_value(data: &[u8]) -> Result + where + Self: Sized, + { + Ok(bytes::Bytes::copy_from_slice(data)) + } +} + +impl ValSer for alloy::primitives::Bytes { + fn encode_value_to(&self, buf: &mut B) + where + B: BufMut + AsMut<[u8]>, + { + buf.put_slice(self); + } + + fn decode_value(data: &[u8]) -> Result + where + Self: Sized, + { + Ok(alloy::primitives::Bytes::copy_from_slice(data)) + } +} diff --git a/crates/storage/src/ser/mod.rs b/crates/storage/src/ser/mod.rs new file mode 100644 index 0000000..2f34e88 --- /dev/null +++ b/crates/storage/src/ser/mod.rs @@ -0,0 +1,9 @@ +mod error; +pub use error::DeserError; + +mod traits; +pub use traits::{KeySer, MAX_KEY_SIZE, ValSer}; + +mod impls; + +mod reth_impls; diff --git a/crates/storage/src/ser/reth_impls.rs b/crates/storage/src/ser/reth_impls.rs new file mode 100644 index 0000000..17f44de --- /dev/null +++ b/crates/storage/src/ser/reth_impls.rs @@ -0,0 +1,105 @@ +use crate::ser::{DeserError, KeySer, MAX_KEY_SIZE, ValSer}; +use alloy::primitives::{Address, B256}; +use reth::primitives::{Account, Bytecode, Header, Log, StorageEntry, TransactionSigned, TxType}; +use reth_db_api::{ + BlockNumberList, + models::{ + AccountBeforeTx, CompactU256, ShardedKey, StoredBlockBodyIndices, + storage_sharded_key::StorageShardedKey, + }, + table::{Compress, Decompress}, +}; + +macro_rules! simple_delegate_compress { + ($ty:ty) => { + impl ValSer for $ty { + fn encode_value_to(&self, buf: &mut B) + where + B: bytes::BufMut + AsMut<[u8]>, + { + self.compress_to_buf(buf); + } + + fn decode_value(data: &[u8]) -> Result + where + Self: Sized, + { + Decompress::decompress(data).map_err(DeserError::from) + } + } + }; + + ($($ty:ty),* $(,)?) => { + $( + simple_delegate_compress!($ty); + )+ + }; +} + +simple_delegate_compress!( + Header, + Account, + Log, + TxType, + StorageEntry, + StoredBlockBodyIndices, + Bytecode, + AccountBeforeTx, + TransactionSigned, + CompactU256, + BlockNumberList, +); + +impl KeySer for ShardedKey { + const SIZE: usize = T::SIZE + u64::SIZE; + + fn encode_key<'a: 'c, 'b: 'c, 'c>(&'a self, buf: &'b mut [u8; MAX_KEY_SIZE]) -> &'c [u8] { + let mut scratch = [0u8; MAX_KEY_SIZE]; + + T::encode_key(&self.key, &mut scratch); + scratch[T::SIZE..Self::SIZE].copy_from_slice(&self.highest_block_number.to_be_bytes()); + *buf = scratch; + + &buf[0..Self::SIZE] + } + + fn decode_key(data: &[u8]) -> Result { + if data.len() < Self::SIZE { + return Err(DeserError::InsufficientData { needed: Self::SIZE, available: data.len() }); + } + + let key = T::decode_key(&data[0..T::SIZE])?; + let highest_block_number = u64::decode_key(&data[T::SIZE..T::SIZE + 8])?; + Ok(Self { key, highest_block_number }) + } +} + +impl KeySer for StorageShardedKey { + const SIZE: usize = Address::SIZE + B256::SIZE + u64::SIZE; + + fn encode_key<'a: 'c, 'b: 'c, 'c>(&'a self, buf: &'b mut [u8; MAX_KEY_SIZE]) -> &'c [u8] { + buf[0..Address::SIZE].copy_from_slice(self.address.as_slice()); + buf[Address::SIZE..Address::SIZE + B256::SIZE] + .copy_from_slice(self.sharded_key.key.as_slice()); + buf[Address::SIZE + B256::SIZE..Self::SIZE] + .copy_from_slice(&self.sharded_key.highest_block_number.to_be_bytes()); + + &buf[0..Self::SIZE] + } + + fn decode_key(mut data: &[u8]) -> Result { + if data.len() < Self::SIZE { + return Err(DeserError::InsufficientData { needed: Self::SIZE, available: data.len() }); + } + + let address = Address::from_slice(&data[0..Address::SIZE]); + data = &data[Address::SIZE..]; + + let storage_key = B256::from_slice(&data[0..B256::SIZE]); + data = &data[B256::SIZE..]; + + let highest_block_number = u64::from_be_bytes(data[0..8].try_into().unwrap()); + + Ok(Self { address, sharded_key: ShardedKey { key: storage_key, highest_block_number } }) + } +} diff --git a/crates/storage/src/ser/traits.rs b/crates/storage/src/ser/traits.rs new file mode 100644 index 0000000..5d4c490 --- /dev/null +++ b/crates/storage/src/ser/traits.rs @@ -0,0 +1,99 @@ +use crate::ser::error::DeserError; +use alloy::primitives::Bytes; + +/// Maximum allowed key size in bytes. +pub const MAX_KEY_SIZE: usize = 64; + +/// Trait for key serialization with fixed-size keys of size no greater than 32 +/// bytes. +pub trait KeySer: Ord + Sized { + /// The fixed size of the serialized key in bytes. + /// Must satisfy `SIZE <= MAX_KEY_SIZE`. + const SIZE: usize; + + /// Compile-time assertion to ensure SIZE is within limits. + #[doc(hidden)] + const ASSERT: () = { + assert!( + Self::SIZE <= MAX_KEY_SIZE, + "KeySer implementations must have SIZE <= MAX_KEY_SIZE" + ); + assert!(Self::SIZE > 0, "KeySer implementations must have SIZE > 0"); + }; + + /// Encode the key into the provided buffer. + /// + /// Writes exactly `SIZE` bytes to `buf[..SIZE]` and returns `SIZE`. + /// The encoding must preserve ordering: for any `k1, k2` where `k1 > k2`, + /// the bytes written by `k1` must be lexicographically greater than those + /// of `k2`. + /// + /// # Returns + /// + /// A slice containing the encoded key. This may be a slice of buf, or may + /// be borrowed from the key itself. This slice must be <= `SIZE` bytes. + fn encode_key<'a: 'c, 'b: 'c, 'c>(&'a self, buf: &'b mut [u8; MAX_KEY_SIZE]) -> &'c [u8]; + + /// Decode a key from a byte slice. + /// + /// # Arguments + /// * `data` - Exactly `SIZE` bytes to decode from. + /// + /// # Errors + /// Returns an error if `data.len() != SIZE` or decoding fails. + fn decode_key(data: &[u8]) -> Result; + + /// Decode an optional key from an optional byte slice. + /// + /// Useful in DB decoding, where the absence of a key is represented by + /// `None`. + fn maybe_decode_key(data: Option<&[u8]>) -> Result, DeserError> { + match data { + Some(d) => Ok(Some(Self::decode_key(d)?)), + None => Ok(None), + } + } +} + +/// Trait for value serialization. +pub trait ValSer { + /// Serialize the value into bytes. + fn encode_value_to(&self, buf: &mut B) + where + B: bytes::BufMut + AsMut<[u8]>; + + /// Serialize the value into bytes and return them. + fn encoded(&self) -> Bytes { + let mut buf = bytes::BytesMut::new(); + self.encode_value_to(&mut buf); + buf.freeze().into() + } + + /// Deserialize the value from bytes, advancing the `data` slice. + fn decode_value(data: &[u8]) -> Result + where + Self: Sized; + + /// Deserialize an optional value from an optional byte slice. + /// + /// Useful in DB decoding, where the absence of a value is represented by + /// `None`. + fn maybe_decode_value(data: Option<&[u8]>) -> Result, DeserError> + where + Self: Sized, + { + match data { + Some(d) => Ok(Some(Self::decode_value(d)?)), + None => Ok(None), + } + } + + /// Deserialize the value from bytes, ensuring all bytes are consumed. + fn decode_value_exact(data: &mut &[u8]) -> Result + where + Self: Sized, + { + let val = Self::decode_value(data)?; + data.is_empty().then_some(val).ok_or(DeserError::InexactDeser { extra_bytes: data.len() }) + } +} diff --git a/crates/storage/src/tables/cold.rs b/crates/storage/src/tables/cold.rs new file mode 100644 index 0000000..e69de29 diff --git a/crates/storage/src/tables/hot.rs b/crates/storage/src/tables/hot.rs new file mode 100644 index 0000000..ec19884 --- /dev/null +++ b/crates/storage/src/tables/hot.rs @@ -0,0 +1,86 @@ +use std::borrow::Cow; + +use crate::tables::Table; +use alloy::primitives::{Address, B256, BlockNumber, FixedBytes, U256}; +use reth::primitives::{Account, Bytecode, Header}; +use reth_db_api::{ + BlockNumberList, + models::{AccountBeforeTx, ShardedKey, storage_sharded_key::StorageShardedKey}, +}; + +tables! { + /// Records recent block Headers, by their number. + Headers, + + /// Records block numbers by hash. + HeaderNumbers, + + /// Records the canonical chain header hashes, by height. + CanonicalHeaders, + + /// Records contract Bytecode, by its hash. + Bytecodes, + + /// Records plain account states, keyed by address. + PlainAccountState, + + /// Records account state change history, keyed by address. + AccountsHistory, BlockNumberList>, + + /// Records storage state change history, keyed by address and storage key. + StorageHistory, + + /// Records account change sets, keyed by block number. + AccountChangeSets, +} + +/// Records plain storage states, keyed by address and storage key. +#[derive(Debug, Copy, Clone, Eq, PartialEq, PartialOrd, Ord)] +pub struct PlainStorageState; + +impl Table for PlainStorageState { + const NAME: &'static str = "PlainStorageState"; + + type Key = FixedBytes<52>; + + type Value = U256; +} + +/// Key for the [`PlainStorageState`] table. +#[derive(Debug, Clone, Eq, PartialEq, PartialOrd, Ord)] +pub struct AccountStorageKey<'a, 'b> { + /// Address of the account. + pub address: Cow<'a, Address>, + /// Storage key. + pub key: Cow<'b, B256>, +} + +impl AccountStorageKey<'static, 'static> { + /// Decode the key from the provided data. + pub fn decode_key(data: &[u8]) -> Result { + if data.len() < Self::SIZE { + return Err(crate::ser::DeserError::InsufficientData { + needed: Self::SIZE, + available: data.len(), + }); + } + + let address = Address::from_slice(&data[0..20]); + let key = B256::from_slice(&data[20..52]); + + Ok(Self { address: Cow::Owned(address), key: Cow::Owned(key) }) + } +} + +impl<'a, 'b> AccountStorageKey<'a, 'b> { + /// Size in bytes. + pub const SIZE: usize = 20 + 32; + + /// Encode the key into the provided buffer. + pub fn encode_key(&self) -> FixedBytes<52> { + let mut buf = [0u8; Self::SIZE]; + buf[0..20].copy_from_slice(self.address.as_slice()); + buf[20..52].copy_from_slice(self.key.as_slice()); + buf.into() + } +} diff --git a/crates/storage/src/tables/macros.rs b/crates/storage/src/tables/macros.rs new file mode 100644 index 0000000..89dc412 --- /dev/null +++ b/crates/storage/src/tables/macros.rs @@ -0,0 +1,21 @@ +macro_rules! tables { + ( + #[doc = $doc:expr] + $name:ident<$key:ty, $value:ty>) => { + #[doc = $doc] + #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] + pub struct $name; + + impl crate::tables::Table for $name { + const NAME: &'static str = stringify!($name); + type Key = $key; + type Value = $value; + } + }; + + ($(#[doc = $doc:expr] $name:ident<$key:ty, $value:ty>),* $(,)?) => { + $( + tables!(#[doc = $doc] $name<$key, $value>); + )* + }; +} diff --git a/crates/storage/src/tables/mod.rs b/crates/storage/src/tables/mod.rs new file mode 100644 index 0000000..bbd6263 --- /dev/null +++ b/crates/storage/src/tables/mod.rs @@ -0,0 +1,21 @@ +#[macro_use] +mod macros; + +/// Tables that are not hot. +pub mod cold; + +/// Tables that are hot, or conditionally hot. +pub mod hot; + +use crate::ser::{KeySer, ValSer}; + +/// Trait for table definitions. +pub trait Table { + /// A Human-readable name for the table. + const NAME: &'static str; + + /// The key type. + type Key: KeySer; + /// The value type. + type Value: ValSer; +} From 671d1624b4c1afc581bd1154bb21f7ac90db2172 Mon Sep 17 00:00:00 2001 From: James Date: Wed, 14 Jan 2026 11:19:38 -0500 Subject: [PATCH 02/16] fix: make ValSer known-size --- crates/storage/src/ser/impls.rs | 379 ++++++- crates/storage/src/ser/reth_impls.rs | 1383 +++++++++++++++++++++++++- crates/storage/src/ser/traits.rs | 13 +- 3 files changed, 1722 insertions(+), 53 deletions(-) diff --git a/crates/storage/src/ser/impls.rs b/crates/storage/src/ser/impls.rs index a5d31ec..5dfa187 100644 --- a/crates/storage/src/ser/impls.rs +++ b/crates/storage/src/ser/impls.rs @@ -4,6 +4,10 @@ use bytes::BufMut; macro_rules! delegate_val_to_key { ($ty:ty) => { impl ValSer for $ty { + fn encoded_size(&self) -> usize { + ::SIZE + } + fn encode_value_to(&self, buf: &mut B) where B: BufMut + AsMut<[u8]>, @@ -110,8 +114,10 @@ ser_be_num!( alloy::primitives::U256, 32; ); -// NB: 52 is for AccountStorageKey which is (20 + 32) -ser_alloy_fixed!(16, 20, 32, 52); +// NB: 52 is for AccountStorageKey which is (20 + 32). +// 65 is for Signature, which is (1 + 32 + 32). +ser_alloy_fixed!(8, 16, 20, 32, 52, 65, 256); +delegate_val_to_key!(alloy::primitives::Address); impl KeySer for alloy::primitives::Address { const SIZE: usize = 20; @@ -130,11 +136,41 @@ impl KeySer for alloy::primitives::Address { } } +impl ValSer for alloy::primitives::Bloom { + fn encoded_size(&self) -> usize { + self.as_slice().len() + } + + fn encode_value_to(&self, buf: &mut B) + where + B: BufMut + AsMut<[u8]>, + { + buf.put_slice(self.as_ref()); + } + + fn decode_value(data: &[u8]) -> Result + where + Self: Sized, + { + if data.len() != 256 { + return Err(DeserError::InsufficientData { needed: 256, available: data.len() }); + } + let mut bloom = Self::default(); + bloom.as_mut_slice().copy_from_slice(data); + Ok(bloom) + } +} + impl ValSer for bytes::Bytes { + fn encoded_size(&self) -> usize { + 2 + self.len() + } + fn encode_value_to(&self, buf: &mut B) where B: BufMut + AsMut<[u8]>, { + buf.put_u16(self.len() as u16); buf.put_slice(self); } @@ -142,22 +178,355 @@ impl ValSer for bytes::Bytes { where Self: Sized, { - Ok(bytes::Bytes::copy_from_slice(data)) + if data.len() < 2 { + return Err(DeserError::InsufficientData { needed: 2, available: data.len() }); + } + let len = u16::from_be_bytes(data[0..2].try_into().unwrap()) as usize; + Ok(bytes::Bytes::copy_from_slice(&data[2..2 + len])) } } impl ValSer for alloy::primitives::Bytes { + fn encoded_size(&self) -> usize { + self.0.encoded_size() + } + fn encode_value_to(&self, buf: &mut B) where B: BufMut + AsMut<[u8]>, { - buf.put_slice(self); + self.0.encode_value_to(buf); } fn decode_value(data: &[u8]) -> Result where Self: Sized, { - Ok(alloy::primitives::Bytes::copy_from_slice(data)) + bytes::Bytes::decode_value(data).map(alloy::primitives::Bytes) + } +} + +impl ValSer for Option +where + T: ValSer, +{ + fn encoded_size(&self) -> usize { + 1 + match self { + Some(inner) => inner.encoded_size(), + None => 0, + } + } + + fn encode_value_to(&self, buf: &mut B) + where + B: bytes::BufMut + AsMut<[u8]>, + { + // Simple presence flag + if let Some(inner) = self { + buf.put_u8(1); + inner.encode_value_to(buf); + } else { + buf.put_u8(0); + } + } + + fn decode_value(data: &[u8]) -> Result + where + Self: Sized, + { + let flag = + data.get(0).ok_or(DeserError::InsufficientData { needed: 1, available: data.len() })?; + match flag { + 0 => Ok(None), + 1 => Ok(Some(T::decode_value(&data[1..])?)), + _ => Err(DeserError::String(format!("Invalid Option flag: {}", flag))), + } + } +} + +impl ValSer for Vec +where + T: ValSer, +{ + fn encoded_size(&self) -> usize { + // 2 bytes for length prefix + 2 + self.iter().map(|item| item.encoded_size()).sum::() + } + + fn encode_value_to(&self, buf: &mut B) + where + B: bytes::BufMut + AsMut<[u8]>, + { + buf.put_u16(self.len() as u16); + self.iter().for_each(|item| item.encode_value_to(buf)); + } + + fn decode_value(mut data: &[u8]) -> Result + where + Self: Sized, + { + if data.len() < 2 { + return Err(DeserError::InsufficientData { needed: 2, available: data.len() }); + } + + let items = u16::from_be_bytes(data[0..2].try_into().unwrap()) as usize; + data = &data[2..]; + + // Preallocate the vector + let mut vec = Vec::with_capacity(items); + + vec.spare_capacity_mut().iter_mut().try_for_each(|slot| { + // Decode the item and advance the data slice + let item = slot.write(T::decode_value(data)?); + // Advance data slice by the size of the decoded item + data = &data[item.encoded_size()..]; + Ok::<_, DeserError>(()) + })?; + + // SAFETY: + // If we did not shortcut return, we have initialized all `items` + // elements. + unsafe { + vec.set_len(items); + } + Ok(vec) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use alloy::primitives::{Address, Bloom, Bytes as AlloBytes, FixedBytes, U160, U256}; + use bytes::Bytes; + + /// Generic roundtrip test for any ValSer type + #[track_caller] + fn test_roundtrip(original: &T) + where + T: ValSer + PartialEq + std::fmt::Debug, + { + // Encode + let mut buf = bytes::BytesMut::new(); + original.encode_value_to(&mut buf); + let encoded = buf.freeze(); + + // Assert that the encoded size matches + assert_eq!( + original.encoded_size(), + encoded.len(), + "Encoded size mismatch: expected {}, got {}", + original.encoded_size(), + encoded.len() + ); + + // Decode + let decoded = T::decode_value(&encoded).expect("Failed to decode value"); + + // Assert equality + assert_eq!(*original, decoded, "Roundtrip failed"); + } + + #[test] + fn test_integer_roundtrips() { + // Test boundary values for all integer types + test_roundtrip(&0u8); + test_roundtrip(&255u8); + test_roundtrip(&127i8); + test_roundtrip(&-128i8); + + test_roundtrip(&0u16); + test_roundtrip(&65535u16); + test_roundtrip(&32767i16); + test_roundtrip(&-32768i16); + + test_roundtrip(&0u32); + test_roundtrip(&4294967295u32); + test_roundtrip(&2147483647i32); + test_roundtrip(&-2147483648i32); + + test_roundtrip(&0u64); + test_roundtrip(&18446744073709551615u64); + test_roundtrip(&9223372036854775807i64); + test_roundtrip(&-9223372036854775808i64); + + test_roundtrip(&0u128); + test_roundtrip(&340282366920938463463374607431768211455u128); + test_roundtrip(&170141183460469231731687303715884105727i128); + test_roundtrip(&-170141183460469231731687303715884105728i128); + + test_roundtrip(&0usize); + test_roundtrip(&usize::MAX); + test_roundtrip(&0isize); + test_roundtrip(&isize::MAX); + test_roundtrip(&isize::MIN); + } + + #[test] + fn test_u256_roundtrips() { + test_roundtrip(&U256::ZERO); + test_roundtrip(&U256::from(1u64)); + test_roundtrip(&U256::from(255u64)); + test_roundtrip(&U256::from(65535u64)); + test_roundtrip(&U256::from(u64::MAX)); + test_roundtrip(&U256::MAX); + } + + #[test] + fn test_u160_roundtrips() { + test_roundtrip(&U160::ZERO); + test_roundtrip(&U160::from(1u64)); + test_roundtrip(&U160::from(u64::MAX)); + // Create a maxed U160 (20 bytes = 160 bits) + let max_u160 = U160::from_be_bytes([0xFF; 20]); + test_roundtrip(&max_u160); + } + + #[test] + fn test_address_roundtrips() { + test_roundtrip(&Address::ZERO); + // Create a test address with known pattern + let test_addr = Address::from([ + 0x01, 0x23, 0x45, 0x67, 0x89, 0xAB, 0xCD, 0xEF, 0x01, 0x23, 0x45, 0x67, 0x89, 0xAB, + 0xCD, 0xEF, 0x01, 0x23, 0x45, 0x67, + ]); + test_roundtrip(&test_addr); + } + + #[test] + fn test_fixedbytes_roundtrips() { + // Test various FixedBytes sizes + test_roundtrip(&FixedBytes::<8>::ZERO); + test_roundtrip(&FixedBytes::<16>::ZERO); + test_roundtrip(&FixedBytes::<20>::ZERO); + test_roundtrip(&FixedBytes::<32>::ZERO); + test_roundtrip(&FixedBytes::<52>::ZERO); + test_roundtrip(&FixedBytes::<65>::ZERO); + test_roundtrip(&FixedBytes::<256>::ZERO); + + // Test with non-zero patterns + let pattern_32 = FixedBytes::<32>::from([ + 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0A, 0x0B, 0x0C, 0x0D, 0x0E, + 0x0F, 0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, 0x18, 0x19, 0x1A, 0x1B, 0x1C, + 0x1D, 0x1E, 0x1F, 0x20, + ]); + test_roundtrip(&pattern_32); + } + + #[test] + fn test_bloom_roundtrips() { + test_roundtrip(&Bloom::ZERO); + // Create a bloom with some bits set + let mut bloom_data = [0u8; 256]; + bloom_data[0] = 0xFF; + bloom_data[127] = 0xAA; + bloom_data[255] = 0x55; + let bloom = Bloom::from(bloom_data); + test_roundtrip(&bloom); + } + + #[test] + fn test_bytes_roundtrips() { + // Test bytes::Bytes + test_roundtrip(&Bytes::new()); + test_roundtrip(&Bytes::from_static(b"hello world")); + test_roundtrip(&Bytes::from(vec![0x00, 0x01, 0x02, 0x03, 0xFF, 0xFE, 0xFD])); + + // Test alloy::primitives::Bytes + test_roundtrip(&AlloBytes::new()); + test_roundtrip(&AlloBytes::from_static(b"hello alloy")); + test_roundtrip(&AlloBytes::copy_from_slice(&[0xDE, 0xAD, 0xBE, 0xEF])); + } + + #[test] + fn test_option_roundtrips() { + // None variants + test_roundtrip(&None::); + test_roundtrip(&None::
); + test_roundtrip(&None::); + + // Some variants + test_roundtrip(&Some(42u32)); + test_roundtrip(&Some(u64::MAX)); + test_roundtrip(&Some(Address::ZERO)); + test_roundtrip(&Some(U256::from(12345u64))); + test_roundtrip(&Some(AlloBytes::from_static(b"test"))); + + // Nested options + test_roundtrip(&Some(Some(123u32))); + test_roundtrip(&Some(None::)); + test_roundtrip(&None::>); + } + + #[test] + fn test_vec_roundtrips() { + // Empty vectors + test_roundtrip(&Vec::::new()); + test_roundtrip(&Vec::
::new()); + + // Single element vectors + test_roundtrip(&vec![42u32]); + test_roundtrip(&vec![Address::ZERO]); + + // Multiple element vectors + test_roundtrip(&vec![1u32, 2, 3, 4, 5]); + test_roundtrip(&vec![U256::ZERO, U256::from(1u64), U256::MAX]); + + // Vector of bytes + test_roundtrip(&vec![ + AlloBytes::from_static(b"first"), + AlloBytes::from_static(b"second"), + AlloBytes::new(), + AlloBytes::from_static(b"last"), + ]); + + // Nested vectors + test_roundtrip(&vec![vec![1u32, 2, 3], vec![], vec![4, 5]]); + + // Vector of options + test_roundtrip(&vec![Some(1u32), None, Some(2u32), Some(3u32), None]); + } + + #[test] + fn test_complex_combinations() { + // Option of Vec + test_roundtrip(&Some(vec![1u32, 2, 3, 4])); + test_roundtrip(&None::>); + + // Vec of Options + test_roundtrip(&vec![Some(Address::ZERO), None, Some(Address::ZERO)]); + + // Option of Option + test_roundtrip(&Some(Some(42u32))); + test_roundtrip(&Some(None::)); + + // Complex nested structure + let complex = vec![ + Some(vec![AlloBytes::from_static(b"hello")]), + None, + Some(vec![ + AlloBytes::from_static(b"world"), + AlloBytes::new(), + AlloBytes::from_static(b"!"), + ]), + ]; + test_roundtrip(&complex); + } + + #[test] + fn test_edge_cases() { + // Maximum values that should still work + test_roundtrip(&vec![0u8; 65535]); // Max length for Vec + + // Large FixedBytes + let large_fixed = FixedBytes::<256>::from([0xFF; 256]); + test_roundtrip(&large_fixed); + + // Very large U256 + let large_u256 = U256::from_str_radix( + "115792089237316195423570985008687907853269984665640564039457584007913129639935", + 10, + ) + .unwrap(); + test_roundtrip(&large_u256); } } diff --git a/crates/storage/src/ser/reth_impls.rs b/crates/storage/src/ser/reth_impls.rs index 17f44de..007c543 100644 --- a/crates/storage/src/ser/reth_impls.rs +++ b/crates/storage/src/ser/reth_impls.rs @@ -1,55 +1,23 @@ use crate::ser::{DeserError, KeySer, MAX_KEY_SIZE, ValSer}; -use alloy::primitives::{Address, B256}; -use reth::primitives::{Account, Bytecode, Header, Log, StorageEntry, TransactionSigned, TxType}; +use alloy::{ + consensus::{EthereumTxEnvelope, Signed, TxEip1559, TxEip2930, TxEip4844, TxEip7702, TxLegacy}, + eips::{ + eip2930::{AccessList, AccessListItem}, + eip7702::{Authorization, SignedAuthorization}, + }, + primitives::{Address, B256, FixedBytes, Signature, TxKind, U256}, +}; +use reth::{ + primitives::{Account, Bytecode, Header, Log, LogData, TransactionSigned, TxType}, + revm::bytecode::{JumpTable, LegacyAnalyzedBytecode, eip7702::Eip7702Bytecode}, +}; use reth_db_api::{ BlockNumberList, models::{ - AccountBeforeTx, CompactU256, ShardedKey, StoredBlockBodyIndices, - storage_sharded_key::StorageShardedKey, + AccountBeforeTx, ShardedKey, StoredBlockBodyIndices, storage_sharded_key::StorageShardedKey, }, - table::{Compress, Decompress}, }; -macro_rules! simple_delegate_compress { - ($ty:ty) => { - impl ValSer for $ty { - fn encode_value_to(&self, buf: &mut B) - where - B: bytes::BufMut + AsMut<[u8]>, - { - self.compress_to_buf(buf); - } - - fn decode_value(data: &[u8]) -> Result - where - Self: Sized, - { - Decompress::decompress(data).map_err(DeserError::from) - } - } - }; - - ($($ty:ty),* $(,)?) => { - $( - simple_delegate_compress!($ty); - )+ - }; -} - -simple_delegate_compress!( - Header, - Account, - Log, - TxType, - StorageEntry, - StoredBlockBodyIndices, - Bytecode, - AccountBeforeTx, - TransactionSigned, - CompactU256, - BlockNumberList, -); - impl KeySer for ShardedKey { const SIZE: usize = T::SIZE + u64::SIZE; @@ -103,3 +71,1328 @@ impl KeySer for StorageShardedKey { Ok(Self { address, sharded_key: ShardedKey { key: storage_key, highest_block_number } }) } } + +macro_rules! by_props { + (@size $($prop:ident),* $(,)?) => { + { + 0 $( + + $prop.encoded_size() + )+ + } + }; + (@encode $buf:ident; $($prop:ident),* $(,)?) => { + { + $( + $prop.encode_value_to($buf); + )+ + } + }; + (@decode $data:ident; $($prop:ident),* $(,)?) => { + { + $( + *$prop = ValSer::decode_value($data)?; + $data = &$data[$prop.encoded_size()..]; + )* + } + }; +} + +impl ValSer for BlockNumberList { + fn encoded_size(&self) -> usize { + 2 + self.serialized_size() + } + + fn encode_value_to(&self, mut buf: &mut B) + where + B: bytes::BufMut + AsMut<[u8]>, + { + use std::io::Write; + let mut writer: bytes::buf::Writer<&mut B> = bytes::BufMut::writer(&mut buf); + + debug_assert!( + self.serialized_size() <= u16::MAX as usize, + "BlockNumberList too large to encode" + ); + + writer.write_all(&(self.serialized_size() as u16).to_be_bytes()).unwrap(); + self.serialize_into(writer).unwrap(); + } + + fn decode_value(data: &[u8]) -> Result + where + Self: Sized, + { + let size = u16::decode_value(&data[..2])? as usize; + BlockNumberList::from_bytes(&data[2..2 + size]) + .map_err(|err| DeserError::String(format!("Failed to decode BlockNumberList {err}"))) + } +} + +impl ValSer for Header { + fn encoded_size(&self) -> usize { + // NB: Destructure to ensure changes are compile errors and mistakes + // are unused var warnings. + let Header { + parent_hash, + ommers_hash, + beneficiary, + state_root, + transactions_root, + receipts_root, + logs_bloom, + difficulty, + number, + gas_limit, + gas_used, + timestamp, + extra_data, + mix_hash, + nonce, + base_fee_per_gas, + withdrawals_root, + blob_gas_used, + excess_blob_gas, + parent_beacon_block_root, + requests_hash, + } = self; + + by_props!( + @size + parent_hash, + ommers_hash, + beneficiary, + state_root, + transactions_root, + receipts_root, + logs_bloom, + difficulty, + number, + gas_limit, + gas_used, + timestamp, + extra_data, + mix_hash, + nonce, + base_fee_per_gas, + withdrawals_root, + blob_gas_used, + excess_blob_gas, + parent_beacon_block_root, + requests_hash, + ) + } + + fn encode_value_to(&self, buf: &mut B) + where + B: bytes::BufMut + AsMut<[u8]>, + { + // NB: Destructure to ensure changes are compile errors and mistakes + // are unused var warnings. + let Header { + parent_hash, + ommers_hash, + beneficiary, + state_root, + transactions_root, + receipts_root, + logs_bloom, + difficulty, + number, + gas_limit, + gas_used, + timestamp, + extra_data, + mix_hash, + nonce, + base_fee_per_gas, + withdrawals_root, + blob_gas_used, + excess_blob_gas, + parent_beacon_block_root, + requests_hash, + } = self; + + by_props!( + @encode buf; + parent_hash, + ommers_hash, + beneficiary, + state_root, + transactions_root, + receipts_root, + logs_bloom, + difficulty, + number, + gas_limit, + gas_used, + timestamp, + extra_data, + mix_hash, + nonce, + base_fee_per_gas, + withdrawals_root, + blob_gas_used, + excess_blob_gas, + parent_beacon_block_root, + requests_hash, + ) + } + + fn decode_value(mut data: &[u8]) -> Result + where + Self: Sized, + { + // NB: Destructure to ensure changes are compile errors and mistakes + // are unused var warnings. + let mut h = Header::default(); + let Header { + parent_hash, + ommers_hash, + beneficiary, + state_root, + transactions_root, + receipts_root, + logs_bloom, + difficulty, + number, + gas_limit, + gas_used, + timestamp, + extra_data, + mix_hash, + nonce, + base_fee_per_gas, + withdrawals_root, + blob_gas_used, + excess_blob_gas, + parent_beacon_block_root, + requests_hash, + } = &mut h; + + by_props!( + @decode data; + parent_hash, + ommers_hash, + beneficiary, + state_root, + transactions_root, + receipts_root, + logs_bloom, + difficulty, + number, + gas_limit, + gas_used, + timestamp, + extra_data, + mix_hash, + nonce, + base_fee_per_gas, + withdrawals_root, + blob_gas_used, + excess_blob_gas, + parent_beacon_block_root, + requests_hash, + ); + Ok(h) + } +} + +impl ValSer for Account { + fn encoded_size(&self) -> usize { + // NB: Destructure to ensure changes are compile errors and mistakes + // are unused var warnings. + let Account { nonce, balance, bytecode_hash } = self; + by_props!( + @size + nonce, + balance, + bytecode_hash, + ) + } + + fn encode_value_to(&self, buf: &mut B) + where + B: bytes::BufMut + AsMut<[u8]>, + { + // NB: Destructure to ensure changes are compile errors and mistakes + // are unused var warnings. + let Account { nonce, balance, bytecode_hash } = self; + by_props!( + @encode buf; + nonce, + balance, + bytecode_hash, + ) + } + + fn decode_value(data: &[u8]) -> Result + where + Self: Sized, + { + // NB: Destructure to ensure changes are compile errors and mistakes + // are unused var warnings. + let mut account = Account::default(); + let Account { nonce, balance, bytecode_hash } = &mut account; + + let mut data = data; + by_props!( + @decode data; + nonce, + balance, + bytecode_hash, + ); + Ok(account) + } +} + +impl ValSer for LogData { + fn encoded_size(&self) -> usize { + let LogData { data, .. } = self; + let topics = self.topics(); + 2 + topics.iter().map(|t| t.encoded_size()).sum::() + data.encoded_size() + } + + fn encode_value_to(&self, buf: &mut B) + where + B: bytes::BufMut + AsMut<[u8]>, + { + let LogData { data, .. } = self; + let topics = self.topics(); + buf.put_u16(topics.len() as u16); + for topic in topics { + topic.encode_value_to(buf); + } + data.encode_value_to(buf); + } + + fn decode_value(data: &[u8]) -> Result + where + Self: Sized, + { + let mut data = data; + let topics_len = u16::decode_value(&data[..2])? as usize; + data = &data[2..]; + + if topics_len > 4 { + return Err(DeserError::String("LogData topics length exceeds maximum of 4".into())); + } + + let mut topics = Vec::with_capacity(topics_len); + for _ in 0..topics_len { + let topic = B256::decode_value(data)?; + data = &data[topic.encoded_size()..]; + topics.push(topic); + } + + let log_data = alloy::primitives::Bytes::decode_value(data)?; + + Ok(LogData::new_unchecked(topics, log_data)) + } +} + +impl ValSer for Log { + fn encoded_size(&self) -> usize { + let Log { address, data } = self; + by_props!( + @size + address, + data, + ) + } + + fn encode_value_to(&self, buf: &mut B) + where + B: bytes::BufMut + AsMut<[u8]>, + { + let Log { address, data } = self; + by_props!( + @encode buf; + address, + data, + ) + } + + fn decode_value(data: &[u8]) -> Result + where + Self: Sized, + { + let mut log = Log::::default(); + let Log { address, data: log_data } = &mut log; + + let mut data = data; + by_props!( + @decode data; + address, + log_data, + ); + Ok(log) + } +} + +impl ValSer for TxType { + fn encoded_size(&self) -> usize { + 1 + } + + fn encode_value_to(&self, buf: &mut B) + where + B: bytes::BufMut + AsMut<[u8]>, + { + buf.put_u8(*self as u8); + } + + fn decode_value(data: &[u8]) -> Result + where + Self: Sized, + { + let byte = u8::decode_value(data)?; + TxType::try_from(byte) + .map_err(|_| DeserError::String(format!("Invalid TxType value: {}", byte))) + } +} + +impl ValSer for StoredBlockBodyIndices { + fn encoded_size(&self) -> usize { + let StoredBlockBodyIndices { first_tx_num, tx_count } = self; + by_props!( + @size + first_tx_num, + tx_count, + ) + } + + fn encode_value_to(&self, buf: &mut B) + where + B: bytes::BufMut + AsMut<[u8]>, + { + let StoredBlockBodyIndices { first_tx_num, tx_count } = self; + by_props!( + @encode buf; + first_tx_num, + tx_count, + ) + } + + fn decode_value(data: &[u8]) -> Result + where + Self: Sized, + { + let mut indices = StoredBlockBodyIndices::default(); + let StoredBlockBodyIndices { first_tx_num, tx_count } = &mut indices; + + let mut data = data; + by_props!( + @decode data; + first_tx_num, + tx_count, + ); + Ok(indices) + } +} + +impl ValSer for Eip7702Bytecode { + fn encoded_size(&self) -> usize { + let Eip7702Bytecode { delegated_address, version, raw } = self; + by_props!( + @size + delegated_address, + version, + raw, + ) + } + + fn encode_value_to(&self, buf: &mut B) + where + B: bytes::BufMut + AsMut<[u8]>, + { + let Eip7702Bytecode { delegated_address, version, raw } = self; + by_props!( + @encode buf; + delegated_address, + version, + raw, + ) + } + + fn decode_value(data: &[u8]) -> Result + where + Self: Sized, + { + let mut eip7702 = Eip7702Bytecode { + delegated_address: Address::ZERO, + version: 0, + raw: alloy::primitives::Bytes::new(), + }; + let Eip7702Bytecode { delegated_address, version, raw } = &mut eip7702; + + let mut data = data; + by_props!( + @decode data; + delegated_address, + version, + raw, + ); + Ok(eip7702) + } +} + +impl ValSer for JumpTable { + fn encoded_size(&self) -> usize { + 2 + 2 + self.as_slice().len() + } + + fn encode_value_to(&self, buf: &mut B) + where + B: bytes::BufMut + AsMut<[u8]>, + { + debug_assert!(self.len() <= u16::MAX as usize, "JumpTable bitlen too large to encode"); + debug_assert!(self.as_slice().len() <= u16::MAX as usize, "JumpTable too large to encode"); + buf.put_u16(self.len() as u16); + buf.put_u16(self.as_slice().len() as u16); + buf.put_slice(self.as_slice()); + } + + fn decode_value(data: &[u8]) -> Result + where + Self: Sized, + { + let bit_len = u16::decode_value(&data[..2])? as usize; + let slice_len = u16::decode_value(&data[2..4])? as usize; + Ok(JumpTable::from_slice(&data[4..4 + slice_len], bit_len)) + } +} + +impl ValSer for LegacyAnalyzedBytecode { + fn encoded_size(&self) -> usize { + let bytecode = self.bytecode(); + let original_len = self.original_len(); + let jump_table = self.jump_table(); + by_props!( + @size + bytecode, + original_len, + jump_table, + ) + } + + fn encode_value_to(&self, buf: &mut B) + where + B: bytes::BufMut + AsMut<[u8]>, + { + let bytecode = self.bytecode(); + let original_len = self.original_len(); + let jump_table = self.jump_table(); + by_props!( + @encode buf; + bytecode, + original_len, + jump_table, + ) + } + + fn decode_value(mut data: &[u8]) -> Result + where + Self: Sized, + { + let mut bytecode = alloy::primitives::Bytes::new(); + let mut original_len = 0usize; + let mut jump_table = JumpTable::default(); + + let bc = &mut bytecode; + let ol = &mut original_len; + let jt = &mut jump_table; + by_props!( + @decode data; + bc, + ol, + jt, + ); + Ok(LegacyAnalyzedBytecode::new(bytecode, original_len, jump_table)) + } +} + +impl ValSer for Bytecode { + fn encoded_size(&self) -> usize { + 1 + match &self.0 { + reth::revm::state::Bytecode::Eip7702(code) => code.encoded_size(), + reth::revm::state::Bytecode::LegacyAnalyzed(code) => code.encoded_size(), + } + } + + fn encode_value_to(&self, buf: &mut B) + where + B: bytes::BufMut + AsMut<[u8]>, + { + match &self.0 { + reth::revm::state::Bytecode::Eip7702(code) => { + buf.put_u8(1); + code.encode_value_to(buf); + } + reth::revm::state::Bytecode::LegacyAnalyzed(code) => { + buf.put_u8(0); + code.encode_value_to(buf); + } + } + } + + fn decode_value(data: &[u8]) -> Result + where + Self: Sized, + { + let ty = u8::decode_value(&data[..1])?; + let data = &data[1..]; + match ty { + 0 => { + let analyzed = LegacyAnalyzedBytecode::decode_value(data)?; + Ok(Bytecode(reth::revm::state::Bytecode::LegacyAnalyzed(analyzed))) + } + 1 => { + let eip7702 = Eip7702Bytecode::decode_value(data)?; + Ok(Bytecode(reth::revm::state::Bytecode::Eip7702(eip7702))) + } + _ => Err(DeserError::String(format!("Invalid Bytecode type value: {}. Max is 1.", ty))), + } + } +} + +impl ValSer for AccountBeforeTx { + fn encoded_size(&self) -> usize { + let AccountBeforeTx { address, info } = self; + by_props!( + @size + address, + info, + ) + } + + fn encode_value_to(&self, buf: &mut B) + where + B: bytes::BufMut + AsMut<[u8]>, + { + let AccountBeforeTx { address, info } = self; + by_props!( + @encode buf; + address, + info, + ) + } + + fn decode_value(mut data: &[u8]) -> Result + where + Self: Sized, + { + let mut abt = AccountBeforeTx::default(); + let AccountBeforeTx { address, info } = &mut abt; + + by_props!( + @decode data; + address, + info, + ); + Ok(abt) + } +} + +impl ValSer for Signature { + fn encoded_size(&self) -> usize { + 65 + } + + fn encode_value_to(&self, buf: &mut B) + where + B: bytes::BufMut + AsMut<[u8]>, + { + FixedBytes(self.as_bytes()).encode_value_to(buf); + } + + fn decode_value(data: &[u8]) -> Result + where + Self: Sized, + { + let bytes = FixedBytes::<65>::decode_value(data)?; + Self::from_raw_array(bytes.as_ref()) + .map_err(|e| DeserError::String(format!("Invalid signature bytes: {}", e))) + } +} + +impl ValSer for TxKind { + fn encoded_size(&self) -> usize { + 1 + match self { + TxKind::Create => 0, + TxKind::Call(_) => 20, + } + } + + fn encode_value_to(&self, buf: &mut B) + where + B: bytes::BufMut + AsMut<[u8]>, + { + match self { + TxKind::Create => { + buf.put_u8(0); + } + TxKind::Call(address) => { + buf.put_u8(1); + address.encode_value_to(buf); + } + } + } + + fn decode_value(data: &[u8]) -> Result + where + Self: Sized, + { + let ty = u8::decode_value(&data[..1])?; + let data = &data[1..]; + match ty { + 0 => Ok(TxKind::Create), + 1 => { + let address = Address::decode_value(data)?; + Ok(TxKind::Call(address)) + } + _ => Err(DeserError::String(format!("Invalid TxKind type value: {}. Max is 1.", ty))), + } + } +} + +impl ValSer for AccessListItem { + fn encoded_size(&self) -> usize { + let AccessListItem { address, storage_keys } = self; + by_props!( + @size + address, + storage_keys, + ) + } + + fn encode_value_to(&self, buf: &mut B) + where + B: bytes::BufMut + AsMut<[u8]>, + { + let AccessListItem { address, storage_keys } = self; + by_props!( + @encode buf; + address, + storage_keys, + ) + } + + fn decode_value(data: &[u8]) -> Result + where + Self: Sized, + { + let mut item = AccessListItem::default(); + let AccessListItem { address, storage_keys } = &mut item; + + let mut data = data; + by_props!( + @decode data; + address, + storage_keys, + ); + Ok(item) + } +} + +impl ValSer for AccessList { + fn encoded_size(&self) -> usize { + self.0.encoded_size() + } + + fn encode_value_to(&self, buf: &mut B) + where + B: bytes::BufMut + AsMut<[u8]>, + { + self.0.encode_value_to(buf); + } + + fn decode_value(data: &[u8]) -> Result + where + Self: Sized, + { + Vec::::decode_value(data).map(AccessList) + } +} + +impl ValSer for Authorization { + fn encoded_size(&self) -> usize { + let Authorization { chain_id, address, nonce } = self; + by_props!( + @size + chain_id, + address, + nonce, + ) + } + + fn encode_value_to(&self, buf: &mut B) + where + B: bytes::BufMut + AsMut<[u8]>, + { + let Authorization { chain_id, address, nonce } = self; + by_props!( + @encode buf; + chain_id, + address, + nonce, + ) + } + + fn decode_value(data: &[u8]) -> Result + where + Self: Sized, + { + let mut auth = Authorization { chain_id: U256::ZERO, address: Address::ZERO, nonce: 0 }; + let Authorization { chain_id, address, nonce } = &mut auth; + + let mut data = data; + by_props!( + @decode data; + chain_id, + address, + nonce, + ); + Ok(auth) + } +} + +impl ValSer for SignedAuthorization { + fn encoded_size(&self) -> usize { + let auth = self.inner(); + let y_parity = self.y_parity(); + let r = self.r(); + let s = self.s(); + by_props!( + @size + auth, + y_parity, + r, + s, + ) + } + + fn encode_value_to(&self, buf: &mut B) + where + B: bytes::BufMut + AsMut<[u8]>, + { + let auth = self.inner(); + let y_parity = &self.y_parity(); + let r = &self.r(); + let s = &self.s(); + by_props!( + @encode buf; + auth, + y_parity, + r, + s, + ) + } + + fn decode_value(mut data: &[u8]) -> Result + where + Self: Sized, + { + let mut auth = Authorization { chain_id: U256::ZERO, address: Address::ZERO, nonce: 0 }; + let mut y_parity = 0u8; + let mut r = U256::ZERO; + let mut s = U256::ZERO; + + let ap = &mut auth; + let yp = &mut y_parity; + let rr = &mut r; + let ss = &mut s; + + by_props!( + @decode data; + ap, + yp, + rr, + ss, + ); + Ok(SignedAuthorization::new_unchecked(auth, y_parity, r, s)) + } +} + +impl ValSer for TxLegacy { + fn encoded_size(&self) -> usize { + let TxLegacy { chain_id, nonce, gas_price, gas_limit, to, value, input } = self; + by_props!( + @size + chain_id, + nonce, + gas_price, + gas_limit, + to, + value, + input, + ) + } + + fn encode_value_to(&self, buf: &mut B) + where + B: bytes::BufMut + AsMut<[u8]>, + { + let TxLegacy { chain_id, nonce, gas_price, gas_limit, to, value, input } = self; + by_props!( + @encode buf; + chain_id, + nonce, + gas_price, + gas_limit, + to, + value, + input, + ) + } + + fn decode_value(mut data: &[u8]) -> Result + where + Self: Sized, + { + let mut tx = TxLegacy::default(); + let TxLegacy { chain_id, nonce, gas_price, gas_limit, to, value, input } = &mut tx; + + by_props!( + @decode data; + chain_id, + nonce, + gas_price, + gas_limit, + to, + value, + input, + ); + Ok(tx) + } +} + +impl ValSer for TxEip2930 { + fn encoded_size(&self) -> usize { + let TxEip2930 { chain_id, nonce, gas_price, gas_limit, to, value, input, access_list } = + self; + by_props!( + @size + chain_id, + nonce, + gas_price, + gas_limit, + to, + value, + input, + access_list, + ) + } + + fn encode_value_to(&self, buf: &mut B) + where + B: bytes::BufMut + AsMut<[u8]>, + { + let TxEip2930 { chain_id, nonce, gas_price, gas_limit, to, value, input, access_list } = + self; + by_props!( + @encode buf; + chain_id, + nonce, + gas_price, + gas_limit, + to, + value, + input, + access_list, + ) + } + + fn decode_value(mut data: &[u8]) -> Result + where + Self: Sized, + { + let mut tx = TxEip2930::default(); + let TxEip2930 { chain_id, nonce, gas_price, gas_limit, to, value, input, access_list } = + &mut tx; + + by_props!( + @decode data; + chain_id, + nonce, + gas_price, + gas_limit, + to, + value, + input, + access_list, + ); + Ok(tx) + } +} + +impl ValSer for TxEip1559 { + fn encoded_size(&self) -> usize { + let TxEip1559 { + chain_id, + nonce, + gas_limit, + max_fee_per_gas, + max_priority_fee_per_gas, + to, + value, + access_list, + input, + } = self; + by_props!( + @size + chain_id, + nonce, + gas_limit, + max_fee_per_gas, + max_priority_fee_per_gas, + to, + value, + access_list, + input + ) + } + + fn encode_value_to(&self, buf: &mut B) + where + B: bytes::BufMut + AsMut<[u8]>, + { + let TxEip1559 { + chain_id, + nonce, + gas_limit, + max_fee_per_gas, + max_priority_fee_per_gas, + to, + value, + access_list, + input, + } = self; + by_props!( + @encode buf; + chain_id, + nonce, + gas_limit, + max_fee_per_gas, + max_priority_fee_per_gas, + to, + value, + access_list, + input + ) + } + + fn decode_value(data: &[u8]) -> Result + where + Self: Sized, + { + let mut tx = TxEip1559::default(); + let TxEip1559 { + chain_id, + nonce, + gas_limit, + max_fee_per_gas, + max_priority_fee_per_gas, + to, + value, + access_list, + input, + } = &mut tx; + + let mut data = data; + by_props!( + @decode data; + chain_id, + nonce, + gas_limit, + max_fee_per_gas, + max_priority_fee_per_gas, + to, + value, + access_list, + input + ); + Ok(tx) + } +} + +impl ValSer for TxEip4844 { + fn encoded_size(&self) -> usize { + let TxEip4844 { + chain_id, + nonce, + gas_limit, + max_fee_per_gas, + max_priority_fee_per_gas, + to, + value, + access_list, + blob_versioned_hashes, + max_fee_per_blob_gas, + input, + } = self; + by_props!( + @size + chain_id, + nonce, + gas_limit, + max_fee_per_gas, + max_priority_fee_per_gas, + to, + value, + access_list, + blob_versioned_hashes, + max_fee_per_blob_gas, + input, + ) + } + + fn encode_value_to(&self, buf: &mut B) + where + B: bytes::BufMut + AsMut<[u8]>, + { + let TxEip4844 { + chain_id, + nonce, + gas_limit, + max_fee_per_gas, + max_priority_fee_per_gas, + to, + value, + access_list, + blob_versioned_hashes, + max_fee_per_blob_gas, + input, + } = self; + by_props!( + @encode buf; + chain_id, + nonce, + gas_limit, + max_fee_per_gas, + max_priority_fee_per_gas, + to, + value, + access_list, + blob_versioned_hashes, + max_fee_per_blob_gas, + input, + ) + } + + fn decode_value(data: &[u8]) -> Result + where + Self: Sized, + { + let mut tx = TxEip4844::default(); + let TxEip4844 { + chain_id, + nonce, + gas_limit, + max_fee_per_gas, + max_priority_fee_per_gas, + to, + value, + access_list, + blob_versioned_hashes, + max_fee_per_blob_gas, + input, + } = &mut tx; + + let mut data = data; + by_props!( + @decode data; + chain_id, + nonce, + gas_limit, + max_fee_per_gas, + max_priority_fee_per_gas, + to, + value, + access_list, + blob_versioned_hashes, + max_fee_per_blob_gas, + input, + ); + Ok(tx) + } +} + +impl ValSer for TxEip7702 { + fn encoded_size(&self) -> usize { + let TxEip7702 { + chain_id, + nonce, + gas_limit, + max_fee_per_gas, + max_priority_fee_per_gas, + to, + value, + access_list, + authorization_list, + input, + } = self; + by_props!( + @size + chain_id, + nonce, + gas_limit, + max_fee_per_gas, + max_priority_fee_per_gas, + to, + value, + access_list, + authorization_list, + input, + ) + } + + fn encode_value_to(&self, buf: &mut B) + where + B: bytes::BufMut + AsMut<[u8]>, + { + let TxEip7702 { + chain_id, + nonce, + gas_limit, + max_fee_per_gas, + max_priority_fee_per_gas, + to, + value, + access_list, + authorization_list, + input, + } = self; + by_props!( + @encode buf; + chain_id, + nonce, + gas_limit, + max_fee_per_gas, + max_priority_fee_per_gas, + to, + value, + access_list, + authorization_list, + input, + ) + } + + fn decode_value(mut data: &[u8]) -> Result + where + Self: Sized, + { + let mut tx = TxEip7702::default(); + let TxEip7702 { + chain_id, + nonce, + gas_limit, + max_fee_per_gas, + max_priority_fee_per_gas, + to, + value, + access_list, + authorization_list, + input, + } = &mut tx; + by_props!( + @decode data; + chain_id, + nonce, + gas_limit, + max_fee_per_gas, + max_priority_fee_per_gas, + to, + value, + access_list, + authorization_list, + input, + ); + Ok(tx) + } +} + +impl ValSer for Signed +where + T: ValSer, + Sig: ValSer, +{ + fn encoded_size(&self) -> usize { + self.signature().encoded_size() + self.tx().encoded_size() + } + + fn encode_value_to(&self, buf: &mut B) + where + B: bytes::BufMut + AsMut<[u8]>, + { + self.signature().encode_value_to(buf); + self.tx().encode_value_to(buf); + } + + fn decode_value(data: &[u8]) -> Result + where + Self: Sized, + { + let mut data = data; + + let signature = Sig::decode_value(data)?; + data = &data[signature.encoded_size()..]; + + let tx = T::decode_value(data)?; + + Ok(Signed::new_unhashed(tx, signature)) + } +} + +impl ValSer for TransactionSigned { + fn encoded_size(&self) -> usize { + self.tx_type().encoded_size() + + match self { + EthereumTxEnvelope::Legacy(signed) => signed.encoded_size(), + EthereumTxEnvelope::Eip2930(signed) => signed.encoded_size(), + EthereumTxEnvelope::Eip1559(signed) => signed.encoded_size(), + EthereumTxEnvelope::Eip4844(signed) => signed.encoded_size(), + EthereumTxEnvelope::Eip7702(signed) => signed.encoded_size(), + } + } + + fn encode_value_to(&self, buf: &mut B) + where + B: bytes::BufMut + AsMut<[u8]>, + { + self.tx_type().encode_value_to(buf); + match self { + EthereumTxEnvelope::Legacy(signed) => { + signed.encode_value_to(buf); + } + EthereumTxEnvelope::Eip2930(signed) => { + signed.encode_value_to(buf); + } + EthereumTxEnvelope::Eip1559(signed) => { + signed.encode_value_to(buf); + } + EthereumTxEnvelope::Eip4844(signed) => { + signed.encode_value_to(buf); + } + EthereumTxEnvelope::Eip7702(signed) => { + signed.encode_value_to(buf); + } + } + } + + fn decode_value(data: &[u8]) -> Result + where + Self: Sized, + { + let ty = TxType::decode_value(data)?; + let data = &data[ty.encoded_size()..]; + match ty { + TxType::Legacy => ValSer::decode_value(data).map(EthereumTxEnvelope::Legacy), + TxType::Eip2930 => ValSer::decode_value(data).map(EthereumTxEnvelope::Eip2930), + TxType::Eip1559 => ValSer::decode_value(data).map(EthereumTxEnvelope::Eip1559), + TxType::Eip4844 => ValSer::decode_value(data).map(EthereumTxEnvelope::Eip4844), + TxType::Eip7702 => ValSer::decode_value(data).map(EthereumTxEnvelope::Eip7702), + } + } +} + + diff --git a/crates/storage/src/ser/traits.rs b/crates/storage/src/ser/traits.rs index 5d4c490..78877af 100644 --- a/crates/storage/src/ser/traits.rs +++ b/crates/storage/src/ser/traits.rs @@ -57,6 +57,11 @@ pub trait KeySer: Ord + Sized { /// Trait for value serialization. pub trait ValSer { + /// The encoded size of the value in bytes. This MUST be accurate, as it is + /// used to allocate buffers for serialization. Inaccurate sizes may result + /// in panics. + fn encoded_size(&self) -> usize; + /// Serialize the value into bytes. fn encode_value_to(&self, buf: &mut B) where @@ -69,7 +74,7 @@ pub trait ValSer { buf.freeze().into() } - /// Deserialize the value from bytes, advancing the `data` slice. + /// Deserialize the value from bytes. fn decode_value(data: &[u8]) -> Result where Self: Sized; @@ -89,11 +94,13 @@ pub trait ValSer { } /// Deserialize the value from bytes, ensuring all bytes are consumed. - fn decode_value_exact(data: &mut &[u8]) -> Result + fn decode_value_exact(data: &[u8]) -> Result where Self: Sized, { let val = Self::decode_value(data)?; - data.is_empty().then_some(val).ok_or(DeserError::InexactDeser { extra_bytes: data.len() }) + (val.encoded_size() == data.len()) + .then_some(val) + .ok_or(DeserError::InexactDeser { extra_bytes: data.len() }) } } From 67c6c16db29fd5c43ed6852d037826aab038175b Mon Sep 17 00:00:00 2001 From: James Date: Wed, 14 Jan 2026 11:41:08 -0500 Subject: [PATCH 03/16] tests: ser tests --- crates/storage/src/ser/impls.rs | 12 +- crates/storage/src/ser/reth_impls.rs | 424 +++++++++++++++++++++++++++ crates/storage/src/ser/traits.rs | 27 +- 3 files changed, 450 insertions(+), 13 deletions(-) diff --git a/crates/storage/src/ser/impls.rs b/crates/storage/src/ser/impls.rs index 5dfa187..96537e8 100644 --- a/crates/storage/src/ser/impls.rs +++ b/crates/storage/src/ser/impls.rs @@ -1,4 +1,5 @@ use crate::ser::{DeserError, KeySer, MAX_KEY_SIZE, ValSer}; +use alloy::primitives::Bloom; use bytes::BufMut; macro_rules! delegate_val_to_key { @@ -136,7 +137,7 @@ impl KeySer for alloy::primitives::Address { } } -impl ValSer for alloy::primitives::Bloom { +impl ValSer for Bloom { fn encoded_size(&self) -> usize { self.as_slice().len() } @@ -152,11 +153,11 @@ impl ValSer for alloy::primitives::Bloom { where Self: Sized, { - if data.len() != 256 { + if data.len() < 256 { return Err(DeserError::InsufficientData { needed: 256, available: data.len() }); } let mut bloom = Self::default(); - bloom.as_mut_slice().copy_from_slice(data); + bloom.as_mut_slice().copy_from_slice(&data[..256]); Ok(bloom) } } @@ -234,8 +235,9 @@ where where Self: Sized, { - let flag = - data.get(0).ok_or(DeserError::InsufficientData { needed: 1, available: data.len() })?; + let flag = data + .first() + .ok_or(DeserError::InsufficientData { needed: 1, available: data.len() })?; match flag { 0 => Ok(None), 1 => Ok(Some(T::decode_value(&data[1..])?)), diff --git a/crates/storage/src/ser/reth_impls.rs b/crates/storage/src/ser/reth_impls.rs index 007c543..2ee25d4 100644 --- a/crates/storage/src/ser/reth_impls.rs +++ b/crates/storage/src/ser/reth_impls.rs @@ -1395,4 +1395,428 @@ impl ValSer for TransactionSigned { } } +#[cfg(test)] +mod tests { + use super::*; + use alloy::primitives::{ + Address, B256, Bloom, Bytes as AlloBytes, Signature, TxKind, U256, keccak256, + }; + use alloy::{ + consensus::{TxEip1559, TxEip2930, TxEip4844, TxEip7702, TxLegacy}, + eips::{ + eip2930::{AccessList, AccessListItem}, + eip7702::{Authorization, SignedAuthorization}, + }, + }; + use reth::primitives::{Account, Header, Log, LogData, TxType}; + use reth::revm::bytecode::JumpTable; + use reth_db_api::{BlockNumberList, models::StoredBlockBodyIndices}; + + /// Generic roundtrip test for any ValSer type + #[track_caller] + fn test_roundtrip(original: &T) + where + T: ValSer + PartialEq + std::fmt::Debug, + { + // Encode + let mut buf = bytes::BytesMut::new(); + original.encode_value_to(&mut buf); + let encoded = buf.freeze(); + + // Assert that the encoded size matches + assert_eq!( + original.encoded_size(), + encoded.len(), + "Encoded size mismatch: expected {}, got {}", + original.encoded_size(), + encoded.len() + ); + + // Decode + let decoded = T::decode_value(&encoded).expect("Failed to decode value"); + + // Assert equality + assert_eq!(*original, decoded, "Roundtrip failed"); + } + + #[test] + fn test_blocknumberlist_roundtrips() { + // Empty list + test_roundtrip(&BlockNumberList::empty()); + + // Single item + let mut single = BlockNumberList::empty(); + single.push(42u64).unwrap(); + test_roundtrip(&single); + + // Multiple items + let mut multiple = BlockNumberList::empty(); + for i in [0, 1, 255, 256, 65535, 65536, u64::MAX] { + multiple.push(i).unwrap(); + } + test_roundtrip(&multiple); + } + + #[test] + fn test_account_roundtrips() { + // Default account + test_roundtrip(&Account::default()); + + // Account with values + let account = Account { + nonce: 42, + balance: U256::from(123456789u64), + bytecode_hash: Some(keccak256(b"hello world")), + }; + test_roundtrip(&account); + + // Account with max values + let max_account = Account { + nonce: u64::MAX, + balance: U256::MAX, + bytecode_hash: Some(B256::from([0xFF; 32])), + }; + test_roundtrip(&max_account); + } + + #[test] + fn test_header_roundtrips() { + // Default header + test_roundtrip(&Header::default()); + + // Header with some values + let header = Header { + number: 12345, + gas_limit: 8000000, + timestamp: 1234567890, + difficulty: U256::from(1000000u64), + ..Default::default() + }; + test_roundtrip(&header); + } + + #[test] + fn test_logdata_roundtrips() { + // Empty log data + test_roundtrip(&LogData::new_unchecked(vec![], AlloBytes::new())); + + // Log data with one topic + test_roundtrip(&LogData::new_unchecked( + vec![B256::from([1; 32])], + AlloBytes::from_static(b"hello"), + )); + + // Log data with multiple topics + test_roundtrip(&LogData::new_unchecked( + vec![ + B256::from([1; 32]), + B256::from([2; 32]), + B256::from([3; 32]), + B256::from([4; 32]), + ], + AlloBytes::from_static(b"world"), + )); + } + + #[test] + fn test_log_roundtrips() { + let log_data = LogData::new_unchecked( + vec![B256::from([1; 32]), B256::from([2; 32])], + AlloBytes::from_static(b"test log data"), + ); + let log = Log { address: Address::from([0x42; 20]), data: log_data }; + test_roundtrip(&log); + } + + #[test] + fn test_txtype_roundtrips() { + test_roundtrip(&TxType::Legacy); + test_roundtrip(&TxType::Eip2930); + test_roundtrip(&TxType::Eip1559); + test_roundtrip(&TxType::Eip4844); + test_roundtrip(&TxType::Eip7702); + } + + #[test] + fn test_stored_block_body_indices_roundtrips() { + test_roundtrip(&StoredBlockBodyIndices { first_tx_num: 0, tx_count: 0 }); + test_roundtrip(&StoredBlockBodyIndices { first_tx_num: 12345, tx_count: 67890 }); + + test_roundtrip(&StoredBlockBodyIndices { first_tx_num: u64::MAX, tx_count: u64::MAX }); + } + + #[test] + fn test_signature_roundtrips() { + test_roundtrip(&Signature::test_signature()); + + // Zero signature + let zero_sig = Signature::new(U256::ZERO, U256::ZERO, false); + test_roundtrip(&zero_sig); + + // Max signature + let max_sig = Signature::new(U256::MAX, U256::MAX, true); + test_roundtrip(&max_sig); + } + + #[test] + fn test_txkind_roundtrips() { + test_roundtrip(&TxKind::Create); + test_roundtrip(&TxKind::Call(Address::ZERO)); + test_roundtrip(&TxKind::Call(Address::from([0xFF; 20]))); + } + + #[test] + fn test_accesslist_roundtrips() { + // Empty access list + test_roundtrip(&AccessList::default()); + + // Access list with one item + let item = AccessListItem { + address: Address::from([0x12; 20]), + storage_keys: vec![B256::from([0x34; 32])], + }; + test_roundtrip(&AccessList(vec![item])); + + // Access list with multiple items and keys + let items = vec![ + AccessListItem { + address: Address::repeat_byte(11), + storage_keys: vec![B256::from([0x22; 32]), B256::from([0x33; 32])], + }, + AccessListItem { address: Address::from([0x44; 20]), storage_keys: vec![] }, + AccessListItem { + address: Address::from([0x55; 20]), + storage_keys: vec![B256::from([0x66; 32])], + }, + ]; + test_roundtrip(&AccessList(items)); + } + + #[test] + fn test_authorization_roundtrips() { + test_roundtrip(&Authorization { + chain_id: U256::from(1u64), + address: Address::repeat_byte(11), + nonce: 0, + }); + + test_roundtrip(&Authorization { + chain_id: U256::MAX, + address: Address::from([0xFF; 20]), + nonce: u64::MAX, + }); + } + + #[test] + fn test_signed_authorization_roundtrips() { + let auth = Authorization { + chain_id: U256::from(1u64), + address: Address::repeat_byte(11), + nonce: 42, + }; + let signed_auth = + SignedAuthorization::new_unchecked(auth, 1, U256::from(12345u64), U256::from(67890u64)); + test_roundtrip(&signed_auth); + } + + #[test] + fn test_tx_legacy_roundtrips() { + test_roundtrip(&TxLegacy::default()); + + let tx = TxLegacy { + chain_id: Some(1), + nonce: 42, + gas_price: 20_000_000_000, + gas_limit: 21000u64, + to: TxKind::Call(Address::repeat_byte(11)), + value: U256::from(1000000000000000000u64), // 1 ETH in wei + input: AlloBytes::from_static(b"hello world"), + }; + test_roundtrip(&tx); + } + + #[test] + fn test_tx_eip2930_roundtrips() { + test_roundtrip(&TxEip2930::default()); + + let access_list = AccessList(vec![AccessListItem { + address: Address::from([0x22; 20]), + storage_keys: vec![B256::from([0x33; 32])], + }]); + + let tx = TxEip2930 { + chain_id: 1, + nonce: 42, + gas_price: 20_000_000_000, + gas_limit: 21000u64, + to: TxKind::Call(Address::repeat_byte(11)), + value: U256::from(1000000000000000000u64), + input: AlloBytes::from_static(b"eip2930 tx"), + access_list, + }; + test_roundtrip(&tx); + } + + #[test] + fn test_tx_eip1559_roundtrips() { + test_roundtrip(&TxEip1559::default()); + + let tx = TxEip1559 { + chain_id: 1, + nonce: 42, + gas_limit: 21000u64, + max_fee_per_gas: 30_000_000_000, + max_priority_fee_per_gas: 2_000_000_000, + to: TxKind::Call(Address::repeat_byte(11)), + value: U256::from(1000000000000000000u64), + input: AlloBytes::from_static(b"eip1559 tx"), + access_list: AccessList::default(), + }; + test_roundtrip(&tx); + } + + #[test] + fn test_tx_eip4844_roundtrips() { + test_roundtrip(&TxEip4844::default()); + + let tx = TxEip4844 { + chain_id: 1, + nonce: 42, + gas_limit: 21000u64, + max_fee_per_gas: 30_000_000_000, + max_priority_fee_per_gas: 2_000_000_000, + to: Address::repeat_byte(11), + value: U256::from(1000000000000000000u64), + input: AlloBytes::from_static(b"eip4844 tx"), + access_list: AccessList::default(), + blob_versioned_hashes: vec![B256::from([0x44; 32])], + max_fee_per_blob_gas: 1_000_000, + }; + test_roundtrip(&tx); + } + + #[test] + fn test_tx_eip7702_roundtrips() { + test_roundtrip(&TxEip7702::default()); + + let auth = SignedAuthorization::new_unchecked( + Authorization { + chain_id: U256::from(1u64), + address: Address::from([0x77; 20]), + nonce: 0, + }, + 1, + U256::from(12345u64), + U256::from(67890u64), + ); + + let tx = TxEip7702 { + chain_id: 1, + nonce: 42, + gas_limit: 21000u64, + max_fee_per_gas: 30_000_000_000, + max_priority_fee_per_gas: 2_000_000_000, + to: Address::repeat_byte(11), + value: U256::from(1000000000000000000u64), + input: AlloBytes::from_static(b"eip7702 tx"), + access_list: AccessList::default(), + authorization_list: vec![auth], + }; + test_roundtrip(&tx); + } + + #[test] + fn test_jump_table_roundtrips() { + // Empty jump table + test_roundtrip(&JumpTable::default()); + + // Jump table with some jumps + let jump_table = JumpTable::from_slice(&[0b10101010, 0b01010101], 16); + test_roundtrip(&jump_table); + } + + #[test] + fn test_complex_combinations() { + // Test a complex Header with all fields populated + let header = Header { + number: 12345, + gas_limit: 8000000, + timestamp: 1234567890, + difficulty: U256::from(1000000u64), + parent_hash: keccak256(b"parent"), + ommers_hash: keccak256(b"ommers"), + beneficiary: Address::from([0xBE; 20]), + state_root: keccak256(b"state"), + transactions_root: keccak256(b"txs"), + receipts_root: keccak256(b"receipts"), + logs_bloom: Bloom::default(), + gas_used: 7999999, + mix_hash: keccak256(b"mix"), + nonce: [0x42; 8].into(), + extra_data: AlloBytes::from_static(b"extra data"), + base_fee_per_gas: Some(1000000000), + withdrawals_root: Some(keccak256(b"withdrawals_root")), + blob_gas_used: Some(500000), + excess_blob_gas: Some(10000), + parent_beacon_block_root: Some(keccak256(b"parent_beacon_block_root")), + requests_hash: Some(keccak256(b"requests_hash")), + }; + test_roundtrip(&header); + + // Test a complex EIP-1559 transaction + let access_list = AccessList(vec![ + AccessListItem { + address: Address::repeat_byte(11), + storage_keys: vec![B256::from([0x22; 32]), B256::from([0x33; 32])], + }, + AccessListItem { address: Address::from([0x44; 20]), storage_keys: vec![] }, + ]); + + let complex_tx = TxEip1559 { + chain_id: 1, + nonce: 123456, + gas_limit: 500000u64, + max_fee_per_gas: 50_000_000_000, + max_priority_fee_per_gas: 3_000_000_000, + to: TxKind::Create, + value: U256::ZERO, + input: AlloBytes::copy_from_slice(&[0xFF; 1000]), // Large input + access_list, + }; + test_roundtrip(&complex_tx); + } + + #[test] + fn test_edge_cases() { + // Very large access list + let large_storage_keys: Vec = + (0..1000).map(|i| B256::from(U256::from(i).to_be_bytes::<32>())).collect(); + let large_access_list = AccessList(vec![AccessListItem { + address: Address::from([0xAA; 20]), + storage_keys: large_storage_keys, + }]); + test_roundtrip(&large_access_list); + + // Transaction with maximum values + let max_tx = TxEip1559 { + chain_id: u64::MAX, + nonce: u64::MAX, + gas_limit: u64::MAX, + max_fee_per_gas: u128::MAX, + max_priority_fee_per_gas: u128::MAX, + to: TxKind::Call(Address::repeat_byte(0xFF)), + value: U256::MAX, + input: AlloBytes::copy_from_slice(&[0xFF; 10000]), // Very large input + access_list: AccessList::default(), + }; + test_roundtrip(&max_tx); + + // BlockNumberList with many numbers + let mut large_list = BlockNumberList::empty(); + for i in 0..10000u64 { + large_list.push(i).unwrap(); + } + test_roundtrip(&large_list); + } +} diff --git a/crates/storage/src/ser/traits.rs b/crates/storage/src/ser/traits.rs index 78877af..88e1fef 100644 --- a/crates/storage/src/ser/traits.rs +++ b/crates/storage/src/ser/traits.rs @@ -6,6 +6,14 @@ pub const MAX_KEY_SIZE: usize = 64; /// Trait for key serialization with fixed-size keys of size no greater than 32 /// bytes. +/// +/// Keys must be FIXED SIZE, of size no greater than `MAX_KEY_SIZE` (64), and +/// no less than 1. The serialization must preserve ordering, i.e., for any two +/// keys `k1` and `k2`, if `k1 > k2`, then the byte representation of `k1` +/// must be lexicographically greater than that of `k2`. +/// +/// In practice, keys are often hashes, addresses, numbers, or composites +/// of these. pub trait KeySer: Ord + Sized { /// The fixed size of the serialized key in bytes. /// Must satisfy `SIZE <= MAX_KEY_SIZE`. @@ -21,16 +29,11 @@ pub trait KeySer: Ord + Sized { assert!(Self::SIZE > 0, "KeySer implementations must have SIZE > 0"); }; - /// Encode the key into the provided buffer. - /// - /// Writes exactly `SIZE` bytes to `buf[..SIZE]` and returns `SIZE`. - /// The encoding must preserve ordering: for any `k1, k2` where `k1 > k2`, - /// the bytes written by `k1` must be lexicographically greater than those - /// of `k2`. + /// Encode the key, optionally using the provided buffer. /// /// # Returns /// - /// A slice containing the encoded key. This may be a slice of buf, or may + /// A slice containing the encoded key. This may be a slice of `buf`, or may /// be borrowed from the key itself. This slice must be <= `SIZE` bytes. fn encode_key<'a: 'c, 'b: 'c, 'c>(&'a self, buf: &'b mut [u8; MAX_KEY_SIZE]) -> &'c [u8]; @@ -56,10 +59,18 @@ pub trait KeySer: Ord + Sized { } /// Trait for value serialization. +/// +/// Values can be of variable size, but must implement accurate size reporting. +/// When serialized, value sizes must be self-describing. I.e. the value must +/// tolerate being deserialized from a byte slice of arbitrary length, consuming +/// only as many bytes as needed. +/// +/// E.g. a correct implementation for an array serializes the length of the +/// array first, so that the deserializer knows how many items to expect. pub trait ValSer { /// The encoded size of the value in bytes. This MUST be accurate, as it is /// used to allocate buffers for serialization. Inaccurate sizes may result - /// in panics. + /// in panics or incorrect behavior. fn encoded_size(&self) -> usize; /// Serialize the value into bytes. From 9173fd22f8b8615ed465da60309e482f0b3132e5 Mon Sep 17 00:00:00 2001 From: James Date: Wed, 14 Jan 2026 14:14:13 -0500 Subject: [PATCH 04/16] feat: hot mem and mdbx --- Cargo.toml | 3 +- crates/storage/Cargo.toml | 1 + crates/storage/src/hot/db.rs | 87 ---- crates/storage/src/hot/error.rs | 16 +- crates/storage/src/hot/mem.rs | 637 ++++++++++++++++++++++++++++ crates/storage/src/hot/mod.rs | 8 +- crates/storage/src/hot/reth_impl.rs | 59 ++- crates/storage/src/hot/traits.rs | 24 +- crates/storage/src/tables/cold.rs | 1 + 9 files changed, 731 insertions(+), 105 deletions(-) delete mode 100644 crates/storage/src/hot/db.rs create mode 100644 crates/storage/src/hot/mem.rs diff --git a/Cargo.toml b/Cargo.toml index 856dba7..1507312 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -5,7 +5,7 @@ resolver = "2" [workspace.package] version = "0.16.0-rc.4" edition = "2024" -rust-version = "1.88" +rust-version = "1.92" authors = ["init4"] license = "MIT OR Apache-2.0" homepage = "https://github.com/init4tech/signet-sdk" @@ -85,6 +85,7 @@ reth-eth-wire-types = { git = "https://github.com/paradigmxyz/reth", tag = "v1.9 reth-evm-ethereum = { git = "https://github.com/paradigmxyz/reth", tag = "v1.9.1" } reth-exex = { git = "https://github.com/paradigmxyz/reth", tag = "v1.9.1" } reth-exex-test-utils = { git = "https://github.com/paradigmxyz/reth", tag = "v1.9.1" } +reth-libmdbx = { git = "https://github.com/paradigmxyz/reth", tag = "v1.9.1" } reth-network-api = { git = "https://github.com/paradigmxyz/reth", tag = "v1.9.1" } reth-network-peers = { git = "https://github.com/paradigmxyz/reth", tag = "v1.9.1" } reth-node-api = { git = "https://github.com/paradigmxyz/reth", tag = "v1.9.1" } diff --git a/crates/storage/Cargo.toml b/crates/storage/Cargo.toml index 0384e1d..312a17f 100644 --- a/crates/storage/Cargo.toml +++ b/crates/storage/Cargo.toml @@ -14,4 +14,5 @@ bytes = "1.11.0" reth.workspace = true reth-db.workspace = true reth-db-api.workspace = true +reth-libmdbx.workspace = true thiserror.workspace = true diff --git a/crates/storage/src/hot/db.rs b/crates/storage/src/hot/db.rs deleted file mode 100644 index a4911da..0000000 --- a/crates/storage/src/hot/db.rs +++ /dev/null @@ -1,87 +0,0 @@ -use crate::hot::{HotKv, HotKvError, HotKvRead, HotKvWrite}; -use std::{ - borrow::Cow, - sync::atomic::{AtomicBool, Ordering}, -}; - -/// Hot database wrapper around a key-value storage. -#[derive(Debug)] -pub struct HotDb { - inner: Inner, - - write_locked: AtomicBool, -} - -impl HotDb { - /// Create a new HotDb wrapping the given inner KV storage. - pub const fn new(inner: Inner) -> Self { - Self { inner, write_locked: AtomicBool::new(false) } - } - - /// Get a read-only handle. - pub fn reader(&self) -> Result - where - Inner: HotKv, - { - self.inner.reader() - } - - /// Get a write handle, if available. If not available, returns - /// [`HotKvError::WriteLocked`]. - pub fn writer(&self) -> Result, HotKvError> - where - Inner: HotKv, - { - if self.write_locked.swap(true, Ordering::AcqRel) { - return Err(HotKvError::WriteLocked); - } - self.inner.writer().map(Some).map(|tx| WriteGuard { tx, db: self }) - } -} - -/// Write guard for a write transaction. -#[derive(Debug)] -pub struct WriteGuard<'a, Inner> -where - Inner: HotKv, -{ - tx: Option, - db: &'a HotDb, -} - -impl Drop for WriteGuard<'_, Inner> -where - Inner: HotKv, -{ - fn drop(&mut self) { - self.db.write_locked.store(false, Ordering::Release); - } -} - -impl HotKvRead for WriteGuard<'_, Inner> -where - Inner: HotKv, -{ - type Error = <::RwTx as HotKvRead>::Error; - - fn get_raw<'a>( - &'a self, - table: &str, - key: &[u8], - ) -> Result>, Self::Error> { - self.tx.as_ref().expect("present until drop").get_raw(table, key) - } -} - -impl HotKvWrite for WriteGuard<'_, Inner> -where - Inner: HotKv, -{ - fn queue_raw_put(&mut self, table: &str, key: &[u8], value: &[u8]) -> Result<(), Self::Error> { - self.tx.as_mut().expect("present until drop").queue_raw_put(table, key, value) - } - - fn raw_commit(mut self) -> Result<(), Self::Error> { - self.tx.take().expect("present until drop").raw_commit() - } -} diff --git a/crates/storage/src/hot/error.rs b/crates/storage/src/hot/error.rs index a5331ce..8dd0924 100644 --- a/crates/storage/src/hot/error.rs +++ b/crates/storage/src/hot/error.rs @@ -1,3 +1,5 @@ +use crate::ser::DeserError; + /// Trait for hot storage read/write errors. #[derive(thiserror::Error, Debug)] pub enum HotKvError { @@ -7,7 +9,7 @@ pub enum HotKvError { /// Deserialization error. Indicates an issue deserializing a key or value. #[error("Deserialization error: {0}")] - DeserError(#[from] crate::ser::DeserError), + Deser(#[from] crate::ser::DeserError), /// Indicates that a write transaction is already in progress. #[error("A write transaction is already in progress")] @@ -24,5 +26,17 @@ impl HotKvError { } } +/// Trait to convert specific read errors into `HotKvError`. +pub trait HotKvReadError: std::error::Error + From + Send + Sync + 'static { + /// Convert the error into a `HotKvError`. + fn into_hot_kv_error(self) -> HotKvError; +} + +impl HotKvReadError for HotKvError { + fn into_hot_kv_error(self) -> HotKvError { + self + } +} + /// Result type for hot storage operations. pub type HotKvResult = Result; diff --git a/crates/storage/src/hot/mem.rs b/crates/storage/src/hot/mem.rs new file mode 100644 index 0000000..ca2ed07 --- /dev/null +++ b/crates/storage/src/hot/mem.rs @@ -0,0 +1,637 @@ +use crate::{ + hot::{HotKv, HotKvError, HotKvRead, HotKvWrite}, + ser::MAX_KEY_SIZE, +}; +use std::{ + borrow::Cow, + collections::BTreeMap, + sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard}, +}; + +type Table = BTreeMap<[u8; MAX_KEY_SIZE], bytes::Bytes>; +type Store = BTreeMap; + +type TableOp = BTreeMap<[u8; MAX_KEY_SIZE], QueuedOp>; +type OpStore = BTreeMap; + +/// A simple in-memory key-value store using a BTreeMap. +/// +/// This implementation supports concurrent multiple concurrent read +/// transactions. Write transactions are exclusive, and cannot overlap +/// with other read or write transactions. +#[derive(Clone)] +pub struct MemKv { + map: Arc>, +} + +impl core::fmt::Debug for MemKv { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + f.debug_struct("MemKv").finish() + } +} + +impl MemKv { + /// Create a new empty in-memory KV store. + pub fn new() -> Self { + Self { map: Arc::new(RwLock::new(BTreeMap::new())) } + } + + #[track_caller] + fn key(k: &[u8]) -> [u8; MAX_KEY_SIZE] { + assert!(k.len() <= MAX_KEY_SIZE, "Key length exceeds MAX_KEY_SIZE"); + let mut buf = [0u8; MAX_KEY_SIZE]; + buf[..k.len()].copy_from_slice(k); + buf + } +} + +impl Default for MemKv { + fn default() -> Self { + Self::new() + } +} + +/// Read-only transaction for MemKv. +pub struct MemKvRoTx { + guard: RwLockReadGuard<'static, Store>, +} + +impl core::fmt::Debug for MemKvRoTx { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + f.debug_struct("MemKvRoTx").finish() + } +} + +// SAFETY: MemKvRoTx holds a read guard which ensures the data remains valid +unsafe impl Send for MemKvRoTx {} +unsafe impl Sync for MemKvRoTx {} + +/// Read-write transaction for MemKv. +pub struct MemKvRwTx { + guard: RwLockWriteGuard<'static, Store>, + queued_ops: OpStore, +} + +impl MemKvRwTx { + fn commit_inner(&mut self) { + let ops = std::mem::take(&mut self.queued_ops); + + for (table, table_ops) in ops.into_iter() { + for (key, op) in table_ops.into_iter() { + match op { + QueuedOp::Put { value } => { + self.guard.entry(table.clone()).or_default().insert(key, value); + } + QueuedOp::Delete => { + if let Some(t) = self.guard.get_mut(&table) { + t.remove(&key); + } + } + } + } + } + } + + /// Downgrade the transaction to a read-only transaction without + /// committing, discarding queued changes. + pub fn downgrade(self) -> MemKvRoTx { + let guard = RwLockWriteGuard::downgrade(self.guard); + + MemKvRoTx { guard } + } + + /// Commit the transaction and downgrade to a read-only transaction. + pub fn commit_downgrade(mut self) -> MemKvRoTx { + self.commit_inner(); + + let guard = RwLockWriteGuard::downgrade(self.guard); + + MemKvRoTx { guard } + } +} + +impl core::fmt::Debug for MemKvRwTx { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + f.debug_struct("MemKvRwTx").finish() + } +} + +#[derive(Debug, Clone)] +enum QueuedOp { + Delete, + Put { value: bytes::Bytes }, +} + +// SAFETY: MemKvRwTx holds a write guard which ensures exclusive access +unsafe impl Send for MemKvRwTx {} + +impl HotKv for MemKv { + type RoTx = MemKvRoTx; + type RwTx = MemKvRwTx; + + fn reader(&self) -> Result { + let guard = self + .map + .try_read() + .map_err(|_| HotKvError::Inner("Failed to acquire read lock".into()))?; + + // SAFETY: We extend the lifetime to 'static, but the guard ensures the + // data remains valid for the lifetime of the transaction + let guard: RwLockReadGuard<'static, Store> = unsafe { std::mem::transmute(guard) }; + + Ok(MemKvRoTx { guard }) + } + + fn writer(&self) -> Result { + let guard = self.map.try_write().map_err(|_| HotKvError::WriteLocked)?; + + // SAFETY: We extend the lifetime to 'static, but the guard ensures exclusive access + let guard: RwLockWriteGuard<'static, Store> = unsafe { std::mem::transmute(guard) }; + + Ok(MemKvRwTx { guard, queued_ops: OpStore::new() }) + } +} + +impl HotKvRead for MemKvRoTx { + type Error = HotKvError; + + fn get_raw<'a>( + &'a self, + table: &str, + key: &[u8], + ) -> Result>, Self::Error> { + // Check queued operations first (read-your-writes consistency) + let key = MemKv::key(key); + + // SAFETY: The guard ensures the map remains valid + + Ok(self + .guard + .get(table) + .and_then(|t| t.get(&key)) + .map(|bytes| Cow::Borrowed(bytes.as_ref()))) + } +} + +impl HotKvRead for MemKvRwTx { + type Error = HotKvError; + + fn get_raw<'a>( + &'a self, + table: &str, + key: &[u8], + ) -> Result>, Self::Error> { + // Check queued operations first (read-your-writes consistency) + let key = MemKv::key(key); + + match self.queued_ops.get(table).and_then(|t| t.get(&key)) { + Some(QueuedOp::Put { value }) => { + return Ok(Some(Cow::Borrowed(value.as_ref()))); + } + Some(QueuedOp::Delete) => { + return Ok(None); + } + None => {} + } + + // If not found in queued ops, check the underlying map + Ok(self + .guard + .get(table) + .and_then(|t| t.get(&key)) + .map(|bytes| Cow::Borrowed(bytes.as_ref()))) + } +} + +impl HotKvWrite for MemKvRwTx { + fn queue_raw_put(&mut self, table: &str, key: &[u8], value: &[u8]) -> Result<(), Self::Error> { + let key = MemKv::key(key); + + let value_bytes = bytes::Bytes::copy_from_slice(value); + + self.queued_ops + .entry(table.to_owned()) + .or_default() + .insert(key, QueuedOp::Put { value: value_bytes }); + Ok(()) + } + + fn queue_raw_delete(&mut self, table: &str, key: &[u8]) -> Result<(), Self::Error> { + let key = MemKv::key(key); + + self.queued_ops.entry(table.to_owned()).or_default().insert(key, QueuedOp::Delete); + Ok(()) + } + + fn raw_commit(mut self) -> Result<(), Self::Error> { + // Apply all queued operations to the map + self.commit_inner(); + + // The write guard is automatically dropped here, releasing the lock + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::tables::Table; + use alloy::primitives::{Address, U256}; + use bytes::Bytes; + + // Test table definitions + #[derive(Debug)] + struct TestTable; + + impl Table for TestTable { + const NAME: &'static str = "test_table"; + type Key = u64; + type Value = Bytes; + } + + #[derive(Debug)] + struct AddressTable; + + impl Table for AddressTable { + const NAME: &'static str = "addresses"; + type Key = Address; + type Value = U256; + } + + #[test] + fn test_new_store() { + let store = MemKv::new(); + let reader = store.reader().unwrap(); + + // Empty store should return None for any key + assert!(reader.get_raw("test", &[1, 2, 3]).unwrap().is_none()); + } + + #[test] + fn test_basic_put_get() { + let store = MemKv::new(); + + // Write some data + { + let mut writer = store.writer().unwrap(); + writer.queue_raw_put("table1", &[1, 2, 3], b"value1").unwrap(); + writer.queue_raw_put("table1", &[4, 5, 6], b"value2").unwrap(); + writer.raw_commit().unwrap(); + } + + // Read the data back + { + let reader = store.reader().unwrap(); + let value1 = reader.get_raw("table1", &[1, 2, 3]).unwrap(); + let value2 = reader.get_raw("table1", &[4, 5, 6]).unwrap(); + let missing = reader.get_raw("table1", &[7, 8, 9]).unwrap(); + + assert_eq!(value1.as_deref(), Some(b"value1" as &[u8])); + assert_eq!(value2.as_deref(), Some(b"value2" as &[u8])); + assert!(missing.is_none()); + } + } + + #[test] + fn test_multiple_tables() { + let store = MemKv::new(); + + // Write to different tables + { + let mut writer = store.writer().unwrap(); + writer.queue_raw_put("table1", &[1], b"table1_value").unwrap(); + writer.queue_raw_put("table2", &[1], b"table2_value").unwrap(); + writer.raw_commit().unwrap(); + } + + // Read from different tables + { + let reader = store.reader().unwrap(); + let value1 = reader.get_raw("table1", &[1]).unwrap(); + let value2 = reader.get_raw("table2", &[1]).unwrap(); + + assert_eq!(value1.as_deref(), Some(b"table1_value" as &[u8])); + assert_eq!(value2.as_deref(), Some(b"table2_value" as &[u8])); + } + } + + #[test] + fn test_overwrite_value() { + let store = MemKv::new(); + + // Write initial value + { + let mut writer = store.writer().unwrap(); + writer.queue_raw_put("table1", &[1], b"original").unwrap(); + writer.raw_commit().unwrap(); + } + + // Overwrite with new value + { + let mut writer = store.writer().unwrap(); + writer.queue_raw_put("table1", &[1], b"updated").unwrap(); + writer.raw_commit().unwrap(); + } + + // Check the value was updated + { + let reader = store.reader().unwrap(); + let value = reader.get_raw("table1", &[1]).unwrap(); + assert_eq!(value.as_deref(), Some(b"updated" as &[u8])); + } + } + + #[test] + fn test_read_your_writes() { + let store = MemKv::new(); + let mut writer = store.writer().unwrap(); + + // Queue some operations but don't commit yet + writer.queue_raw_put("table1", &[1], b"queued_value").unwrap(); + + // Should be able to read the queued value + let value = writer.get_raw("table1", &[1]).unwrap(); + assert_eq!(value.as_deref(), Some(b"queued_value" as &[u8])); + + writer.raw_commit().unwrap(); + + // After commit, other readers should see it + { + let reader = store.reader().unwrap(); + let value = reader.get_raw("table1", &[1]).unwrap(); + assert_eq!(value.as_deref(), Some(b"queued_value" as &[u8])); + } + } + + #[test] + fn test_typed_operations() { + let store = MemKv::new(); + + // Write using typed interface + { + let mut writer = store.writer().unwrap(); + writer.queue_put::(&42u64, &Bytes::from_static(b"hello world")).unwrap(); + writer.queue_put::(&100u64, &Bytes::from_static(b"another value")).unwrap(); + writer.raw_commit().unwrap(); + } + + // Read using typed interface + { + let reader = store.reader().unwrap(); + let value1 = reader.get::(&42u64).unwrap(); + let value2 = reader.get::(&100u64).unwrap(); + let missing = reader.get::(&999u64).unwrap(); + + assert_eq!(value1, Some(Bytes::from_static(b"hello world"))); + assert_eq!(value2, Some(Bytes::from_static(b"another value"))); + assert!(missing.is_none()); + } + } + + #[test] + fn test_address_table() { + let store = MemKv::new(); + + let addr1 = Address::from([0x11; 20]); + let addr2 = Address::from([0x22; 20]); + let balance1 = U256::from(1000u64); + let balance2 = U256::from(2000u64); + + // Write address data + { + let mut writer = store.writer().unwrap(); + writer.queue_put::(&addr1, &balance1).unwrap(); + writer.queue_put::(&addr2, &balance2).unwrap(); + writer.raw_commit().unwrap(); + } + + // Read address data + { + let reader = store.reader().unwrap(); + let bal1 = reader.get::(&addr1).unwrap(); + let bal2 = reader.get::(&addr2).unwrap(); + + assert_eq!(bal1, Some(balance1)); + assert_eq!(bal2, Some(balance2)); + } + } + + #[test] + fn test_batch_operations() { + let store = MemKv::new(); + + let entries = [ + (1u64, Bytes::from_static(b"first")), + (2u64, Bytes::from_static(b"second")), + (3u64, Bytes::from_static(b"third")), + ]; + + // Write batch + { + let mut writer = store.writer().unwrap(); + let entry_refs: Vec<_> = entries.iter().map(|(k, v)| (k, v)).collect(); + writer.queue_put_many::(entry_refs).unwrap(); + writer.raw_commit().unwrap(); + } + + // Read batch + { + let reader = store.reader().unwrap(); + let keys: Vec<_> = entries.iter().map(|(k, _)| k).collect(); + let values = reader.get_many::(keys).unwrap(); + + assert_eq!(values.len(), 3); + assert_eq!(values[0], Some(Bytes::from_static(b"first"))); + assert_eq!(values[1], Some(Bytes::from_static(b"second"))); + assert_eq!(values[2], Some(Bytes::from_static(b"third"))); + } + } + + #[test] + fn test_concurrent_readers() { + let store = MemKv::new(); + + // Write some initial data + { + let mut writer = store.writer().unwrap(); + writer.queue_raw_put("table1", &[1], b"value1").unwrap(); + writer.raw_commit().unwrap(); + } + + // Multiple readers should be able to read concurrently + let reader1 = store.reader().unwrap(); + let reader2 = store.reader().unwrap(); + + let value1 = reader1.get_raw("table1", &[1]).unwrap(); + let value2 = reader2.get_raw("table1", &[1]).unwrap(); + + assert_eq!(value1.as_deref(), Some(b"value1" as &[u8])); + assert_eq!(value2.as_deref(), Some(b"value1" as &[u8])); + } + + #[test] + fn test_write_lock_exclusivity() { + let store = MemKv::new(); + + // Get a writer + let _writer1 = store.writer().unwrap(); + + // Second writer should fail + match store.writer() { + Err(HotKvError::WriteLocked) => {} // Expected + Ok(_) => panic!("Should not be able to get second writer"), + Err(e) => panic!("Unexpected error: {:?}", e), + } + } + + #[test] + fn test_empty_values() { + let store = MemKv::new(); + + { + let mut writer = store.writer().unwrap(); + writer.queue_raw_put("table1", &[1], b"").unwrap(); + writer.raw_commit().unwrap(); + } + + { + let reader = store.reader().unwrap(); + let value = reader.get_raw("table1", &[1]).unwrap(); + assert_eq!(value.as_deref(), Some(b"" as &[u8])); + } + } + + #[test] + fn test_multiple_operations_same_transaction() { + let store = MemKv::new(); + + { + let mut writer = store.writer().unwrap(); + + // Multiple operations on same key - last one should win + writer.queue_raw_put("table1", &[1], b"first").unwrap(); + writer.queue_raw_put("table1", &[1], b"second").unwrap(); + writer.queue_raw_put("table1", &[1], b"third").unwrap(); + + // Read-your-writes should return the latest value + let value = writer.get_raw("table1", &[1]).unwrap(); + assert_eq!(value.as_deref(), Some(b"third" as &[u8])); + + writer.raw_commit().unwrap(); + } + + { + let reader = store.reader().unwrap(); + let value = reader.get_raw("table1", &[1]).unwrap(); + assert_eq!(value.as_deref(), Some(b"third" as &[u8])); + } + } + + #[test] + fn test_isolation() { + let store = MemKv::new(); + + // Write initial value + { + let mut writer = store.writer().unwrap(); + writer.queue_raw_put("table1", &[1], b"original").unwrap(); + writer.raw_commit().unwrap(); + } + + // Start a read transaction + { + let reader = store.reader().unwrap(); + let original_value = reader.get_raw("table1", &[1]).unwrap(); + assert_eq!(original_value.as_deref(), Some(b"original" as &[u8])); + } + + // Update the value in a separate transaction + { + let mut writer = store.writer().unwrap(); + writer.queue_raw_put("table1", &[1], b"updated").unwrap(); + writer.raw_commit().unwrap(); + } + + // The value should now be latest for new readers + { + // New reader should see the updated value + let new_reader = store.reader().unwrap(); + let updated_value = new_reader.get_raw("table1", &[1]).unwrap(); + assert_eq!(updated_value.as_deref(), Some(b"updated" as &[u8])); + } + } + + #[test] + fn test_rollback_on_drop() { + let store = MemKv::new(); + + { + let mut writer = store.writer().unwrap(); + writer.queue_raw_put("table1", &[1], b"should_not_persist").unwrap(); + // Drop without committing + } + + // Value should not be persisted + { + let reader = store.reader().unwrap(); + let value = reader.get_raw("table1", &[1]).unwrap(); + assert!(value.is_none()); + } + } + + #[test] + fn write_two_tables() { + let store = MemKv::new(); + + { + let mut writer = store.writer().unwrap(); + writer.queue_raw_put("table1", &[1], b"value1").unwrap(); + writer.queue_raw_put("table2", &[2], b"value2").unwrap(); + writer.raw_commit().unwrap(); + } + + { + let reader = store.reader().unwrap(); + let value1 = reader.get_raw("table1", &[1]).unwrap(); + let value2 = reader.get_raw("table2", &[2]).unwrap(); + + assert_eq!(value1.as_deref(), Some(b"value1" as &[u8])); + assert_eq!(value2.as_deref(), Some(b"value2" as &[u8])); + } + } + + #[test] + fn test_downgrades() { + let store = MemKv::new(); + { + // Write some data + // Start a read-write transaction + let mut rw_tx = store.writer().unwrap(); + rw_tx.queue_raw_put("table1", &[1, 2, 3], b"value1").unwrap(); + rw_tx.queue_raw_put("table1", &[4, 5, 6], b"value2").unwrap(); + + let ro_tx = rw_tx.commit_downgrade(); + + // Read the data back + let value1 = ro_tx.get_raw("table1", &[1, 2, 3]).unwrap(); + let value2 = ro_tx.get_raw("table1", &[4, 5, 6]).unwrap(); + + assert_eq!(value1.as_deref(), Some(b"value1" as &[u8])); + assert_eq!(value2.as_deref(), Some(b"value2" as &[u8])); + } + + { + // Start another read-write transaction + let mut rw_tx = store.writer().unwrap(); + rw_tx.queue_raw_put("table2", &[7, 8, 9], b"value3").unwrap(); + + // Value should not be set + let ro_tx = rw_tx.downgrade(); + + // Read the data back + let value3 = ro_tx.get_raw("table2", &[7, 8, 9]).unwrap(); + + assert!(value3.is_none()); + } + } +} diff --git a/crates/storage/src/hot/mod.rs b/crates/storage/src/hot/mod.rs index 2055543..8ef53fb 100644 --- a/crates/storage/src/hot/mod.rs +++ b/crates/storage/src/hot/mod.rs @@ -1,11 +1,11 @@ -mod db; -pub use db::{HotDb, WriteGuard}; - mod db_traits; pub use db_traits::{HotDbReader, HotDbWriter}; mod error; -pub use error::{HotKvError, HotKvResult}; +pub use error::{HotKvError, HotKvReadError, HotKvResult}; + +mod mem; +pub use mem::{MemKv, MemKvRoTx, MemKvRwTx}; mod reth_impl; diff --git a/crates/storage/src/hot/reth_impl.rs b/crates/storage/src/hot/reth_impl.rs index 706fad6..f085aba 100644 --- a/crates/storage/src/hot/reth_impl.rs +++ b/crates/storage/src/hot/reth_impl.rs @@ -1,8 +1,30 @@ +use crate::{ + hot::{HotKvRead, HotKvReadError, HotKvWrite}, + ser::DeserError, +}; +use reth_db::mdbx::{RW, TransactionKind, WriteFlags, tx::Tx}; +use reth_db_api::DatabaseError; use std::borrow::Cow; -use crate::{hot::HotKvRead, ser::DeserError}; -use reth_db::mdbx::{TransactionKind, tx::Tx}; -use reth_db_api::DatabaseError; +/// Error type for reth-libmdbx based hot storage. +#[derive(Debug, thiserror::Error)] +pub enum MdbxError { + /// Inner error + #[error(transparent)] + Mdbx(#[from] reth_libmdbx::Error), + /// Deser. + #[error(transparent)] + Deser(#[from] DeserError), +} + +impl HotKvReadError for MdbxError { + fn into_hot_kv_error(self) -> super::HotKvError { + match self { + MdbxError::Mdbx(e) => super::HotKvError::from_err(e), + MdbxError::Deser(e) => super::HotKvError::Deser(e), + } + } +} impl From for DatabaseError { fn from(value: DeserError) -> Self { @@ -14,19 +36,36 @@ impl HotKvRead for Tx where K: TransactionKind, { - type Error = DatabaseError; + type Error = MdbxError; fn get_raw<'a>( &'a self, table: &str, key: &[u8], ) -> Result>, Self::Error> { - let dbi = self - .inner - .open_db(Some(table)) - .map(|db| db.dbi()) - .map_err(|e| DatabaseError::Open(e.into()))?; + let dbi = self.inner.open_db(Some(table)).map(|db| db.dbi())?; + + self.inner.get(dbi, key.as_ref()).map_err(MdbxError::Mdbx) + } +} + +impl HotKvWrite for Tx { + fn queue_raw_put(&mut self, table: &str, key: &[u8], value: &[u8]) -> Result<(), Self::Error> { + let dbi = self.inner.open_db(Some(table)).map(|db| db.dbi())?; + + self.inner.put(dbi, key, value, WriteFlags::UPSERT).map(|_| ()).map_err(MdbxError::Mdbx) + } + + fn queue_raw_delete(&mut self, table: &str, key: &[u8]) -> Result<(), Self::Error> { + let dbi = self.inner.open_db(Some(table)).map(|db| db.dbi())?; + + self.inner.del(dbi, key, None).map(|_| ()).map_err(MdbxError::Mdbx) + } + + fn raw_commit(self) -> Result<(), Self::Error> { + // when committing, mdbx returns true on failure + let res = self.inner.commit()?; - self.inner.get(dbi, key.as_ref()).map_err(|err| DatabaseError::Read(err.into())) + if res.0 { Err(reth_libmdbx::Error::Other(1).into()) } else { Ok(()) } } } diff --git a/crates/storage/src/hot/traits.rs b/crates/storage/src/hot/traits.rs index 11f03c8..e6f09d8 100644 --- a/crates/storage/src/hot/traits.rs +++ b/crates/storage/src/hot/traits.rs @@ -1,7 +1,7 @@ use std::borrow::Cow; use crate::{ - hot::HotKvError, + hot::{HotKvError, HotKvReadError}, ser::{KeySer, MAX_KEY_SIZE, ValSer}, tables::Table, }; @@ -34,9 +34,12 @@ pub trait HotKv { /// Trait for hot storage read transactions. pub trait HotKvRead { /// Error type for read operations. - type Error: std::error::Error + From + Send + Sync + 'static; + type Error: HotKvReadError; /// Get a raw value from a specific table. + /// + /// The `key` buf must be <= [`MAX_KEY_SIZE`] bytes. Implementations are + /// allowed to panic if this is not the case. fn get_raw<'a>(&'a self, table: &str, key: &[u8]) -> Result>, Self::Error>; @@ -91,8 +94,17 @@ pub trait HotKvRead { /// Trait for hot storage write transactions. pub trait HotKvWrite: HotKvRead { /// Queue a raw put operation. + /// + /// The `key` buf must be <= [`MAX_KEY_SIZE`] bytes. Implementations are + /// allowed to panic if this is not the case. fn queue_raw_put(&mut self, table: &str, key: &[u8], value: &[u8]) -> Result<(), Self::Error>; + /// Queue a raw delete operation. + /// + /// The `key` buf must be <= [`MAX_KEY_SIZE`] bytes. Implementations are + /// allowed to panic if this is not the case. + fn queue_raw_delete(&mut self, table: &str, key: &[u8]) -> Result<(), Self::Error>; + /// Queue a put operation for a specific table. fn queue_put(&mut self, key: &T::Key, value: &T::Value) -> Result<(), Self::Error> { let mut key_buf = [0u8; MAX_KEY_SIZE]; @@ -102,6 +114,14 @@ pub trait HotKvWrite: HotKvRead { self.queue_raw_put(T::NAME, key_bytes, &value_bytes) } + /// Queue a delete operation for a specific table. + fn queue_delete(&mut self, key: &T::Key) -> Result<(), Self::Error> { + let mut key_buf = [0u8; MAX_KEY_SIZE]; + let key_bytes = key.encode_key(&mut key_buf); + + self.queue_raw_delete(T::NAME, key_bytes) + } + /// Queue many put operations for a specific table. fn queue_put_many<'a, 'b, T, I>(&mut self, entries: I) -> Result<(), Self::Error> where diff --git a/crates/storage/src/tables/cold.rs b/crates/storage/src/tables/cold.rs index e69de29..8b13789 100644 --- a/crates/storage/src/tables/cold.rs +++ b/crates/storage/src/tables/cold.rs @@ -0,0 +1 @@ + From 2a72a74f9d3c793713b75cd0d244bfeb21a7c019 Mon Sep 17 00:00:00 2001 From: James Date: Thu, 15 Jan 2026 08:45:41 -0500 Subject: [PATCH 05/16] fix: mem cache stuff mostly --- .../storage/src/hot/{reth_impl.rs => mdbx.rs} | 44 +++- crates/storage/src/hot/mem.rs | 204 +++++++++++++++--- crates/storage/src/hot/mod.rs | 2 +- crates/storage/src/hot/traits.rs | 22 ++ crates/storage/src/ser/reth_impls.rs | 58 +++-- crates/storage/src/tables/macros.rs | 1 + 6 files changed, 274 insertions(+), 57 deletions(-) rename crates/storage/src/hot/{reth_impl.rs => mdbx.rs} (56%) diff --git a/crates/storage/src/hot/reth_impl.rs b/crates/storage/src/hot/mdbx.rs similarity index 56% rename from crates/storage/src/hot/reth_impl.rs rename to crates/storage/src/hot/mdbx.rs index f085aba..b7704b0 100644 --- a/crates/storage/src/hot/reth_impl.rs +++ b/crates/storage/src/hot/mdbx.rs @@ -1,9 +1,13 @@ use crate::{ - hot::{HotKvRead, HotKvReadError, HotKvWrite}, + hot::{HotKv, HotKvError, HotKvRead, HotKvReadError, HotKvWrite}, ser::DeserError, }; -use reth_db::mdbx::{RW, TransactionKind, WriteFlags, tx::Tx}; +use reth_db::{ + Database, DatabaseEnv, + mdbx::{RW, TransactionKind, WriteFlags, tx::Tx}, +}; use reth_db_api::DatabaseError; +use reth_libmdbx::RO; use std::borrow::Cow; /// Error type for reth-libmdbx based hot storage. @@ -12,16 +16,22 @@ pub enum MdbxError { /// Inner error #[error(transparent)] Mdbx(#[from] reth_libmdbx::Error), + + /// Reth error. + #[error(transparent)] + Reth(#[from] DatabaseError), + /// Deser. #[error(transparent)] Deser(#[from] DeserError), } impl HotKvReadError for MdbxError { - fn into_hot_kv_error(self) -> super::HotKvError { + fn into_hot_kv_error(self) -> HotKvError { match self { - MdbxError::Mdbx(e) => super::HotKvError::from_err(e), - MdbxError::Deser(e) => super::HotKvError::Deser(e), + MdbxError::Mdbx(e) => HotKvError::from_err(e), + MdbxError::Deser(e) => HotKvError::Deser(e), + MdbxError::Reth(e) => HotKvError::from_err(e), } } } @@ -32,6 +42,19 @@ impl From for DatabaseError { } } +impl HotKv for DatabaseEnv { + type RoTx = Tx; + type RwTx = Tx; + + fn reader(&self) -> Result { + self.tx().map_err(HotKvError::from_err) + } + + fn writer(&self) -> Result { + self.tx_mut().map_err(HotKvError::from_err) + } +} + impl HotKvRead for Tx where K: TransactionKind, @@ -62,6 +85,17 @@ impl HotKvWrite for Tx { self.inner.del(dbi, key, None).map(|_| ()).map_err(MdbxError::Mdbx) } + fn queue_raw_clear(&mut self, table: &str) -> Result<(), Self::Error> { + // Future: port more of reth's db env with dbi caching to avoid + // repeated open_db calls + let dbi = self.inner.open_db(Some(table)).map(|db| db.dbi())?; + self.inner.clear_db(dbi).map(|_| ()).map_err(MdbxError::Mdbx) + } + + fn queue_raw_create(&mut self, table: &str) -> Result<(), Self::Error> { + self.inner.create_db(Some(table), Default::default()).map(|_| ()).map_err(MdbxError::Mdbx) + } + fn raw_commit(self) -> Result<(), Self::Error> { // when committing, mdbx returns true on failure let res = self.inner.commit()?; diff --git a/crates/storage/src/hot/mem.rs b/crates/storage/src/hot/mem.rs index ca2ed07..0a6fdb7 100644 --- a/crates/storage/src/hot/mem.rs +++ b/crates/storage/src/hot/mem.rs @@ -11,8 +11,8 @@ use std::{ type Table = BTreeMap<[u8; MAX_KEY_SIZE], bytes::Bytes>; type Store = BTreeMap; -type TableOp = BTreeMap<[u8; MAX_KEY_SIZE], QueuedOp>; -type OpStore = BTreeMap; +type TableOp = BTreeMap<[u8; MAX_KEY_SIZE], QueuedKvOp>; +type OpStore = BTreeMap; /// A simple in-memory key-value store using a BTreeMap. /// @@ -54,6 +54,9 @@ impl Default for MemKv { /// Read-only transaction for MemKv. pub struct MemKvRoTx { guard: RwLockReadGuard<'static, Store>, + + // Keep the store alive while the transaction exists + _store: Arc>, } impl core::fmt::Debug for MemKvRoTx { @@ -70,25 +73,17 @@ unsafe impl Sync for MemKvRoTx {} pub struct MemKvRwTx { guard: RwLockWriteGuard<'static, Store>, queued_ops: OpStore, + + // Keep the store alive while the transaction exists + _store: Arc>, } impl MemKvRwTx { fn commit_inner(&mut self) { let ops = std::mem::take(&mut self.queued_ops); - for (table, table_ops) in ops.into_iter() { - for (key, op) in table_ops.into_iter() { - match op { - QueuedOp::Put { value } => { - self.guard.entry(table.clone()).or_default().insert(key, value); - } - QueuedOp::Delete => { - if let Some(t) = self.guard.get_mut(&table) { - t.remove(&key); - } - } - } - } + for (table, table_op) in ops.into_iter() { + table_op.apply(&table, &mut self.guard); } } @@ -97,7 +92,7 @@ impl MemKvRwTx { pub fn downgrade(self) -> MemKvRoTx { let guard = RwLockWriteGuard::downgrade(self.guard); - MemKvRoTx { guard } + MemKvRoTx { guard, _store: self._store } } /// Commit the transaction and downgrade to a read-only transaction. @@ -106,7 +101,7 @@ impl MemKvRwTx { let guard = RwLockWriteGuard::downgrade(self.guard); - MemKvRoTx { guard } + MemKvRoTx { guard, _store: self._store } } } @@ -116,12 +111,90 @@ impl core::fmt::Debug for MemKvRwTx { } } +/// Queued key-value operation #[derive(Debug, Clone)] -enum QueuedOp { +enum QueuedKvOp { Delete, Put { value: bytes::Bytes }, } +impl QueuedKvOp { + /// Apply the op to a table + fn apply(self, table: &mut Table, key: [u8; MAX_KEY_SIZE]) { + match self { + QueuedKvOp::Put { value } => { + table.insert(key, value); + } + QueuedKvOp::Delete => { + table.remove(&key); + } + } + } +} + +/// Queued table operation +#[derive(Debug)] +enum QueuedTableOp { + Modify { ops: TableOp }, + Clear { new_table: TableOp }, +} + +impl Default for QueuedTableOp { + fn default() -> Self { + QueuedTableOp::Modify { ops: TableOp::new() } + } +} + +impl QueuedTableOp { + const fn is_clear(&self) -> bool { + matches!(self, QueuedTableOp::Clear { .. }) + } + + fn get(&self, key: &[u8; MAX_KEY_SIZE]) -> Option<&QueuedKvOp> { + match self { + QueuedTableOp::Modify { ops } => ops.get(key), + QueuedTableOp::Clear { new_table } => new_table.get(key), + } + } + + fn put(&mut self, key: [u8; MAX_KEY_SIZE], op: QueuedKvOp) { + match self { + QueuedTableOp::Modify { ops } | QueuedTableOp::Clear { new_table: ops } => { + ops.insert(key, op); + } + } + } + + fn delete(&mut self, key: [u8; MAX_KEY_SIZE]) { + match self { + QueuedTableOp::Modify { ops } | QueuedTableOp::Clear { new_table: ops } => { + ops.insert(key, QueuedKvOp::Delete); + } + } + } + + /// Get mutable reference to the inner ops if applicable + fn apply(self, key: &str, store: &mut Store) { + match self { + QueuedTableOp::Modify { ops } => { + let table = store.entry(key.to_owned()).or_default(); + for (key, op) in ops { + op.apply(table, key); + } + } + QueuedTableOp::Clear { new_table } => { + let mut table = Table::new(); + for (k, op) in new_table { + op.apply(&mut table, k); + } + + // replace the table entirely + store.insert(key.to_owned(), table); + } + } + } +} + // SAFETY: MemKvRwTx holds a write guard which ensures exclusive access unsafe impl Send for MemKvRwTx {} @@ -135,20 +208,21 @@ impl HotKv for MemKv { .try_read() .map_err(|_| HotKvError::Inner("Failed to acquire read lock".into()))?; - // SAFETY: We extend the lifetime to 'static, but the guard ensures the - // data remains valid for the lifetime of the transaction + // SAFETY: This is safe-ish, as we ensure the map is not dropped until + // the guard is also dropped. let guard: RwLockReadGuard<'static, Store> = unsafe { std::mem::transmute(guard) }; - Ok(MemKvRoTx { guard }) + Ok(MemKvRoTx { guard, _store: self.map.clone() }) } fn writer(&self) -> Result { let guard = self.map.try_write().map_err(|_| HotKvError::WriteLocked)?; - // SAFETY: We extend the lifetime to 'static, but the guard ensures exclusive access + // SAFETY: This is safe-ish, as we ensure the map is not dropped until + // the guard is also dropped. let guard: RwLockWriteGuard<'static, Store> = unsafe { std::mem::transmute(guard) }; - Ok(MemKvRwTx { guard, queued_ops: OpStore::new() }) + Ok(MemKvRwTx { guard, _store: self.map.clone(), queued_ops: OpStore::new() }) } } @@ -184,14 +258,20 @@ impl HotKvRead for MemKvRwTx { // Check queued operations first (read-your-writes consistency) let key = MemKv::key(key); - match self.queued_ops.get(table).and_then(|t| t.get(&key)) { - Some(QueuedOp::Put { value }) => { - return Ok(Some(Cow::Borrowed(value.as_ref()))); - } - Some(QueuedOp::Delete) => { + if let Some(table) = self.queued_ops.get(table) { + if table.is_clear() { return Ok(None); } - None => {} + + match table.get(&key) { + Some(QueuedKvOp::Put { value }) => { + return Ok(Some(Cow::Borrowed(value.as_ref()))); + } + Some(QueuedKvOp::Delete) => { + return Ok(None); + } + None => {} + } } // If not found in queued ops, check the underlying map @@ -212,14 +292,24 @@ impl HotKvWrite for MemKvRwTx { self.queued_ops .entry(table.to_owned()) .or_default() - .insert(key, QueuedOp::Put { value: value_bytes }); + .put(key, QueuedKvOp::Put { value: value_bytes }); Ok(()) } fn queue_raw_delete(&mut self, table: &str, key: &[u8]) -> Result<(), Self::Error> { let key = MemKv::key(key); - self.queued_ops.entry(table.to_owned()).or_default().insert(key, QueuedOp::Delete); + self.queued_ops.entry(table.to_owned()).or_default().delete(key); + Ok(()) + } + + fn queue_raw_clear(&mut self, table: &str) -> Result<(), Self::Error> { + self.queued_ops + .insert(table.to_owned(), QueuedTableOp::Clear { new_table: TableOp::new() }); + Ok(()) + } + + fn queue_raw_create(&mut self, _table: &str) -> Result<(), Self::Error> { Ok(()) } @@ -245,6 +335,7 @@ mod tests { impl Table for TestTable { const NAME: &'static str = "test_table"; + type Key = u64; type Value = Bytes; } @@ -634,4 +725,55 @@ mod tests { assert!(value3.is_none()); } } + + #[test] + fn test_clear_table() { + let store = MemKv::new(); + + { + let mut writer = store.writer().unwrap(); + writer.queue_raw_put("table1", &[1], b"value1").unwrap(); + writer.queue_raw_put("table1", &[2], b"value2").unwrap(); + writer.raw_commit().unwrap(); + } + + { + let reader = store.reader().unwrap(); + + let value1 = reader.get_raw("table1", &[1]).unwrap(); + let value2 = reader.get_raw("table1", &[2]).unwrap(); + + assert_eq!(value1.as_deref(), Some(b"value1" as &[u8])); + assert_eq!(value2.as_deref(), Some(b"value2" as &[u8])); + } + + { + let mut writer = store.writer().unwrap(); + + let value1 = writer.get_raw("table1", &[1]).unwrap(); + let value2 = writer.get_raw("table1", &[2]).unwrap(); + + assert_eq!(value1.as_deref(), Some(b"value1" as &[u8])); + assert_eq!(value2.as_deref(), Some(b"value2" as &[u8])); + + writer.queue_raw_clear("table1").unwrap(); + + let value1 = writer.get_raw("table1", &[1]).unwrap(); + let value2 = writer.get_raw("table1", &[2]).unwrap(); + + assert!(value1.is_none()); + assert!(value2.is_none()); + + writer.raw_commit().unwrap(); + } + + { + let reader = store.reader().unwrap(); + let value1 = reader.get_raw("table1", &[1]).unwrap(); + let value2 = reader.get_raw("table1", &[2]).unwrap(); + + assert!(value1.is_none()); + assert!(value2.is_none()); + } + } } diff --git a/crates/storage/src/hot/mod.rs b/crates/storage/src/hot/mod.rs index 8ef53fb..dfec724 100644 --- a/crates/storage/src/hot/mod.rs +++ b/crates/storage/src/hot/mod.rs @@ -7,7 +7,7 @@ pub use error::{HotKvError, HotKvReadError, HotKvResult}; mod mem; pub use mem::{MemKv, MemKvRoTx, MemKvRwTx}; -mod reth_impl; +mod mdbx; mod traits; pub use traits::{HotKv, HotKvRead, HotKvWrite}; diff --git a/crates/storage/src/hot/traits.rs b/crates/storage/src/hot/traits.rs index e6f09d8..05c35c8 100644 --- a/crates/storage/src/hot/traits.rs +++ b/crates/storage/src/hot/traits.rs @@ -105,6 +105,12 @@ pub trait HotKvWrite: HotKvRead { /// allowed to panic if this is not the case. fn queue_raw_delete(&mut self, table: &str, key: &[u8]) -> Result<(), Self::Error>; + /// Queue a raw clear operation for a specific table. + fn queue_raw_clear(&mut self, table: &str) -> Result<(), Self::Error>; + + /// Queue a raw create operation for a specific table. + fn queue_raw_create(&mut self, table: &str) -> Result<(), Self::Error>; + /// Queue a put operation for a specific table. fn queue_put(&mut self, key: &T::Key, value: &T::Value) -> Result<(), Self::Error> { let mut key_buf = [0u8; MAX_KEY_SIZE]; @@ -142,6 +148,22 @@ pub trait HotKvWrite: HotKvRead { Ok(()) } + /// Queue creation of a specific table. + fn queue_create(&mut self) -> Result<(), Self::Error> + where + T: Table, + { + self.queue_raw_create(T::NAME) + } + + /// Queue clearing all entries in a specific table. + fn queue_clear(&mut self) -> Result<(), Self::Error> + where + T: Table, + { + self.queue_raw_clear(T::NAME) + } + /// Commit the queued operations. fn raw_commit(self) -> Result<(), Self::Error>; } diff --git a/crates/storage/src/ser/reth_impls.rs b/crates/storage/src/ser/reth_impls.rs index 2ee25d4..33c20bc 100644 --- a/crates/storage/src/ser/reth_impls.rs +++ b/crates/storage/src/ser/reth_impls.rs @@ -18,30 +18,48 @@ use reth_db_api::{ }, }; -impl KeySer for ShardedKey { - const SIZE: usize = T::SIZE + u64::SIZE; - - fn encode_key<'a: 'c, 'b: 'c, 'c>(&'a self, buf: &'b mut [u8; MAX_KEY_SIZE]) -> &'c [u8] { - let mut scratch = [0u8; MAX_KEY_SIZE]; - - T::encode_key(&self.key, &mut scratch); - scratch[T::SIZE..Self::SIZE].copy_from_slice(&self.highest_block_number.to_be_bytes()); - *buf = scratch; - - &buf[0..Self::SIZE] - } +// Specialized impls for the sharded key types. This was implemented +// generically, but there are only 2 types, and we can skip pushing a scratch +// buffer to the stack this way. +macro_rules! sharded_key { + ($ty:ty) => { + impl KeySer for ShardedKey<$ty> { + const SIZE: usize = <$ty as KeySer>::SIZE + u64::SIZE; + + fn encode_key<'a: 'c, 'b: 'c, 'c>( + &'a self, + buf: &'b mut [u8; MAX_KEY_SIZE], + ) -> &'c [u8] { + const SIZE: usize = <$ty as KeySer>::SIZE; + + let prefix = self.key.as_slice(); + + buf[0..SIZE].copy_from_slice(prefix); + buf[SIZE..Self::SIZE].copy_from_slice(&self.highest_block_number.to_be_bytes()); + + &buf[0..Self::SIZE] + } - fn decode_key(data: &[u8]) -> Result { - if data.len() < Self::SIZE { - return Err(DeserError::InsufficientData { needed: Self::SIZE, available: data.len() }); + fn decode_key(data: &[u8]) -> Result { + const SIZE: usize = <$ty as KeySer>::SIZE; + if data.len() < Self::SIZE { + return Err(DeserError::InsufficientData { + needed: Self::SIZE, + available: data.len(), + }); + } + + let key = <$ty as KeySer>::decode_key(&data[0..SIZE])?; + let highest_block_number = u64::decode_key(&data[SIZE..SIZE + 8])?; + Ok(Self { key, highest_block_number }) + } } - - let key = T::decode_key(&data[0..T::SIZE])?; - let highest_block_number = u64::decode_key(&data[T::SIZE..T::SIZE + 8])?; - Ok(Self { key, highest_block_number }) - } + }; } +sharded_key!(B256); +sharded_key!(Address); + impl KeySer for StorageShardedKey { const SIZE: usize = Address::SIZE + B256::SIZE + u64::SIZE; diff --git a/crates/storage/src/tables/macros.rs b/crates/storage/src/tables/macros.rs index 89dc412..70c52e8 100644 --- a/crates/storage/src/tables/macros.rs +++ b/crates/storage/src/tables/macros.rs @@ -8,6 +8,7 @@ macro_rules! tables { impl crate::tables::Table for $name { const NAME: &'static str = stringify!($name); + type Key = $key; type Value = $value; } From 299d76d4d37182d28ebd685a1e994250810c61b0 Mon Sep 17 00:00:00 2001 From: James Date: Thu, 15 Jan 2026 08:50:36 -0500 Subject: [PATCH 06/16] docs: more of them --- crates/storage/src/hot/mem.rs | 12 ++++++++---- crates/storage/src/hot/traits.rs | 3 +++ 2 files changed, 11 insertions(+), 4 deletions(-) diff --git a/crates/storage/src/hot/mem.rs b/crates/storage/src/hot/mem.rs index 0a6fdb7..f5e2353 100644 --- a/crates/storage/src/hot/mem.rs +++ b/crates/storage/src/hot/mem.rs @@ -14,11 +14,15 @@ type Store = BTreeMap; type TableOp = BTreeMap<[u8; MAX_KEY_SIZE], QueuedKvOp>; type OpStore = BTreeMap; -/// A simple in-memory key-value store using a BTreeMap. +/// A simple in-memory key-value store using [`BTreeMap`]s. /// -/// This implementation supports concurrent multiple concurrent read -/// transactions. Write transactions are exclusive, and cannot overlap -/// with other read or write transactions. +/// The store is backed by an [`RwLock`]. As a result, this implementation +/// supports concurrent multiple concurrent read transactions, but write +/// transactions are exclusive, and cannot overlap with other read or write +/// transactions. +/// +/// This implementation is primarily intended for testing and +/// development purposes. #[derive(Clone)] pub struct MemKv { map: Arc>, diff --git a/crates/storage/src/hot/traits.rs b/crates/storage/src/hot/traits.rs index 05c35c8..42c0e15 100644 --- a/crates/storage/src/hot/traits.rs +++ b/crates/storage/src/hot/traits.rs @@ -18,6 +18,9 @@ pub trait HotKv { /// Create a read-write transaction. /// + /// This is allowed to fail with [`Err(HotKvError::WriteLocked)`] if + /// multiple write transactions are not supported concurrently. + /// /// # Returns /// /// - `Ok(Some(tx))` if the write transaction was created successfully. From b4b63ae31f175b3c55db0737ef5adeed89436a3b Mon Sep 17 00:00:00 2001 From: James Date: Thu, 15 Jan 2026 09:00:17 -0500 Subject: [PATCH 07/16] tests: keyser for reth impls --- crates/storage/src/ser/reth_impls.rs | 279 ++++++++++++++++++++++++++- 1 file changed, 278 insertions(+), 1 deletion(-) diff --git a/crates/storage/src/ser/reth_impls.rs b/crates/storage/src/ser/reth_impls.rs index 33c20bc..ccdb2b2 100644 --- a/crates/storage/src/ser/reth_impls.rs +++ b/crates/storage/src/ser/reth_impls.rs @@ -20,7 +20,8 @@ use reth_db_api::{ // Specialized impls for the sharded key types. This was implemented // generically, but there are only 2 types, and we can skip pushing a scratch -// buffer to the stack this way. +// buffer, because we know the 2 types involved are already fixed-size byte +// arrays. macro_rules! sharded_key { ($ty:ty) => { impl KeySer for ShardedKey<$ty> { @@ -1837,4 +1838,280 @@ mod tests { } test_roundtrip(&large_list); } + + // KeySer Tests + // ============ + + /// Generic roundtrip test for any KeySer type + #[track_caller] + fn test_key_roundtrip(original: &T) + where + T: KeySer + PartialEq + std::fmt::Debug, + { + let mut buf = [0u8; MAX_KEY_SIZE]; + let encoded = original.encode_key(&mut buf); + + // Assert that the encoded size matches the const SIZE + assert_eq!( + encoded.len(), + T::SIZE, + "Encoded key length mismatch: expected {}, got {}", + T::SIZE, + encoded.len() + ); + + // Decode and verify + let decoded = T::decode_key(encoded).expect("Failed to decode key"); + assert_eq!(*original, decoded, "Key roundtrip failed"); + } + + /// Test ordering preservation for KeySer types + #[track_caller] + fn test_key_ordering(keys: &[T]) + where + T: KeySer + Ord + std::fmt::Debug + Clone, + { + let mut sorted_keys = keys.to_vec(); + sorted_keys.sort(); + + let mut encoded_keys: Vec<_> = sorted_keys + .iter() + .map(|k| { + let mut buf = [0u8; MAX_KEY_SIZE]; + let encoded = k.encode_key(&mut buf); + encoded.to_vec() + }) + .collect(); + + // Check that encoded keys are also sorted lexicographically + let original_encoded = encoded_keys.clone(); + encoded_keys.sort(); + + assert_eq!(original_encoded, encoded_keys, "Key encoding does not preserve ordering"); + } + + #[test] + fn test_sharded_key_b256_roundtrips() { + // Test with B256 keys + let key1 = ShardedKey { key: B256::ZERO, highest_block_number: 0 }; + test_key_roundtrip(&key1); + + let key2 = ShardedKey { key: B256::repeat_byte(0xFF), highest_block_number: u64::MAX }; + test_key_roundtrip(&key2); + + let key3 = ShardedKey { + key: B256::from([ + 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0A, 0x0B, 0x0C, 0x0D, 0x0E, + 0x0F, 0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, 0x18, 0x19, 0x1A, 0x1B, 0x1C, + 0x1D, 0x1E, 0x1F, 0x20, + ]), + highest_block_number: 12345, + }; + test_key_roundtrip(&key3); + } + + #[test] + fn test_sharded_key_address_roundtrips() { + // Test with Address keys + let key1 = ShardedKey { key: Address::ZERO, highest_block_number: 0 }; + test_key_roundtrip(&key1); + + let key2 = ShardedKey { key: Address::repeat_byte(0xFF), highest_block_number: u64::MAX }; + test_key_roundtrip(&key2); + + let key3 = ShardedKey { + key: Address::from([ + 0x12, 0x34, 0x56, 0x78, 0x9A, 0xBC, 0xDE, 0xF0, 0x12, 0x34, 0x56, 0x78, 0x9A, 0xBC, + 0xDE, 0xF0, 0x12, 0x34, 0x56, 0x78, + ]), + highest_block_number: 9876543210, + }; + test_key_roundtrip(&key3); + } + + #[test] + fn test_storage_sharded_key_roundtrips() { + // Test basic cases + let key1 = StorageShardedKey { + address: Address::ZERO, + sharded_key: ShardedKey { key: B256::ZERO, highest_block_number: 0 }, + }; + test_key_roundtrip(&key1); + + // Test with maximum values + let key2 = StorageShardedKey { + address: Address::repeat_byte(0xFF), + sharded_key: ShardedKey { + key: B256::repeat_byte(0xFF), + highest_block_number: u64::MAX, + }, + }; + test_key_roundtrip(&key2); + + // Test with realistic values + let key3 = StorageShardedKey { + address: Address::from([ + 0xd8, 0xdA, 0x6B, 0xF2, 0x69, 0x64, 0xaF, 0x9D, 0x7e, 0xEd, 0x9e, 0x03, 0xE5, 0x34, + 0x15, 0xD3, 0x7a, 0xA9, 0x60, 0x45, + ]), // Example address + sharded_key: ShardedKey { + key: keccak256(b"storage_slot"), + highest_block_number: 1000000, + }, + }; + test_key_roundtrip(&key3); + } + + #[test] + fn test_sharded_key_b256_ordering() { + let keys = vec![ + ShardedKey { key: B256::ZERO, highest_block_number: 0 }, + ShardedKey { key: B256::ZERO, highest_block_number: 1 }, + ShardedKey { key: B256::ZERO, highest_block_number: u64::MAX }, + ShardedKey { key: B256::from([0x01; 32]), highest_block_number: 0 }, + ShardedKey { key: B256::from([0x01; 32]), highest_block_number: 1 }, + ShardedKey { key: B256::repeat_byte(0xFF), highest_block_number: 0 }, + ShardedKey { key: B256::repeat_byte(0xFF), highest_block_number: u64::MAX }, + ]; + test_key_ordering(&keys); + } + + #[test] + fn test_sharded_key_address_ordering() { + let keys = vec![ + ShardedKey { key: Address::ZERO, highest_block_number: 0 }, + ShardedKey { key: Address::ZERO, highest_block_number: 1 }, + ShardedKey { key: Address::ZERO, highest_block_number: u64::MAX }, + ShardedKey { key: Address::from([0x01; 20]), highest_block_number: 0 }, + ShardedKey { key: Address::from([0x01; 20]), highest_block_number: 1 }, + ShardedKey { key: Address::repeat_byte(0xFF), highest_block_number: 0 }, + ShardedKey { key: Address::repeat_byte(0xFF), highest_block_number: u64::MAX }, + ]; + test_key_ordering(&keys); + } + + #[test] + fn test_storage_sharded_key_ordering() { + let keys = vec![ + StorageShardedKey { + address: Address::ZERO, + sharded_key: ShardedKey { key: B256::ZERO, highest_block_number: 0 }, + }, + StorageShardedKey { + address: Address::ZERO, + sharded_key: ShardedKey { key: B256::ZERO, highest_block_number: 1 }, + }, + StorageShardedKey { + address: Address::ZERO, + sharded_key: ShardedKey { key: B256::from([0x01; 32]), highest_block_number: 0 }, + }, + StorageShardedKey { + address: Address::from([0x01; 20]), + sharded_key: ShardedKey { key: B256::ZERO, highest_block_number: 0 }, + }, + StorageShardedKey { + address: Address::repeat_byte(0xFF), + sharded_key: ShardedKey { + key: B256::repeat_byte(0xFF), + highest_block_number: u64::MAX, + }, + }, + ]; + test_key_ordering(&keys); + } + + #[test] + fn test_key_size_constants() { + // Verify the SIZE constants are correct + assert_eq!(ShardedKey::::SIZE, B256::SIZE + u64::SIZE); + assert_eq!(ShardedKey::
::SIZE, Address::SIZE + u64::SIZE); + assert_eq!(StorageShardedKey::SIZE, Address::SIZE + B256::SIZE + u64::SIZE); + + // Verify sizes are within MAX_KEY_SIZE + assert!(ShardedKey::::SIZE <= MAX_KEY_SIZE); + assert!(ShardedKey::
::SIZE <= MAX_KEY_SIZE); + assert!(StorageShardedKey::SIZE <= MAX_KEY_SIZE); + } + + #[test] + fn test_key_decode_insufficient_data() { + // Test ShardedKey with insufficient data + let short_data = [0u8; 10]; // Much smaller than required + + match ShardedKey::::decode_key(&short_data) { + Err(DeserError::InsufficientData { needed, available }) => { + assert_eq!(needed, ShardedKey::::SIZE); + assert_eq!(available, 10); + } + other => panic!("Expected InsufficientData error, got: {:?}", other), + } + + // Test ShardedKey
with insufficient data + match ShardedKey::
::decode_key(&short_data) { + Err(DeserError::InsufficientData { needed, available }) => { + assert_eq!(needed, ShardedKey::
::SIZE); + assert_eq!(available, 10); + } + other => panic!("Expected InsufficientData error, got: {:?}", other), + } + + // Test StorageShardedKey with insufficient data + match StorageShardedKey::decode_key(&short_data) { + Err(DeserError::InsufficientData { needed, available }) => { + assert_eq!(needed, StorageShardedKey::SIZE); + assert_eq!(available, 10); + } + other => panic!("Expected InsufficientData error, got: {:?}", other), + } + } + + #[test] + fn test_key_encode_decode_boundary_values() { + // Test boundary values for block numbers + let boundary_keys = vec![ + ShardedKey { key: B256::ZERO, highest_block_number: 0 }, + ShardedKey { key: B256::ZERO, highest_block_number: 1 }, + ShardedKey { key: B256::ZERO, highest_block_number: u64::MAX - 1 }, + ShardedKey { key: B256::ZERO, highest_block_number: u64::MAX }, + ]; + + for key in &boundary_keys { + test_key_roundtrip(key); + } + + // Test that ordering is preserved across boundaries + test_key_ordering(&boundary_keys); + } + + #[test] + fn test_storage_sharded_key_component_ordering() { + // Test that address takes precedence in ordering + let addr1 = Address::from([0x01; 20]); + let addr2 = Address::from([0x02; 20]); + + let key1 = StorageShardedKey { + address: addr1, + sharded_key: ShardedKey { + key: B256::repeat_byte(0xFF), + highest_block_number: u64::MAX, + }, + }; + + let key2 = StorageShardedKey { + address: addr2, + sharded_key: ShardedKey { key: B256::ZERO, highest_block_number: 0 }, + }; + + // key1 should be less than key2 because addr1 < addr2 + assert!(key1 < key2); + + // This should be reflected in the encoding + let mut buf1 = [0u8; MAX_KEY_SIZE]; + let mut buf2 = [0u8; MAX_KEY_SIZE]; + + let encoded1 = key1.encode_key(&mut buf1); + let encoded2 = key2.encode_key(&mut buf2); + + assert!(encoded1 < encoded2, "Encoding should preserve ordering"); + } } From 6c43b777b9a8b9557f2450c1ee597a566f4fd801 Mon Sep 17 00:00:00 2001 From: James Date: Thu, 15 Jan 2026 10:12:31 -0500 Subject: [PATCH 08/16] feat: revm databasing --- crates/storage/Cargo.toml | 2 + crates/storage/src/hot/mod.rs | 3 + crates/storage/src/hot/revm.rs | 630 +++++++++++++++++++++++++++ crates/storage/src/hot/traits.rs | 27 +- crates/storage/src/ser/reth_impls.rs | 17 +- 5 files changed, 663 insertions(+), 16 deletions(-) create mode 100644 crates/storage/src/hot/revm.rs diff --git a/crates/storage/Cargo.toml b/crates/storage/Cargo.toml index 312a17f..7d18346 100644 --- a/crates/storage/Cargo.toml +++ b/crates/storage/Cargo.toml @@ -10,9 +10,11 @@ repository.workspace = true [dependencies] alloy.workspace = true +auto_impl = "1.3.0" bytes = "1.11.0" reth.workspace = true reth-db.workspace = true reth-db-api.workspace = true reth-libmdbx.workspace = true thiserror.workspace = true +trevm.workspace = true diff --git a/crates/storage/src/hot/mod.rs b/crates/storage/src/hot/mod.rs index dfec724..2720015 100644 --- a/crates/storage/src/hot/mod.rs +++ b/crates/storage/src/hot/mod.rs @@ -9,5 +9,8 @@ pub use mem::{MemKv, MemKvRoTx, MemKvRwTx}; mod mdbx; +mod revm; +pub use revm::{RevmRead, RevmWrite}; + mod traits; pub use traits::{HotKv, HotKvRead, HotKvWrite}; diff --git a/crates/storage/src/hot/revm.rs b/crates/storage/src/hot/revm.rs new file mode 100644 index 0000000..fea9266 --- /dev/null +++ b/crates/storage/src/hot/revm.rs @@ -0,0 +1,630 @@ +use crate::{ + hot::{HotKvError, HotKvRead, HotKvWrite}, + tables::hot::{AccountStorageKey, Bytecodes, PlainAccountState}, +}; +use alloy::primitives::{Address, B256, KECCAK256_EMPTY}; +use core::fmt; +use reth::primitives::Account; +use trevm::revm::{ + database::{DBErrorMarker, Database, DatabaseRef, TryDatabaseCommit}, + primitives::{HashMap, StorageKey, StorageValue}, + state::{self, AccountInfo, Bytecode as RevmBytecode}, +}; + +// Error marker implementation +impl DBErrorMarker for HotKvError {} + +/// Read-only [`Database`] and [`DatabaseRef`] adapter. +pub struct RevmRead { + reader: T, +} + +impl RevmRead { + /// Create a new read adapter + pub const fn new(reader: T) -> Self { + Self { reader } + } +} + +impl fmt::Debug for RevmRead { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("RevmRead").finish() + } +} + +// HotKvRead implementation for RevmRead +impl HotKvRead for RevmRead { + type Error = T::Error; + + fn get_raw<'a>( + &'a self, + table: &str, + key: &[u8], + ) -> Result>, Self::Error> { + self.reader.get_raw(table, key) + } +} + +/// Read-write REVM database adapter. This adapter allows committing changes. +/// Despite the naming of [`TryDatabaseCommit::try_commit`], the changes are +/// only persisted when [`Self::persist`] is called. This is because of a +/// mismatch in semantics between the two systems. +pub struct RevmWrite { + writer: T, +} +impl RevmWrite { + /// Create a new write adapter + pub const fn new(writer: T) -> Self { + Self { writer } + } + + /// Persist the changes made in this write transaction. + pub fn persist(self) -> Result<(), T::Error> { + self.writer.raw_commit() + } +} + +impl fmt::Debug for RevmWrite { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("RevmWrite").finish() + } +} + +// HotKvWrite implementation for RevmWrite +impl HotKvRead for RevmWrite { + type Error = T::Error; + + fn get_raw<'a>( + &'a self, + table: &str, + key: &[u8], + ) -> Result>, Self::Error> { + self.writer.get_raw(table, key) + } +} + +impl HotKvWrite for RevmWrite { + fn queue_raw_put(&mut self, table: &str, key: &[u8], value: &[u8]) -> Result<(), Self::Error> { + self.writer.queue_raw_put(table, key, value) + } + + fn queue_raw_delete(&mut self, table: &str, key: &[u8]) -> Result<(), Self::Error> { + self.writer.queue_raw_delete(table, key) + } + + fn queue_raw_clear(&mut self, table: &str) -> Result<(), Self::Error> { + self.writer.queue_raw_clear(table) + } + + fn queue_raw_create(&mut self, table: &str) -> Result<(), Self::Error> { + self.writer.queue_raw_create(table) + } + + fn raw_commit(self) -> Result<(), Self::Error> { + self.writer.raw_commit() + } +} + +// DatabaseRef implementation for RevmRead +impl DatabaseRef for RevmRead +where + T::Error: DBErrorMarker, +{ + type Error = T::Error; + + fn basic_ref(&self, address: Address) -> Result, Self::Error> { + let account_opt = self.reader.get::(&address)?; + + let Some(account) = account_opt else { + return Ok(None); + }; + + let code_hash = account.bytecode_hash.unwrap_or(KECCAK256_EMPTY); + let code = if code_hash != KECCAK256_EMPTY { + self.reader.get::(&code_hash)?.map(|b| b.0) + } else { + None + }; + + Ok(Some(AccountInfo { balance: account.balance, nonce: account.nonce, code_hash, code })) + } + + fn code_by_hash_ref(&self, code_hash: B256) -> Result { + Ok(self.reader.get::(&code_hash)?.map(|bytecode| bytecode.0).unwrap_or_default()) + } + + fn storage_ref( + &self, + address: Address, + index: StorageKey, + ) -> Result { + let storage_key = AccountStorageKey { + address: std::borrow::Cow::Borrowed(&address), + key: std::borrow::Cow::Owned(B256::from_slice(&index.to_be_bytes::<32>())), + } + .encode_key(); + + Ok(self + .reader + .get::(&storage_key)? + .unwrap_or_default()) + } + + fn block_hash_ref(&self, _number: u64) -> Result { + // This would need to be implemented based on your block hash storage + // For now, return zero hash + Ok(B256::ZERO) + } +} + +// Database implementation for RevmRead +impl Database for RevmRead +where + T::Error: DBErrorMarker, +{ + type Error = T::Error; + + fn basic(&mut self, address: Address) -> Result, Self::Error> { + self.basic_ref(address) + } + + fn code_by_hash(&mut self, code_hash: B256) -> Result { + self.code_by_hash_ref(code_hash) + } + + fn storage( + &mut self, + address: Address, + index: StorageKey, + ) -> Result { + self.storage_ref(address, index) + } + + fn block_hash(&mut self, number: u64) -> Result { + self.block_hash_ref(number) + } +} + +// DatabaseRef implementation for RevmWrite (delegates to read operations) +impl DatabaseRef for RevmWrite +where + T::Error: DBErrorMarker, +{ + type Error = T::Error; + + fn basic_ref(&self, address: Address) -> Result, Self::Error> { + let account_opt = self.writer.get::(&address)?; + + let Some(account) = account_opt else { + return Ok(None); + }; + + let code_hash = account.bytecode_hash.unwrap_or(KECCAK256_EMPTY); + let code = if code_hash != KECCAK256_EMPTY { + self.writer.get::(&code_hash)?.map(|b| b.0) + } else { + None + }; + + Ok(Some(AccountInfo { balance: account.balance, nonce: account.nonce, code_hash, code })) + } + + fn code_by_hash_ref(&self, code_hash: B256) -> Result { + Ok(self.writer.get::(&code_hash)?.map(|bytecode| bytecode.0).unwrap_or_default()) + } + + fn storage_ref( + &self, + address: Address, + index: StorageKey, + ) -> Result { + let storage_key = AccountStorageKey { + address: std::borrow::Cow::Borrowed(&address), + key: std::borrow::Cow::Owned(B256::from_slice(&index.to_be_bytes::<32>())), + } + .encode_key(); + + Ok(self + .writer + .get::(&storage_key)? + .unwrap_or_default()) + } + + fn block_hash_ref(&self, _number: u64) -> Result { + // This would need to be implemented based on your block hash storage + // For now, return zero hash + Ok(B256::ZERO) + } +} + +// Database implementation for RevmWrite +impl Database for RevmWrite +where + T::Error: DBErrorMarker, +{ + type Error = T::Error; + + fn basic(&mut self, address: Address) -> Result, Self::Error> { + self.basic_ref(address) + } + + fn code_by_hash(&mut self, code_hash: B256) -> Result { + self.code_by_hash_ref(code_hash) + } + + fn storage( + &mut self, + address: Address, + index: StorageKey, + ) -> Result { + self.storage_ref(address, index) + } + + fn block_hash(&mut self, number: u64) -> Result { + self.block_hash_ref(number) + } +} + +// TryDatabaseCommit implementation for RevmWrite +impl TryDatabaseCommit for RevmWrite +where + T::Error: DBErrorMarker, +{ + type Error = T::Error; + + fn try_commit(&mut self, changes: HashMap) -> Result<(), Self::Error> { + for (address, account) in changes { + // Handle account info changes + let account_data = Account { + nonce: account.info.nonce, + balance: account.info.balance, + bytecode_hash: (account.info.code_hash != KECCAK256_EMPTY) + .then_some(account.info.code_hash), + }; + self.writer.queue_put::(&address, &account_data)?; + + // Handle storage changes + for (key, value) in account.storage { + let storage_key = AccountStorageKey { + address: std::borrow::Cow::Borrowed(&address), + key: std::borrow::Cow::Owned(B256::from_slice(&key.to_be_bytes::<32>())), + } + .encode_key(); + self.writer.queue_put::( + &storage_key, + &value.present_value(), + )?; + } + } + + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{ + hot::{HotKv, HotKvRead, HotKvWrite, MemKv}, + tables::hot::{Bytecodes, PlainAccountState, PlainStorageState}, + }; + use alloy::primitives::{Address, B256, U256}; + use reth::primitives::{Account, Bytecode}; + use trevm::revm::{ + database::{Database, DatabaseRef, TryDatabaseCommit}, + primitives::{HashMap, StorageKey, StorageValue}, + state::{Account as RevmAccount, AccountInfo, Bytecode as RevmBytecode}, + }; + + /// Create a test account with some data + fn create_test_account() -> (Address, Account) { + let address = Address::from_slice(&[0x1; 20]); + let account = Account { + nonce: 42, + balance: U256::from(1000u64), + bytecode_hash: Some(B256::from_slice(&[0x2; 32])), + }; + (address, account) + } + + /// Create test bytecode + fn create_test_bytecode() -> (B256, Bytecode) { + let hash = B256::from_slice(&[0x2; 32]); + let code = RevmBytecode::new_raw(vec![0x60, 0x80, 0x60, 0x40].into()); + let bytecode = Bytecode(code); + (hash, bytecode) + } + + #[test] + fn test_database_ref_traits() -> Result<(), Box> { + let mem_kv = MemKv::default(); + + let (address, account) = create_test_account(); + let (hash, bytecode) = create_test_bytecode(); + + { + // Setup data using HotKv + let mut writer = mem_kv.revm_writer()?; + writer.queue_put::(&address, &account)?; + writer.queue_put::(&hash, &bytecode)?; + writer.persist()?; + } + + { + // Read using REVM DatabaseRef traits + let reader = mem_kv.revm_reader()?; + + // Test basic_ref + let account_info = reader.basic_ref(address)?; + assert!(account_info.is_some()); + let info = account_info.unwrap(); + assert_eq!(info.nonce, 42); + assert_eq!(info.balance, U256::from(1000u64)); + assert_eq!(info.code_hash, hash); + + // Test code_by_hash_ref + let retrieved_code = reader.code_by_hash_ref(hash)?; + assert_eq!(retrieved_code, bytecode.0); + + // Test storage_ref (should be zero for non-existent storage) + let storage_val = reader.storage_ref(address, StorageKey::from(U256::from(123u64)))?; + assert_eq!(storage_val, U256::ZERO); + + // Test block_hash_ref + let block_hash = reader.block_hash_ref(123)?; + assert_eq!(block_hash, B256::ZERO); + } + + Ok(()) + } + + #[test] + fn test_database_mutable_traits() -> Result<(), Box> { + let mem_kv = MemKv::default(); + + let (address, account) = create_test_account(); + let (hash, bytecode) = create_test_bytecode(); + + { + // Setup data using HotKv + let mut writer = mem_kv.revm_writer()?; + writer.queue_put::(&address, &account)?; + writer.queue_put::(&hash, &bytecode)?; + writer.persist()?; + } + + { + // Read using mutable REVM Database traits + let mut reader = mem_kv.revm_reader()?; + + // Test basic + let account_info = reader.basic(address)?; + assert!(account_info.is_some()); + let info = account_info.unwrap(); + assert_eq!(info.nonce, 42); + assert_eq!(info.balance, U256::from(1000u64)); + + // Test code_by_hash + let retrieved_code = reader.code_by_hash(hash)?; + assert_eq!(retrieved_code, bytecode.0); + + // Test storage + let storage_val = reader.storage(address, StorageKey::from(U256::from(123u64)))?; + assert_eq!(storage_val, U256::ZERO); + + // Test block_hash + let block_hash = reader.block_hash(123)?; + assert_eq!(block_hash, B256::ZERO); + } + + Ok(()) + } + + #[test] + fn test_write_database_traits() -> Result<(), Box> { + let mem_kv = MemKv::default(); + + let (address, account) = create_test_account(); + let (hash, bytecode) = create_test_bytecode(); + + { + // Setup initial data + let mut writer = mem_kv.revm_writer()?; + writer.queue_put::(&address, &account)?; + writer.queue_put::(&hash, &bytecode)?; + writer.persist()?; + } + + { + // Test write operations using DatabaseRef and Database traits + let mut writer = mem_kv.revm_writer()?; + + // Test read operations work on writer + let account_info = writer.basic_ref(address)?; + assert!(account_info.is_some()); + + let account_info_mut = writer.basic(address)?; + assert!(account_info_mut.is_some()); + + let code = writer.code_by_hash_ref(hash)?; + assert_eq!(code, bytecode.0); + + let code_mut = writer.code_by_hash(hash)?; + assert_eq!(code_mut, bytecode.0); + + // Don't persist this writer to test that reads work + } + + Ok(()) + } + + #[test] + fn test_try_database_commit() -> Result<(), Box> { + let mem_kv = MemKv::default(); + + let address = Address::from_slice(&[0x1; 20]); + + { + let mut writer = mem_kv.revm_writer()?; + + // Create REVM state changes + let mut changes = HashMap::default(); + let account_info = AccountInfo { + nonce: 55, + balance: U256::from(2000u64), + code_hash: KECCAK256_EMPTY, + code: None, + }; + + let mut storage = HashMap::default(); + storage.insert( + StorageKey::from(U256::from(100u64)), + trevm::revm::state::EvmStorageSlot::new(U256::from(200u64), 0), + ); + + let revm_account = RevmAccount { + info: account_info, + storage, + status: trevm::revm::state::AccountStatus::Touched, + transaction_id: 0, + }; + + changes.insert(address, revm_account); + + // Commit changes using REVM trait + writer.try_commit(changes)?; + writer.persist()?; + } + + { + // Verify changes were persisted using HotKv traits + let reader = mem_kv.revm_reader()?; + + let account: Option = reader.get::(&address)?; + assert!(account.is_some()); + let acc = account.unwrap(); + assert_eq!(acc.nonce, 55); + assert_eq!(acc.balance, U256::from(2000u64)); + assert_eq!(acc.bytecode_hash, None); + + // Check storage was written + let storage_key = AccountStorageKey { + address: std::borrow::Cow::Borrowed(&address), + key: std::borrow::Cow::Owned(B256::from_slice( + &U256::from(100u64).to_be_bytes::<32>(), + )), + } + .encode_key(); + + let storage_val: Option = + reader.get::(&storage_key)?; + assert_eq!(storage_val, Some(U256::from(200u64))); + } + + Ok(()) + } + + #[test] + fn test_mixed_usage_patterns() -> Result<(), Box> { + let mem_kv = MemKv::default(); + + let address1 = Address::from_slice(&[0x1; 20]); + let address2 = Address::from_slice(&[0x2; 20]); + + // Write some data using HotKv + { + let mut writer = mem_kv.revm_writer()?; + let account = Account { nonce: 10, balance: U256::from(500u64), bytecode_hash: None }; + writer.queue_put::(&address1, &account)?; + writer.persist()?; + } + + // Write more data using REVM traits + { + let mut writer = mem_kv.revm_writer()?; + let mut changes = HashMap::default(); + let revm_account = RevmAccount { + info: AccountInfo { + nonce: 20, + balance: U256::from(1500u64), + code_hash: KECCAK256_EMPTY, + code: None, + }, + storage: HashMap::default(), + status: trevm::revm::state::AccountStatus::Touched, + transaction_id: 0, + }; + changes.insert(address2, revm_account); + writer.try_commit(changes)?; + writer.persist()?; + } + + // Read using mixed approaches + { + let reader = mem_kv.revm_reader()?; + + // Read address1 using HotKv + let account1: Option = reader.get::(&address1)?; + assert!(account1.is_some()); + assert_eq!(account1.unwrap().nonce, 10); + + // Read address2 using REVM DatabaseRef + let account2_info = reader.basic_ref(address2)?; + assert!(account2_info.is_some()); + assert_eq!(account2_info.unwrap().nonce, 20); + } + + Ok(()) + } + + #[test] + fn test_error_handling() -> Result<(), Box> { + let mem_kv = MemKv::default(); + + let address = Address::from_slice(&[0x1; 20]); + let hash = B256::from_slice(&[0x99; 32]); + + let reader = mem_kv.revm_reader()?; + + // Test reading non-existent account + let account_info = reader.basic_ref(address)?; + assert!(account_info.is_none()); + + // Test reading non-existent code + let code = reader.code_by_hash_ref(hash)?; + assert!(code.is_empty()); + + // Test reading non-existent storage + let storage = reader.storage_ref(address, StorageKey::from(U256::from(123u64)))?; + assert_eq!(storage, U256::ZERO); + + Ok(()) + } + + #[test] + fn test_concurrent_readers() -> Result<(), Box> { + let mem_kv = MemKv::default(); + + let (address, account) = create_test_account(); + + // Setup data + { + let mut writer = mem_kv.revm_writer()?; + writer.queue_put::(&address, &account)?; + writer.persist()?; + } + + // Create multiple readers + let reader1 = mem_kv.revm_reader()?; + let reader2 = mem_kv.revm_reader()?; + + // Both should read the same data + let account1 = reader1.basic_ref(address)?; + let account2 = reader2.basic_ref(address)?; + + assert_eq!(account1, account2); + assert!(account1.is_some()); + + Ok(()) + } +} diff --git a/crates/storage/src/hot/traits.rs b/crates/storage/src/hot/traits.rs index 42c0e15..b23aed1 100644 --- a/crates/storage/src/hot/traits.rs +++ b/crates/storage/src/hot/traits.rs @@ -1,12 +1,16 @@ use std::borrow::Cow; use crate::{ - hot::{HotKvError, HotKvReadError}, + hot::{ + HotKvError, HotKvReadError, + revm::{RevmRead, RevmWrite}, + }, ser::{KeySer, MAX_KEY_SIZE, ValSer}, tables::Table, }; /// Trait for hot storage. This is a KV store with read/write transactions. +#[auto_impl::auto_impl(&, Arc, Box)] pub trait HotKv { /// The read-only transaction type. type RoTx: HotKvRead; @@ -16,6 +20,15 @@ pub trait HotKv { /// Create a read-only transaction. fn reader(&self) -> Result; + /// Create a read-only transaction, and wrap it in an adapter for the + /// revm [`DatabaseRef`] trait. The resulting reader can be used directly + /// with [`trevm`] and [`revm`]. + /// + /// [`DatabaseRef`]: trevm::revm::database::DatabaseRef + fn revm_reader(&self) -> Result, HotKvError> { + self.reader().map(RevmRead::new) + } + /// Create a read-write transaction. /// /// This is allowed to fail with [`Err(HotKvError::WriteLocked)`] if @@ -32,9 +45,21 @@ pub trait HotKv { /// [`Err(HotKvError::Inner)`]: HotKvError::Inner /// [`Err(HotKvError::WriteLocked)`]: HotKvError::WriteLocked fn writer(&self) -> Result; + + /// Create a read-write transaction, and wrap it in an adapter for the + /// revm [`TryDatabaseCommit`] trait. The resulting writer can be used + /// directly with [`trevm`] and [`revm`]. + /// + /// + /// [`revm`]: trevm::revm + /// [`TryDatabaseCommit`]: trevm::revm::database::TryDatabaseCommit + fn revm_writer(&self) -> Result, HotKvError> { + self.writer().map(RevmWrite::new) + } } /// Trait for hot storage read transactions. +#[auto_impl::auto_impl(&, Arc, Box)] pub trait HotKvRead { /// Error type for read operations. type Error: HotKvReadError; diff --git a/crates/storage/src/ser/reth_impls.rs b/crates/storage/src/ser/reth_impls.rs index ccdb2b2..ef19bfc 100644 --- a/crates/storage/src/ser/reth_impls.rs +++ b/crates/storage/src/ser/reth_impls.rs @@ -1951,8 +1951,8 @@ mod tests { // Test with realistic values let key3 = StorageShardedKey { address: Address::from([ - 0xd8, 0xdA, 0x6B, 0xF2, 0x69, 0x64, 0xaF, 0x9D, 0x7e, 0xEd, 0x9e, 0x03, 0xE5, 0x34, - 0x15, 0xD3, 0x7a, 0xA9, 0x60, 0x45, + 0xd8, 0xda, 0x6b, 0xf2, 0x69, 0x64, 0xaf, 0x9d, 0x7e, 0xed, 0x9e, 0x03, 0xe5, 0x34, + 0x15, 0xd3, 0x7a, 0xa9, 0x60, 0x45, ]), // Example address sharded_key: ShardedKey { key: keccak256(b"storage_slot"), @@ -2020,19 +2020,6 @@ mod tests { test_key_ordering(&keys); } - #[test] - fn test_key_size_constants() { - // Verify the SIZE constants are correct - assert_eq!(ShardedKey::::SIZE, B256::SIZE + u64::SIZE); - assert_eq!(ShardedKey::
::SIZE, Address::SIZE + u64::SIZE); - assert_eq!(StorageShardedKey::SIZE, Address::SIZE + B256::SIZE + u64::SIZE); - - // Verify sizes are within MAX_KEY_SIZE - assert!(ShardedKey::::SIZE <= MAX_KEY_SIZE); - assert!(ShardedKey::
::SIZE <= MAX_KEY_SIZE); - assert!(StorageShardedKey::SIZE <= MAX_KEY_SIZE); - } - #[test] fn test_key_decode_insufficient_data() { // Test ShardedKey with insufficient data From 3e1a423a001dab546a56478b5da97c94e7ccf82c Mon Sep 17 00:00:00 2001 From: James Date: Thu, 15 Jan 2026 14:26:04 -0500 Subject: [PATCH 09/16] feat: fixed and dupsort --- crates/storage/Cargo.toml | 3 + crates/storage/src/hot/db_traits.rs | 14 +- crates/storage/src/hot/mdbx.rs | 562 +++++++++++++++++++++++++++- crates/storage/src/hot/mem.rs | 111 ++++-- crates/storage/src/hot/revm.rs | 205 +++++++--- crates/storage/src/hot/traits.rs | 86 ++++- crates/storage/src/ser/impls.rs | 47 ++- crates/storage/src/tables/hot.rs | 78 +--- crates/storage/src/tables/macros.rs | 33 +- crates/storage/src/tables/mod.rs | 46 ++- 10 files changed, 1007 insertions(+), 178 deletions(-) diff --git a/crates/storage/Cargo.toml b/crates/storage/Cargo.toml index 7d18346..a1cc1b0 100644 --- a/crates/storage/Cargo.toml +++ b/crates/storage/Cargo.toml @@ -18,3 +18,6 @@ reth-db-api.workspace = true reth-libmdbx.workspace = true thiserror.workspace = true trevm.workspace = true + +[dev-dependencies] +tempfile.workspace = true diff --git a/crates/storage/src/hot/db_traits.rs b/crates/storage/src/hot/db_traits.rs index b71b637..31b9cdf 100644 --- a/crates/storage/src/hot/db_traits.rs +++ b/crates/storage/src/hot/db_traits.rs @@ -1,10 +1,9 @@ use crate::{ hot::{HotKvRead, HotKvWrite}, - tables::hot::{self as tables, AccountStorageKey}, + tables::hot::{self as tables}, }; use alloy::primitives::{Address, B256, U256}; use reth::primitives::{Account, Bytecode, Header, SealedHeader, StorageEntry}; -use std::borrow::Cow; /// Trait for database read operations. pub trait HotDbReader: sealed::Sealed { @@ -67,12 +66,7 @@ where } fn get_storage(&self, address: &Address, key: &B256) -> Result, Self::Error> { - let storage_key = AccountStorageKey { - address: std::borrow::Cow::Borrowed(address), - key: std::borrow::Cow::Borrowed(key), - }; - let key = storage_key.encode_key(); - self.get::(&key) + self.get_dual::(address, key) } } @@ -147,9 +141,7 @@ where key: &B256, entry: &U256, ) -> Result<(), Self::Error> { - let storage_key = - AccountStorageKey { address: Cow::Borrowed(address), key: Cow::Borrowed(key) }; - self.queue_put::(&storage_key.encode_key(), entry) + self.queue_put_dual::(address, key, entry) } fn commit(self) -> Result<(), Self::Error> { diff --git a/crates/storage/src/hot/mdbx.rs b/crates/storage/src/hot/mdbx.rs index b7704b0..dd9289f 100644 --- a/crates/storage/src/hot/mdbx.rs +++ b/crates/storage/src/hot/mdbx.rs @@ -1,7 +1,9 @@ use crate::{ hot::{HotKv, HotKvError, HotKvRead, HotKvReadError, HotKvWrite}, - ser::DeserError, + ser::{DeserError, KeySer, MAX_KEY_SIZE, ValSer}, + tables::DualKeyed, }; +use bytes::{BufMut, BytesMut}; use reth_db::{ Database, DatabaseEnv, mdbx::{RW, TransactionKind, WriteFlags, tx::Tx}, @@ -61,7 +63,7 @@ where { type Error = MdbxError; - fn get_raw<'a>( + fn raw_get<'a>( &'a self, table: &str, key: &[u8], @@ -70,6 +72,52 @@ where self.inner.get(dbi, key.as_ref()).map_err(MdbxError::Mdbx) } + + fn raw_get_dual<'a>( + &'a self, + _table: &str, + _key1: &[u8], + _key2: &[u8], + ) -> Result>, Self::Error> { + unimplemented!("Not implemented: raw_get_dual. Use get_dual instead."); + } + + fn get_dual( + &self, + key1: &T::K1, + key2: &T::K2, + ) -> Result, Self::Error> { + let mut key1_buf = [0u8; MAX_KEY_SIZE]; + let key1_bytes = key1.encode_key(&mut key1_buf); + + // K2 slice must be EXACTLY the size of the fixed value size, if the + // table has one. This is a bit ugly, and results in an extra + // allocation for fixed-size values. This could be avoided using + // max value size. + let value_bytes = if let Some(size) = T::FIXED_VALUE_SIZE { + let buf = vec![0u8; size]; + let _ = key2.encode_key(&mut buf[..MAX_KEY_SIZE].try_into().unwrap()); + + let db = self.inner.open_db(Some(T::NAME))?; + let mut cursor = self.inner.cursor(&db).map_err(MdbxError::Mdbx)?; + cursor.get_both_range(key1_bytes, &buf).map_err(MdbxError::Mdbx) + } else { + let mut buf = [0u8; MAX_KEY_SIZE]; + let encoded = key2.encode_key(&mut buf); + + let db = self.inner.open_db(Some(T::NAME))?; + let mut cursor = self.inner.cursor(&db).map_err(MdbxError::Mdbx)?; + cursor.get_both_range::>(key1_bytes, encoded).map_err(MdbxError::Mdbx) + }; + + let Some(value_bytes) = value_bytes? else { + return Ok(None); + }; + // we need to strip the key2 prefix from the value bytes before decoding + let value_bytes = &value_bytes[<::K2 as KeySer>::SIZE..]; + + T::Value::decode_value(&value_bytes).map(Some).map_err(Into::into) + } } impl HotKvWrite for Tx { @@ -79,6 +127,43 @@ impl HotKvWrite for Tx { self.inner.put(dbi, key, value, WriteFlags::UPSERT).map(|_| ()).map_err(MdbxError::Mdbx) } + fn queue_raw_put_dual( + &mut self, + _table: &str, + _key1: &[u8], + _key2: &[u8], + _value: &[u8], + ) -> Result<(), Self::Error> { + unimplemented!("Not implemented: queue_raw_put_dual. Use queue_put_dual instead."); + } + + // Specialized put for dual-keyed tables. + fn queue_put_dual( + &mut self, + key1: &T::K1, + key2: &T::K2, + value: &T::Value, + ) -> Result<(), Self::Error> { + let k2_size = ::SIZE; + let mut scratch = [0u8; MAX_KEY_SIZE]; + + // This will be the total length of key2 + value, reserved in mdbx + let encoded_len = k2_size + value.encoded_size(); + + // Prepend the value with k2. + let mut buf = BytesMut::with_capacity(encoded_len); + let encoded_k2 = key2.encode_key(&mut scratch); + buf.put_slice(encoded_k2); + value.encode_value_to(&mut buf); + + let encoded_k1 = key1.encode_key(&mut scratch); + // NB: DUPSORT and RESERVE are incompatible :( + let db = self.inner.open_db(Some(T::NAME))?; + self.inner.put(db.dbi(), encoded_k1, &buf, Default::default())?; + + Ok(()) + } + fn queue_raw_delete(&mut self, table: &str, key: &[u8]) -> Result<(), Self::Error> { let dbi = self.inner.open_db(Some(table)).map(|db| db.dbi())?; @@ -92,8 +177,22 @@ impl HotKvWrite for Tx { self.inner.clear_db(dbi).map(|_| ()).map_err(MdbxError::Mdbx) } - fn queue_raw_create(&mut self, table: &str) -> Result<(), Self::Error> { - self.inner.create_db(Some(table), Default::default()).map(|_| ()).map_err(MdbxError::Mdbx) + fn queue_raw_create( + &mut self, + table: &str, + dual_key: bool, + fixed_val: bool, + ) -> Result<(), Self::Error> { + let mut flags = Default::default(); + + if dual_key { + flags |= reth_libmdbx::DatabaseFlags::DUP_SORT; + if fixed_val { + flags |= reth_libmdbx::DatabaseFlags::DUP_FIXED; + } + } + + self.inner.create_db(Some(table), flags).map(|_| ()).map_err(MdbxError::Mdbx) } fn raw_commit(self) -> Result<(), Self::Error> { @@ -103,3 +202,458 @@ impl HotKvWrite for Tx { if res.0 { Err(reth_libmdbx::Error::Other(1).into()) } else { Ok(()) } } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::{ + hot::{HotDbWriter, HotKv, HotKvRead, HotKvWrite}, + tables::hot, + }; + use alloy::primitives::{Address, B256, BlockNumber, U256}; + use reth::primitives::{Account, Bytecode, Header}; + use reth_db::DatabaseEnv; + + /// A test database wrapper that automatically cleans up on drop + struct TestDb { + db: DatabaseEnv, + #[allow(dead_code)] + temp_dir: tempfile::TempDir, + } + + impl std::ops::Deref for TestDb { + type Target = DatabaseEnv; + + fn deref(&self) -> &Self::Target { + &self.db + } + } + + /// Create a temporary MDBX database for testing that will be automatically cleaned up + fn create_test_db() -> TestDb { + let temp_dir = tempfile::tempdir().expect("Failed to create temp directory"); + + // Create the database + let db = reth_db::create_db(&temp_dir, Default::default()).unwrap(); + + // Create tables from the `crate::tables::hot` module + let mut writer = db.writer().unwrap(); + + writer.queue_create::().unwrap(); + writer.queue_create::().unwrap(); + writer.queue_create::().unwrap(); + writer.queue_create::().unwrap(); + writer.queue_create::().unwrap(); + writer.queue_create::().unwrap(); + writer.queue_create::().unwrap(); + writer.queue_create::().unwrap(); + writer.queue_create::().unwrap(); + writer.queue_create::().unwrap(); + + writer.commit().expect("Failed to commit table creation"); + + TestDb { db, temp_dir } + } + + /// Create test data + fn create_test_account() -> (Address, Account) { + let address = Address::from_slice(&[0x1; 20]); + let account = Account { + nonce: 42, + balance: U256::from(1000u64), + bytecode_hash: Some(B256::from_slice(&[0x2; 32])), + }; + (address, account) + } + + fn create_test_bytecode() -> (B256, Bytecode) { + let hash = B256::from_slice(&[0x2; 32]); + let code = reth::primitives::Bytecode::new_raw(vec![0x60, 0x80, 0x60, 0x40].into()); + (hash, code) + } + + fn create_test_header() -> (BlockNumber, Header) { + let block_number = 12345; + let header = Header { + number: block_number, + gas_limit: 8000000, + gas_used: 100000, + timestamp: 1640995200, + parent_hash: B256::from_slice(&[0x3; 32]), + state_root: B256::from_slice(&[0x4; 32]), + ..Default::default() + }; + (block_number, header) + } + + #[test] + fn test_hotkv_basic_operations() { + let db = create_test_db(); + let (address, account) = create_test_account(); + let (hash, bytecode) = create_test_bytecode(); + + // Test HotKv::writer() and basic write operations + { + let mut writer: Tx = db.writer().unwrap(); + + // Create tables first + writer.queue_create::().unwrap(); + + // Write account data + writer.queue_put::(&address, &account).unwrap(); + writer.queue_put::(&hash, &bytecode).unwrap(); + + // Commit the transaction + writer.raw_commit().unwrap(); + } + + // Test HotKv::reader() and basic read operations + { + let reader: Tx = db.reader().unwrap(); + + // Read account data + let read_account: Option = + reader.get::(&address).unwrap(); + assert_eq!(read_account, Some(account)); + + // Read bytecode + let read_bytecode: Option = reader.get::(&hash).unwrap(); + assert_eq!(read_bytecode, Some(bytecode)); + + // Test non-existent data + let nonexistent_addr = Address::from_slice(&[0xff; 20]); + let nonexistent_account: Option = + reader.get::(&nonexistent_addr).unwrap(); + assert_eq!(nonexistent_account, None); + } + } + + #[test] + fn test_raw_operations() { + let db = create_test_db(); + + let table_name = "test_table"; + let key = b"test_key"; + let value = b"test_value"; + + // Test raw write operations + { + let mut writer: Tx = db.writer().unwrap(); + + // Create table + writer.queue_raw_create(table_name, false, false).unwrap(); + + // Put raw data + writer.queue_raw_put(table_name, key, value).unwrap(); + + writer.raw_commit().unwrap(); + } + + // Test raw read operations + { + let reader: Tx = db.reader().unwrap(); + + let read_value = reader.raw_get(table_name, key).unwrap(); + assert_eq!(read_value.as_deref(), Some(value.as_slice())); + + // Test non-existent key + let nonexistent = reader.raw_get(table_name, b"nonexistent").unwrap(); + assert_eq!(nonexistent, None); + } + + // Test raw delete + { + let mut writer: Tx = db.writer().unwrap(); + + writer.queue_raw_delete(table_name, key).unwrap(); + writer.raw_commit().unwrap(); + } + + // Verify deletion + { + let reader: Tx = db.reader().unwrap(); + let deleted_value = reader.raw_get(table_name, key).unwrap(); + assert_eq!(deleted_value, None); + } + } + + #[test] + fn test_dual_keyed_operations() { + let db = create_test_db(); + + let address = Address::from_slice(&[0x1; 20]); + let storage_key = B256::from_slice(&[0x5; 32]); + let storage_value = U256::from(999u64); + + // Test dual-keyed table operations + { + let mut writer: Tx = db.writer().unwrap(); + + // Put storage data using dual keys + writer + .queue_put_dual::(&address, &storage_key, &storage_value) + .unwrap(); + + writer.raw_commit().unwrap(); + } + + // Test reading dual-keyed data + { + let reader: Tx = db.reader().unwrap(); + + // Read storage using dual key lookup + let read_value = + reader.get_dual::(&address, &storage_key).unwrap().unwrap(); + + assert_eq!(read_value, storage_value); + } + } + + #[test] + fn test_table_management() { + let db = create_test_db(); + + // Add some data + let (block_number, header) = create_test_header(); + { + let mut writer: Tx = db.writer().unwrap(); + writer.queue_put::(&block_number, &header).unwrap(); + writer.raw_commit().unwrap(); + } + + // Verify data exists + { + let reader: Tx = db.reader().unwrap(); + let read_header: Option
= reader.get::(&block_number).unwrap(); + assert_eq!(read_header, Some(header.clone())); + } + + // Clear the table + { + let mut writer: Tx = db.writer().unwrap(); + writer.queue_clear::().unwrap(); + writer.raw_commit().unwrap(); + } + + // Verify table is empty + { + let reader: Tx = db.reader().unwrap(); + let read_header: Option
= reader.get::(&block_number).unwrap(); + assert_eq!(read_header, None); + } + } + + #[test] + fn test_batch_operations() { + let db = create_test_db(); + + // Create test data + let accounts: Vec<(Address, Account)> = (0..10) + .map(|i| { + let mut addr_bytes = [0u8; 20]; + addr_bytes[19] = i; + let address = Address::from_slice(&addr_bytes); + let account = Account { + nonce: i.into(), + balance: U256::from((i as u64) * 100), + bytecode_hash: None, + }; + (address, account) + }) + .collect(); + + // Test batch writes + { + let mut writer: Tx = db.writer().unwrap(); + + // Write multiple accounts + for (address, account) in &accounts { + writer.queue_put::(address, account).unwrap(); + } + + writer.raw_commit().unwrap(); + } + + // Test batch reads + { + let reader: Tx = db.reader().unwrap(); + + for (address, expected_account) in &accounts { + let read_account: Option = + reader.get::(address).unwrap(); + assert_eq!(read_account.as_ref(), Some(expected_account)); + } + } + + // Test batch get_many + { + let reader: Tx = db.reader().unwrap(); + let addresses: Vec
= accounts.iter().map(|(addr, _)| *addr).collect(); + let read_accounts: Vec> = + reader.get_many::(addresses.iter()).unwrap(); + + for (i, (_, expected_account)) in accounts.iter().enumerate() { + assert_eq!(read_accounts[i].as_ref(), Some(expected_account)); + } + } + } + + #[test] + fn test_transaction_isolation() { + let db = create_test_db(); + let (address, account) = create_test_account(); + + // Setup initial data + { + let mut writer: Tx = db.writer().unwrap(); + writer.queue_put::(&address, &account).unwrap(); + writer.raw_commit().unwrap(); + } + + // Start a reader transaction + let reader: Tx = db.reader().unwrap(); + + // Modify data in a writer transaction + { + let mut writer: Tx = db.writer().unwrap(); + let modified_account = + Account { nonce: 999, balance: U256::from(9999u64), bytecode_hash: None }; + writer.queue_put::(&address, &modified_account).unwrap(); + writer.raw_commit().unwrap(); + } + + // Reader should still see original data (snapshot isolation) + { + let read_account: Option = + reader.get::(&address).unwrap(); + assert_eq!(read_account, Some(account)); + } + + // New reader should see modified data + { + let new_reader: Tx = db.reader().unwrap(); + let read_account: Option = + new_reader.get::(&address).unwrap(); + assert_eq!(read_account.unwrap().nonce, 999); + } + } + + #[test] + fn test_multiple_readers() { + let db = create_test_db(); + let (address, account) = create_test_account(); + + // Setup data + { + let mut writer: Tx = db.writer().unwrap(); + writer.queue_put::(&address, &account).unwrap(); + writer.raw_commit().unwrap(); + } + + // Create multiple readers + let reader1: Tx = db.reader().unwrap(); + let reader2: Tx = db.reader().unwrap(); + let reader3: Tx = db.reader().unwrap(); + + // All readers should see the same data + let account1: Option = reader1.get::(&address).unwrap(); + let account2: Option = reader2.get::(&address).unwrap(); + let account3: Option = reader3.get::(&address).unwrap(); + + assert_eq!(account1, Some(account)); + assert_eq!(account2, Some(account)); + assert_eq!(account3, Some(account)); + } + + #[test] + fn test_error_handling() { + let db = create_test_db(); + + // Test reading from non-existent table + { + let reader: Tx = db.reader().unwrap(); + let result = reader.raw_get("nonexistent_table", b"key"); + + // Should handle gracefully (may return None or error depending on MDBX behavior) + match result { + Ok(None) => {} // This is fine + Err(_) => {} // This is also acceptable for non-existent table + Ok(Some(_)) => panic!("Should not return data for non-existent table"), + } + } + + // Test writing to a table without creating it first + { + let mut writer: Tx = db.writer().unwrap(); + let (address, account) = create_test_account(); + + // This should handle the case where table doesn't exist + let result = writer.queue_put::(&address, &account); + match result { + Ok(_) => { + // If it succeeds, commit should work + writer.raw_commit().unwrap(); + } + Err(_) => { + // If it fails, that's expected behavior + } + } + } + } + + #[test] + fn test_serialization_roundtrip() { + let db = create_test_db(); + + // Test various data types + let (block_number, header) = create_test_header(); + let canonical_hash = B256::from_slice(&[0x7; 32]); + + { + let mut writer: Tx = db.writer().unwrap(); + + // Create tables + + // Write different types + writer.queue_put::(&block_number, &header).unwrap(); + writer.queue_put::(&block_number, &canonical_hash).unwrap(); + + writer.raw_commit().unwrap(); + } + + { + let reader: Tx = db.reader().unwrap(); + + // Read and verify + let read_header: Option
= reader.get::(&block_number).unwrap(); + assert_eq!(read_header, Some(header)); + + let read_hash: Option = + reader.get::(&block_number).unwrap(); + assert_eq!(read_hash, Some(canonical_hash)); + } + } + + #[test] + fn test_large_data() { + let db = create_test_db(); + + // Create a large bytecode + let hash = B256::from_slice(&[0x8; 32]); + let large_code_vec: Vec = (0..10000).map(|i| (i % 256) as u8).collect(); + let large_bytecode = Bytecode::new_raw(large_code_vec.clone().into()); + + { + let mut writer: Tx = db.writer().unwrap(); + writer.queue_create::().unwrap(); + writer.queue_put::(&hash, &large_bytecode).unwrap(); + writer.raw_commit().unwrap(); + } + + { + let reader: Tx = db.reader().unwrap(); + let read_bytecode: Option = reader.get::(&hash).unwrap(); + assert_eq!(read_bytecode, Some(large_bytecode)); + } + } +} diff --git a/crates/storage/src/hot/mem.rs b/crates/storage/src/hot/mem.rs index f5e2353..8917e78 100644 --- a/crates/storage/src/hot/mem.rs +++ b/crates/storage/src/hot/mem.rs @@ -47,6 +47,15 @@ impl MemKv { buf[..k.len()].copy_from_slice(k); buf } + + #[track_caller] + fn dual_key(k1: &[u8], k2: &[u8]) -> [u8; MAX_KEY_SIZE] { + assert!(k1.len() + k2.len() <= MAX_KEY_SIZE, "Combined key length exceeds MAX_KEY_SIZE"); + let mut buf = [0u8; MAX_KEY_SIZE]; + buf[..k1.len()].copy_from_slice(k1); + buf[k1.len()..k1.len() + k2.len()].copy_from_slice(k2); + buf + } } impl Default for MemKv { @@ -233,7 +242,7 @@ impl HotKv for MemKv { impl HotKvRead for MemKvRoTx { type Error = HotKvError; - fn get_raw<'a>( + fn raw_get<'a>( &'a self, table: &str, key: &[u8], @@ -249,12 +258,22 @@ impl HotKvRead for MemKvRoTx { .and_then(|t| t.get(&key)) .map(|bytes| Cow::Borrowed(bytes.as_ref()))) } + + fn raw_get_dual<'a>( + &'a self, + table: &str, + key1: &[u8], + key2: &[u8], + ) -> Result>, Self::Error> { + let key = MemKv::dual_key(key1, key2); + self.raw_get(table, &key) + } } impl HotKvRead for MemKvRwTx { type Error = HotKvError; - fn get_raw<'a>( + fn raw_get<'a>( &'a self, table: &str, key: &[u8], @@ -285,6 +304,16 @@ impl HotKvRead for MemKvRwTx { .and_then(|t| t.get(&key)) .map(|bytes| Cow::Borrowed(bytes.as_ref()))) } + + fn raw_get_dual<'a>( + &'a self, + table: &str, + key1: &[u8], + key2: &[u8], + ) -> Result>, Self::Error> { + let key = MemKv::dual_key(key1, key2); + self.raw_get(table, &key) + } } impl HotKvWrite for MemKvRwTx { @@ -300,6 +329,17 @@ impl HotKvWrite for MemKvRwTx { Ok(()) } + fn queue_raw_put_dual( + &mut self, + table: &str, + key1: &[u8], + key2: &[u8], + value: &[u8], + ) -> Result<(), Self::Error> { + let key = MemKv::dual_key(key1, key2); + self.queue_raw_put(table, &key, value) + } + fn queue_raw_delete(&mut self, table: &str, key: &[u8]) -> Result<(), Self::Error> { let key = MemKv::key(key); @@ -313,7 +353,12 @@ impl HotKvWrite for MemKvRwTx { Ok(()) } - fn queue_raw_create(&mut self, _table: &str) -> Result<(), Self::Error> { + fn queue_raw_create( + &mut self, + _table: &str, + _dual_key: bool, + _dual_fixed: bool, + ) -> Result<(), Self::Error> { Ok(()) } @@ -359,7 +404,7 @@ mod tests { let reader = store.reader().unwrap(); // Empty store should return None for any key - assert!(reader.get_raw("test", &[1, 2, 3]).unwrap().is_none()); + assert!(reader.raw_get("test", &[1, 2, 3]).unwrap().is_none()); } #[test] @@ -377,9 +422,9 @@ mod tests { // Read the data back { let reader = store.reader().unwrap(); - let value1 = reader.get_raw("table1", &[1, 2, 3]).unwrap(); - let value2 = reader.get_raw("table1", &[4, 5, 6]).unwrap(); - let missing = reader.get_raw("table1", &[7, 8, 9]).unwrap(); + let value1 = reader.raw_get("table1", &[1, 2, 3]).unwrap(); + let value2 = reader.raw_get("table1", &[4, 5, 6]).unwrap(); + let missing = reader.raw_get("table1", &[7, 8, 9]).unwrap(); assert_eq!(value1.as_deref(), Some(b"value1" as &[u8])); assert_eq!(value2.as_deref(), Some(b"value2" as &[u8])); @@ -402,8 +447,8 @@ mod tests { // Read from different tables { let reader = store.reader().unwrap(); - let value1 = reader.get_raw("table1", &[1]).unwrap(); - let value2 = reader.get_raw("table2", &[1]).unwrap(); + let value1 = reader.raw_get("table1", &[1]).unwrap(); + let value2 = reader.raw_get("table2", &[1]).unwrap(); assert_eq!(value1.as_deref(), Some(b"table1_value" as &[u8])); assert_eq!(value2.as_deref(), Some(b"table2_value" as &[u8])); @@ -431,7 +476,7 @@ mod tests { // Check the value was updated { let reader = store.reader().unwrap(); - let value = reader.get_raw("table1", &[1]).unwrap(); + let value = reader.raw_get("table1", &[1]).unwrap(); assert_eq!(value.as_deref(), Some(b"updated" as &[u8])); } } @@ -445,7 +490,7 @@ mod tests { writer.queue_raw_put("table1", &[1], b"queued_value").unwrap(); // Should be able to read the queued value - let value = writer.get_raw("table1", &[1]).unwrap(); + let value = writer.raw_get("table1", &[1]).unwrap(); assert_eq!(value.as_deref(), Some(b"queued_value" as &[u8])); writer.raw_commit().unwrap(); @@ -453,7 +498,7 @@ mod tests { // After commit, other readers should see it { let reader = store.reader().unwrap(); - let value = reader.get_raw("table1", &[1]).unwrap(); + let value = reader.raw_get("table1", &[1]).unwrap(); assert_eq!(value.as_deref(), Some(b"queued_value" as &[u8])); } } @@ -557,8 +602,8 @@ mod tests { let reader1 = store.reader().unwrap(); let reader2 = store.reader().unwrap(); - let value1 = reader1.get_raw("table1", &[1]).unwrap(); - let value2 = reader2.get_raw("table1", &[1]).unwrap(); + let value1 = reader1.raw_get("table1", &[1]).unwrap(); + let value2 = reader2.raw_get("table1", &[1]).unwrap(); assert_eq!(value1.as_deref(), Some(b"value1" as &[u8])); assert_eq!(value2.as_deref(), Some(b"value1" as &[u8])); @@ -591,7 +636,7 @@ mod tests { { let reader = store.reader().unwrap(); - let value = reader.get_raw("table1", &[1]).unwrap(); + let value = reader.raw_get("table1", &[1]).unwrap(); assert_eq!(value.as_deref(), Some(b"" as &[u8])); } } @@ -609,7 +654,7 @@ mod tests { writer.queue_raw_put("table1", &[1], b"third").unwrap(); // Read-your-writes should return the latest value - let value = writer.get_raw("table1", &[1]).unwrap(); + let value = writer.raw_get("table1", &[1]).unwrap(); assert_eq!(value.as_deref(), Some(b"third" as &[u8])); writer.raw_commit().unwrap(); @@ -617,7 +662,7 @@ mod tests { { let reader = store.reader().unwrap(); - let value = reader.get_raw("table1", &[1]).unwrap(); + let value = reader.raw_get("table1", &[1]).unwrap(); assert_eq!(value.as_deref(), Some(b"third" as &[u8])); } } @@ -636,7 +681,7 @@ mod tests { // Start a read transaction { let reader = store.reader().unwrap(); - let original_value = reader.get_raw("table1", &[1]).unwrap(); + let original_value = reader.raw_get("table1", &[1]).unwrap(); assert_eq!(original_value.as_deref(), Some(b"original" as &[u8])); } @@ -651,7 +696,7 @@ mod tests { { // New reader should see the updated value let new_reader = store.reader().unwrap(); - let updated_value = new_reader.get_raw("table1", &[1]).unwrap(); + let updated_value = new_reader.raw_get("table1", &[1]).unwrap(); assert_eq!(updated_value.as_deref(), Some(b"updated" as &[u8])); } } @@ -669,7 +714,7 @@ mod tests { // Value should not be persisted { let reader = store.reader().unwrap(); - let value = reader.get_raw("table1", &[1]).unwrap(); + let value = reader.raw_get("table1", &[1]).unwrap(); assert!(value.is_none()); } } @@ -687,8 +732,8 @@ mod tests { { let reader = store.reader().unwrap(); - let value1 = reader.get_raw("table1", &[1]).unwrap(); - let value2 = reader.get_raw("table2", &[2]).unwrap(); + let value1 = reader.raw_get("table1", &[1]).unwrap(); + let value2 = reader.raw_get("table2", &[2]).unwrap(); assert_eq!(value1.as_deref(), Some(b"value1" as &[u8])); assert_eq!(value2.as_deref(), Some(b"value2" as &[u8])); @@ -708,8 +753,8 @@ mod tests { let ro_tx = rw_tx.commit_downgrade(); // Read the data back - let value1 = ro_tx.get_raw("table1", &[1, 2, 3]).unwrap(); - let value2 = ro_tx.get_raw("table1", &[4, 5, 6]).unwrap(); + let value1 = ro_tx.raw_get("table1", &[1, 2, 3]).unwrap(); + let value2 = ro_tx.raw_get("table1", &[4, 5, 6]).unwrap(); assert_eq!(value1.as_deref(), Some(b"value1" as &[u8])); assert_eq!(value2.as_deref(), Some(b"value2" as &[u8])); @@ -724,7 +769,7 @@ mod tests { let ro_tx = rw_tx.downgrade(); // Read the data back - let value3 = ro_tx.get_raw("table2", &[7, 8, 9]).unwrap(); + let value3 = ro_tx.raw_get("table2", &[7, 8, 9]).unwrap(); assert!(value3.is_none()); } @@ -744,8 +789,8 @@ mod tests { { let reader = store.reader().unwrap(); - let value1 = reader.get_raw("table1", &[1]).unwrap(); - let value2 = reader.get_raw("table1", &[2]).unwrap(); + let value1 = reader.raw_get("table1", &[1]).unwrap(); + let value2 = reader.raw_get("table1", &[2]).unwrap(); assert_eq!(value1.as_deref(), Some(b"value1" as &[u8])); assert_eq!(value2.as_deref(), Some(b"value2" as &[u8])); @@ -754,16 +799,16 @@ mod tests { { let mut writer = store.writer().unwrap(); - let value1 = writer.get_raw("table1", &[1]).unwrap(); - let value2 = writer.get_raw("table1", &[2]).unwrap(); + let value1 = writer.raw_get("table1", &[1]).unwrap(); + let value2 = writer.raw_get("table1", &[2]).unwrap(); assert_eq!(value1.as_deref(), Some(b"value1" as &[u8])); assert_eq!(value2.as_deref(), Some(b"value2" as &[u8])); writer.queue_raw_clear("table1").unwrap(); - let value1 = writer.get_raw("table1", &[1]).unwrap(); - let value2 = writer.get_raw("table1", &[2]).unwrap(); + let value1 = writer.raw_get("table1", &[1]).unwrap(); + let value2 = writer.raw_get("table1", &[2]).unwrap(); assert!(value1.is_none()); assert!(value2.is_none()); @@ -773,8 +818,8 @@ mod tests { { let reader = store.reader().unwrap(); - let value1 = reader.get_raw("table1", &[1]).unwrap(); - let value2 = reader.get_raw("table1", &[2]).unwrap(); + let value1 = reader.raw_get("table1", &[1]).unwrap(); + let value2 = reader.raw_get("table1", &[2]).unwrap(); assert!(value1.is_none()); assert!(value2.is_none()); diff --git a/crates/storage/src/hot/revm.rs b/crates/storage/src/hot/revm.rs index fea9266..1f06969 100644 --- a/crates/storage/src/hot/revm.rs +++ b/crates/storage/src/hot/revm.rs @@ -1,10 +1,14 @@ use crate::{ hot::{HotKvError, HotKvRead, HotKvWrite}, - tables::hot::{AccountStorageKey, Bytecodes, PlainAccountState}, + tables::{ + DualKeyed, Table, + hot::{self, Bytecodes, PlainAccountState}, + }, }; use alloy::primitives::{Address, B256, KECCAK256_EMPTY}; use core::fmt; use reth::primitives::Account; +use std::borrow::Cow; use trevm::revm::{ database::{DBErrorMarker, Database, DatabaseRef, TryDatabaseCommit}, primitives::{HashMap, StorageKey, StorageValue}, @@ -33,15 +37,45 @@ impl fmt::Debug for RevmRead { } // HotKvRead implementation for RevmRead -impl HotKvRead for RevmRead { - type Error = T::Error; +impl HotKvRead for RevmRead { + type Error = U::Error; - fn get_raw<'a>( + fn raw_get<'a>( &'a self, table: &str, key: &[u8], ) -> Result>, Self::Error> { - self.reader.get_raw(table, key) + self.reader.raw_get(table, key) + } + + fn raw_get_dual<'a>( + &'a self, + table: &str, + key1: &[u8], + key2: &[u8], + ) -> Result>, Self::Error> { + self.reader.raw_get_dual(table, key1, key2) + } + + fn get(&self, key: &T::Key) -> Result, Self::Error> { + self.reader.get::(key) + } + + fn get_dual( + &self, + key1: &T::K1, + key2: &T::K2, + ) -> Result, Self::Error> { + self.reader.get_dual::(key1, key2) + } + + fn get_many<'a, T, I>(&self, keys: I) -> Result>, Self::Error> + where + T::Key: 'a, + T: Table, + I: IntoIterator, + { + self.reader.get_many::(keys) } } @@ -49,45 +83,86 @@ impl HotKvRead for RevmRead { /// Despite the naming of [`TryDatabaseCommit::try_commit`], the changes are /// only persisted when [`Self::persist`] is called. This is because of a /// mismatch in semantics between the two systems. -pub struct RevmWrite { - writer: T, +pub struct RevmWrite { + writer: U, } -impl RevmWrite { + +impl RevmWrite { /// Create a new write adapter - pub const fn new(writer: T) -> Self { + pub const fn new(writer: U) -> Self { Self { writer } } /// Persist the changes made in this write transaction. - pub fn persist(self) -> Result<(), T::Error> { + pub fn persist(self) -> Result<(), U::Error> { self.writer.raw_commit() } } -impl fmt::Debug for RevmWrite { +impl fmt::Debug for RevmWrite { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("RevmWrite").finish() } } // HotKvWrite implementation for RevmWrite -impl HotKvRead for RevmWrite { - type Error = T::Error; +impl HotKvRead for RevmWrite { + type Error = U::Error; - fn get_raw<'a>( + fn raw_get<'a>( &'a self, table: &str, key: &[u8], ) -> Result>, Self::Error> { - self.writer.get_raw(table, key) + self.writer.raw_get(table, key) + } + + fn raw_get_dual<'a>( + &'a self, + table: &str, + key1: &[u8], + key2: &[u8], + ) -> Result>, Self::Error> { + self.writer.raw_get_dual(table, key1, key2) + } + + fn get(&self, key: &T::Key) -> Result, Self::Error> { + self.writer.get::(key) + } + + fn get_dual( + &self, + key1: &T::K1, + key2: &T::K2, + ) -> Result, Self::Error> { + self.writer.get_dual::(key1, key2) + } + + fn get_many<'a, T, I>(&self, keys: I) -> Result>, Self::Error> + where + T::Key: 'a, + T: Table, + I: IntoIterator, + { + self.writer.get_many::(keys) } } -impl HotKvWrite for RevmWrite { +impl HotKvWrite for RevmWrite { fn queue_raw_put(&mut self, table: &str, key: &[u8], value: &[u8]) -> Result<(), Self::Error> { self.writer.queue_raw_put(table, key, value) } + fn queue_raw_put_dual( + &mut self, + table: &str, + key1: &[u8], + key2: &[u8], + value: &[u8], + ) -> Result<(), Self::Error> { + self.writer.queue_raw_put_dual(table, key1, key2, value) + } + fn queue_raw_delete(&mut self, table: &str, key: &[u8]) -> Result<(), Self::Error> { self.writer.queue_raw_delete(table, key) } @@ -96,13 +171,59 @@ impl HotKvWrite for RevmWrite { self.writer.queue_raw_clear(table) } - fn queue_raw_create(&mut self, table: &str) -> Result<(), Self::Error> { - self.writer.queue_raw_create(table) + fn queue_raw_create( + &mut self, + table: &str, + dual_key: bool, + dual_fixed: bool, + ) -> Result<(), Self::Error> { + self.writer.queue_raw_create(table, dual_key, dual_fixed) } fn raw_commit(self) -> Result<(), Self::Error> { self.writer.raw_commit() } + + fn queue_put(&mut self, key: &T::Key, value: &T::Value) -> Result<(), Self::Error> { + self.writer.queue_put::(key, value) + } + + fn queue_put_dual( + &mut self, + key1: &T::K1, + key2: &T::K2, + value: &T::Value, + ) -> Result<(), Self::Error> { + self.writer.queue_put_dual::(key1, key2, value) + } + + fn queue_delete(&mut self, key: &T::Key) -> Result<(), Self::Error> { + self.writer.queue_delete::(key) + } + + fn queue_put_many<'a, 'b, T, I>(&mut self, entries: I) -> Result<(), Self::Error> + where + T: Table, + T::Key: 'a, + T::Value: 'b, + I: IntoIterator, + { + self.writer.queue_put_many::(entries) + } + + fn queue_create(&mut self) -> Result<(), Self::Error> + where + T: Table, + { + self.writer.queue_create::() + } + + fn queue_clear(&mut self) -> Result<(), Self::Error> + where + T: Table, + { + self.writer.queue_clear::() + } } // DatabaseRef implementation for RevmRead @@ -138,16 +259,9 @@ where address: Address, index: StorageKey, ) -> Result { - let storage_key = AccountStorageKey { - address: std::borrow::Cow::Borrowed(&address), - key: std::borrow::Cow::Owned(B256::from_slice(&index.to_be_bytes::<32>())), - } - .encode_key(); + let key = B256::from_slice(&index.to_be_bytes::<32>()); - Ok(self - .reader - .get::(&storage_key)? - .unwrap_or_default()) + Ok(self.reader.get_dual::(&address, &key)?.unwrap_or_default()) } fn block_hash_ref(&self, _number: u64) -> Result { @@ -218,16 +332,8 @@ where address: Address, index: StorageKey, ) -> Result { - let storage_key = AccountStorageKey { - address: std::borrow::Cow::Borrowed(&address), - key: std::borrow::Cow::Owned(B256::from_slice(&index.to_be_bytes::<32>())), - } - .encode_key(); - - Ok(self - .writer - .get::(&storage_key)? - .unwrap_or_default()) + let key = B256::from_slice(&index.to_be_bytes::<32>()); + Ok(self.writer.get_dual::(&address, &key)?.unwrap_or_default()) } fn block_hash_ref(&self, _number: u64) -> Result { @@ -285,13 +391,10 @@ where // Handle storage changes for (key, value) in account.storage { - let storage_key = AccountStorageKey { - address: std::borrow::Cow::Borrowed(&address), - key: std::borrow::Cow::Owned(B256::from_slice(&key.to_be_bytes::<32>())), - } - .encode_key(); - self.writer.queue_put::( - &storage_key, + let key = B256::from_slice(&key.to_be_bytes::<32>()); + self.writer.queue_put_dual::( + &address, + &key, &value.present_value(), )?; } @@ -306,7 +409,7 @@ mod tests { use super::*; use crate::{ hot::{HotKv, HotKvRead, HotKvWrite, MemKv}, - tables::hot::{Bytecodes, PlainAccountState, PlainStorageState}, + tables::hot::{Bytecodes, PlainAccountState}, }; use alloy::primitives::{Address, B256, U256}; use reth::primitives::{Account, Bytecode}; @@ -507,17 +610,9 @@ mod tests { assert_eq!(acc.balance, U256::from(2000u64)); assert_eq!(acc.bytecode_hash, None); - // Check storage was written - let storage_key = AccountStorageKey { - address: std::borrow::Cow::Borrowed(&address), - key: std::borrow::Cow::Owned(B256::from_slice( - &U256::from(100u64).to_be_bytes::<32>(), - )), - } - .encode_key(); - + let key = B256::with_last_byte(100); let storage_val: Option = - reader.get::(&storage_key)?; + reader.get_dual::(&address, &key)?; assert_eq!(storage_val, Some(U256::from(200u64))); } diff --git a/crates/storage/src/hot/traits.rs b/crates/storage/src/hot/traits.rs index b23aed1..7483338 100644 --- a/crates/storage/src/hot/traits.rs +++ b/crates/storage/src/hot/traits.rs @@ -6,7 +6,7 @@ use crate::{ revm::{RevmRead, RevmWrite}, }, ser::{KeySer, MAX_KEY_SIZE, ValSer}, - tables::Table, + tables::{DualKeyed, Table}, }; /// Trait for hot storage. This is a KV store with read/write transactions. @@ -68,9 +68,22 @@ pub trait HotKvRead { /// /// The `key` buf must be <= [`MAX_KEY_SIZE`] bytes. Implementations are /// allowed to panic if this is not the case. - fn get_raw<'a>(&'a self, table: &str, key: &[u8]) + /// + /// If the table is dual-keyed, the output may be implementation-defined. + fn raw_get<'a>(&'a self, table: &str, key: &[u8]) -> Result>, Self::Error>; + /// Get a raw value from a specific table with dual keys. + /// + /// If the table is not dual-keyed, the output may be + /// implementation-defined. + fn raw_get_dual<'a>( + &'a self, + table: &str, + key1: &[u8], + key2: &[u8], + ) -> Result>, Self::Error>; + /// Get a value from a specific table. fn get(&self, key: &T::Key) -> Result, Self::Error> { let mut key_buf = [0u8; MAX_KEY_SIZE]; @@ -80,11 +93,28 @@ pub trait HotKvRead { "Encoded key length does not match expected size" ); - let Some(value_bytes) = self.get_raw(T::NAME, key_bytes)? else { + let Some(value_bytes) = self.raw_get(T::NAME, key_bytes)? else { + return Ok(None); + }; + T::Value::decode_value(&value_bytes).map(Some).map_err(Into::into) + } + + /// Get a value from a specific dual-keyed table. + fn get_dual( + &self, + key1: &T::K1, + key2: &T::K2, + ) -> Result, Self::Error> { + let mut key1_buf = [0u8; MAX_KEY_SIZE]; + let mut key2_buf = [0u8; MAX_KEY_SIZE]; + + let key1_bytes = key1.encode_key(&mut key1_buf); + let key2_bytes = key2.encode_key(&mut key2_buf); + + let Some(value_bytes) = self.raw_get_dual(T::NAME, key1_bytes, key2_bytes)? else { return Ok(None); }; - let data = &value_bytes[..]; - T::Value::decode_value(data).map(Some).map_err(Into::into) + T::Value::decode_value(&value_bytes).map(Some).map_err(Into::into) } /// Get many values from a specific table. @@ -110,7 +140,7 @@ pub trait HotKvRead { let mut key_buf = [0u8; MAX_KEY_SIZE]; keys.into_iter() - .map(|key| self.get_raw(T::NAME, key.encode_key(&mut key_buf))) + .map(|key| self.raw_get(T::NAME, key.encode_key(&mut key_buf))) .map(|maybe_val| { maybe_val .and_then(|val| ValSer::maybe_decode_value(val.as_deref()).map_err(Into::into)) @@ -127,6 +157,18 @@ pub trait HotKvWrite: HotKvRead { /// allowed to panic if this is not the case. fn queue_raw_put(&mut self, table: &str, key: &[u8], value: &[u8]) -> Result<(), Self::Error>; + /// Queue a raw put operation for a dual-keyed table. + //// + /// The `key1` and `key2` buf must be <= [`MAX_KEY_SIZE`] bytes. + /// Implementations are allowed to panic if this is not the case. + fn queue_raw_put_dual( + &mut self, + table: &str, + key1: &[u8], + key2: &[u8], + value: &[u8], + ) -> Result<(), Self::Error>; + /// Queue a raw delete operation. /// /// The `key` buf must be <= [`MAX_KEY_SIZE`] bytes. Implementations are @@ -137,7 +179,19 @@ pub trait HotKvWrite: HotKvRead { fn queue_raw_clear(&mut self, table: &str) -> Result<(), Self::Error>; /// Queue a raw create operation for a specific table. - fn queue_raw_create(&mut self, table: &str) -> Result<(), Self::Error>; + /// + /// This abstraction supports two table specializations: + /// 1. `dual_key`: whether the table uses dual keys (interior maps, called + /// `DUPSORT` in LMDB/MDBX). + /// 2. `fixed_val`: whether the table has fixed-size values. + /// + /// Database implementations can use this information for optimizations. + fn queue_raw_create( + &mut self, + table: &str, + dual_key: bool, + fixed_val: bool, + ) -> Result<(), Self::Error>; /// Queue a put operation for a specific table. fn queue_put(&mut self, key: &T::Key, value: &T::Value) -> Result<(), Self::Error> { @@ -148,6 +202,22 @@ pub trait HotKvWrite: HotKvRead { self.queue_raw_put(T::NAME, key_bytes, &value_bytes) } + /// Queue a put operation for a specific dual-keyed table. + fn queue_put_dual( + &mut self, + key1: &T::K1, + key2: &T::K2, + value: &T::Value, + ) -> Result<(), Self::Error> { + let mut key1_buf = [0u8; MAX_KEY_SIZE]; + let mut key2_buf = [0u8; MAX_KEY_SIZE]; + let key1_bytes = key1.encode_key(&mut key1_buf); + let key2_bytes = key2.encode_key(&mut key2_buf); + let value_bytes = value.encoded(); + + self.queue_raw_put_dual(T::NAME, key1_bytes, key2_bytes, &value_bytes) + } + /// Queue a delete operation for a specific table. fn queue_delete(&mut self, key: &T::Key) -> Result<(), Self::Error> { let mut key_buf = [0u8; MAX_KEY_SIZE]; @@ -181,7 +251,7 @@ pub trait HotKvWrite: HotKvRead { where T: Table, { - self.queue_raw_create(T::NAME) + self.queue_raw_create(T::NAME, T::DUAL_KEY, T::DUAL_FIXED_VAL) } /// Queue clearing all entries in a specific table. diff --git a/crates/storage/src/ser/impls.rs b/crates/storage/src/ser/impls.rs index 96537e8..bcfb819 100644 --- a/crates/storage/src/ser/impls.rs +++ b/crates/storage/src/ser/impls.rs @@ -1,6 +1,8 @@ use crate::ser::{DeserError, KeySer, MAX_KEY_SIZE, ValSer}; -use alloy::primitives::Bloom; +use alloy::primitives::{Address, B256, Bloom}; use bytes::BufMut; +use reth::primitives::StorageEntry; +use reth_db::models::BlockNumberAddress; macro_rules! delegate_val_to_key { ($ty:ty) => { @@ -120,7 +122,7 @@ ser_be_num!( ser_alloy_fixed!(8, 16, 20, 32, 52, 65, 256); delegate_val_to_key!(alloy::primitives::Address); -impl KeySer for alloy::primitives::Address { +impl KeySer for Address { const SIZE: usize = 20; fn encode_key<'a: 'c, 'b: 'c, 'c>(&'a self, _buf: &'b mut [u8; MAX_KEY_SIZE]) -> &'c [u8] { @@ -295,6 +297,47 @@ where } } +impl KeySer for BlockNumberAddress { + const SIZE: usize = u64::SIZE + Address::SIZE; + + fn encode_key<'a: 'c, 'b: 'c, 'c>(&'a self, buf: &'b mut [u8; MAX_KEY_SIZE]) -> &'c [u8] { + buf.copy_from_slice(&self.0.0.to_be_bytes()); + buf[8..28].copy_from_slice(self.0.1.as_ref()); + &buf[..Self::SIZE] + } + + fn decode_key(data: &[u8]) -> Result { + if data.len() < Self::SIZE { + return Err(DeserError::InsufficientData { needed: Self::SIZE, available: data.len() }); + } + let number = u64::from_be_bytes(data[0..8].try_into().unwrap()); + let address = Address::from_slice(&data[8..28]); + Ok(BlockNumberAddress((number, address))) + } +} + +impl ValSer for StorageEntry { + fn encoded_size(&self) -> usize { + self.key.encoded_size() + self.value.encoded_size() + } + + fn encode_value_to(&self, buf: &mut B) + where + B: bytes::BufMut + AsMut<[u8]>, + { + self.key.encode_value_to(buf); + self.value.encode_value_to(buf); + } + + fn decode_value(data: &[u8]) -> Result + where + Self: Sized, + { + let key: B256 = ValSer::decode_value(data)?; + let value = ValSer::decode_value(&data[key.encoded_size()..])?; + Ok(StorageEntry { key, value }) + } +} #[cfg(test)] mod tests { use super::*; diff --git a/crates/storage/src/tables/hot.rs b/crates/storage/src/tables/hot.rs index ec19884..6c0fabb 100644 --- a/crates/storage/src/tables/hot.rs +++ b/crates/storage/src/tables/hot.rs @@ -1,8 +1,6 @@ -use std::borrow::Cow; - -use crate::tables::Table; -use alloy::primitives::{Address, B256, BlockNumber, FixedBytes, U256}; -use reth::primitives::{Account, Bytecode, Header}; +use alloy::primitives::{Address, B256, BlockNumber, U256}; +use reth::primitives::{Account, Bytecode, Header, StorageEntry}; +use reth_db::models::BlockNumberAddress; use reth_db_api::{ BlockNumberList, models::{AccountBeforeTx, ShardedKey, storage_sharded_key::StorageShardedKey}, @@ -10,77 +8,35 @@ use reth_db_api::{ tables! { /// Records recent block Headers, by their number. - Headers, + Headers Header>, /// Records block numbers by hash. - HeaderNumbers, + HeaderNumbers BlockNumber>, /// Records the canonical chain header hashes, by height. - CanonicalHeaders, + CanonicalHeaders B256>, /// Records contract Bytecode, by its hash. - Bytecodes, + Bytecodes Bytecode>, /// Records plain account states, keyed by address. - PlainAccountState, + PlainAccountState
Account>, /// Records account state change history, keyed by address. - AccountsHistory, BlockNumberList>, + AccountsHistory => BlockNumberList>, /// Records storage state change history, keyed by address and storage key. - StorageHistory, - - /// Records account change sets, keyed by block number. - AccountChangeSets, -} - -/// Records plain storage states, keyed by address and storage key. -#[derive(Debug, Copy, Clone, Eq, PartialEq, PartialOrd, Ord)] -pub struct PlainStorageState; - -impl Table for PlainStorageState { - const NAME: &'static str = "PlainStorageState"; - - type Key = FixedBytes<52>; + StorageHistory BlockNumberList>, - type Value = U256; } -/// Key for the [`PlainStorageState`] table. -#[derive(Debug, Clone, Eq, PartialEq, PartialOrd, Ord)] -pub struct AccountStorageKey<'a, 'b> { - /// Address of the account. - pub address: Cow<'a, Address>, - /// Storage key. - pub key: Cow<'b, B256>, -} - -impl AccountStorageKey<'static, 'static> { - /// Decode the key from the provided data. - pub fn decode_key(data: &[u8]) -> Result { - if data.len() < Self::SIZE { - return Err(crate::ser::DeserError::InsufficientData { - needed: Self::SIZE, - available: data.len(), - }); - } - - let address = Address::from_slice(&data[0..20]); - let key = B256::from_slice(&data[20..52]); - - Ok(Self { address: Cow::Owned(address), key: Cow::Owned(key) }) - } -} +tables! { + /// Records plain storage states, keyed by address and storage key. + PlainStorageState
B256 => U256> size: Some(32 + 32), -impl<'a, 'b> AccountStorageKey<'a, 'b> { - /// Size in bytes. - pub const SIZE: usize = 20 + 32; + /// Records account states before transactions, keyed by (address, block number). + StorageChangeSets B256 => StorageEntry> size: Some(32 + 32 + 32), - /// Encode the key into the provided buffer. - pub fn encode_key(&self) -> FixedBytes<52> { - let mut buf = [0u8; Self::SIZE]; - buf[0..20].copy_from_slice(self.address.as_slice()); - buf[20..52].copy_from_slice(self.key.as_slice()); - buf.into() - } + /// Records account states before transactions, keyed by (address, block number). + AccountChangeSets Address => AccountBeforeTx> size: None, } diff --git a/crates/storage/src/tables/macros.rs b/crates/storage/src/tables/macros.rs index 70c52e8..c17c13d 100644 --- a/crates/storage/src/tables/macros.rs +++ b/crates/storage/src/tables/macros.rs @@ -1,7 +1,8 @@ macro_rules! tables { ( #[doc = $doc:expr] - $name:ident<$key:ty, $value:ty>) => { + $name:ident<$key:ty => $value:ty> + ) => { #[doc = $doc] #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] pub struct $name; @@ -14,9 +15,35 @@ macro_rules! tables { } }; - ($(#[doc = $doc:expr] $name:ident<$key:ty, $value:ty>),* $(,)?) => { + ( + #[doc = $doc:expr] + $name:ident<$key:ty => $subkey:ty => $value:ty> size: $fixed:expr + ) => { + #[doc = $doc] + #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] + pub struct $name; + + impl crate::tables::DualKeyed for $name { + const NAME: &'static str = stringify!($name); + + const FIXED_VALUE_SIZE: Option = { $fixed }; + + type K1 = $key; + type K2 = $subkey; + + type Value = $value; + } + }; + + ($(#[doc = $doc:expr] $name:ident<$key:ty => $value:ty>),* $(,)?) => { + $( + tables!(#[doc = $doc] $name<$key => $value>); + )* + }; + + ($(#[doc = $doc:expr] $name:ident<$key:ty => $subkey:ty => $value:ty> size: $fixed:expr),* $(,)?) => { $( - tables!(#[doc = $doc] $name<$key, $value>); + tables!(#[doc = $doc] $name<$key => $subkey => $value> size: $fixed); )* }; } diff --git a/crates/storage/src/tables/mod.rs b/crates/storage/src/tables/mod.rs index bbd6263..90187aa 100644 --- a/crates/storage/src/tables/mod.rs +++ b/crates/storage/src/tables/mod.rs @@ -11,11 +11,55 @@ use crate::ser::{KeySer, ValSer}; /// Trait for table definitions. pub trait Table { - /// A Human-readable name for the table. + /// A short, human-readable name for the table. const NAME: &'static str; + /// Indicates that this table uses dual keys. + const DUAL_KEY: bool = false; + + /// Indicates that this table has fixed-size values. + const DUAL_FIXED_VAL: bool = false; + /// The key type. type Key: KeySer; /// The value type. type Value: ValSer; } + +/// Trait for tables with two keys. +/// +/// This trait aims to capture tables that use a composite key made up of two +/// distinct parts. This is useful for representing (e.g.) dupsort or other +/// nested map optimizations. +pub trait DualKeyed { + /// A short, human-readable name for the table. + const NAME: &'static str; + + /// If the value size is fixed, `Some(size)`. Otherwise, `None`. + const FIXED_VALUE_SIZE: Option = None; + + /// The first key type. + type K1: KeySer; + + /// The second key type. + type K2: KeySer; + + /// The value type. + type Value: ValSer; +} + +impl Table for T +where + T: DualKeyed, +{ + const NAME: &'static str = T::NAME; + + /// Indicates that this table uses dual keys. + const DUAL_KEY: bool = true; + + /// Indicates that this table has fixed-size values. + const DUAL_FIXED_VAL: bool = T::FIXED_VALUE_SIZE.is_some(); + + type Key = T::K1; + type Value = T::Value; +} From a4707314bdddaf2fbcc18e9e0bb280fa2caca74d Mon Sep 17 00:00:00 2001 From: James Date: Thu, 15 Jan 2026 16:26:35 -0500 Subject: [PATCH 10/16] refactor: expand table, contract dualkeyed --- crates/storage/src/hot/mdbx.rs | 18 ++++---- crates/storage/src/hot/revm.rs | 12 +++--- crates/storage/src/hot/traits.rs | 10 ++--- crates/storage/src/tables/hot.rs | 48 +++++++++++++-------- crates/storage/src/tables/macros.rs | 66 +++++++++++++++++++---------- crates/storage/src/tables/mod.rs | 54 +++++++++++------------ 6 files changed, 118 insertions(+), 90 deletions(-) diff --git a/crates/storage/src/hot/mdbx.rs b/crates/storage/src/hot/mdbx.rs index dd9289f..c1f31de 100644 --- a/crates/storage/src/hot/mdbx.rs +++ b/crates/storage/src/hot/mdbx.rs @@ -1,7 +1,7 @@ use crate::{ hot::{HotKv, HotKvError, HotKvRead, HotKvReadError, HotKvWrite}, ser::{DeserError, KeySer, MAX_KEY_SIZE, ValSer}, - tables::DualKeyed, + tables::{DualKeyed, Table}, }; use bytes::{BufMut, BytesMut}; use reth_db::{ @@ -84,8 +84,8 @@ where fn get_dual( &self, - key1: &T::K1, - key2: &T::K2, + key1: &T::Key, + key2: &T::Key2, ) -> Result, Self::Error> { let mut key1_buf = [0u8; MAX_KEY_SIZE]; let key1_bytes = key1.encode_key(&mut key1_buf); @@ -94,7 +94,7 @@ where // table has one. This is a bit ugly, and results in an extra // allocation for fixed-size values. This could be avoided using // max value size. - let value_bytes = if let Some(size) = T::FIXED_VALUE_SIZE { + let value_bytes = if let Some(size) = ::FIXED_VAL_SIZE { let buf = vec![0u8; size]; let _ = key2.encode_key(&mut buf[..MAX_KEY_SIZE].try_into().unwrap()); @@ -114,9 +114,9 @@ where return Ok(None); }; // we need to strip the key2 prefix from the value bytes before decoding - let value_bytes = &value_bytes[<::K2 as KeySer>::SIZE..]; + let value_bytes = &value_bytes[<::Key2 as KeySer>::SIZE..]; - T::Value::decode_value(&value_bytes).map(Some).map_err(Into::into) + T::Value::decode_value(value_bytes).map(Some).map_err(Into::into) } } @@ -140,11 +140,11 @@ impl HotKvWrite for Tx { // Specialized put for dual-keyed tables. fn queue_put_dual( &mut self, - key1: &T::K1, - key2: &T::K2, + key1: &T::Key, + key2: &T::Key2, value: &T::Value, ) -> Result<(), Self::Error> { - let k2_size = ::SIZE; + let k2_size = ::SIZE; let mut scratch = [0u8; MAX_KEY_SIZE]; // This will be the total length of key2 + value, reserved in mdbx diff --git a/crates/storage/src/hot/revm.rs b/crates/storage/src/hot/revm.rs index 1f06969..5f43641 100644 --- a/crates/storage/src/hot/revm.rs +++ b/crates/storage/src/hot/revm.rs @@ -63,8 +63,8 @@ impl HotKvRead for RevmRead { fn get_dual( &self, - key1: &T::K1, - key2: &T::K2, + key1: &T::Key, + key2: &T::Key2, ) -> Result, Self::Error> { self.reader.get_dual::(key1, key2) } @@ -132,8 +132,8 @@ impl HotKvRead for RevmWrite { fn get_dual( &self, - key1: &T::K1, - key2: &T::K2, + key1: &T::Key, + key2: &T::Key2, ) -> Result, Self::Error> { self.writer.get_dual::(key1, key2) } @@ -190,8 +190,8 @@ impl HotKvWrite for RevmWrite { fn queue_put_dual( &mut self, - key1: &T::K1, - key2: &T::K2, + key1: &T::Key, + key2: &T::Key2, value: &T::Value, ) -> Result<(), Self::Error> { self.writer.queue_put_dual::(key1, key2, value) diff --git a/crates/storage/src/hot/traits.rs b/crates/storage/src/hot/traits.rs index 7483338..44f65c7 100644 --- a/crates/storage/src/hot/traits.rs +++ b/crates/storage/src/hot/traits.rs @@ -102,8 +102,8 @@ pub trait HotKvRead { /// Get a value from a specific dual-keyed table. fn get_dual( &self, - key1: &T::K1, - key2: &T::K2, + key1: &T::Key, + key2: &T::Key2, ) -> Result, Self::Error> { let mut key1_buf = [0u8; MAX_KEY_SIZE]; let mut key2_buf = [0u8; MAX_KEY_SIZE]; @@ -205,8 +205,8 @@ pub trait HotKvWrite: HotKvRead { /// Queue a put operation for a specific dual-keyed table. fn queue_put_dual( &mut self, - key1: &T::K1, - key2: &T::K2, + key1: &T::Key, + key2: &T::Key2, value: &T::Value, ) -> Result<(), Self::Error> { let mut key1_buf = [0u8; MAX_KEY_SIZE]; @@ -251,7 +251,7 @@ pub trait HotKvWrite: HotKvRead { where T: Table, { - self.queue_raw_create(T::NAME, T::DUAL_KEY, T::DUAL_FIXED_VAL) + self.queue_raw_create(T::NAME, T::DUAL_KEY, T::IS_FIXED_VAL) } /// Queue clearing all entries in a specific table. diff --git a/crates/storage/src/tables/hot.rs b/crates/storage/src/tables/hot.rs index 6c0fabb..af20489 100644 --- a/crates/storage/src/tables/hot.rs +++ b/crates/storage/src/tables/hot.rs @@ -6,37 +6,49 @@ use reth_db_api::{ models::{AccountBeforeTx, ShardedKey, storage_sharded_key::StorageShardedKey}, }; -tables! { +table! { /// Records recent block Headers, by their number. - Headers Header>, + Headers Header> +} +table! { /// Records block numbers by hash. - HeaderNumbers BlockNumber>, + HeaderNumbers BlockNumber> +} - /// Records the canonical chain header hashes, by height. - CanonicalHeaders B256>, +table! { + /// Records the canonical chain header hashes, by height. + CanonicalHeaders B256> +} +table! { /// Records contract Bytecode, by its hash. - Bytecodes Bytecode>, - - /// Records plain account states, keyed by address. - PlainAccountState
Account>, - + Bytecodes Bytecode> +} +table! { + /// Records plain account states, keyed by address. + PlainAccountState
Account> +} +table! { /// Records account state change history, keyed by address. - AccountsHistory => BlockNumberList>, - + AccountsHistory => BlockNumberList> +} +table! { /// Records storage state change history, keyed by address and storage key. - StorageHistory BlockNumberList>, - + StorageHistory BlockNumberList> } -tables! { +table! { /// Records plain storage states, keyed by address and storage key. - PlainStorageState
B256 => U256> size: Some(32 + 32), + PlainStorageState
B256 => U256> is 32 + 32 +} +table! { /// Records account states before transactions, keyed by (address, block number). - StorageChangeSets B256 => StorageEntry> size: Some(32 + 32 + 32), + StorageChangeSets B256 => StorageEntry> is 32 + 32 + 32 +} +table! { /// Records account states before transactions, keyed by (address, block number). - AccountChangeSets Address => AccountBeforeTx> size: None, + AccountChangeSets Address => AccountBeforeTx> } diff --git a/crates/storage/src/tables/macros.rs b/crates/storage/src/tables/macros.rs index c17c13d..85e479f 100644 --- a/crates/storage/src/tables/macros.rs +++ b/crates/storage/src/tables/macros.rs @@ -1,7 +1,8 @@ -macro_rules! tables { +macro_rules! table { ( + @implement #[doc = $doc:expr] - $name:ident<$key:ty => $value:ty> + $name:ident, $key:ty, $value:ty, $dual:expr, $fixed:expr ) => { #[doc = $doc] #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] @@ -9,6 +10,8 @@ macro_rules! tables { impl crate::tables::Table for $name { const NAME: &'static str = stringify!($name); + const DUAL_KEY: bool = $dual; + const FIXED_VAL_SIZE: Option = $fixed; type Key = $key; type Value = $value; @@ -17,33 +20,52 @@ macro_rules! tables { ( #[doc = $doc:expr] - $name:ident<$key:ty => $subkey:ty => $value:ty> size: $fixed:expr + $name:ident<$key:ty => $value:ty> ) => { - #[doc = $doc] - #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] - pub struct $name; - - impl crate::tables::DualKeyed for $name { - const NAME: &'static str = stringify!($name); + table!(@implement + #[doc = $doc] + $name, + $key, + $value, + false, + None + ); + }; - const FIXED_VALUE_SIZE: Option = { $fixed }; - type K1 = $key; - type K2 = $subkey; + ( + #[doc = $doc:expr] + $name:ident<$key:ty => $subkey:ty => $value:ty> + ) => { + table!(@implement + #[doc = $doc] + $name, + $key, + $value, + true, + None + ); - type Value = $value; + impl crate::tables::DualKeyed for $name { + type Key2 = $subkey; } }; - ($(#[doc = $doc:expr] $name:ident<$key:ty => $value:ty>),* $(,)?) => { - $( - tables!(#[doc = $doc] $name<$key => $value>); - )* - }; + ( + #[doc = $doc:expr] + $name:ident<$key:ty => $subkey:ty => $value:ty> is $fixed:expr + ) => { + table!(@implement + #[doc = $doc] + $name, + $key, + $value, + true, + Some($fixed) + ); - ($(#[doc = $doc:expr] $name:ident<$key:ty => $subkey:ty => $value:ty> size: $fixed:expr),* $(,)?) => { - $( - tables!(#[doc = $doc] $name<$key => $subkey => $value> size: $fixed); - )* + impl crate::tables::DualKeyed for $name { + type Key2 = $subkey; + } }; } diff --git a/crates/storage/src/tables/mod.rs b/crates/storage/src/tables/mod.rs index 90187aa..efd8ffa 100644 --- a/crates/storage/src/tables/mod.rs +++ b/crates/storage/src/tables/mod.rs @@ -9,6 +9,9 @@ pub mod hot; use crate::ser::{KeySer, ValSer}; +/// The maximum size of a dual key (in bytes). +pub const MAX_FIXED_VAL_SIZE: usize = 96; + /// Trait for table definitions. pub trait Table { /// A short, human-readable name for the table. @@ -17,8 +20,21 @@ pub trait Table { /// Indicates that this table uses dual keys. const DUAL_KEY: bool = false; + /// True if the table is guaranteed to have fixed-size values, false + /// otherwise. + const FIXED_VAL_SIZE: Option = None; + /// Indicates that this table has fixed-size values. - const DUAL_FIXED_VAL: bool = false; + const IS_FIXED_VAL: bool = Self::FIXED_VAL_SIZE.is_some(); + + /// Compile-time assertions for the table. + #[doc(hidden)] + const ASSERT: () = { + // Ensure that fixed-size values do not exceed the maximum allowed size. + if let Some(size) = Self::FIXED_VAL_SIZE { + assert!(size <= MAX_FIXED_VAL_SIZE, "Fixed value size exceeds maximum allowed size"); + } + }; /// The key type. type Key: KeySer; @@ -31,35 +47,13 @@ pub trait Table { /// This trait aims to capture tables that use a composite key made up of two /// distinct parts. This is useful for representing (e.g.) dupsort or other /// nested map optimizations. -pub trait DualKeyed { - /// A short, human-readable name for the table. - const NAME: &'static str; - - /// If the value size is fixed, `Some(size)`. Otherwise, `None`. - const FIXED_VALUE_SIZE: Option = None; - - /// The first key type. - type K1: KeySer; - +pub trait DualKeyed: Table { /// The second key type. - type K2: KeySer; - - /// The value type. - type Value: ValSer; -} - -impl Table for T -where - T: DualKeyed, -{ - const NAME: &'static str = T::NAME; - - /// Indicates that this table uses dual keys. - const DUAL_KEY: bool = true; - - /// Indicates that this table has fixed-size values. - const DUAL_FIXED_VAL: bool = T::FIXED_VALUE_SIZE.is_some(); + type Key2: KeySer; - type Key = T::K1; - type Value = T::Value; + /// Compile-time assertions for the dual-keyed table. + #[doc(hidden)] + const ASSERT: () = { + assert!(Self::DUAL_KEY, "DualKeyed tables must have DUAL_KEY = true"); + }; } From 8b6ba3e8ebbaa02f7ac8a5ad5c31943fdb435339 Mon Sep 17 00:00:00 2001 From: James Date: Thu, 15 Jan 2026 17:15:54 -0500 Subject: [PATCH 11/16] refactor: get many reordering allowed --- crates/storage/src/hot/mdbx.rs | 4 ++-- crates/storage/src/hot/mem.rs | 6 +++--- crates/storage/src/hot/revm.rs | 10 ++++++++-- crates/storage/src/hot/traits.rs | 21 +++++++++++++++------ 4 files changed, 28 insertions(+), 13 deletions(-) diff --git a/crates/storage/src/hot/mdbx.rs b/crates/storage/src/hot/mdbx.rs index c1f31de..a28f836 100644 --- a/crates/storage/src/hot/mdbx.rs +++ b/crates/storage/src/hot/mdbx.rs @@ -489,11 +489,11 @@ mod tests { { let reader: Tx = db.reader().unwrap(); let addresses: Vec
= accounts.iter().map(|(addr, _)| *addr).collect(); - let read_accounts: Vec> = + let read_accounts: Vec<(_, Option)> = reader.get_many::(addresses.iter()).unwrap(); for (i, (_, expected_account)) in accounts.iter().enumerate() { - assert_eq!(read_accounts[i].as_ref(), Some(expected_account)); + assert_eq!(read_accounts[i].1.as_ref(), Some(expected_account)); } } } diff --git a/crates/storage/src/hot/mem.rs b/crates/storage/src/hot/mem.rs index 8917e78..b7e84d2 100644 --- a/crates/storage/src/hot/mem.rs +++ b/crates/storage/src/hot/mem.rs @@ -581,9 +581,9 @@ mod tests { let values = reader.get_many::(keys).unwrap(); assert_eq!(values.len(), 3); - assert_eq!(values[0], Some(Bytes::from_static(b"first"))); - assert_eq!(values[1], Some(Bytes::from_static(b"second"))); - assert_eq!(values[2], Some(Bytes::from_static(b"third"))); + assert_eq!(values[0], (&1u64, Some(Bytes::from_static(b"first")))); + assert_eq!(values[1], (&2u64, Some(Bytes::from_static(b"second")))); + assert_eq!(values[2], (&3u64, Some(Bytes::from_static(b"third")))); } } diff --git a/crates/storage/src/hot/revm.rs b/crates/storage/src/hot/revm.rs index 5f43641..6eb5228 100644 --- a/crates/storage/src/hot/revm.rs +++ b/crates/storage/src/hot/revm.rs @@ -69,7 +69,10 @@ impl HotKvRead for RevmRead { self.reader.get_dual::(key1, key2) } - fn get_many<'a, T, I>(&self, keys: I) -> Result>, Self::Error> + fn get_many<'a, T, I>( + &self, + keys: I, + ) -> Result)>, Self::Error> where T::Key: 'a, T: Table, @@ -138,7 +141,10 @@ impl HotKvRead for RevmWrite { self.writer.get_dual::(key1, key2) } - fn get_many<'a, T, I>(&self, keys: I) -> Result>, Self::Error> + fn get_many<'a, T, I>( + &self, + keys: I, + ) -> Result)>, Self::Error> where T::Key: 'a, T: Table, diff --git a/crates/storage/src/hot/traits.rs b/crates/storage/src/hot/traits.rs index 44f65c7..d20616a 100644 --- a/crates/storage/src/hot/traits.rs +++ b/crates/storage/src/hot/traits.rs @@ -125,13 +125,21 @@ pub trait HotKvRead { /// /// # Returns /// - /// A vector of `Option`, where each element corresponds to the - /// value for the respective key in the input iterator. If a key does not - /// exist in the table, the corresponding element will be `None`. + /// A vector of `(&'a T::Key, Option)`, where each element + /// corresponds to the value for the respective key in the input iterator. + /// If a key does not exist in the table, the corresponding element will be + /// `None`. + /// + /// Implementations ARE NOT required to preserve the order of the input + /// keys in the output vector. Users should not rely on any specific + /// ordering. /// /// If any error occurs during retrieval or deserialization, the entire /// operation will return an error. - fn get_many<'a, T, I>(&self, keys: I) -> Result>, Self::Error> + fn get_many<'a, T, I>( + &self, + keys: I, + ) -> Result)>, Self::Error> where T::Key: 'a, T: Table, @@ -140,10 +148,11 @@ pub trait HotKvRead { let mut key_buf = [0u8; MAX_KEY_SIZE]; keys.into_iter() - .map(|key| self.raw_get(T::NAME, key.encode_key(&mut key_buf))) - .map(|maybe_val| { + .map(|key| (key, self.raw_get(T::NAME, key.encode_key(&mut key_buf)))) + .map(|(key, maybe_val)| { maybe_val .and_then(|val| ValSer::maybe_decode_value(val.as_deref()).map_err(Into::into)) + .map(|res| (key, res)) }) .collect() } From bb575a5e1f22d466623178f33c5f019915fe5f6e Mon Sep 17 00:00:00 2001 From: James Date: Fri, 16 Jan 2026 08:57:35 -0500 Subject: [PATCH 12/16] fix: get_dual no alloc --- crates/storage/src/hot/mdbx.rs | 10 ++++++---- crates/storage/src/hot/mod.rs | 2 +- crates/storage/src/hot/revm.rs | 12 +++--------- crates/storage/src/hot/traits.rs | 10 +++++----- crates/storage/src/tables/hot.rs | 4 ++-- crates/storage/src/tables/mod.rs | 2 +- 6 files changed, 18 insertions(+), 22 deletions(-) diff --git a/crates/storage/src/hot/mdbx.rs b/crates/storage/src/hot/mdbx.rs index a28f836..67f8610 100644 --- a/crates/storage/src/hot/mdbx.rs +++ b/crates/storage/src/hot/mdbx.rs @@ -1,7 +1,7 @@ use crate::{ hot::{HotKv, HotKvError, HotKvRead, HotKvReadError, HotKvWrite}, ser::{DeserError, KeySer, MAX_KEY_SIZE, ValSer}, - tables::{DualKeyed, Table}, + tables::{DualKeyed, MAX_FIXED_VAL_SIZE}, }; use bytes::{BufMut, BytesMut}; use reth_db::{ @@ -94,13 +94,15 @@ where // table has one. This is a bit ugly, and results in an extra // allocation for fixed-size values. This could be avoided using // max value size. - let value_bytes = if let Some(size) = ::FIXED_VAL_SIZE { - let buf = vec![0u8; size]; + let value_bytes = if T::IS_FIXED_VAL { + let buf = [0u8; MAX_KEY_SIZE + MAX_FIXED_VAL_SIZE]; let _ = key2.encode_key(&mut buf[..MAX_KEY_SIZE].try_into().unwrap()); + let kv_size = ::SIZE + T::FIXED_VAL_SIZE.unwrap(); + let db = self.inner.open_db(Some(T::NAME))?; let mut cursor = self.inner.cursor(&db).map_err(MdbxError::Mdbx)?; - cursor.get_both_range(key1_bytes, &buf).map_err(MdbxError::Mdbx) + cursor.get_both_range(key1_bytes, &buf[..kv_size]).map_err(MdbxError::Mdbx) } else { let mut buf = [0u8; MAX_KEY_SIZE]; let encoded = key2.encode_key(&mut buf); diff --git a/crates/storage/src/hot/mod.rs b/crates/storage/src/hot/mod.rs index 2720015..bbf1836 100644 --- a/crates/storage/src/hot/mod.rs +++ b/crates/storage/src/hot/mod.rs @@ -13,4 +13,4 @@ mod revm; pub use revm::{RevmRead, RevmWrite}; mod traits; -pub use traits::{HotKv, HotKvRead, HotKvWrite}; +pub use traits::{HotKv, HotKvRead, HotKvWrite, KeyValue}; diff --git a/crates/storage/src/hot/revm.rs b/crates/storage/src/hot/revm.rs index 6eb5228..7b7bcbb 100644 --- a/crates/storage/src/hot/revm.rs +++ b/crates/storage/src/hot/revm.rs @@ -1,5 +1,5 @@ use crate::{ - hot::{HotKvError, HotKvRead, HotKvWrite}, + hot::{HotKvError, HotKvRead, HotKvWrite, KeyValue}, tables::{ DualKeyed, Table, hot::{self, Bytecodes, PlainAccountState}, @@ -69,10 +69,7 @@ impl HotKvRead for RevmRead { self.reader.get_dual::(key1, key2) } - fn get_many<'a, T, I>( - &self, - keys: I, - ) -> Result)>, Self::Error> + fn get_many<'a, T, I>(&self, keys: I) -> Result>, Self::Error> where T::Key: 'a, T: Table, @@ -141,10 +138,7 @@ impl HotKvRead for RevmWrite { self.writer.get_dual::(key1, key2) } - fn get_many<'a, T, I>( - &self, - keys: I, - ) -> Result)>, Self::Error> + fn get_many<'a, T, I>(&self, keys: I) -> Result>, Self::Error> where T::Key: 'a, T: Table, diff --git a/crates/storage/src/hot/traits.rs b/crates/storage/src/hot/traits.rs index d20616a..23eb2ba 100644 --- a/crates/storage/src/hot/traits.rs +++ b/crates/storage/src/hot/traits.rs @@ -9,6 +9,9 @@ use crate::{ tables::{DualKeyed, Table}, }; +/// A key-value pair from a table. +pub type KeyValue<'a, T> = (&'a ::Key, Option<::Value>); + /// Trait for hot storage. This is a KV store with read/write transactions. #[auto_impl::auto_impl(&, Arc, Box)] pub trait HotKv { @@ -125,7 +128,7 @@ pub trait HotKvRead { /// /// # Returns /// - /// A vector of `(&'a T::Key, Option)`, where each element + /// A vector of [`KeyValue`] where each element /// corresponds to the value for the respective key in the input iterator. /// If a key does not exist in the table, the corresponding element will be /// `None`. @@ -136,10 +139,7 @@ pub trait HotKvRead { /// /// If any error occurs during retrieval or deserialization, the entire /// operation will return an error. - fn get_many<'a, T, I>( - &self, - keys: I, - ) -> Result)>, Self::Error> + fn get_many<'a, T, I>(&self, keys: I) -> Result>, Self::Error> where T::Key: 'a, T: Table, diff --git a/crates/storage/src/tables/hot.rs b/crates/storage/src/tables/hot.rs index af20489..94736a7 100644 --- a/crates/storage/src/tables/hot.rs +++ b/crates/storage/src/tables/hot.rs @@ -40,12 +40,12 @@ table! { table! { /// Records plain storage states, keyed by address and storage key. - PlainStorageState
B256 => U256> is 32 + 32 + PlainStorageState
B256 => U256> is 32 } table! { /// Records account states before transactions, keyed by (address, block number). - StorageChangeSets B256 => StorageEntry> is 32 + 32 + 32 + StorageChangeSets B256 => StorageEntry> is 32 + 32 } table! { diff --git a/crates/storage/src/tables/mod.rs b/crates/storage/src/tables/mod.rs index efd8ffa..572daac 100644 --- a/crates/storage/src/tables/mod.rs +++ b/crates/storage/src/tables/mod.rs @@ -10,7 +10,7 @@ pub mod hot; use crate::ser::{KeySer, ValSer}; /// The maximum size of a dual key (in bytes). -pub const MAX_FIXED_VAL_SIZE: usize = 96; +pub const MAX_FIXED_VAL_SIZE: usize = 64; /// Trait for table definitions. pub trait Table { From 93c0afef33a24584a2a8b9857a7aa69bbfa31d3c Mon Sep 17 00:00:00 2001 From: James Date: Fri, 16 Jan 2026 09:37:26 -0500 Subject: [PATCH 13/16] feat: put_header --- crates/storage/src/hot/db_traits.rs | 59 ++++++++++++++++------------- crates/storage/src/hot/mdbx.rs | 19 +++++----- crates/storage/src/tables/hot.rs | 8 ++-- 3 files changed, 44 insertions(+), 42 deletions(-) diff --git a/crates/storage/src/hot/db_traits.rs b/crates/storage/src/hot/db_traits.rs index 31b9cdf..0786578 100644 --- a/crates/storage/src/hot/db_traits.rs +++ b/crates/storage/src/hot/db_traits.rs @@ -16,9 +16,6 @@ pub trait HotDbReader: sealed::Sealed { /// Read a block number by its hash. fn get_header_number(&self, hash: &B256) -> Result, Self::Error>; - /// Read the canonical hash by block number. - fn get_canonical_hash(&self, number: u64) -> Result, Self::Error>; - /// Read contract Bytecode by its hash. fn get_bytecode(&self, code_hash: &B256) -> Result, Self::Error>; @@ -37,6 +34,14 @@ pub trait HotDbReader: sealed::Sealed { let opt = self.get_storage(address, key)?; Ok(opt.map(|value| StorageEntry { key: *key, value })) } + + /// Read a block header by its hash. + fn header_by_hash(&self, hash: &B256) -> Result, Self::Error> { + let Some(number) = self.get_header_number(hash)? else { + return Ok(None); + }; + self.get_header(number) + } } impl HotDbReader for T @@ -53,10 +58,6 @@ where self.get::(hash) } - fn get_canonical_hash(&self, number: u64) -> Result, Self::Error> { - self.get::(&number) - } - fn get_bytecode(&self, code_hash: &B256) -> Result, Self::Error> { self.get::(code_hash) } @@ -75,14 +76,19 @@ pub trait HotDbWriter: sealed::Sealed { /// The error type for write operations type Error: std::error::Error + Send + Sync + 'static + From; - /// Read the latest block header. - fn put_header(&mut self, header: &Header) -> Result<(), Self::Error>; - - /// Write a block number by its hash. - fn put_header_number(&mut self, hash: &B256, number: u64) -> Result<(), Self::Error>; + /// Write a block header. This will leave the DB in an inconsistent state + /// until the corresponding header number is also written. Users should + /// prefer [`Self::put_header`] instead. + fn put_header_inconsistent(&mut self, header: &Header) -> Result<(), Self::Error>; - /// Write the canonical hash by block number. - fn put_canonical_hash(&mut self, number: u64, hash: &B256) -> Result<(), Self::Error>; + /// Write a block number by its hash. This will leave the DB in an + /// inconsistent state until the corresponding header is also written. + /// Users should prefer [`Self::put_header`] instead. + fn put_header_number_inconsistent( + &mut self, + hash: &B256, + number: u64, + ) -> Result<(), Self::Error>; /// Write contract Bytecode by its hash. fn put_bytecode(&mut self, code_hash: &B256, bytecode: &Bytecode) -> Result<(), Self::Error>; @@ -98,15 +104,14 @@ pub trait HotDbWriter: sealed::Sealed { entry: &U256, ) -> Result<(), Self::Error>; + /// Write a sealed block header (header + number). + fn put_header(&mut self, header: &SealedHeader) -> Result<(), Self::Error> { + self.put_header_inconsistent(header.header()) + .and_then(|_| self.put_header_number_inconsistent(&header.hash(), header.number)) + } + /// Commit the write transaction. fn commit(self) -> Result<(), Self::Error>; - - /// Write a canonical header (header, number mapping, and canonical hash). - fn put_canonical(&mut self, header: &SealedHeader) -> Result<(), Self::Error> { - self.put_header(header)?; - self.put_header_number(&header.hash(), header.number)?; - self.put_canonical_hash(header.number, &header.hash()) - } } impl HotDbWriter for T @@ -115,18 +120,18 @@ where { type Error = ::Error; - fn put_header(&mut self, header: &Header) -> Result<(), Self::Error> { + fn put_header_inconsistent(&mut self, header: &Header) -> Result<(), Self::Error> { self.queue_put::(&header.number, header) } - fn put_header_number(&mut self, hash: &B256, number: u64) -> Result<(), Self::Error> { + fn put_header_number_inconsistent( + &mut self, + hash: &B256, + number: u64, + ) -> Result<(), Self::Error> { self.queue_put::(hash, &number) } - fn put_canonical_hash(&mut self, number: u64, hash: &B256) -> Result<(), Self::Error> { - self.queue_put::(&number, hash) - } - fn put_bytecode(&mut self, code_hash: &B256, bytecode: &Bytecode) -> Result<(), Self::Error> { self.queue_put::(code_hash, bytecode) } diff --git a/crates/storage/src/hot/mdbx.rs b/crates/storage/src/hot/mdbx.rs index 67f8610..0f3a2bf 100644 --- a/crates/storage/src/hot/mdbx.rs +++ b/crates/storage/src/hot/mdbx.rs @@ -212,7 +212,10 @@ mod tests { hot::{HotDbWriter, HotKv, HotKvRead, HotKvWrite}, tables::hot, }; - use alloy::primitives::{Address, B256, BlockNumber, U256}; + use alloy::{ + consensus::Sealed, + primitives::{Address, B256, BlockNumber, U256}, + }; use reth::primitives::{Account, Bytecode, Header}; use reth_db::DatabaseEnv; @@ -243,7 +246,6 @@ mod tests { writer.queue_create::().unwrap(); writer.queue_create::().unwrap(); - writer.queue_create::().unwrap(); writer.queue_create::().unwrap(); writer.queue_create::().unwrap(); writer.queue_create::().unwrap(); @@ -609,16 +611,14 @@ mod tests { // Test various data types let (block_number, header) = create_test_header(); - let canonical_hash = B256::from_slice(&[0x7; 32]); + let header = Sealed::new(header); { let mut writer: Tx = db.writer().unwrap(); - // Create tables - // Write different types writer.queue_put::(&block_number, &header).unwrap(); - writer.queue_put::(&block_number, &canonical_hash).unwrap(); + writer.queue_put::(&header.hash(), &block_number).unwrap(); writer.raw_commit().unwrap(); } @@ -628,11 +628,10 @@ mod tests { // Read and verify let read_header: Option
= reader.get::(&block_number).unwrap(); - assert_eq!(read_header, Some(header)); + assert_eq!(read_header.as_ref(), Some(header.inner())); - let read_hash: Option = - reader.get::(&block_number).unwrap(); - assert_eq!(read_hash, Some(canonical_hash)); + let read_hash: Option = reader.get::(&header.hash()).unwrap(); + assert_eq!(read_hash, Some(header.number)); } } diff --git a/crates/storage/src/tables/hot.rs b/crates/storage/src/tables/hot.rs index 94736a7..ee801e9 100644 --- a/crates/storage/src/tables/hot.rs +++ b/crates/storage/src/tables/hot.rs @@ -16,23 +16,21 @@ table! { HeaderNumbers BlockNumber> } -table! { - /// Records the canonical chain header hashes, by height. - CanonicalHeaders B256> -} - table! { /// Records contract Bytecode, by its hash. Bytecodes Bytecode> } + table! { /// Records plain account states, keyed by address. PlainAccountState
Account> } + table! { /// Records account state change history, keyed by address. AccountsHistory => BlockNumberList> } + table! { /// Records storage state change history, keyed by address and storage key. StorageHistory BlockNumberList> From 41c662feb5b356e732e8ca63d9db705533033325 Mon Sep 17 00:00:00 2001 From: James Date: Fri, 16 Jan 2026 11:36:17 -0500 Subject: [PATCH 14/16] feat: history read/write traits --- crates/storage/src/hot/db_traits.rs | 133 ++++++++++++++++++++++++--- crates/storage/src/hot/mdbx.rs | 2 +- crates/storage/src/hot/mod.rs | 2 +- crates/storage/src/hot/traits.rs | 17 +++- crates/storage/src/ser/reth_impls.rs | 42 +++++---- crates/storage/src/tables/hot.rs | 25 +++-- 6 files changed, 169 insertions(+), 52 deletions(-) diff --git a/crates/storage/src/hot/db_traits.rs b/crates/storage/src/hot/db_traits.rs index 0786578..b0006d4 100644 --- a/crates/storage/src/hot/db_traits.rs +++ b/crates/storage/src/hot/db_traits.rs @@ -4,12 +4,11 @@ use crate::{ }; use alloy::primitives::{Address, B256, U256}; use reth::primitives::{Account, Bytecode, Header, SealedHeader, StorageEntry}; +use reth_db::{BlockNumberList, models::BlockNumberAddress}; +use reth_db_api::models::ShardedKey; /// Trait for database read operations. -pub trait HotDbReader: sealed::Sealed { - /// The error type for read operations - type Error: std::error::Error + Send + Sync + 'static + From; - +pub trait HotDbRead: HotKvRead + sealed::Sealed { /// Read a block header by its number. fn get_header(&self, number: u64) -> Result, Self::Error>; @@ -44,12 +43,10 @@ pub trait HotDbReader: sealed::Sealed { } } -impl HotDbReader for T +impl HotDbRead for T where T: HotKvRead, { - type Error = ::Error; - fn get_header(&self, number: u64) -> Result, Self::Error> { self.get::(&number) } @@ -72,10 +69,7 @@ where } /// Trait for database write operations. -pub trait HotDbWriter: sealed::Sealed { - /// The error type for write operations - type Error: std::error::Error + Send + Sync + 'static + From; - +pub trait HotDbWrite: HotKvWrite + sealed::Sealed { /// Write a block header. This will leave the DB in an inconsistent state /// until the corresponding header number is also written. Users should /// prefer [`Self::put_header`] instead. @@ -114,12 +108,10 @@ pub trait HotDbWriter: sealed::Sealed { fn commit(self) -> Result<(), Self::Error>; } -impl HotDbWriter for T +impl HotDbWrite for T where T: HotKvWrite, { - type Error = ::Error; - fn put_header_inconsistent(&mut self, header: &Header) -> Result<(), Self::Error> { self.queue_put::(&header.number, header) } @@ -154,6 +146,119 @@ where } } +/// Trait for history read operations. +pub trait HotHistoryRead: HotDbRead { + /// Get the list of block numbers where an account was touched. + /// Get the list of block numbers where an account was touched. + fn get_account_history( + &self, + address: &Address, + latest_height: u64, + ) -> Result, Self::Error> { + self.get_dual::(address, &latest_height) + } + /// Get the account change (pre-state) for an account at a specific block. + /// + /// If the return value is `None`, the account was not changed in that + /// block. + fn get_account_change( + &self, + block_number: u64, + address: &Address, + ) -> Result, Self::Error> { + self.get_dual::(&block_number, address) + } + + /// Get the storage history for an account and storage slot. The returned + /// list will contain block numbers where the storage slot was changed. + fn get_storage_history( + &self, + address: &Address, + slot: B256, + highest_block_number: u64, + ) -> Result, Self::Error> { + let sharded_key = ShardedKey::new(slot, highest_block_number); + self.get_dual::(address, &sharded_key) + } + + /// Get the storage change (before state) for a specific storage slot at a + /// specific block. + /// + /// If the return value is `None`, the storage slot was not changed in that + /// block. If the return value is `Some(value)`, the value is the pre-state + /// of the storage slot before the change in that block. If the value is + /// `U256::ZERO`, that indicates that the storage slot was not set before + /// the change. + fn get_storage_change( + &self, + block_number: u64, + address: &Address, + slot: &B256, + ) -> Result, Self::Error> { + let block_number_address = BlockNumberAddress((block_number, *address)); + self.get_dual::(&block_number_address, slot) + } +} + +impl HotHistoryRead for T where T: HotDbRead {} + +/// Trait for history write operations. +pub trait HotHistoryWrite: HotDbWrite { + /// Maintain a list of block numbers where an account was touched. + /// + /// Accounts are keyed + fn write_account_history( + &mut self, + address: &Address, + latest_height: u64, + touched: &BlockNumberList, + ) -> Result<(), Self::Error> { + self.queue_put_dual::(address, &latest_height, touched) + } + + /// Write a storage change (before state) for an account at a specific + /// block. + fn write_storage_change( + &mut self, + block_number: u64, + address: Address, + slot: &B256, + value: &U256, + ) -> Result<(), Self::Error> { + let block_number_address = BlockNumberAddress((block_number, address)); + self.queue_put_dual::(&block_number_address, slot, value) + } + + /// Write an account change (pre-state) for an account at a specific + /// block. + + /// Write storage history, by highest block number and touched block + /// numbers. + fn write_storage_history( + &mut self, + address: &Address, + slot: B256, + highest_block_number: u64, + touched: &BlockNumberList, + ) -> Result<(), Self::Error> { + let sharded_key = ShardedKey::new(slot, highest_block_number); + self.queue_put_dual::(address, &sharded_key, touched) + } + + /// Write a storage change (before state) for an account at a specific + /// block. + fn write_account_change( + &mut self, + block_number: u64, + address: Address, + pre_state: &Account, + ) -> Result<(), Self::Error> { + self.queue_put_dual::(&block_number, &address, pre_state) + } +} + +impl HotHistoryWrite for T where T: HotDbWrite + HotKvWrite {} + mod sealed { use crate::hot::HotKvRead; diff --git a/crates/storage/src/hot/mdbx.rs b/crates/storage/src/hot/mdbx.rs index 0f3a2bf..3b23e54 100644 --- a/crates/storage/src/hot/mdbx.rs +++ b/crates/storage/src/hot/mdbx.rs @@ -209,7 +209,7 @@ impl HotKvWrite for Tx { mod tests { use super::*; use crate::{ - hot::{HotDbWriter, HotKv, HotKvRead, HotKvWrite}, + hot::{HotDbWrite, HotKv, HotKvRead, HotKvWrite}, tables::hot, }; use alloy::{ diff --git a/crates/storage/src/hot/mod.rs b/crates/storage/src/hot/mod.rs index bbf1836..6f4786c 100644 --- a/crates/storage/src/hot/mod.rs +++ b/crates/storage/src/hot/mod.rs @@ -1,5 +1,5 @@ mod db_traits; -pub use db_traits::{HotDbReader, HotDbWriter}; +pub use db_traits::{HotDbRead, HotDbWrite, HotHistoryRead, HotHistoryWrite}; mod error; pub use error::{HotKvError, HotKvReadError, HotKvResult}; diff --git a/crates/storage/src/hot/traits.rs b/crates/storage/src/hot/traits.rs index 23eb2ba..567ea5d 100644 --- a/crates/storage/src/hot/traits.rs +++ b/crates/storage/src/hot/traits.rs @@ -72,13 +72,18 @@ pub trait HotKvRead { /// The `key` buf must be <= [`MAX_KEY_SIZE`] bytes. Implementations are /// allowed to panic if this is not the case. /// - /// If the table is dual-keyed, the output may be implementation-defined. + /// If the table is dual-keyed, the output MAY be implementation-defined. fn raw_get<'a>(&'a self, table: &str, key: &[u8]) -> Result>, Self::Error>; /// Get a raw value from a specific table with dual keys. /// - /// If the table is not dual-keyed, the output may be + /// If `key1` is present, but `key2` is not in the table, the output is + /// implementation-defined. For sorted databases, it SHOULD return the value + /// of the NEXT populated key. It MAY also return `None`, even if other + /// subkeys are populated. + /// + /// If the table is not dual-keyed, the output MAY be /// implementation-defined. fn raw_get_dual<'a>( &'a self, @@ -103,6 +108,14 @@ pub trait HotKvRead { } /// Get a value from a specific dual-keyed table. + /// + /// If `key1` is present, but `key2` is not in the table, the output is + /// implementation-defined. For sorted databases, it SHOULD return the value + /// of the NEXT populated key. It MAY also return `None`, even if other + /// subkeys are populated. + /// + /// If the table is not dual-keyed, the output MAY be + /// implementation-defined. fn get_dual( &self, key1: &T::Key, diff --git a/crates/storage/src/ser/reth_impls.rs b/crates/storage/src/ser/reth_impls.rs index ef19bfc..66ce986 100644 --- a/crates/storage/src/ser/reth_impls.rs +++ b/crates/storage/src/ser/reth_impls.rs @@ -5,7 +5,7 @@ use alloy::{ eip2930::{AccessList, AccessListItem}, eip7702::{Authorization, SignedAuthorization}, }, - primitives::{Address, B256, FixedBytes, Signature, TxKind, U256}, + primitives::{Address, B256, FixedBytes, KECCAK256_EMPTY, Signature, TxKind, U256}, }; use reth::{ primitives::{Account, Bytecode, Header, Log, LogData, TransactionSigned, TxType}, @@ -320,13 +320,9 @@ impl ValSer for Account { fn encoded_size(&self) -> usize { // NB: Destructure to ensure changes are compile errors and mistakes // are unused var warnings. - let Account { nonce, balance, bytecode_hash } = self; - by_props!( - @size - nonce, - balance, - bytecode_hash, - ) + let Account { nonce, balance, bytecode_hash: _ } = self; + + nonce.encoded_size() + balance.encoded_size() + 32 } fn encode_value_to(&self, buf: &mut B) @@ -336,12 +332,11 @@ impl ValSer for Account { // NB: Destructure to ensure changes are compile errors and mistakes // are unused var warnings. let Account { nonce, balance, bytecode_hash } = self; - by_props!( - @encode buf; - nonce, - balance, - bytecode_hash, - ) + { + nonce.encode_value_to(buf); + balance.encode_value_to(buf); + bytecode_hash.unwrap_or(KECCAK256_EMPTY).encode_value_to(buf); + } } fn decode_value(data: &[u8]) -> Result @@ -354,12 +349,19 @@ impl ValSer for Account { let Account { nonce, balance, bytecode_hash } = &mut account; let mut data = data; - by_props!( - @decode data; - nonce, - balance, - bytecode_hash, - ); + { + *nonce = ValSer::decode_value(data)?; + data = &data[nonce.encoded_size()..]; + *balance = ValSer::decode_value(data)?; + data = &data[balance.encoded_size()..]; + + let bch: B256 = ValSer::decode_value(data)?; + if bch == KECCAK256_EMPTY { + *bytecode_hash = None; + } else { + *bytecode_hash = Some(bch); + } + }; Ok(account) } } diff --git a/crates/storage/src/tables/hot.rs b/crates/storage/src/tables/hot.rs index ee801e9..7e306ea 100644 --- a/crates/storage/src/tables/hot.rs +++ b/crates/storage/src/tables/hot.rs @@ -1,10 +1,7 @@ use alloy::primitives::{Address, B256, BlockNumber, U256}; -use reth::primitives::{Account, Bytecode, Header, StorageEntry}; +use reth::primitives::{Account, Bytecode, Header}; use reth_db::models::BlockNumberAddress; -use reth_db_api::{ - BlockNumberList, - models::{AccountBeforeTx, ShardedKey, storage_sharded_key::StorageShardedKey}, -}; +use reth_db_api::{BlockNumberList, models::ShardedKey}; table! { /// Records recent block Headers, by their number. @@ -27,26 +24,26 @@ table! { } table! { - /// Records account state change history, keyed by address. - AccountsHistory => BlockNumberList> + /// Records plain storage states, keyed by address and storage key. + PlainStorageState
B256 => U256> is 32 } table! { - /// Records storage state change history, keyed by address and storage key. - StorageHistory BlockNumberList> + /// Records account state change history, keyed by address. + AccountsHistory
u64 => BlockNumberList> } table! { - /// Records plain storage states, keyed by address and storage key. - PlainStorageState
B256 => U256> is 32 + /// Records account states before transactions, keyed by (block_number, address). + AccountChangeSets Address => Account> is 96 } table! { - /// Records account states before transactions, keyed by (address, block number). - StorageChangeSets B256 => StorageEntry> is 32 + 32 + /// Records storage state change history, keyed by address and storage key. + StorageHistory
ShardedKey => BlockNumberList> } table! { /// Records account states before transactions, keyed by (address, block number). - AccountChangeSets Address => AccountBeforeTx> + StorageChangeSets B256 => U256> is 32 } From 0ff2bde24a3c86637b54300d99bb8d3c8363ffca Mon Sep 17 00:00:00 2001 From: James Date: Fri, 16 Jan 2026 11:41:30 -0500 Subject: [PATCH 15/16] refactor: roll impls up --- crates/storage/src/hot/db_traits.rs | 104 ++++++++++------------------ 1 file changed, 35 insertions(+), 69 deletions(-) diff --git a/crates/storage/src/hot/db_traits.rs b/crates/storage/src/hot/db_traits.rs index b0006d4..25ef573 100644 --- a/crates/storage/src/hot/db_traits.rs +++ b/crates/storage/src/hot/db_traits.rs @@ -10,19 +10,29 @@ use reth_db_api::models::ShardedKey; /// Trait for database read operations. pub trait HotDbRead: HotKvRead + sealed::Sealed { /// Read a block header by its number. - fn get_header(&self, number: u64) -> Result, Self::Error>; + fn get_header(&self, number: u64) -> Result, Self::Error> { + self.get::(&number) + } /// Read a block number by its hash. - fn get_header_number(&self, hash: &B256) -> Result, Self::Error>; + fn get_header_number(&self, hash: &B256) -> Result, Self::Error> { + self.get::(hash) + } /// Read contract Bytecode by its hash. - fn get_bytecode(&self, code_hash: &B256) -> Result, Self::Error>; + fn get_bytecode(&self, code_hash: &B256) -> Result, Self::Error> { + self.get::(code_hash) + } /// Read an account by its address. - fn get_account(&self, address: &Address) -> Result, Self::Error>; + fn get_account(&self, address: &Address) -> Result, Self::Error> { + self.get::(address) + } /// Read a storage slot by its address and key. - fn get_storage(&self, address: &Address, key: &B256) -> Result, Self::Error>; + fn get_storage(&self, address: &Address, key: &B256) -> Result, Self::Error> { + self.get_dual::(address, key) + } /// Read a [`StorageEntry`] by its address and key. fn get_storage_entry( @@ -43,79 +53,20 @@ pub trait HotDbRead: HotKvRead + sealed::Sealed { } } -impl HotDbRead for T -where - T: HotKvRead, -{ - fn get_header(&self, number: u64) -> Result, Self::Error> { - self.get::(&number) - } - - fn get_header_number(&self, hash: &B256) -> Result, Self::Error> { - self.get::(hash) - } - - fn get_bytecode(&self, code_hash: &B256) -> Result, Self::Error> { - self.get::(code_hash) - } - - fn get_account(&self, address: &Address) -> Result, Self::Error> { - self.get::(address) - } - - fn get_storage(&self, address: &Address, key: &B256) -> Result, Self::Error> { - self.get_dual::(address, key) - } -} +impl HotDbRead for T where T: HotKvRead {} /// Trait for database write operations. pub trait HotDbWrite: HotKvWrite + sealed::Sealed { /// Write a block header. This will leave the DB in an inconsistent state /// until the corresponding header number is also written. Users should /// prefer [`Self::put_header`] instead. - fn put_header_inconsistent(&mut self, header: &Header) -> Result<(), Self::Error>; - - /// Write a block number by its hash. This will leave the DB in an - /// inconsistent state until the corresponding header is also written. - /// Users should prefer [`Self::put_header`] instead. - fn put_header_number_inconsistent( - &mut self, - hash: &B256, - number: u64, - ) -> Result<(), Self::Error>; - - /// Write contract Bytecode by its hash. - fn put_bytecode(&mut self, code_hash: &B256, bytecode: &Bytecode) -> Result<(), Self::Error>; - - /// Write an account by its address. - fn put_account(&mut self, address: &Address, account: &Account) -> Result<(), Self::Error>; - - /// Write a storage entry by its address and key. - fn put_storage( - &mut self, - address: &Address, - key: &B256, - entry: &U256, - ) -> Result<(), Self::Error>; - - /// Write a sealed block header (header + number). - fn put_header(&mut self, header: &SealedHeader) -> Result<(), Self::Error> { - self.put_header_inconsistent(header.header()) - .and_then(|_| self.put_header_number_inconsistent(&header.hash(), header.number)) - } - - /// Commit the write transaction. - fn commit(self) -> Result<(), Self::Error>; -} - -impl HotDbWrite for T -where - T: HotKvWrite, -{ fn put_header_inconsistent(&mut self, header: &Header) -> Result<(), Self::Error> { self.queue_put::(&header.number, header) } + /// Write a block number by its hash. This will leave the DB in an + /// inconsistent state until the corresponding header is also written. + /// Users should prefer [`Self::put_header`] instead. fn put_header_number_inconsistent( &mut self, hash: &B256, @@ -124,14 +75,17 @@ where self.queue_put::(hash, &number) } + /// Write contract Bytecode by its hash. fn put_bytecode(&mut self, code_hash: &B256, bytecode: &Bytecode) -> Result<(), Self::Error> { self.queue_put::(code_hash, bytecode) } + /// Write an account by its address. fn put_account(&mut self, address: &Address, account: &Account) -> Result<(), Self::Error> { self.queue_put::(address, account) } + /// Write a storage entry by its address and key. fn put_storage( &mut self, address: &Address, @@ -141,11 +95,23 @@ where self.queue_put_dual::(address, key, entry) } - fn commit(self) -> Result<(), Self::Error> { + /// Write a sealed block header (header + number). + fn put_header(&mut self, header: &SealedHeader) -> Result<(), Self::Error> { + self.put_header_inconsistent(header.header()) + .and_then(|_| self.put_header_number_inconsistent(&header.hash(), header.number)) + } + + /// Commit the write transaction. + fn commit(self) -> Result<(), Self::Error> + where + Self: Sized, + { HotKvWrite::raw_commit(self) } } +impl HotDbWrite for T where T: HotKvWrite {} + /// Trait for history read operations. pub trait HotHistoryRead: HotDbRead { /// Get the list of block numbers where an account was touched. From d57ef087d70ba881e7734719c02d6d05128047ef Mon Sep 17 00:00:00 2001 From: James Date: Fri, 16 Jan 2026 23:07:49 -0500 Subject: [PATCH 16/16] feat: cursors --- crates/storage/Cargo.toml | 3 +- crates/storage/src/hot/impls/mdbx.rs | 1356 ++++++++++++ crates/storage/src/hot/impls/mem.rs | 1869 +++++++++++++++++ crates/storage/src/hot/impls/mod.rs | 309 +++ crates/storage/src/hot/mdbx.rs | 660 ------ crates/storage/src/hot/mem.rs | 828 -------- crates/storage/src/hot/mod.rs | 19 +- .../storage/src/hot/{ => model}/db_traits.rs | 30 +- crates/storage/src/hot/{ => model}/error.rs | 0 crates/storage/src/hot/model/mod.rs | 38 + crates/storage/src/hot/{ => model}/revm.rs | 61 +- crates/storage/src/hot/{ => model}/traits.rs | 84 +- crates/storage/src/hot/model/traverse.rs | 399 ++++ crates/storage/src/ser/impls.rs | 2 +- crates/storage/src/ser/reth_impls.rs | 138 +- crates/storage/src/ser/traits.rs | 2 +- crates/storage/src/tables/macros.rs | 3 + crates/storage/src/tables/mod.rs | 76 +- 18 files changed, 4192 insertions(+), 1685 deletions(-) create mode 100644 crates/storage/src/hot/impls/mdbx.rs create mode 100644 crates/storage/src/hot/impls/mem.rs create mode 100644 crates/storage/src/hot/impls/mod.rs delete mode 100644 crates/storage/src/hot/mdbx.rs delete mode 100644 crates/storage/src/hot/mem.rs rename crates/storage/src/hot/{ => model}/db_traits.rs (96%) rename crates/storage/src/hot/{ => model}/error.rs (100%) create mode 100644 crates/storage/src/hot/model/mod.rs rename crates/storage/src/hot/{ => model}/revm.rs (93%) rename crates/storage/src/hot/{ => model}/traits.rs (78%) create mode 100644 crates/storage/src/hot/model/traverse.rs diff --git a/crates/storage/Cargo.toml b/crates/storage/Cargo.toml index a1cc1b0..8b25ecd 100644 --- a/crates/storage/Cargo.toml +++ b/crates/storage/Cargo.toml @@ -13,11 +13,12 @@ alloy.workspace = true auto_impl = "1.3.0" bytes = "1.11.0" reth.workspace = true -reth-db.workspace = true +reth-db = { workspace = true, features = ["test-utils"] } reth-db-api.workspace = true reth-libmdbx.workspace = true thiserror.workspace = true trevm.workspace = true [dev-dependencies] +serial_test = "3.3.1" tempfile.workspace = true diff --git a/crates/storage/src/hot/impls/mdbx.rs b/crates/storage/src/hot/impls/mdbx.rs new file mode 100644 index 0000000..ecdfbf1 --- /dev/null +++ b/crates/storage/src/hot/impls/mdbx.rs @@ -0,0 +1,1356 @@ +use crate::{ + hot::model::{ + DualKeyValue, DualKeyedTraverse, DualTableTraverse, HotKv, HotKvError, HotKvRead, + HotKvReadError, HotKvWrite, KvTraverse, KvTraverseMut, RawDualKeyValue, RawKeyValue, + RawValue, + }, + ser::{DeserError, KeySer, MAX_KEY_SIZE, ValSer}, + tables::{DualKeyed, MAX_FIXED_VAL_SIZE}, +}; +use bytes::{BufMut, BytesMut}; +use reth_db::{ + Database, DatabaseEnv, + mdbx::{RW, TransactionKind, WriteFlags, tx::Tx}, +}; +use reth_db_api::DatabaseError; +use reth_libmdbx::{Cursor, RO}; +use std::borrow::Cow; + +/// Error type for reth-libmdbx based hot storage. +#[derive(Debug, thiserror::Error)] +pub enum MdbxError { + /// Inner error + #[error(transparent)] + Mdbx(#[from] reth_libmdbx::Error), + + /// Reth error. + #[error(transparent)] + Reth(#[from] DatabaseError), + + /// Deser. + #[error(transparent)] + Deser(#[from] DeserError), +} + +impl HotKvReadError for MdbxError { + fn into_hot_kv_error(self) -> HotKvError { + match self { + MdbxError::Mdbx(e) => HotKvError::from_err(e), + MdbxError::Deser(e) => HotKvError::Deser(e), + MdbxError::Reth(e) => HotKvError::from_err(e), + } + } +} + +impl From for DatabaseError { + fn from(value: DeserError) -> Self { + DatabaseError::Other(value.to_string()) + } +} + +impl HotKv for DatabaseEnv { + type RoTx = Tx; + type RwTx = Tx; + + fn reader(&self) -> Result { + self.tx().map_err(HotKvError::from_err) + } + + fn writer(&self) -> Result { + self.tx_mut().map_err(HotKvError::from_err) + } +} + +impl HotKvRead for Tx +where + K: TransactionKind, +{ + type Error = MdbxError; + + type Traverse<'a> = Cursor; + + fn raw_traverse<'a>(&'a self, table: &str) -> Result, Self::Error> { + let db = self.inner.open_db(Some(table))?; + let cursor = self.inner.cursor(&db)?; + + Ok(cursor) + } + + fn raw_get<'a>( + &'a self, + table: &str, + key: &[u8], + ) -> Result>, Self::Error> { + let dbi = self.inner.open_db(Some(table)).map(|db| db.dbi())?; + + self.inner.get(dbi, key.as_ref()).map_err(MdbxError::Mdbx) + } + + fn raw_get_dual<'a>( + &'a self, + _table: &str, + _key1: &[u8], + _key2: &[u8], + ) -> Result>, Self::Error> { + unimplemented!("Not implemented: raw_get_dual. Use get_dual instead."); + } + + fn get_dual( + &self, + key1: &T::Key, + key2: &T::Key2, + ) -> Result, Self::Error> { + let db = self.inner.open_db(Some(T::NAME))?; + let mut cursor = self.inner.cursor(&db)?; + + DualTableTraverse::::exact_dual(&mut cursor, key1, key2) + } +} + +impl HotKvWrite for Tx { + type TraverseMut<'a> = Cursor; + + fn raw_traverse_mut<'a>( + &'a mut self, + table: &str, + ) -> Result, Self::Error> { + let db = self.inner.open_db(Some(table))?; + let cursor = self.inner.cursor(&db)?; + + Ok(cursor) + } + + fn queue_raw_put(&mut self, table: &str, key: &[u8], value: &[u8]) -> Result<(), Self::Error> { + let dbi = self.inner.open_db(Some(table)).map(|db| db.dbi())?; + + self.inner.put(dbi, key, value, WriteFlags::UPSERT).map(|_| ()).map_err(MdbxError::Mdbx) + } + + fn queue_raw_put_dual( + &mut self, + _table: &str, + _key1: &[u8], + _key2: &[u8], + _value: &[u8], + ) -> Result<(), Self::Error> { + unimplemented!("Not implemented: queue_raw_put_dual. Use queue_put_dual instead."); + } + + // Specialized put for dual-keyed tables. + fn queue_put_dual( + &mut self, + key1: &T::Key, + key2: &T::Key2, + value: &T::Value, + ) -> Result<(), Self::Error> { + let k2_size = ::SIZE; + let mut scratch = [0u8; MAX_KEY_SIZE]; + + // This will be the total length of key2 + value, reserved in mdbx + let encoded_len = k2_size + value.encoded_size(); + + // Prepend the value with k2. + let mut buf = BytesMut::with_capacity(encoded_len); + let encoded_k2 = key2.encode_key(&mut scratch); + buf.put_slice(encoded_k2); + value.encode_value_to(&mut buf); + + let encoded_k1 = key1.encode_key(&mut scratch); + // NB: DUPSORT and RESERVE are incompatible :( + let db = self.inner.open_db(Some(T::NAME))?; + self.inner.put(db.dbi(), encoded_k1, &buf, Default::default())?; + + Ok(()) + } + + fn queue_raw_delete(&mut self, table: &str, key: &[u8]) -> Result<(), Self::Error> { + let dbi = self.inner.open_db(Some(table)).map(|db| db.dbi())?; + + self.inner.del(dbi, key, None).map(|_| ()).map_err(MdbxError::Mdbx) + } + + fn queue_raw_clear(&mut self, table: &str) -> Result<(), Self::Error> { + // Future: port more of reth's db env with dbi caching to avoid + // repeated open_db calls + let dbi = self.inner.open_db(Some(table)).map(|db| db.dbi())?; + self.inner.clear_db(dbi).map(|_| ()).map_err(MdbxError::Mdbx) + } + + fn queue_raw_create( + &mut self, + table: &str, + dual_key: bool, + fixed_val: bool, + ) -> Result<(), Self::Error> { + let mut flags = Default::default(); + + if dual_key { + flags |= reth_libmdbx::DatabaseFlags::DUP_SORT; + if fixed_val { + flags |= reth_libmdbx::DatabaseFlags::DUP_FIXED; + } + } + + self.inner.create_db(Some(table), flags).map(|_| ()).map_err(MdbxError::Mdbx) + } + + fn raw_commit(self) -> Result<(), Self::Error> { + // when committing, mdbx returns true on failure + let res = self.inner.commit()?; + + if res.0 { Err(reth_libmdbx::Error::Other(1).into()) } else { Ok(()) } + } +} + +impl KvTraverse for Cursor +where + K: TransactionKind, +{ + fn first<'a>(&'a mut self) -> Result>, MdbxError> { + Cursor::first(self).map_err(MdbxError::Mdbx) + } + + fn last<'a>(&'a mut self) -> Result>, MdbxError> { + Cursor::last(self).map_err(MdbxError::Mdbx) + } + + fn exact<'a>(&'a mut self, key: &[u8]) -> Result>, MdbxError> { + Cursor::set(self, key).map_err(MdbxError::Mdbx) + } + + fn lower_bound<'a>(&'a mut self, key: &[u8]) -> Result>, MdbxError> { + Cursor::set_range(self, key).map_err(MdbxError::Mdbx) + } + + fn read_next<'a>(&'a mut self) -> Result>, MdbxError> { + Cursor::next(self).map_err(MdbxError::Mdbx) + } + + fn read_prev<'a>(&'a mut self) -> Result>, MdbxError> { + Cursor::prev(self).map_err(MdbxError::Mdbx) + } +} + +impl KvTraverseMut for Cursor { + fn delete_current(&mut self) -> Result<(), MdbxError> { + Cursor::del(self, Default::default()).map_err(MdbxError::Mdbx) + } +} + +impl DualKeyedTraverse for Cursor +where + K: TransactionKind, +{ + fn exact_dual<'a>( + &'a mut self, + _key1: &[u8], + _key2: &[u8], + ) -> Result>, MdbxError> { + unimplemented!("Use DualTableTraverse for exact_dual"); + } + + fn next_dual_above<'a>( + &'a mut self, + _key1: &[u8], + _key2: &[u8], + ) -> Result>, MdbxError> { + unimplemented!("Use DualTableTraverse for next_dual_above"); + } + + fn next_k1<'a>(&'a mut self) -> Result>, MdbxError> { + unimplemented!("Use DualTableTraverse for next_k1"); + } + + fn next_k2<'a>(&'a mut self) -> Result>, MdbxError> { + unimplemented!("Use DualTableTraverse for next_k2"); + } +} + +impl DualTableTraverse for Cursor +where + T: DualKeyed, + K: TransactionKind, +{ + fn next_dual_above( + &mut self, + key1: &T::Key, + key2: &T::Key2, + ) -> Result>, MdbxError> { + Ok(get_both_range_helper::(self, key1, key2)? + .map(T::decode_prepended_value) + .transpose()? + .map(|(k2, v)| (key1.clone(), k2, v))) + } + + fn next_k1(&mut self) -> Result>, MdbxError> { + let Some((k, v)) = self.next_nodup::, Cow<'_, [u8]>>()? else { + return Ok(None); + }; + + let k1 = T::Key::decode_key(&k)?; + let (k2, v) = T::decode_prepended_value(v)?; + + Ok(Some((k1, k2, v))) + } + + fn next_k2(&mut self) -> Result>, MdbxError> { + let Some((k, v)) = self.next_dup::, Cow<'_, [u8]>>()? else { + return Ok(None); + }; + + let k = T::Key::decode_key(&k)?; + let (k2, v) = T::decode_prepended_value(v)?; + + Ok(Some((k, k2, v))) + } +} + +/// Helper to handle dup fixed value tables +fn dup_fixed_helper( + cursor: &mut Cursor, + key1: &T::Key, + key2: &T::Key2, + f: impl FnOnce(&mut Cursor, &[u8], &[u8]) -> Result, +) -> Result +where + T: DualKeyed, + K: TransactionKind, +{ + let mut key1_buf = [0u8; MAX_KEY_SIZE]; + let mut key2_buf = [0u8; MAX_KEY_SIZE]; + let key1_bytes = key1.encode_key(&mut key1_buf); + let key2_bytes = key2.encode_key(&mut key2_buf); + + // K2 slice must be EXACTLY the size of the fixed value size, if the + // table has one. This is a bit ugly, and results in an extra + // allocation for fixed-size values. This could be avoided using + // max value size. + if T::IS_FIXED_VAL { + let mut buf = [0u8; MAX_KEY_SIZE + MAX_FIXED_VAL_SIZE]; + buf[..::SIZE].copy_from_slice(key2_bytes); + + let kvs: usize = ::SIZE + T::FIXED_VAL_SIZE.unwrap(); + + f(cursor, key1_bytes, &buf[..kvs]) + } else { + f(cursor, key1_bytes, key2_bytes) + } +} + +// Helper to call get_both_range with dup fixed handling +fn get_both_range_helper<'a, T, K>( + cursor: &'a mut Cursor, + key1: &T::Key, + key2: &T::Key2, +) -> Result>, MdbxError> +where + T: DualKeyed, + K: TransactionKind, +{ + dup_fixed_helper::>>( + cursor, + key1, + key2, + |cursor, key1_bytes, key2_bytes| { + cursor.get_both_range(key1_bytes, key2_bytes).map_err(MdbxError::Mdbx) + }, + ) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{ + hot::model::{HotDbWrite, HotKv, HotKvRead, HotKvWrite, TableTraverse, TableTraverseMut}, + tables::{SingleKey, Table, hot}, + }; + use alloy::primitives::{Address, B256, BlockNumber, Bytes, U256}; + use reth::primitives::{Account, Bytecode, Header, SealedHeader}; + use reth_db::DatabaseEnv; + use serial_test::serial; + + // Test table definitions for traversal tests + #[derive(Debug)] + struct TestTable; + + impl Table for TestTable { + const NAME: &'static str = "mdbx_test_table"; + type Key = u64; + type Value = Bytes; + } + + impl SingleKey for TestTable {} + + /// Create a temporary MDBX database for testing that will be automatically cleaned up + fn run_test(f: F) { + let db = reth_db::test_utils::create_test_rw_db(); + + // Create tables from the `crate::tables::hot` module + let mut writer = db.db().writer().unwrap(); + + writer.queue_create::().unwrap(); + writer.queue_create::().unwrap(); + writer.queue_create::().unwrap(); + writer.queue_create::().unwrap(); + writer.queue_create::().unwrap(); + writer.queue_create::().unwrap(); + writer.queue_create::().unwrap(); + writer.queue_create::().unwrap(); + writer.queue_create::().unwrap(); + writer.queue_create::().unwrap(); + + writer.commit().expect("Failed to commit table creation"); + + f(db.db()); + } + + /// Create test data + fn create_test_account() -> (Address, Account) { + let address = Address::from_slice(&[0x1; 20]); + let account = Account { + nonce: 42, + balance: U256::from(1000u64), + bytecode_hash: Some(B256::from_slice(&[0x2; 32])), + }; + (address, account) + } + + fn create_test_bytecode() -> (B256, Bytecode) { + let hash = B256::from_slice(&[0x2; 32]); + let code = reth::primitives::Bytecode::new_raw(vec![0x60, 0x80, 0x60, 0x40].into()); + (hash, code) + } + + fn create_test_header() -> (BlockNumber, Header) { + let block_number = 12345; + let header = Header { + number: block_number, + gas_limit: 8000000, + gas_used: 100000, + timestamp: 1640995200, + parent_hash: B256::from_slice(&[0x3; 32]), + state_root: B256::from_slice(&[0x4; 32]), + ..Default::default() + }; + (block_number, header) + } + + #[test] + #[serial] + fn test_hotkv_basic_operations() { + run_test(test_hotkv_basic_operations_inner); + } + + fn test_hotkv_basic_operations_inner(db: &DatabaseEnv) { + let (address, account) = create_test_account(); + let (hash, bytecode) = create_test_bytecode(); + + // Test HotKv::writer() and basic write operations + { + let mut writer: Tx = db.writer().unwrap(); + + // Create tables first + writer.queue_create::().unwrap(); + + // Write account data + writer.queue_put::(&address, &account).unwrap(); + writer.queue_put::(&hash, &bytecode).unwrap(); + + // Commit the transaction + writer.raw_commit().unwrap(); + } + + // Test HotKv::reader() and basic read operations + { + let reader: Tx = db.reader().unwrap(); + + // Read account data + let read_account: Option = + reader.get::(&address).unwrap(); + assert_eq!(read_account, Some(account)); + + // Read bytecode + let read_bytecode: Option = reader.get::(&hash).unwrap(); + assert_eq!(read_bytecode, Some(bytecode)); + + // Test non-existent data + let nonexistent_addr = Address::from_slice(&[0xff; 20]); + let nonexistent_account: Option = + reader.get::(&nonexistent_addr).unwrap(); + assert_eq!(nonexistent_account, None); + } + } + + #[test] + #[serial] + fn test_raw_operations() { + run_test(test_raw_operations_inner) + } + + fn test_raw_operations_inner(db: &DatabaseEnv) { + let table_name = "test_table"; + let key = b"test_key"; + let value = b"test_value"; + + // Test raw write operations + { + let mut writer: Tx = db.writer().unwrap(); + + // Create table + writer.queue_raw_create(table_name, false, false).unwrap(); + + // Put raw data + writer.queue_raw_put(table_name, key, value).unwrap(); + + writer.raw_commit().unwrap(); + } + + // Test raw read operations + { + let reader: Tx = db.reader().unwrap(); + + let read_value = reader.raw_get(table_name, key).unwrap(); + assert_eq!(read_value.as_deref(), Some(value.as_slice())); + + // Test non-existent key + let nonexistent = reader.raw_get(table_name, b"nonexistent").unwrap(); + assert_eq!(nonexistent, None); + } + + // Test raw delete + { + let mut writer: Tx = db.writer().unwrap(); + + writer.queue_raw_delete(table_name, key).unwrap(); + writer.raw_commit().unwrap(); + } + + // Verify deletion + { + let reader: Tx = db.reader().unwrap(); + let deleted_value = reader.raw_get(table_name, key).unwrap(); + assert_eq!(deleted_value, None); + } + } + + #[test] + #[serial] + fn test_dual_keyed_operations() { + run_test(test_dual_keyed_operations_inner) + } + + fn test_dual_keyed_operations_inner(db: &DatabaseEnv) { + let address = Address::from_slice(&[0x1; 20]); + let storage_key = B256::from_slice(&[0x5; 32]); + let storage_value = U256::from(999u64); + + // Test dual-keyed table operations + { + let mut writer: Tx = db.writer().unwrap(); + + // Put storage data using dual keys + writer + .queue_put_dual::(&address, &storage_key, &storage_value) + .unwrap(); + + writer.raw_commit().unwrap(); + } + + // Test reading dual-keyed data + { + let reader: Tx = db.reader().unwrap(); + + // Read storage using dual key lookup + let read_value = + reader.get_dual::(&address, &storage_key).unwrap().unwrap(); + + assert_eq!(read_value, storage_value); + } + } + + #[test] + #[serial] + fn test_table_management() { + run_test(test_table_management_inner) + } + + fn test_table_management_inner(db: &DatabaseEnv) { + // Add some data + let (block_number, header) = create_test_header(); + { + let mut writer: Tx = db.writer().unwrap(); + writer.queue_put::(&block_number, &header).unwrap(); + writer.raw_commit().unwrap(); + } + + // Verify data exists + { + let reader: Tx = db.reader().unwrap(); + let read_header: Option
= reader.get::(&block_number).unwrap(); + assert_eq!(read_header, Some(header.clone())); + } + + // Clear the table + { + let mut writer: Tx = db.writer().unwrap(); + writer.queue_clear::().unwrap(); + writer.raw_commit().unwrap(); + } + + // Verify table is empty + { + let reader: Tx = db.reader().unwrap(); + let read_header: Option
= reader.get::(&block_number).unwrap(); + assert_eq!(read_header, None); + } + } + + #[test] + fn test_batch_operations() { + run_test(test_batch_operations_inner) + } + + fn test_batch_operations_inner(db: &DatabaseEnv) { + // Create test data + let accounts: Vec<(Address, Account)> = (0..10) + .map(|i| { + let mut addr_bytes = [0u8; 20]; + addr_bytes[19] = i; + let address = Address::from_slice(&addr_bytes); + let account = Account { + nonce: i.into(), + balance: U256::from((i as u64) * 100), + bytecode_hash: None, + }; + (address, account) + }) + .collect(); + + // Test batch writes + { + let mut writer: Tx = db.writer().unwrap(); + + // Write multiple accounts + for (address, account) in &accounts { + writer.queue_put::(address, account).unwrap(); + } + + writer.raw_commit().unwrap(); + } + + // Test batch reads + { + let reader: Tx = db.reader().unwrap(); + + for (address, expected_account) in &accounts { + let read_account: Option = + reader.get::(address).unwrap(); + assert_eq!(read_account.as_ref(), Some(expected_account)); + } + } + + // Test batch get_many + { + let reader: Tx = db.reader().unwrap(); + let addresses: Vec
= accounts.iter().map(|(addr, _)| *addr).collect(); + let read_accounts: Vec<(_, Option)> = + reader.get_many::(addresses.iter()).unwrap(); + + for (i, (_, expected_account)) in accounts.iter().enumerate() { + assert_eq!(read_accounts[i].1.as_ref(), Some(expected_account)); + } + } + } + + #[test] + fn test_transaction_isolation() { + run_test(test_transaction_isolation_inner) + } + + fn test_transaction_isolation_inner(db: &DatabaseEnv) { + let (address, account) = create_test_account(); + + // Setup initial data + { + let mut writer: Tx = db.writer().unwrap(); + writer.queue_put::(&address, &account).unwrap(); + writer.raw_commit().unwrap(); + } + + // Start a reader transaction + let reader: Tx = db.reader().unwrap(); + + // Modify data in a writer transaction + { + let mut writer: Tx = db.writer().unwrap(); + let modified_account = + Account { nonce: 999, balance: U256::from(9999u64), bytecode_hash: None }; + writer.queue_put::(&address, &modified_account).unwrap(); + writer.raw_commit().unwrap(); + } + + // Reader should still see original data (snapshot isolation) + { + let read_account: Option = + reader.get::(&address).unwrap(); + assert_eq!(read_account, Some(account)); + } + + // New reader should see modified data + { + let new_reader: Tx = db.reader().unwrap(); + let read_account: Option = + new_reader.get::(&address).unwrap(); + assert_eq!(read_account.unwrap().nonce, 999); + } + } + + #[test] + fn test_multiple_readers() { + run_test(test_multiple_readers_inner) + } + + fn test_multiple_readers_inner(db: &DatabaseEnv) { + let (address, account) = create_test_account(); + + // Setup data + { + let mut writer: Tx = db.writer().unwrap(); + writer.queue_put::(&address, &account).unwrap(); + writer.raw_commit().unwrap(); + } + + // Create multiple readers + let reader1: Tx = db.reader().unwrap(); + let reader2: Tx = db.reader().unwrap(); + let reader3: Tx = db.reader().unwrap(); + + // All readers should see the same data + let account1: Option = reader1.get::(&address).unwrap(); + let account2: Option = reader2.get::(&address).unwrap(); + let account3: Option = reader3.get::(&address).unwrap(); + + assert_eq!(account1, Some(account)); + assert_eq!(account2, Some(account)); + assert_eq!(account3, Some(account)); + } + + #[test] + fn test_error_handling() { + run_test(test_error_handling_inner) + } + + fn test_error_handling_inner(db: &DatabaseEnv) { + // Test reading from non-existent table + { + let reader: Tx = db.reader().unwrap(); + let result = reader.raw_get("nonexistent_table", b"key"); + + // Should handle gracefully (may return None or error depending on MDBX behavior) + match result { + Ok(None) => {} // This is fine + Err(_) => {} // This is also acceptable for non-existent table + Ok(Some(_)) => panic!("Should not return data for non-existent table"), + } + } + + // Test writing to a table without creating it first + { + let mut writer: Tx = db.writer().unwrap(); + let (address, account) = create_test_account(); + + // This should handle the case where table doesn't exist + let result = writer.queue_put::(&address, &account); + match result { + Ok(_) => { + // If it succeeds, commit should work + writer.raw_commit().unwrap(); + } + Err(_) => { + // If it fails, that's expected behavior + } + } + } + } + + #[test] + fn test_serialization_roundtrip() { + run_test(test_serialization_roundtrip_inner) + } + + fn test_serialization_roundtrip_inner(db: &DatabaseEnv) { + // Test various data types + let (block_number, header) = create_test_header(); + let header = SealedHeader::new_unhashed(header); + + { + let mut writer: Tx = db.writer().unwrap(); + + // Write different types + writer.put_header(&header).unwrap(); + + writer.raw_commit().unwrap(); + } + + { + let reader: Tx = db.reader().unwrap(); + + // Read and verify + let read_header: Option
= reader.get::(&block_number).unwrap(); + assert_eq!(read_header.as_ref(), Some(header.header())); + + let read_hash: Option = reader.get::(&header.hash()).unwrap(); + assert_eq!(read_hash, Some(header.number)); + } + } + + #[test] + fn test_large_data() { + run_test(test_large_data_inner) + } + + fn test_large_data_inner(db: &DatabaseEnv) { + // Create a large bytecode + let hash = B256::from_slice(&[0x8; 32]); + let large_code_vec: Vec = (0..10000).map(|i| (i % 256) as u8).collect(); + let large_bytecode = Bytecode::new_raw(large_code_vec.clone().into()); + + { + let mut writer: Tx = db.writer().unwrap(); + writer.queue_create::().unwrap(); + writer.queue_put::(&hash, &large_bytecode).unwrap(); + writer.raw_commit().unwrap(); + } + + { + let reader: Tx = db.reader().unwrap(); + let read_bytecode: Option = reader.get::(&hash).unwrap(); + assert_eq!(read_bytecode, Some(large_bytecode)); + } + } + + // ======================================================================== + // Cursor Traversal Tests + // ======================================================================== + + #[test] + fn test_table_traverse_basic_navigation() { + run_test(test_table_traverse_basic_navigation_inner) + } + + fn test_table_traverse_basic_navigation_inner(db: &DatabaseEnv) { + // Setup test data with multiple entries + let test_data: Vec<(u64, Bytes)> = vec![ + (1, Bytes::from_static(b"value_001")), + (2, Bytes::from_static(b"value_002")), + (3, Bytes::from_static(b"value_003")), + (10, Bytes::from_static(b"value_010")), + (20, Bytes::from_static(b"value_020")), + ]; + + // Insert test data + { + let mut writer: Tx = db.writer().unwrap(); + for (key, value) in &test_data { + writer.queue_put::(key, value).unwrap(); + } + writer.raw_commit().unwrap(); + } + + // Test cursor traversal + { + let tx: Tx = db.reader().unwrap(); + let db_handle = tx.inner.open_db(Some(TestTable::NAME)).unwrap(); + let mut cursor = tx.inner.cursor(&db_handle).unwrap(); + + // Test first() + let first_result = TableTraverse::::first(&mut cursor).unwrap(); + assert!(first_result.is_some()); + let (key, value) = first_result.unwrap(); + assert_eq!(key, test_data[0].0); + assert_eq!(value, test_data[0].1); + + // Test last() + let last_result = TableTraverse::::last(&mut cursor).unwrap(); + assert!(last_result.is_some()); + let (key, value) = last_result.unwrap(); + assert_eq!(key, test_data.last().unwrap().0); + assert_eq!(value, test_data.last().unwrap().1); + + // Test exact lookup + let exact_result = TableTraverse::::exact(&mut cursor, &2u64).unwrap(); + assert!(exact_result.is_some()); + assert_eq!(exact_result.unwrap(), test_data[1].1); + + // Test exact lookup for non-existent key + let missing_result = + TableTraverse::::exact(&mut cursor, &999u64).unwrap(); + assert!(missing_result.is_none()); + + // Test next_above (range lookup) + let range_result = + TableTraverse::::lower_bound(&mut cursor, &5u64).unwrap(); + assert!(range_result.is_some()); + let (key, value) = range_result.unwrap(); + assert_eq!(key, test_data[3].0); // key 10 + assert_eq!(value, test_data[3].1); + } + } + + #[test] + fn test_table_traverse_sequential_navigation() { + run_test(test_table_traverse_sequential_navigation_inner) + } + + fn test_table_traverse_sequential_navigation_inner(db: &DatabaseEnv) { + // Setup sequential test data + let test_data: Vec<(u64, Bytes)> = (1..=10) + .map(|i| { + let s = format!("value_{:03}", i); + let s = s.as_bytes(); + let value = Bytes::copy_from_slice(s); + (i, value) + }) + .collect(); + + // Insert test data + { + let mut writer: Tx = db.writer().unwrap(); + for (key, value) in &test_data { + writer.queue_put::(key, value).unwrap(); + } + writer.raw_commit().unwrap(); + } + + // Test sequential navigation + { + let tx: Tx = db.reader().unwrap(); + let db_handle = tx.inner.open_db(Some(TestTable::NAME)).unwrap(); + let mut cursor = tx.inner.cursor(&db_handle).unwrap(); + + // Start from first and traverse forward + let mut current_idx = 0; + let first_result = TableTraverse::::first(&mut cursor).unwrap(); + assert!(first_result.is_some()); + + let (key, value) = first_result.unwrap(); + assert_eq!(key, test_data[current_idx].0); + assert_eq!(value, test_data[current_idx].1); + + // Navigate forward through all entries + while current_idx < test_data.len() - 1 { + let next_result = TableTraverse::::read_next(&mut cursor).unwrap(); + assert!(next_result.is_some()); + + current_idx += 1; + let (key, value) = next_result.unwrap(); + assert_eq!(key, test_data[current_idx].0); + assert_eq!(value, test_data[current_idx].1); + } + + // Next should return None at the end + let beyond_end = TableTraverse::::read_next(&mut cursor).unwrap(); + assert!(beyond_end.is_none()); + + // Navigate backward + while current_idx > 0 { + let prev_result = TableTraverse::::read_prev(&mut cursor).unwrap(); + assert!(prev_result.is_some()); + + current_idx -= 1; + let (key, value) = prev_result.unwrap(); + assert_eq!(key, test_data[current_idx].0); + assert_eq!(value, test_data[current_idx].1); + } + + // Previous should return None at the beginning + let before_start = TableTraverse::::read_prev(&mut cursor).unwrap(); + assert!(before_start.is_none()); + } + } + + #[test] + fn test_table_traverse_mut_delete() { + run_test(test_table_traverse_mut_delete_inner) + } + + fn test_table_traverse_mut_delete_inner(db: &DatabaseEnv) { + let test_data: Vec<(u64, Bytes)> = vec![ + (1, Bytes::from_static(b"delete_value_1")), + (2, Bytes::from_static(b"delete_value_2")), + (3, Bytes::from_static(b"delete_value_3")), + ]; + + // Insert test data + { + let mut writer: Tx = db.writer().unwrap(); + for (key, value) in &test_data { + writer.queue_put::(key, value).unwrap(); + } + writer.raw_commit().unwrap(); + } + // Test cursor deletion + { + let tx: Tx = db.writer().unwrap(); + let db_handle = tx.inner.open_db(Some(TestTable::NAME)).unwrap(); + let mut cursor = tx.inner.cursor(&db_handle).unwrap(); + + // Navigate to middle entry + let first = TableTraverse::::first(&mut cursor).unwrap().unwrap(); + assert_eq!(first.0, test_data[0].0); + + let next = TableTraverse::::read_next(&mut cursor).unwrap().unwrap(); + assert_eq!(next.0, test_data[1].0); + + // Delete current entry (key 2) + TableTraverseMut::::delete_current(&mut cursor).unwrap(); + + tx.raw_commit().unwrap(); + } + + // Verify deletion + { + let tx: Tx = db.reader().unwrap(); + let db_handle = tx.inner.open_db(Some(TestTable::NAME)).unwrap(); + let mut cursor = tx.inner.cursor(&db_handle).unwrap(); + + // Should only have first and third entries + let first = TableTraverse::::first(&mut cursor).unwrap().unwrap(); + assert_eq!(first.0, test_data[0].0); + + let second = TableTraverse::::read_next(&mut cursor).unwrap().unwrap(); + assert_eq!(second.0, test_data[2].0); + + // Should be no more entries + let none = TableTraverse::::read_next(&mut cursor).unwrap(); + assert!(none.is_none()); + + // Verify deleted key is gone + let missing = + TableTraverse::::exact(&mut cursor, &test_data[1].0).unwrap(); + assert!(missing.is_none()); + } + } + + #[test] + fn test_table_traverse_accounts() { + run_test(test_table_traverse_accounts_inner) + } + + fn test_table_traverse_accounts_inner(db: &DatabaseEnv) { + // Setup test accounts + let test_accounts: Vec<(Address, Account)> = (0..5) + .map(|i| { + let mut addr_bytes = [0u8; 20]; + addr_bytes[19] = i; + let address = Address::from_slice(&addr_bytes); + let account = Account { + nonce: (i as u64) * 10, + balance: U256::from((i as u64) * 1000), + bytecode_hash: if i % 2 == 0 { Some(B256::from_slice(&[i; 32])) } else { None }, + }; + (address, account) + }) + .collect(); + + // Insert test data + { + let mut writer: Tx = db.writer().unwrap(); + + for (address, account) in &test_accounts { + writer.queue_put::(address, account).unwrap(); + } + + writer.raw_commit().unwrap(); + } + + // Test typed table traversal + { + let tx: Tx = db.reader().unwrap(); + let db_handle = tx.inner.open_db(Some(hot::PlainAccountState::NAME)).unwrap(); + let mut cursor = tx.inner.cursor(&db_handle).unwrap(); + + // Test first with type-safe operations + let first_raw = TableTraverse::::first(&mut cursor).unwrap(); + assert!(first_raw.is_some()); + let (first_key, first_account) = first_raw.unwrap(); + assert_eq!(first_key, test_accounts[0].0); + assert_eq!(first_account, test_accounts[0].1); + + // Test last + let last_raw = TableTraverse::::last(&mut cursor).unwrap(); + assert!(last_raw.is_some()); + let (last_key, last_account) = last_raw.unwrap(); + assert_eq!(last_key, test_accounts.last().unwrap().0); + assert_eq!(last_account, test_accounts.last().unwrap().1); + + // Test exact lookup + let target_address = &test_accounts[2].0; + let exact_account = + TableTraverse::::exact(&mut cursor, target_address) + .unwrap(); + assert!(exact_account.is_some()); + assert_eq!(exact_account.unwrap(), test_accounts[2].1); + + // Test range lookup + let mut partial_addr = [0u8; 20]; + partial_addr[19] = 3; // Between entries 2 and 3 + let range_addr = Address::from_slice(&partial_addr); + + let range_result = + TableTraverse::::lower_bound(&mut cursor, &range_addr) + .unwrap(); + assert!(range_result.is_some()); + let (found_addr, found_account) = range_result.unwrap(); + assert_eq!(found_addr, test_accounts[3].0); + assert_eq!(found_account, test_accounts[3].1); + } + } + + #[test] + fn test_dual_table_traverse() { + run_test(test_dual_table_traverse_inner) + } + + fn test_dual_table_traverse_inner(db: &DatabaseEnv) { + let one_addr = Address::repeat_byte(0x01); + let two_addr = Address::repeat_byte(0x02); + + let one_slot = B256::with_last_byte(0x01); + let two_slot = B256::with_last_byte(0x06); + let three_slot = B256::with_last_byte(0x09); + + let one_value = U256::from(0x100); + let two_value = U256::from(0x200); + let three_value = U256::from(0x300); + let four_value = U256::from(0x400); + let five_value = U256::from(0x500); + + // Setup test storage data + let test_storage: Vec<(Address, B256, U256)> = vec![ + (one_addr, one_slot, one_value), + (one_addr, two_slot, two_value), + (one_addr, three_slot, three_value), + (two_addr, one_slot, four_value), + (two_addr, two_slot, five_value), + ]; + + // Insert test data + { + let mut writer: Tx = db.writer().unwrap(); + + for (address, storage_key, value) in &test_storage { + writer + .queue_put_dual::(address, storage_key, value) + .unwrap(); + } + + writer.raw_commit().unwrap(); + } + + // Test dual-keyed traversal + { + let tx: Tx = db.reader().unwrap(); + let db_handle = tx.inner.open_db(Some(hot::PlainStorageState::NAME)).unwrap(); + let mut cursor = tx.inner.cursor(&db_handle).unwrap(); + + // Test exact dual lookup + let address = &test_storage[1].0; + let storage_key = &test_storage[1].1; + let expected_value = &test_storage[1].2; + + let exact_result = DualTableTraverse::::exact_dual( + &mut cursor, + address, + storage_key, + ) + .unwrap() + .unwrap(); + assert_eq!(exact_result, *expected_value); + + // Test range lookup for dual keys + let search_key = B256::with_last_byte(0x02); + let range_result = DualTableTraverse::::next_dual_above( + &mut cursor, + &test_storage[0].0, // Address 0x01 + &search_key, + ) + .unwrap() + .unwrap(); + + let (found_addr, found_key, found_value) = range_result; + assert_eq!(found_addr, test_storage[1].0); // Same address + assert_eq!(found_key, test_storage[1].1); // Next storage key (0x02) + assert_eq!(found_value, test_storage[1].2); // Corresponding value + + // Test next_k1 (move to next primary key) + // First position cursor at first entry of first address + DualTableTraverse::::exact_dual( + &mut cursor, + &test_storage[0].0, + &test_storage[0].1, + ) + .unwrap(); + + // Move to next primary key (different address) + let next_k1_result = + DualTableTraverse::::next_k1(&mut cursor).unwrap(); + assert!(next_k1_result.is_some()); + let (next_addr, next_storage_key, next_value) = next_k1_result.unwrap(); + assert_eq!(next_addr, test_storage[3].0); // Address 0x02 + assert_eq!(next_storage_key, test_storage[3].1); // First storage key for new address + assert_eq!(next_value, test_storage[3].2); + } + } + + #[test] + fn test_dual_table_traverse_empty_results() { + run_test(test_dual_table_traverse_empty_results_inner) + } + + fn test_dual_table_traverse_empty_results_inner(db: &DatabaseEnv) { + // Setup minimal test data + let address = Address::from_slice(&[0x01; 20]); + let storage_key = B256::from_slice(&[0x01; 32]); + let value = U256::from(100); + + { + let mut writer: Tx = db.writer().unwrap(); + writer + .queue_put_dual::(&address, &storage_key, &value) + .unwrap(); + writer.raw_commit().unwrap(); + } + + { + let tx: Tx = db.reader().unwrap(); + let db_handle = tx.inner.open_db(Some(hot::PlainStorageState::NAME)).unwrap(); + let mut cursor = tx.inner.cursor(&db_handle).unwrap(); + + // Test exact lookup for non-existent dual key + let missing_addr = Address::from_slice(&[0xFF; 20]); + let missing_key = B256::from_slice(&[0xFF; 32]); + + let exact_missing = DualTableTraverse::::exact_dual( + &mut cursor, + &missing_addr, + &missing_key, + ) + .unwrap(); + assert!(exact_missing.is_none()); + + // Test range lookup beyond all data + let beyond_key = B256::from_slice(&[0xFF; 32]); + let range_missing = DualTableTraverse::::next_dual_above( + &mut cursor, + &address, + &beyond_key, + ) + .unwrap(); + assert!(range_missing.is_none()); + + // Position at the only entry, then try next_k1 + DualTableTraverse::::exact_dual( + &mut cursor, + &address, + &storage_key, + ) + .unwrap(); + + let next_k1_missing = + DualTableTraverse::::next_k1(&mut cursor).unwrap(); + assert!(next_k1_missing.is_none()); + } + } + + #[test] + fn test_table_traverse_empty_table() { + run_test(test_table_traverse_empty_table_inner) + } + + fn test_table_traverse_empty_table_inner(db: &DatabaseEnv) { + // TestTable is already created but empty + { + let tx: Tx = db.reader().unwrap(); + let db_handle = tx.inner.open_db(Some(TestTable::NAME)).unwrap(); + let mut cursor = tx.inner.cursor(&db_handle).unwrap(); + + // All operations should return None on empty table + assert!(TableTraverse::::first(&mut cursor).unwrap().is_none()); + assert!(TableTraverse::::last(&mut cursor).unwrap().is_none()); + assert!(TableTraverse::::exact(&mut cursor, &42u64).unwrap().is_none()); + assert!( + TableTraverse::::lower_bound(&mut cursor, &42u64).unwrap().is_none() + ); + assert!(TableTraverse::::read_next(&mut cursor).unwrap().is_none()); + assert!(TableTraverse::::read_prev(&mut cursor).unwrap().is_none()); + } + } + + #[test] + fn test_table_traverse_state_management() { + run_test(test_table_traverse_state_management_inner) + } + + fn test_table_traverse_state_management_inner(db: &DatabaseEnv) { + let test_data: Vec<(u64, Bytes)> = vec![ + (1, Bytes::from_static(b"state_value_1")), + (2, Bytes::from_static(b"state_value_2")), + (3, Bytes::from_static(b"state_value_3")), + ]; + + { + let mut writer: Tx = db.writer().unwrap(); + for (key, value) in &test_data { + writer.queue_put::(key, value).unwrap(); + } + writer.raw_commit().unwrap(); + } + + { + let tx: Tx = db.reader().unwrap(); + let db_handle = tx.inner.open_db(Some(TestTable::NAME)).unwrap(); + let mut cursor = tx.inner.cursor(&db_handle).unwrap(); + + // Test that cursor operations maintain state correctly + + // Start at first + let first = TableTraverse::::first(&mut cursor).unwrap().unwrap(); + assert_eq!(first.0, test_data[0].0); + + // Move to second via next + let second = TableTraverse::::read_next(&mut cursor).unwrap().unwrap(); + assert_eq!(second.0, test_data[1].0); + + // Jump to last + let last = TableTraverse::::last(&mut cursor).unwrap().unwrap(); + assert_eq!(last.0, test_data[2].0); + + // Move back via prev + let back_to_second = + TableTraverse::::read_prev(&mut cursor).unwrap().unwrap(); + assert_eq!(back_to_second.0, test_data[1].0); + + // Use exact to jump to specific position + let exact_first = + TableTraverse::::exact(&mut cursor, &test_data[0].0).unwrap(); + assert!(exact_first.is_some()); + assert_eq!(exact_first.unwrap(), test_data[0].1); + + // Verify cursor is now positioned at first entry + let next_from_first = + TableTraverse::::read_next(&mut cursor).unwrap().unwrap(); + assert_eq!(next_from_first.0, test_data[1].0); + + // Use range lookup - look for key >= 1, should find key 1 + let range_lookup = + TableTraverse::::lower_bound(&mut cursor, &1u64).unwrap().unwrap(); + assert_eq!(range_lookup.0, test_data[0].0); // Should find key 1 + + // Verify we can continue navigation from range position + let next_after_range = + TableTraverse::::read_next(&mut cursor).unwrap().unwrap(); + assert_eq!(next_after_range.0, test_data[1].0); + } + } +} diff --git a/crates/storage/src/hot/impls/mem.rs b/crates/storage/src/hot/impls/mem.rs new file mode 100644 index 0000000..ed43cd2 --- /dev/null +++ b/crates/storage/src/hot/impls/mem.rs @@ -0,0 +1,1869 @@ +use crate::{ + hot::model::{ + DualKeyValue, DualKeyedTraverse, DualTableTraverse, HotKv, HotKvError, HotKvRead, + HotKvReadError, HotKvWrite, KvTraverse, KvTraverseMut, RawDualKeyValue, RawKeyValue, + RawValue, + }, + ser::{DeserError, KeySer, MAX_KEY_SIZE}, + tables::DualKeyed, +}; +use bytes::Bytes; +use std::{ + borrow::Cow, + collections::BTreeMap, + sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard}, +}; + +// Type aliases for store structure +type MemStoreKey = [u8; MAX_KEY_SIZE * 2]; +type StoreTable = BTreeMap; +type Store = BTreeMap; + +// Type aliases for queued operations +type TableOp = BTreeMap; +type OpStore = BTreeMap; + +/// A simple in-memory key-value store using [`BTreeMap`]s. +/// +/// The store is backed by an [`RwLock`]. As a result, this implementation +/// supports concurrent multiple concurrent read transactions, but write +/// transactions are exclusive, and cannot overlap with other read or write +/// transactions. +/// +/// This implementation is primarily intended for testing and +/// development purposes. +#[derive(Clone)] +pub struct MemKv { + map: Arc>, +} + +impl core::fmt::Debug for MemKv { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + f.debug_struct("MemKv").finish() + } +} + +impl MemKv { + /// Create a new empty in-memory KV store. + pub fn new() -> Self { + Self { map: Arc::new(RwLock::new(BTreeMap::new())) } + } + + #[track_caller] + fn key(k: &[u8]) -> MemStoreKey { + assert!(k.len() <= MAX_KEY_SIZE * 2, "Key length exceeds MAX_KEY_SIZE"); + let mut buf = [0u8; MAX_KEY_SIZE * 2]; + buf[..k.len()].copy_from_slice(k); + buf + } + + #[track_caller] + fn dual_key(k1: &[u8], k2: &[u8]) -> MemStoreKey { + assert!( + k1.len() + k2.len() <= MAX_KEY_SIZE * 2, + "Combined key length exceeds MAX_KEY_SIZE" + ); + let mut buf = [0u8; MAX_KEY_SIZE * 2]; + buf[..MAX_KEY_SIZE.min(k1.len())].copy_from_slice(k1); + buf[MAX_KEY_SIZE..MAX_KEY_SIZE + k2.len()].copy_from_slice(k2); + buf + } + + /// SAFETY: + /// Caller must ensure that `key` lives long enough. + #[track_caller] + fn split_dual_key<'a>(key: &[u8]) -> (Cow<'a, [u8]>, Cow<'a, [u8]>) { + assert_eq!(key.len(), MAX_KEY_SIZE * 2, "Key length does not match expected dual key size"); + let k1 = &key[..MAX_KEY_SIZE]; + let k2 = &key[MAX_KEY_SIZE..]; + + unsafe { std::mem::transmute((Cow::Borrowed(k1), Cow::Borrowed(k2))) } + } +} + +impl Default for MemKv { + fn default() -> Self { + Self::new() + } +} + +/// Read-only transaction for MemKv. +pub struct MemKvRoTx { + guard: RwLockReadGuard<'static, Store>, + + // Keep the store alive while the transaction exists + _store: Arc>, +} + +impl core::fmt::Debug for MemKvRoTx { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + f.debug_struct("MemKvRoTx").finish() + } +} + +// SAFETY: MemKvRoTx holds a read guard which ensures the data remains valid +unsafe impl Send for MemKvRoTx {} +unsafe impl Sync for MemKvRoTx {} + +/// Read-write transaction for MemKv. +pub struct MemKvRwTx { + guard: RwLockWriteGuard<'static, Store>, + queued_ops: OpStore, + + // Keep the store alive while the transaction exists + _store: Arc>, +} + +impl MemKvRwTx { + fn commit_inner(&mut self) { + let ops = std::mem::take(&mut self.queued_ops); + + for (table, table_op) in ops.into_iter() { + table_op.apply(&table, &mut self.guard); + } + } + + /// Downgrade the transaction to a read-only transaction without + /// committing, discarding queued changes. + pub fn downgrade(self) -> MemKvRoTx { + let guard = RwLockWriteGuard::downgrade(self.guard); + + MemKvRoTx { guard, _store: self._store } + } + + /// Commit the transaction and downgrade to a read-only transaction. + pub fn commit_downgrade(mut self) -> MemKvRoTx { + self.commit_inner(); + + let guard = RwLockWriteGuard::downgrade(self.guard); + + MemKvRoTx { guard, _store: self._store } + } +} + +impl core::fmt::Debug for MemKvRwTx { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + f.debug_struct("MemKvRwTx").finish() + } +} + +/// Queued key-value operation +#[derive(Debug, Clone)] +enum QueuedKvOp { + Delete, + Put { value: Bytes }, +} + +impl QueuedKvOp { + /// Apply the op to a table + fn apply(self, table: &mut StoreTable, key: MemStoreKey) { + match self { + QueuedKvOp::Put { value } => { + table.insert(key, value); + } + QueuedKvOp::Delete => { + table.remove(&key); + } + } + } +} + +/// Queued table operation +#[derive(Debug)] +enum QueuedTableOp { + Modify { ops: TableOp }, + Clear { new_table: TableOp }, +} + +impl Default for QueuedTableOp { + fn default() -> Self { + QueuedTableOp::Modify { ops: TableOp::new() } + } +} + +impl QueuedTableOp { + const fn is_clear(&self) -> bool { + matches!(self, QueuedTableOp::Clear { .. }) + } + + fn get(&self, key: &MemStoreKey) -> Option<&QueuedKvOp> { + match self { + QueuedTableOp::Modify { ops } => ops.get(key), + QueuedTableOp::Clear { new_table } => new_table.get(key), + } + } + + fn put(&mut self, key: MemStoreKey, op: QueuedKvOp) { + match self { + QueuedTableOp::Modify { ops } | QueuedTableOp::Clear { new_table: ops } => { + ops.insert(key, op); + } + } + } + + fn delete(&mut self, key: MemStoreKey) { + match self { + QueuedTableOp::Modify { ops } | QueuedTableOp::Clear { new_table: ops } => { + ops.insert(key, QueuedKvOp::Delete); + } + } + } + + /// Get mutable reference to the inner ops if applicable + fn apply(self, key: &str, store: &mut Store) { + match self { + QueuedTableOp::Modify { ops } => { + let table = store.entry(key.to_owned()).or_default(); + for (key, op) in ops { + op.apply(table, key); + } + } + QueuedTableOp::Clear { new_table } => { + let mut table = StoreTable::new(); + for (k, op) in new_table { + op.apply(&mut table, k); + } + + // replace the table entirely + store.insert(key.to_owned(), table); + } + } + } +} + +// SAFETY: MemKvRwTx holds a write guard which ensures exclusive access +unsafe impl Send for MemKvRwTx {} + +/// Error type for MemKv operations +#[derive(Debug, thiserror::Error)] +pub enum MemKvError { + /// Hot KV error + #[error(transparent)] + HotKv(#[from] HotKvError), + + /// Serialization error + #[error(transparent)] + Deser(#[from] DeserError), +} + +impl trevm::revm::database::DBErrorMarker for MemKvError {} + +impl HotKvReadError for MemKvError { + fn into_hot_kv_error(self) -> HotKvError { + match self { + MemKvError::HotKv(e) => e, + MemKvError::Deser(e) => HotKvError::Deser(e), + } + } +} + +/// Memory cursor for traversing a BTreeMap +pub struct MemKvCursor<'a> { + table: &'a StoreTable, + current_key: Option, +} + +impl core::fmt::Debug for MemKvCursor<'_> { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + f.debug_struct("MemKvCursor").finish() + } +} + +impl<'a> MemKvCursor<'a> { + /// Create a new cursor for the given table + pub const fn new(table: &'a StoreTable) -> Self { + Self { table, current_key: None } + } + + /// Get the current key the cursor is positioned at + pub fn current_key(&self) -> MemStoreKey { + self.current_key.unwrap_or([0u8; MAX_KEY_SIZE * 2]) + } + + /// Set the current key the cursor is positioned at + pub const fn set_current_key(&mut self, key: MemStoreKey) { + self.current_key = Some(key); + } + + /// Clear the current key the cursor is positioned at + pub const fn clear_current_key(&mut self) { + self.current_key = None; + } + + /// Get the current k1 the cursor is positioned at + fn current_k1(&self) -> [u8; MAX_KEY_SIZE] { + self.current_key + .map(|key| key[..MAX_KEY_SIZE].try_into().unwrap()) + .unwrap_or([0u8; MAX_KEY_SIZE]) + } +} + +impl<'a> KvTraverse for MemKvCursor<'a> { + fn first<'b>(&'b mut self) -> Result>, MemKvError> { + let Some((key, value)) = self.table.first_key_value() else { + self.clear_current_key(); + return Ok(None); + }; + self.current_key = Some(*key); + Ok(Some((Cow::Borrowed(key), Cow::Borrowed(value.as_ref())))) + } + + fn last<'b>(&'b mut self) -> Result>, MemKvError> { + let Some((key, value)) = self.table.last_key_value() else { + self.clear_current_key(); + return Ok(None); + }; + self.current_key = Some(*key); + Ok(Some((Cow::Borrowed(key), Cow::Borrowed(value.as_ref())))) + } + + fn exact<'b>(&'b mut self, key: &[u8]) -> Result>, MemKvError> { + let search_key = MemKv::key(key); + self.set_current_key(search_key); + if let Some(value) = self.table.get(&search_key) { + Ok(Some(Cow::Borrowed(value.as_ref()))) + } else { + Ok(None) + } + } + + fn lower_bound<'b>(&'b mut self, key: &[u8]) -> Result>, MemKvError> { + let search_key = MemKv::key(key); + + // Use range to find the first key >= search_key + if let Some((found_key, value)) = self.table.range(search_key..).next() { + self.set_current_key(*found_key); + Ok(Some((Cow::Borrowed(found_key), Cow::Borrowed(value.as_ref())))) + } else { + self.current_key = self.table.last_key_value().map(|(k, _)| *k); + Ok(None) + } + } + + fn read_next<'b>(&'b mut self) -> Result>, MemKvError> { + use core::ops::Bound; + let current = self.current_key(); + // Use Excluded bound to find strictly greater than current key + let Some((found_key, value)) = + self.table.range((Bound::Excluded(current), Bound::Unbounded)).next() + else { + return Ok(None); + }; + self.set_current_key(*found_key); + Ok(Some((Cow::Borrowed(found_key), Cow::Borrowed(value.as_ref())))) + } + + fn read_prev<'b>(&'b mut self) -> Result>, MemKvError> { + let current = self.current_key(); + let Some((k, v)) = self.table.range(..current).next_back() else { + self.clear_current_key(); + return Ok(None); + }; + self.set_current_key(*k); + Ok(Some((Cow::Borrowed(k), Cow::Borrowed(v.as_ref())))) + } +} + +// Implement DualKeyedTraverse (basic implementation - delegates to raw methods) +impl<'a> DualKeyedTraverse for MemKvCursor<'a> { + fn exact_dual<'b>( + &'b mut self, + key1: &[u8], + key2: &[u8], + ) -> Result>, MemKvError> { + let combined_key = MemKv::dual_key(key1, key2); + KvTraverse::exact(self, &combined_key) + } + + fn next_dual_above<'b>( + &'b mut self, + key1: &[u8], + key2: &[u8], + ) -> Result>, MemKvError> { + let combined_key = MemKv::dual_key(key1, key2); + let Some((found_key, value)) = KvTraverse::lower_bound(self, &combined_key)? else { + return Ok(None); + }; + let (k1, k2) = MemKv::split_dual_key(found_key.as_ref()); + Ok(Some((k1, k2, value))) + } + + fn next_k1<'b>(&'b mut self) -> Result>, MemKvError> { + // scan forward until finding a new k1 + let last_k1 = self.current_k1(); + + DualKeyedTraverse::next_dual_above(self, &last_k1, &[0xffu8; MAX_KEY_SIZE]) + } + + fn next_k2<'b>(&'b mut self) -> Result>, MemKvError> { + let current_key = self.current_key(); + let (current_k1, current_k2) = MemKv::split_dual_key(¤t_key); + + // scan forward until finding a new k2 for the same k1 + DualKeyedTraverse::next_dual_above(self, ¤t_k1, ¤t_k2) + } +} + +// Implement DualTableTraverse for typed dual-keyed table access +impl<'a, T> DualTableTraverse for MemKvCursor<'a> +where + T: DualKeyed, +{ + fn next_dual_above( + &mut self, + key1: &T::Key, + key2: &T::Key2, + ) -> Result>, MemKvError> { + let mut key1_buf = [0u8; MAX_KEY_SIZE]; + let mut key2_buf = [0u8; MAX_KEY_SIZE]; + let key1_bytes = key1.encode_key(&mut key1_buf); + let key2_bytes = key2.encode_key(&mut key2_buf); + + DualKeyedTraverse::next_dual_above(self, key1_bytes, key2_bytes)? + .map(T::decode_kkv_tuple) + .transpose() + .map_err(Into::into) + } + + fn next_k1(&mut self) -> Result>, MemKvError> { + DualKeyedTraverse::next_k1(self)?.map(T::decode_kkv_tuple).transpose().map_err(Into::into) + } + + fn next_k2(&mut self) -> Result>, MemKvError> { + DualKeyedTraverse::next_k2(self)?.map(T::decode_kkv_tuple).transpose().map_err(Into::into) + } +} + +/// Memory cursor for read-write operations +pub struct MemKvCursorMut<'a> { + table: &'a StoreTable, + queued_ops: &'a mut TableOp, + is_cleared: bool, + current_key: Option, +} + +impl core::fmt::Debug for MemKvCursorMut<'_> { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + f.debug_struct("MemKvCursorMut").field("is_cleared", &self.is_cleared).finish() + } +} + +impl<'a> MemKvCursorMut<'a> { + /// Create a new mutable cursor for the given table and queued operations + const fn new(table: &'a StoreTable, queued_ops: &'a mut TableOp, is_cleared: bool) -> Self { + Self { table, queued_ops, is_cleared, current_key: None } + } + + /// Get the current key the cursor is positioned at + pub fn current_key(&self) -> MemStoreKey { + self.current_key.unwrap_or([0u8; MAX_KEY_SIZE * 2]) + } + + /// Set the current key the cursor is positioned at + pub const fn set_current_key(&mut self, key: MemStoreKey) { + self.current_key = Some(key); + } + + /// Clear the current key the cursor is positioned at + pub const fn clear_current_key(&mut self) { + self.current_key = None; + } + + /// Get the current k1 the cursor is positioned at + fn current_k1(&self) -> [u8; MAX_KEY_SIZE] { + self.current_key + .map(|key| key[..MAX_KEY_SIZE].try_into().unwrap()) + .unwrap_or([0u8; MAX_KEY_SIZE]) + } + + /// Get value for a key, returning owned bytes + fn get_owned(&self, key: &MemStoreKey) -> Option { + if let Some(op) = self.queued_ops.get(key) { + match op { + QueuedKvOp::Put { value } => Some(value.clone()), + QueuedKvOp::Delete => None, + } + } else if !self.is_cleared { + self.table.get(key).cloned() + } else { + None + } + } + + /// Get the first key-value pair >= key, returning owned data + fn get_range_owned(&self, key: &MemStoreKey) -> Option<(MemStoreKey, Bytes)> { + let q = self.queued_ops.range(*key..).next(); + let c = if !self.is_cleared { self.table.range(*key..).next() } else { None }; + + match (q, c) { + (None, None) => None, + (Some((qk, queued)), Some((ck, current))) => { + if qk <= ck { + // Queued operation takes precedence + match queued { + QueuedKvOp::Put { value } => Some((*qk, value.clone())), + QueuedKvOp::Delete => { + // Skip deleted entry and look for next + let mut next_key = *qk; + for i in (0..next_key.len()).rev() { + if next_key[i] < u8::MAX { + next_key[i] += 1; + break; + } + next_key[i] = 0; + } + self.get_range_owned(&next_key) + } + } + } else { + Some((*ck, current.clone())) + } + } + (Some((qk, queued)), None) => match queued { + QueuedKvOp::Put { value } => Some((*qk, value.clone())), + QueuedKvOp::Delete => { + let mut next_key = *qk; + for i in (0..next_key.len()).rev() { + if next_key[i] < u8::MAX { + next_key[i] += 1; + break; + } + next_key[i] = 0; + } + self.get_range_owned(&next_key) + } + }, + (None, Some((ck, current))) => Some((*ck, current.clone())), + } + } + + /// Get the first key-value pair > key (strictly greater), returning owned data + fn get_range_exclusive_owned(&self, key: &MemStoreKey) -> Option<(MemStoreKey, Bytes)> { + use core::ops::Bound; + + let q = self.queued_ops.range((Bound::Excluded(*key), Bound::Unbounded)).next(); + let c = if !self.is_cleared { + self.table.range((Bound::Excluded(*key), Bound::Unbounded)).next() + } else { + None + }; + + match (q, c) { + (None, None) => None, + (Some((qk, queued)), Some((ck, current))) => { + if qk <= ck { + // Queued operation takes precedence + match queued { + QueuedKvOp::Put { value } => Some((*qk, value.clone())), + QueuedKvOp::Delete => { + // This key is deleted, recurse to find the next one + self.get_range_exclusive_owned(qk) + } + } + } else { + // Check if the current key has a delete queued + if let Some(QueuedKvOp::Delete) = self.queued_ops.get(ck) { + self.get_range_exclusive_owned(ck) + } else { + Some((*ck, current.clone())) + } + } + } + (Some((qk, queued)), None) => match queued { + QueuedKvOp::Put { value } => Some((*qk, value.clone())), + QueuedKvOp::Delete => self.get_range_exclusive_owned(qk), + }, + (None, Some((ck, current))) => { + // Check if the current key has a delete queued + if let Some(QueuedKvOp::Delete) = self.queued_ops.get(ck) { + self.get_range_exclusive_owned(ck) + } else { + Some((*ck, current.clone())) + } + } + } + } + + /// Get the last key-value pair < key, returning owned data + fn get_range_reverse_owned(&self, key: &MemStoreKey) -> Option<(MemStoreKey, Bytes)> { + let q = self.queued_ops.range(..*key).next_back(); + let c = if !self.is_cleared { self.table.range(..*key).next_back() } else { None }; + + match (q, c) { + (None, None) => None, + (Some((qk, queued)), Some((ck, current))) => { + if qk >= ck { + // Queued operation takes precedence + match queued { + QueuedKvOp::Put { value } => Some((*qk, value.clone())), + QueuedKvOp::Delete => self.get_range_reverse_owned(qk), + } + } else { + Some((*ck, current.clone())) + } + } + (Some((qk, queued)), None) => match queued { + QueuedKvOp::Put { value } => Some((*qk, value.clone())), + QueuedKvOp::Delete => self.get_range_reverse_owned(qk), + }, + (None, Some((ck, current))) => Some((*ck, current.clone())), + } + } +} + +impl<'a> KvTraverse for MemKvCursorMut<'a> { + fn first<'b>(&'b mut self) -> Result>, MemKvError> { + let start_key = [0u8; MAX_KEY_SIZE * 2]; + + // Get the first effective key-value pair + if let Some((key, value)) = self.get_range_owned(&start_key) { + self.current_key = Some(key); + Ok(Some((Cow::Owned(key.to_vec()), Cow::Owned(value.to_vec())))) + } else { + self.current_key = None; + Ok(None) + } + } + + fn last<'b>(&'b mut self) -> Result>, MemKvError> { + let end_key = [0xffu8; MAX_KEY_SIZE * 2]; + + if let Some((key, value)) = self.get_range_reverse_owned(&end_key) { + self.current_key = Some(key); + Ok(Some((Cow::Owned(key.to_vec()), Cow::Owned(value.to_vec())))) + } else { + self.current_key = None; + Ok(None) + } + } + + fn exact<'b>(&'b mut self, key: &[u8]) -> Result>, MemKvError> { + let search_key = MemKv::key(key); + self.current_key = Some(search_key); + + if let Some(value) = self.get_owned(&search_key) { + Ok(Some(Cow::Owned(value.to_vec()))) + } else { + Ok(None) + } + } + + fn lower_bound<'b>(&'b mut self, key: &[u8]) -> Result>, MemKvError> { + let search_key = MemKv::key(key); + + if let Some((found_key, value)) = self.get_range_owned(&search_key) { + self.current_key = Some(found_key); + Ok(Some((Cow::Owned(found_key.to_vec()), Cow::Owned(value.to_vec())))) + } else { + self.current_key = None; + Ok(None) + } + } + + fn read_next<'b>(&'b mut self) -> Result>, MemKvError> { + let current = self.current_key(); + + // Use exclusive range to find strictly greater than current key + if let Some((found_key, value)) = self.get_range_exclusive_owned(¤t) { + self.current_key = Some(found_key); + Ok(Some((Cow::Owned(found_key.to_vec()), Cow::Owned(value.to_vec())))) + } else { + self.current_key = None; + Ok(None) + } + } + + fn read_prev<'b>(&'b mut self) -> Result>, MemKvError> { + let current = self.current_key(); + + if let Some((found_key, value)) = self.get_range_reverse_owned(¤t) { + self.current_key = Some(found_key); + Ok(Some((Cow::Owned(found_key.to_vec()), Cow::Owned(value.to_vec())))) + } else { + self.current_key = None; + Ok(None) + } + } +} + +impl<'a> KvTraverseMut for MemKvCursorMut<'a> { + fn delete_current(&mut self) -> Result<(), MemKvError> { + if let Some(key) = self.current_key { + // Queue a delete operation + self.queued_ops.insert(key, QueuedKvOp::Delete); + Ok(()) + } else { + Err(MemKvError::HotKv(HotKvError::Inner("No current key to delete".into()))) + } + } +} + +impl<'a> DualKeyedTraverse for MemKvCursorMut<'a> { + fn exact_dual<'b>( + &'b mut self, + key1: &[u8], + key2: &[u8], + ) -> Result>, MemKvError> { + let combined_key = MemKv::dual_key(key1, key2); + KvTraverse::exact(self, &combined_key) + } + + fn next_dual_above<'b>( + &'b mut self, + key1: &[u8], + key2: &[u8], + ) -> Result>, MemKvError> { + let combined_key = MemKv::dual_key(key1, key2); + let Some((found_key, value)) = KvTraverse::lower_bound(self, &combined_key)? else { + return Ok(None); + }; + + let (key1, key2) = MemKv::split_dual_key(found_key.as_ref()); + Ok(Some((key1, key2, value))) + } + + fn next_k1<'b>(&'b mut self) -> Result>, MemKvError> { + // scan forward until finding a new k1 + let last_k1 = self.current_k1(); + + DualKeyedTraverse::next_dual_above(self, &last_k1, &[0xffu8; MAX_KEY_SIZE]) + } + + fn next_k2<'b>(&'b mut self) -> Result>, MemKvError> { + let current_key = self.current_key(); + let (current_k1, current_k2) = MemKv::split_dual_key(¤t_key); + + // scan forward until finding a new k2 for the same k1 + DualKeyedTraverse::next_dual_above(self, ¤t_k1, ¤t_k2) + } +} + +// Implement DualTableTraverse for typed dual-keyed table access +impl<'a, T> DualTableTraverse for MemKvCursorMut<'a> +where + T: DualKeyed, +{ + fn next_dual_above( + &mut self, + key1: &T::Key, + key2: &T::Key2, + ) -> Result>, MemKvError> { + let mut key1_buf = [0u8; MAX_KEY_SIZE]; + let mut key2_buf = [0u8; MAX_KEY_SIZE]; + let key1_bytes = key1.encode_key(&mut key1_buf); + let key2_bytes = key2.encode_key(&mut key2_buf); + + DualKeyedTraverse::next_dual_above(self, key1_bytes, key2_bytes)? + .map(T::decode_kkv_tuple) + .transpose() + .map_err(Into::into) + } + + fn next_k1(&mut self) -> Result>, MemKvError> { + DualKeyedTraverse::next_k1(self)?.map(T::decode_kkv_tuple).transpose().map_err(Into::into) + } + + fn next_k2(&mut self) -> Result>, MemKvError> { + DualKeyedTraverse::next_k2(self)?.map(T::decode_kkv_tuple).transpose().map_err(Into::into) + } +} + +impl HotKv for MemKv { + type RoTx = MemKvRoTx; + type RwTx = MemKvRwTx; + + fn reader(&self) -> Result { + let guard = self + .map + .try_read() + .map_err(|_| HotKvError::Inner("Failed to acquire read lock".into()))?; + + // SAFETY: This is safe-ish, as we ensure the map is not dropped until + // the guard is also dropped. + let guard: RwLockReadGuard<'static, Store> = unsafe { std::mem::transmute(guard) }; + + Ok(MemKvRoTx { guard, _store: self.map.clone() }) + } + + fn writer(&self) -> Result { + let guard = self.map.try_write().map_err(|_| HotKvError::WriteLocked)?; + + // SAFETY: This is safe-ish, as we ensure the map is not dropped until + // the guard is also dropped. + let guard: RwLockWriteGuard<'static, Store> = unsafe { std::mem::transmute(guard) }; + + Ok(MemKvRwTx { guard, _store: self.map.clone(), queued_ops: OpStore::new() }) + } +} + +impl HotKvRead for MemKvRoTx { + type Error = MemKvError; + + type Traverse<'a> = MemKvCursor<'a>; + + fn raw_traverse<'a>(&'a self, table: &str) -> Result, Self::Error> { + let table_data = self.guard.get(table).unwrap_or(&EMPTY_TABLE); + Ok(MemKvCursor::new(table_data)) + } + + fn raw_get<'a>( + &'a self, + table: &str, + key: &[u8], + ) -> Result>, Self::Error> { + // Check queued operations first (read-your-writes consistency) + let key = MemKv::key(key); + + // SAFETY: The guard ensures the map remains valid + + Ok(self + .guard + .get(table) + .and_then(|t| t.get(&key)) + .map(|bytes| Cow::Borrowed(bytes.as_ref()))) + } + + fn raw_get_dual<'a>( + &'a self, + table: &str, + key1: &[u8], + key2: &[u8], + ) -> Result>, Self::Error> { + let key = MemKv::dual_key(key1, key2); + self.raw_get(table, &key) + } +} + +static EMPTY_TABLE: StoreTable = BTreeMap::new(); + +impl MemKvRoTx { + /// Get a cursor for the specified table + pub fn cursor<'a>(&'a self, table: &str) -> Result, MemKvError> { + let table_data = self.guard.get(table).unwrap_or(&EMPTY_TABLE); + Ok(MemKvCursor::new(table_data)) + } +} + +impl HotKvRead for MemKvRwTx { + type Error = MemKvError; + + type Traverse<'a> = MemKvCursor<'a>; + + fn raw_traverse<'a>(&'a self, table: &str) -> Result, Self::Error> { + let table_data = self.guard.get(table).unwrap_or(&EMPTY_TABLE); + Ok(MemKvCursor::new(table_data)) + } + + fn raw_get<'a>( + &'a self, + table: &str, + key: &[u8], + ) -> Result>, Self::Error> { + // Check queued operations first (read-your-writes consistency) + let key = MemKv::key(key); + + if let Some(table) = self.queued_ops.get(table) { + if table.is_clear() { + return Ok(None); + } + + match table.get(&key) { + Some(QueuedKvOp::Put { value }) => { + return Ok(Some(Cow::Borrowed(value.as_ref()))); + } + Some(QueuedKvOp::Delete) => { + return Ok(None); + } + None => {} + } + } + + // If not found in queued ops, check the underlying map + Ok(self + .guard + .get(table) + .and_then(|t| t.get(&key)) + .map(|bytes| Cow::Borrowed(bytes.as_ref()))) + } + + fn raw_get_dual<'a>( + &'a self, + table: &str, + key1: &[u8], + key2: &[u8], + ) -> Result>, Self::Error> { + let key = MemKv::dual_key(key1, key2); + self.raw_get(table, &key) + } +} + +impl MemKvRwTx { + /// Get a read-only cursor for the specified table + /// Note: This cursor will NOT see pending writes from this transaction + pub fn cursor<'a>(&'a self, table: &str) -> Result, MemKvError> { + if let Some(table_data) = self.guard.get(table) { + Ok(MemKvCursor::new(table_data)) + } else { + Err(MemKvError::HotKv(HotKvError::Inner(format!("Table '{}' not found", table).into()))) + } + } + + /// Get a mutable cursor for the specified table + /// This cursor will see both committed data and pending writes from this transaction + pub fn cursor_mut<'a>(&'a mut self, table: &str) -> Result, MemKvError> { + // Get or create the table data + let table_data = self.guard.entry(table.to_owned()).or_default(); + + // Get or create the queued operations for this table + let table_ops = self.queued_ops.entry(table.to_owned()).or_default(); + + let is_cleared = table_ops.is_clear(); + + // Extract the inner TableOp from QueuedTableOp + let ops = match table_ops { + QueuedTableOp::Modify { ops } => ops, + QueuedTableOp::Clear { new_table } => new_table, + }; + + Ok(MemKvCursorMut::new(table_data, ops, is_cleared)) + } +} + +impl HotKvWrite for MemKvRwTx { + type TraverseMut<'a> + = MemKvCursorMut<'a> + where + Self: 'a; + + fn raw_traverse_mut<'a>( + &'a mut self, + table: &str, + ) -> Result, Self::Error> { + self.cursor_mut(table) + } + + fn queue_raw_put(&mut self, table: &str, key: &[u8], value: &[u8]) -> Result<(), Self::Error> { + let key = MemKv::key(key); + + let value_bytes = Bytes::copy_from_slice(value); + + self.queued_ops + .entry(table.to_owned()) + .or_default() + .put(key, QueuedKvOp::Put { value: value_bytes }); + Ok(()) + } + + fn queue_raw_put_dual( + &mut self, + table: &str, + key1: &[u8], + key2: &[u8], + value: &[u8], + ) -> Result<(), Self::Error> { + let key = MemKv::dual_key(key1, key2); + self.queue_raw_put(table, &key, value) + } + + fn queue_raw_delete(&mut self, table: &str, key: &[u8]) -> Result<(), Self::Error> { + let key = MemKv::key(key); + + self.queued_ops.entry(table.to_owned()).or_default().delete(key); + Ok(()) + } + + fn queue_raw_clear(&mut self, table: &str) -> Result<(), Self::Error> { + self.queued_ops + .insert(table.to_owned(), QueuedTableOp::Clear { new_table: TableOp::new() }); + Ok(()) + } + + fn queue_raw_create( + &mut self, + _table: &str, + _dual_key: bool, + _dual_fixed: bool, + ) -> Result<(), Self::Error> { + Ok(()) + } + + fn raw_commit(mut self) -> Result<(), Self::Error> { + // Apply all queued operations to the map + self.commit_inner(); + + // The write guard is automatically dropped here, releasing the lock + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{ + hot::model::{DualTableTraverse, TableTraverse, TableTraverseMut}, + tables::{SingleKey, Table}, + }; + use alloy::primitives::{Address, U256}; + use bytes::Bytes; + + // Test table definitions + #[derive(Debug)] + struct TestTable; + + impl SingleKey for TestTable {} + + impl Table for TestTable { + const NAME: &'static str = "test_table"; + + type Key = u64; + type Value = Bytes; + } + + #[derive(Debug)] + struct AddressTable; + + impl Table for AddressTable { + const NAME: &'static str = "addresses"; + type Key = Address; + type Value = U256; + } + + impl SingleKey for AddressTable {} + + #[derive(Debug)] + struct DualTestTable; + + impl Table for DualTestTable { + const NAME: &'static str = "dual_test_table"; + type Key = u64; + type Value = Bytes; + } + + impl crate::tables::DualKeyed for DualTestTable { + type Key2 = u32; + } + + #[test] + fn test_new_store() { + let store = MemKv::new(); + let reader = store.reader().unwrap(); + + // Empty store should return None for any key + assert!(reader.raw_get("test", &[1, 2, 3]).unwrap().is_none()); + } + + #[test] + fn test_basic_put_get() { + let store = MemKv::new(); + + // Write some data + { + let mut writer = store.writer().unwrap(); + writer.queue_raw_put("table1", &[1, 2, 3], b"value1").unwrap(); + writer.queue_raw_put("table1", &[4, 5, 6], b"value2").unwrap(); + writer.raw_commit().unwrap(); + } + + // Read the data back + { + let reader = store.reader().unwrap(); + let value1 = reader.raw_get("table1", &[1, 2, 3]).unwrap(); + let value2 = reader.raw_get("table1", &[4, 5, 6]).unwrap(); + let missing = reader.raw_get("table1", &[7, 8, 9]).unwrap(); + + assert_eq!(value1.as_deref(), Some(b"value1" as &[u8])); + assert_eq!(value2.as_deref(), Some(b"value2" as &[u8])); + assert!(missing.is_none()); + } + } + + #[test] + fn test_multiple_tables() { + let store = MemKv::new(); + + // Write to different tables + { + let mut writer = store.writer().unwrap(); + writer.queue_raw_put("table1", &[1], b"table1_value").unwrap(); + writer.queue_raw_put("table2", &[1], b"table2_value").unwrap(); + writer.raw_commit().unwrap(); + } + + // Read from different tables + { + let reader = store.reader().unwrap(); + let value1 = reader.raw_get("table1", &[1]).unwrap(); + let value2 = reader.raw_get("table2", &[1]).unwrap(); + + assert_eq!(value1.as_deref(), Some(b"table1_value" as &[u8])); + assert_eq!(value2.as_deref(), Some(b"table2_value" as &[u8])); + } + } + + #[test] + fn test_overwrite_value() { + let store = MemKv::new(); + + // Write initial value + { + let mut writer = store.writer().unwrap(); + writer.queue_raw_put("table1", &[1], b"original").unwrap(); + writer.raw_commit().unwrap(); + } + + // Overwrite with new value + { + let mut writer = store.writer().unwrap(); + writer.queue_raw_put("table1", &[1], b"updated").unwrap(); + writer.raw_commit().unwrap(); + } + + // Check the value was updated + { + let reader = store.reader().unwrap(); + let value = reader.raw_get("table1", &[1]).unwrap(); + assert_eq!(value.as_deref(), Some(b"updated" as &[u8])); + } + } + + #[test] + fn test_read_your_writes() { + let store = MemKv::new(); + let mut writer = store.writer().unwrap(); + + // Queue some operations but don't commit yet + writer.queue_raw_put("table1", &[1], b"queued_value").unwrap(); + + // Should be able to read the queued value + let value = writer.raw_get("table1", &[1]).unwrap(); + assert_eq!(value.as_deref(), Some(b"queued_value" as &[u8])); + + writer.raw_commit().unwrap(); + + // After commit, other readers should see it + { + let reader = store.reader().unwrap(); + let value = reader.raw_get("table1", &[1]).unwrap(); + assert_eq!(value.as_deref(), Some(b"queued_value" as &[u8])); + } + } + + #[test] + fn test_typed_operations() { + let store = MemKv::new(); + + // Write using typed interface + { + let mut writer = store.writer().unwrap(); + writer.queue_put::(&42u64, &Bytes::from_static(b"hello world")).unwrap(); + writer.queue_put::(&100u64, &Bytes::from_static(b"another value")).unwrap(); + writer.raw_commit().unwrap(); + } + + // Read using typed interface + { + let reader = store.reader().unwrap(); + let value1 = reader.get::(&42u64).unwrap(); + let value2 = reader.get::(&100u64).unwrap(); + let missing = reader.get::(&999u64).unwrap(); + + assert_eq!(value1, Some(Bytes::from_static(b"hello world"))); + assert_eq!(value2, Some(Bytes::from_static(b"another value"))); + assert!(missing.is_none()); + } + } + + #[test] + fn test_address_table() { + let store = MemKv::new(); + + let addr1 = Address::from([0x11; 20]); + let addr2 = Address::from([0x22; 20]); + let balance1 = U256::from(1000u64); + let balance2 = U256::from(2000u64); + + // Write address data + { + let mut writer = store.writer().unwrap(); + writer.queue_put::(&addr1, &balance1).unwrap(); + writer.queue_put::(&addr2, &balance2).unwrap(); + writer.raw_commit().unwrap(); + } + + // Read address data + { + let reader = store.reader().unwrap(); + let bal1 = reader.get::(&addr1).unwrap(); + let bal2 = reader.get::(&addr2).unwrap(); + + assert_eq!(bal1, Some(balance1)); + assert_eq!(bal2, Some(balance2)); + } + } + + #[test] + fn test_batch_operations() { + let store = MemKv::new(); + + let entries = [ + (1u64, Bytes::from_static(b"first")), + (2u64, Bytes::from_static(b"second")), + (3u64, Bytes::from_static(b"third")), + ]; + + // Write batch + { + let mut writer = store.writer().unwrap(); + let entry_refs: Vec<_> = entries.iter().map(|(k, v)| (k, v)).collect(); + writer.queue_put_many::(entry_refs).unwrap(); + writer.raw_commit().unwrap(); + } + + // Read batch + { + let reader = store.reader().unwrap(); + let keys: Vec<_> = entries.iter().map(|(k, _)| k).collect(); + let values = reader.get_many::(keys).unwrap(); + + assert_eq!(values.len(), 3); + assert_eq!(values[0], (&1u64, Some(Bytes::from_static(b"first")))); + assert_eq!(values[1], (&2u64, Some(Bytes::from_static(b"second")))); + assert_eq!(values[2], (&3u64, Some(Bytes::from_static(b"third")))); + } + } + + #[test] + fn test_concurrent_readers() { + let store = MemKv::new(); + + // Write some initial data + { + let mut writer = store.writer().unwrap(); + writer.queue_raw_put("table1", &[1], b"value1").unwrap(); + writer.raw_commit().unwrap(); + } + + // Multiple readers should be able to read concurrently + let reader1 = store.reader().unwrap(); + let reader2 = store.reader().unwrap(); + + let value1 = reader1.raw_get("table1", &[1]).unwrap(); + let value2 = reader2.raw_get("table1", &[1]).unwrap(); + + assert_eq!(value1.as_deref(), Some(b"value1" as &[u8])); + assert_eq!(value2.as_deref(), Some(b"value1" as &[u8])); + } + + #[test] + fn test_write_lock_exclusivity() { + let store = MemKv::new(); + + // Get a writer + let _writer1 = store.writer().unwrap(); + + // Second writer should fail + match store.writer() { + Err(HotKvError::WriteLocked) => {} // Expected + Ok(_) => panic!("Should not be able to get second writer"), + Err(e) => panic!("Unexpected error: {:?}", e), + } + } + + #[test] + fn test_empty_values() { + let store = MemKv::new(); + + { + let mut writer = store.writer().unwrap(); + writer.queue_raw_put("table1", &[1], b"").unwrap(); + writer.raw_commit().unwrap(); + } + + { + let reader = store.reader().unwrap(); + let value = reader.raw_get("table1", &[1]).unwrap(); + assert_eq!(value.as_deref(), Some(b"" as &[u8])); + } + } + + #[test] + fn test_multiple_operations_same_transaction() { + let store = MemKv::new(); + + { + let mut writer = store.writer().unwrap(); + + // Multiple operations on same key - last one should win + writer.queue_raw_put("table1", &[1], b"first").unwrap(); + writer.queue_raw_put("table1", &[1], b"second").unwrap(); + writer.queue_raw_put("table1", &[1], b"third").unwrap(); + + // Read-your-writes should return the latest value + let value = writer.raw_get("table1", &[1]).unwrap(); + assert_eq!(value.as_deref(), Some(b"third" as &[u8])); + + writer.raw_commit().unwrap(); + } + + { + let reader = store.reader().unwrap(); + let value = reader.raw_get("table1", &[1]).unwrap(); + assert_eq!(value.as_deref(), Some(b"third" as &[u8])); + } + } + + #[test] + fn test_isolation() { + let store = MemKv::new(); + + // Write initial value + { + let mut writer = store.writer().unwrap(); + writer.queue_raw_put("table1", &[1], b"original").unwrap(); + writer.raw_commit().unwrap(); + } + + // Start a read transaction + { + let reader = store.reader().unwrap(); + let original_value = reader.raw_get("table1", &[1]).unwrap(); + assert_eq!(original_value.as_deref(), Some(b"original" as &[u8])); + } + + // Update the value in a separate transaction + { + let mut writer = store.writer().unwrap(); + writer.queue_raw_put("table1", &[1], b"updated").unwrap(); + writer.raw_commit().unwrap(); + } + + // The value should now be latest for new readers + { + // New reader should see the updated value + let new_reader = store.reader().unwrap(); + let updated_value = new_reader.raw_get("table1", &[1]).unwrap(); + assert_eq!(updated_value.as_deref(), Some(b"updated" as &[u8])); + } + } + + #[test] + fn test_rollback_on_drop() { + let store = MemKv::new(); + + { + let mut writer = store.writer().unwrap(); + writer.queue_raw_put("table1", &[1], b"should_not_persist").unwrap(); + // Drop without committing + } + + // Value should not be persisted + { + let reader = store.reader().unwrap(); + let value = reader.raw_get("table1", &[1]).unwrap(); + assert!(value.is_none()); + } + } + + #[test] + fn write_two_tables() { + let store = MemKv::new(); + + { + let mut writer = store.writer().unwrap(); + writer.queue_raw_put("table1", &[1], b"value1").unwrap(); + writer.queue_raw_put("table2", &[2], b"value2").unwrap(); + writer.raw_commit().unwrap(); + } + + { + let reader = store.reader().unwrap(); + let value1 = reader.raw_get("table1", &[1]).unwrap(); + let value2 = reader.raw_get("table2", &[2]).unwrap(); + + assert_eq!(value1.as_deref(), Some(b"value1" as &[u8])); + assert_eq!(value2.as_deref(), Some(b"value2" as &[u8])); + } + } + + #[test] + fn test_downgrades() { + let store = MemKv::new(); + { + // Write some data + // Start a read-write transaction + let mut rw_tx = store.writer().unwrap(); + rw_tx.queue_raw_put("table1", &[1, 2, 3], b"value1").unwrap(); + rw_tx.queue_raw_put("table1", &[4, 5, 6], b"value2").unwrap(); + + let ro_tx = rw_tx.commit_downgrade(); + + // Read the data back + let value1 = ro_tx.raw_get("table1", &[1, 2, 3]).unwrap(); + let value2 = ro_tx.raw_get("table1", &[4, 5, 6]).unwrap(); + + assert_eq!(value1.as_deref(), Some(b"value1" as &[u8])); + assert_eq!(value2.as_deref(), Some(b"value2" as &[u8])); + } + + { + // Start another read-write transaction + let mut rw_tx = store.writer().unwrap(); + rw_tx.queue_raw_put("table2", &[7, 8, 9], b"value3").unwrap(); + + // Value should not be set + let ro_tx = rw_tx.downgrade(); + + // Read the data back + let value3 = ro_tx.raw_get("table2", &[7, 8, 9]).unwrap(); + + assert!(value3.is_none()); + } + } + + #[test] + fn test_clear_table() { + let store = MemKv::new(); + + { + let mut writer = store.writer().unwrap(); + writer.queue_raw_put("table1", &[1], b"value1").unwrap(); + writer.queue_raw_put("table1", &[2], b"value2").unwrap(); + writer.raw_commit().unwrap(); + } + + { + let reader = store.reader().unwrap(); + + let value1 = reader.raw_get("table1", &[1]).unwrap(); + let value2 = reader.raw_get("table1", &[2]).unwrap(); + + assert_eq!(value1.as_deref(), Some(b"value1" as &[u8])); + assert_eq!(value2.as_deref(), Some(b"value2" as &[u8])); + } + + { + let mut writer = store.writer().unwrap(); + + let value1 = writer.raw_get("table1", &[1]).unwrap(); + let value2 = writer.raw_get("table1", &[2]).unwrap(); + + assert_eq!(value1.as_deref(), Some(b"value1" as &[u8])); + assert_eq!(value2.as_deref(), Some(b"value2" as &[u8])); + + writer.queue_raw_clear("table1").unwrap(); + + let value1 = writer.raw_get("table1", &[1]).unwrap(); + let value2 = writer.raw_get("table1", &[2]).unwrap(); + + assert!(value1.is_none()); + assert!(value2.is_none()); + + writer.raw_commit().unwrap(); + } + + { + let reader = store.reader().unwrap(); + let value1 = reader.raw_get("table1", &[1]).unwrap(); + let value2 = reader.raw_get("table1", &[2]).unwrap(); + + assert!(value1.is_none()); + assert!(value2.is_none()); + } + } + + // ======================================================================== + // Cursor Traversal Tests + // ======================================================================== + + #[test] + fn test_cursor_basic_navigation() { + let store = MemKv::new(); + + // Setup test data using TestTable + let test_data = vec![ + (1u64, Bytes::from_static(b"value_001")), + (2u64, Bytes::from_static(b"value_002")), + (3u64, Bytes::from_static(b"value_003")), + (10u64, Bytes::from_static(b"value_010")), + (20u64, Bytes::from_static(b"value_020")), + ]; + + // Insert data + { + let mut writer = store.writer().unwrap(); + for (key, value) in &test_data { + writer.queue_put::(key, value).unwrap(); + } + writer.raw_commit().unwrap(); + } + + // Test cursor navigation + { + let reader = store.reader().unwrap(); + let mut cursor = reader.cursor(TestTable::NAME).unwrap(); + + // Test first() + let (key, value) = TableTraverse::::first(&mut cursor).unwrap().unwrap(); + assert_eq!(key, test_data[0].0); + assert_eq!(value, test_data[0].1); + + // Test last() + let last_result = TableTraverse::::last(&mut cursor).unwrap(); + assert!(last_result.is_some()); + let (key, value) = last_result.unwrap(); + assert_eq!(key, test_data.last().unwrap().0); + assert_eq!(value, test_data.last().unwrap().1); + + // Test exact lookup + let exact_result = TableTraverse::::exact(&mut cursor, &2u64).unwrap(); + assert!(exact_result.is_some()); + assert_eq!(exact_result.unwrap(), test_data[1].1); + + // Test next_above (range lookup) + let range_result = + TableTraverse::::lower_bound(&mut cursor, &5u64).unwrap(); + assert!(range_result.is_some()); + let (key, value) = range_result.unwrap(); + assert_eq!(key, test_data[3].0); // 10u64 + assert_eq!(value, test_data[3].1); + } + } + + #[test] + fn test_cursor_sequential_navigation() { + let store = MemKv::new(); + + // Setup sequential test data using TestTable + let test_data: Vec<(u64, Bytes)> = (1..=5) + .map(|i| { + let key = i; + let value = Bytes::from(format!("value_{:03}", i)); + (key, value) + }) + .collect(); + + // Insert data + { + let mut writer = store.writer().unwrap(); + for (key, value) in &test_data { + writer.queue_put::(key, value).unwrap(); + } + writer.raw_commit().unwrap(); + } + + // Test sequential navigation + { + let reader = store.reader().unwrap(); + let mut cursor = reader.cursor(TestTable::NAME).unwrap(); + + // Start from first and traverse forward + let mut current_idx = 0; + let first_result = TableTraverse::::first(&mut cursor).unwrap(); + assert!(first_result.is_some()); + + let (key, value) = first_result.unwrap(); + assert_eq!(key, test_data[current_idx].0); + assert_eq!(value, test_data[current_idx].1); + + // Navigate forward through all entries + while current_idx < test_data.len() - 1 { + let next_result = TableTraverse::::read_next(&mut cursor).unwrap(); + assert!(next_result.is_some()); + + current_idx += 1; + let (key, value) = next_result.unwrap(); + assert_eq!(key, test_data[current_idx].0); + assert_eq!(value, test_data[current_idx].1); + } + + // Next should return None at the end + let beyond_end = TableTraverse::::read_next(&mut cursor).unwrap(); + assert!(beyond_end.is_none()); + + // Navigate backward + while current_idx > 0 { + let prev_result = TableTraverse::::read_prev(&mut cursor).unwrap(); + + assert!(prev_result.is_some()); + + current_idx -= 1; + let (key, value) = prev_result.unwrap(); + assert_eq!(key, test_data[current_idx].0); + assert_eq!(value, test_data[current_idx].1); + } + + // Previous should return None at the beginning + let before_start = TableTraverse::::read_prev(&mut cursor).unwrap(); + assert!(before_start.is_none()); + } + } + + #[test] + fn test_cursor_mut_operations() { + let store = MemKv::new(); + + let test_data = vec![ + (1u64, Bytes::from_static(b"delete_value_1")), + (2u64, Bytes::from_static(b"delete_value_2")), + (3u64, Bytes::from_static(b"delete_value_3")), + ]; + + // Insert initial data + { + let mut writer = store.writer().unwrap(); + for (key, value) in &test_data { + writer.queue_put::(key, value).unwrap(); + } + writer.raw_commit().unwrap(); + } + + // Test mutable cursor operations + { + let mut writer = store.writer().unwrap(); + let mut cursor = writer.cursor_mut(TestTable::NAME).unwrap(); + + // Navigate to middle entry + let first = TableTraverse::::first(&mut cursor).unwrap().unwrap(); + assert_eq!(first.0, test_data[0].0); + + let next = TableTraverse::::read_next(&mut cursor).unwrap().unwrap(); + assert_eq!(next.0, test_data[1].0); + + // Delete current entry (key 2) + TableTraverseMut::::delete_current(&mut cursor).unwrap(); + + writer.raw_commit().unwrap(); + } + + // Verify deletion + { + let reader = store.reader().unwrap(); + let mut cursor = reader.cursor(TestTable::NAME).unwrap(); + + // Should only have first and third entries + let first = TableTraverse::::first(&mut cursor).unwrap().unwrap(); + assert_eq!(first.0, test_data[0].0); + + let second = TableTraverse::::read_next(&mut cursor).unwrap().unwrap(); + assert_eq!(second.0, test_data[2].0); + + // Should be no more entries + let none = TableTraverse::::read_next(&mut cursor).unwrap(); + assert!(none.is_none()); + + // Verify deleted key is gone + let missing = + TableTraverse::::exact(&mut cursor, &test_data[1].0).unwrap(); + assert!(missing.is_none()); + } + } + + #[test] + fn test_table_traverse_typed() { + let store = MemKv::new(); + + // Setup test data using the test table + let test_data: Vec<(u64, bytes::Bytes)> = (0..5) + .map(|i| { + let key = i * 10; + let value = bytes::Bytes::from(format!("test_value_{}", i)); + (key, value) + }) + .collect(); + + // Insert data + { + let mut writer = store.writer().unwrap(); + for (key, value) in &test_data { + writer.queue_put::(key, value).unwrap(); + } + writer.raw_commit().unwrap(); + } + + // Test typed table traversal + { + let reader = store.reader().unwrap(); + let mut cursor = reader.cursor(TestTable::NAME).unwrap(); + + // Test first with type-safe operations + let first_raw = TableTraverse::::first(&mut cursor).unwrap(); + assert!(first_raw.is_some()); + let (first_key, first_value) = first_raw.unwrap(); + assert_eq!(first_key, test_data[0].0); + assert_eq!(first_value, test_data[0].1); + + // Test last + let last_raw = TableTraverse::::last(&mut cursor).unwrap(); + assert!(last_raw.is_some()); + let (last_key, last_value) = last_raw.unwrap(); + assert_eq!(last_key, test_data.last().unwrap().0); + assert_eq!(last_value, test_data.last().unwrap().1); + + // Test exact lookup + let target_key = &test_data[2].0; + let exact_value = + TableTraverse::::exact(&mut cursor, target_key).unwrap(); + assert!(exact_value.is_some()); + assert_eq!(exact_value.unwrap(), test_data[2].1); + + // Test range lookup + let range_key = 15u64; // Between entries 1 and 2 + + let range_result = + TableTraverse::::lower_bound(&mut cursor, &range_key).unwrap(); + assert!(range_result.is_some()); + let (found_key, found_value) = range_result.unwrap(); + assert_eq!(found_key, test_data[2].0); // key 20 + assert_eq!(found_value, test_data[2].1); + } + } + + #[test] + fn test_cursor_empty_table() { + let store = MemKv::new(); + + // Create an empty table first + { + let mut writer = store.writer().unwrap(); + writer.queue_raw_create(TestTable::NAME, false, false).unwrap(); + writer.raw_commit().unwrap(); + } + + { + let reader = store.reader().unwrap(); + let mut cursor = reader.cursor(TestTable::NAME).unwrap(); + + // All operations should return None on empty table + assert!(TableTraverse::::first(&mut cursor).unwrap().is_none()); + assert!(TableTraverse::::last(&mut cursor).unwrap().is_none()); + assert!(TableTraverse::::exact(&mut cursor, &42u64).unwrap().is_none()); + assert!( + TableTraverse::::lower_bound(&mut cursor, &42u64).unwrap().is_none() + ); + assert!(TableTraverse::::read_next(&mut cursor).unwrap().is_none()); + assert!(TableTraverse::::read_prev(&mut cursor).unwrap().is_none()); + } + } + + #[test] + fn test_cursor_state_management() { + let store = MemKv::new(); + + let test_data = vec![ + (1u64, Bytes::from_static(b"state_value_1")), + (2u64, Bytes::from_static(b"state_value_2")), + (3u64, Bytes::from_static(b"state_value_3")), + ]; + + { + let mut writer = store.writer().unwrap(); + for (key, value) in &test_data { + writer.queue_put::(key, value).unwrap(); + } + writer.raw_commit().unwrap(); + } + + { + let reader = store.reader().unwrap(); + let mut cursor = reader.cursor(TestTable::NAME).unwrap(); + + // Test that cursor operations maintain state correctly + + // Start at first + let first = TableTraverse::::first(&mut cursor).unwrap().unwrap(); + assert_eq!(first.0, test_data[0].0); + + // Move to second via next + let second = TableTraverse::::read_next(&mut cursor).unwrap().unwrap(); + assert_eq!(second.0, test_data[1].0); + + // Jump to last + let last = TableTraverse::::last(&mut cursor).unwrap().unwrap(); + assert_eq!(last.0, test_data[2].0); + + // Move back via prev + let back_to_second = + TableTraverse::::read_prev(&mut cursor).unwrap().unwrap(); + assert_eq!(back_to_second.0, test_data[1].0); + + // Use exact to jump to specific position + let exact_first = + TableTraverse::::exact(&mut cursor, &test_data[0].0).unwrap(); + assert!(exact_first.is_some()); + assert_eq!(exact_first.unwrap(), test_data[0].1); + + // Verify cursor is now positioned at first entry + let next_from_first = + TableTraverse::::read_next(&mut cursor).unwrap().unwrap(); + assert_eq!(next_from_first.0, test_data[1].0); + + // Use range lookup + let range_lookup = + TableTraverse::::lower_bound(&mut cursor, &1u64).unwrap().unwrap(); // Should find key 1 + assert_eq!(range_lookup.0, test_data[0].0); + + // Verify we can continue navigation from range position + let next_after_range = + TableTraverse::::read_next(&mut cursor).unwrap().unwrap(); + assert_eq!(next_after_range.0, test_data[1].0); + } + } + + #[test] + fn test_dual_key_operations() { + let store = MemKv::new(); + + // Test dual key storage and retrieval using DualTestTable + let dual_data = vec![ + (1u64, 100u32, Bytes::from_static(b"value1")), + (1u64, 200u32, Bytes::from_static(b"value2")), + (2u64, 100u32, Bytes::from_static(b"value3")), + ]; + + { + let mut writer = store.writer().unwrap(); + for (key1, key2, value) in &dual_data { + writer.queue_put_dual::(key1, key2, value).unwrap(); + } + writer.raw_commit().unwrap(); + } + + // Test dual key traversal + { + let reader = store.reader().unwrap(); + let mut cursor = reader.cursor(DualTestTable::NAME).unwrap(); + + // Test exact dual lookup + let exact_result = + DualTableTraverse::::exact_dual(&mut cursor, &1u64, &200u32) + .unwrap(); + assert!(exact_result.is_some()); + assert_eq!(exact_result.unwrap(), Bytes::from_static(b"value2")); + + // Test missing dual key + let missing_result = + DualTableTraverse::::exact_dual(&mut cursor, &3u64, &100u32) + .unwrap(); + assert!(missing_result.is_none()); + + // Test next_dual_above + let range_result = + DualTableTraverse::::next_dual_above(&mut cursor, &1u64, &150u32) + .unwrap(); + assert!(range_result.is_some()); + let (k, k2, value) = range_result.unwrap(); + assert_eq!(k, 1u64); + assert_eq!(k2, 200u32); + assert_eq!(value, Bytes::from_static(b"value2")); + + // Test next_k1 to find next different first key + let next_k1_result = + DualTableTraverse::::next_k1(&mut cursor).unwrap(); + assert!(next_k1_result.is_some()); + let (k, k2, value) = next_k1_result.unwrap(); + assert_eq!(k, 2u64); + assert_eq!(k2, 100u32); + assert_eq!(value, Bytes::from_static(b"value3")); + } + } +} diff --git a/crates/storage/src/hot/impls/mod.rs b/crates/storage/src/hot/impls/mod.rs new file mode 100644 index 0000000..7119d0a --- /dev/null +++ b/crates/storage/src/hot/impls/mod.rs @@ -0,0 +1,309 @@ +/// An in-memory key-value store implementation. +pub mod mem; + +/// MDBX-backed key-value store implementation. +pub mod mdbx; + +#[cfg(test)] +mod test { + use crate::{ + hot::{ + mem, + model::{HotDbRead, HotDbWrite, HotHistoryRead, HotHistoryWrite, HotKv, HotKvWrite}, + }, + tables::hot, + }; + use alloy::primitives::{B256, Bytes, U256, address, b256}; + use reth::primitives::{Account, Bytecode, Header, SealedHeader}; + use reth_db::BlockNumberList; + + #[test] + fn mem_conformance() { + let hot_kv = mem::MemKv::new(); + conformance(&hot_kv); + } + + #[test] + fn mdbx_conformance() { + let db = reth_db::test_utils::create_test_rw_db(); + + // Create tables from the `crate::tables::hot` module + let mut writer = db.db().writer().unwrap(); + + writer.queue_create::().unwrap(); + writer.queue_create::().unwrap(); + writer.queue_create::().unwrap(); + writer.queue_create::().unwrap(); + writer.queue_create::().unwrap(); + writer.queue_create::().unwrap(); + writer.queue_create::().unwrap(); + writer.queue_create::().unwrap(); + writer.queue_create::().unwrap(); + + writer.commit().expect("Failed to commit table creation"); + + conformance(db.db()); + } + + fn conformance(hot_kv: &T) { + test_header_roundtrip(hot_kv); + test_account_roundtrip(hot_kv); + test_storage_roundtrip(hot_kv); + test_bytecode_roundtrip(hot_kv); + test_account_history(hot_kv); + test_storage_history(hot_kv); + test_account_changes(hot_kv); + test_storage_changes(hot_kv); + test_missing_reads(hot_kv); + } + + /// Test writing and reading headers via HotDbWrite/HotDbRead + fn test_header_roundtrip(hot_kv: &T) { + let header = Header { number: 42, gas_limit: 1_000_000, ..Default::default() }; + let sealed = SealedHeader::seal_slow(header.clone()); + let hash = sealed.hash(); + + // Write header + { + let mut writer = hot_kv.writer().unwrap(); + writer.put_header(&sealed).unwrap(); + writer.commit().unwrap(); + } + + // Read header by number + { + let reader = hot_kv.reader().unwrap(); + let read_header = reader.get_header(42).unwrap(); + assert!(read_header.is_some()); + assert_eq!(read_header.unwrap().number, 42); + } + + // Read header number by hash + { + let reader = hot_kv.reader().unwrap(); + let read_number = reader.get_header_number(&hash).unwrap(); + assert!(read_number.is_some()); + assert_eq!(read_number.unwrap(), 42); + } + + // Read header by hash + { + let reader = hot_kv.reader().unwrap(); + let read_header = reader.header_by_hash(&hash).unwrap(); + assert!(read_header.is_some()); + assert_eq!(read_header.unwrap().number, 42); + } + } + + /// Test writing and reading accounts via HotDbWrite/HotDbRead + fn test_account_roundtrip(hot_kv: &T) { + let addr = address!("0x1234567890123456789012345678901234567890"); + let account = + Account { nonce: 5, balance: U256::from(1000), bytecode_hash: Some(B256::ZERO) }; + + // Write account + { + let mut writer = hot_kv.writer().unwrap(); + writer.put_account(&addr, &account).unwrap(); + writer.commit().unwrap(); + } + + // Read account + { + let reader = hot_kv.reader().unwrap(); + let read_account = reader.get_account(&addr).unwrap(); + assert!(read_account.is_some()); + let read_account = read_account.unwrap(); + assert_eq!(read_account.nonce, 5); + assert_eq!(read_account.balance, U256::from(1000)); + } + } + + /// Test writing and reading storage via HotDbWrite/HotDbRead + fn test_storage_roundtrip(hot_kv: &T) { + let addr = address!("0xabcdef0123456789abcdef0123456789abcdef01"); + let slot = b256!("0x0000000000000000000000000000000000000000000000000000000000000001"); + let value = U256::from(999); + + // Write storage + { + let mut writer = hot_kv.writer().unwrap(); + writer.put_storage(&addr, &slot, &value).unwrap(); + writer.commit().unwrap(); + } + + // Read storage + { + let reader = hot_kv.reader().unwrap(); + let read_value = reader.get_storage(&addr, &slot).unwrap(); + assert!(read_value.is_some()); + assert_eq!(read_value.unwrap(), U256::from(999)); + } + + // Read storage entry + { + let reader = hot_kv.reader().unwrap(); + let read_entry = reader.get_storage_entry(&addr, &slot).unwrap(); + assert!(read_entry.is_some()); + let entry = read_entry.unwrap(); + assert_eq!(entry.key, slot); + assert_eq!(entry.value, U256::from(999)); + } + } + + /// Test writing and reading bytecode via HotDbWrite/HotDbRead + fn test_bytecode_roundtrip(hot_kv: &T) { + let code = Bytes::from_static(&[0x60, 0x00, 0x60, 0x00, 0xf3]); // Simple EVM bytecode + let bytecode = Bytecode::new_raw(code); + let code_hash = bytecode.hash_slow(); + + // Write bytecode + { + let mut writer = hot_kv.writer().unwrap(); + writer.put_bytecode(&code_hash, &bytecode).unwrap(); + writer.commit().unwrap(); + } + + // Read bytecode + { + let reader = hot_kv.reader().unwrap(); + let read_bytecode = reader.get_bytecode(&code_hash).unwrap(); + assert!(read_bytecode.is_some()); + } + } + + /// Test account history via HotHistoryWrite/HotHistoryRead + fn test_account_history(hot_kv: &T) { + let addr = address!("0x1111111111111111111111111111111111111111"); + let touched_blocks = BlockNumberList::new([10, 20, 30]).unwrap(); + let latest_height = 100u64; + + // Write account history + { + let mut writer = hot_kv.writer().unwrap(); + writer.write_account_history(&addr, latest_height, &touched_blocks).unwrap(); + writer.commit().unwrap(); + } + + // Read account history + { + let reader = hot_kv.reader().unwrap(); + let read_history = reader.get_account_history(&addr, latest_height).unwrap(); + assert!(read_history.is_some()); + let history = read_history.unwrap(); + assert_eq!(history.iter().collect::>(), vec![10, 20, 30]); + } + } + + /// Test storage history via HotHistoryWrite/HotHistoryRead + fn test_storage_history(hot_kv: &T) { + let addr = address!("0x2222222222222222222222222222222222222222"); + let slot = b256!("0x0000000000000000000000000000000000000000000000000000000000000042"); + let touched_blocks = BlockNumberList::new([5, 15, 25]).unwrap(); + let highest_block = 50u64; + + // Write storage history + { + let mut writer = hot_kv.writer().unwrap(); + writer.write_storage_history(&addr, slot, highest_block, &touched_blocks).unwrap(); + writer.commit().unwrap(); + } + + // Read storage history + { + let reader = hot_kv.reader().unwrap(); + let read_history = reader.get_storage_history(&addr, slot, highest_block).unwrap(); + assert!(read_history.is_some()); + let history = read_history.unwrap(); + assert_eq!(history.iter().collect::>(), vec![5, 15, 25]); + } + } + + /// Test account change sets via HotHistoryWrite/HotHistoryRead + fn test_account_changes(hot_kv: &T) { + let addr = address!("0x3333333333333333333333333333333333333333"); + let pre_state = Account { nonce: 10, balance: U256::from(5000), bytecode_hash: None }; + let block_number = 100u64; + + // Write account change + { + let mut writer = hot_kv.writer().unwrap(); + writer.write_account_change(block_number, addr, &pre_state).unwrap(); + writer.commit().unwrap(); + } + + // Read account change + { + let reader = hot_kv.reader().unwrap(); + let read_change = reader.get_account_change(block_number, &addr).unwrap(); + assert!(read_change.is_some()); + let change = read_change.unwrap(); + assert_eq!(change.nonce, 10); + assert_eq!(change.balance, U256::from(5000)); + } + } + + /// Test storage change sets via HotHistoryWrite/HotHistoryRead + fn test_storage_changes(hot_kv: &T) { + let addr = address!("0x4444444444444444444444444444444444444444"); + let slot = b256!("0x0000000000000000000000000000000000000000000000000000000000000099"); + let pre_value = U256::from(12345); + let block_number = 200u64; + + // Write storage change + { + let mut writer = hot_kv.writer().unwrap(); + writer.write_storage_change(block_number, addr, &slot, &pre_value).unwrap(); + writer.commit().unwrap(); + } + + // Read storage change + { + let reader = hot_kv.reader().unwrap(); + let read_change = reader.get_storage_change(block_number, &addr, &slot).unwrap(); + assert!(read_change.is_some()); + assert_eq!(read_change.unwrap(), U256::from(12345)); + } + } + + /// Test that missing reads return None + fn test_missing_reads(hot_kv: &T) { + let missing_addr = address!("0x9999999999999999999999999999999999999999"); + let missing_hash = + b256!("0xffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff"); + let missing_slot = + b256!("0xeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee"); + + let reader = hot_kv.reader().unwrap(); + + // Missing header + assert!(reader.get_header(999999).unwrap().is_none()); + + // Missing header number + assert!(reader.get_header_number(&missing_hash).unwrap().is_none()); + + // Missing account + assert!(reader.get_account(&missing_addr).unwrap().is_none()); + + // Missing storage + assert!(reader.get_storage(&missing_addr, &missing_slot).unwrap().is_none()); + + // Missing bytecode + assert!(reader.get_bytecode(&missing_hash).unwrap().is_none()); + + // Missing header by hash + assert!(reader.header_by_hash(&missing_hash).unwrap().is_none()); + + // Missing account history + assert!(reader.get_account_history(&missing_addr, 1000).unwrap().is_none()); + + // Missing storage history + assert!(reader.get_storage_history(&missing_addr, missing_slot, 1000).unwrap().is_none()); + + // Missing account change + assert!(reader.get_account_change(999999, &missing_addr).unwrap().is_none()); + + // Missing storage change + assert!(reader.get_storage_change(999999, &missing_addr, &missing_slot).unwrap().is_none()); + } +} diff --git a/crates/storage/src/hot/mdbx.rs b/crates/storage/src/hot/mdbx.rs deleted file mode 100644 index 3b23e54..0000000 --- a/crates/storage/src/hot/mdbx.rs +++ /dev/null @@ -1,660 +0,0 @@ -use crate::{ - hot::{HotKv, HotKvError, HotKvRead, HotKvReadError, HotKvWrite}, - ser::{DeserError, KeySer, MAX_KEY_SIZE, ValSer}, - tables::{DualKeyed, MAX_FIXED_VAL_SIZE}, -}; -use bytes::{BufMut, BytesMut}; -use reth_db::{ - Database, DatabaseEnv, - mdbx::{RW, TransactionKind, WriteFlags, tx::Tx}, -}; -use reth_db_api::DatabaseError; -use reth_libmdbx::RO; -use std::borrow::Cow; - -/// Error type for reth-libmdbx based hot storage. -#[derive(Debug, thiserror::Error)] -pub enum MdbxError { - /// Inner error - #[error(transparent)] - Mdbx(#[from] reth_libmdbx::Error), - - /// Reth error. - #[error(transparent)] - Reth(#[from] DatabaseError), - - /// Deser. - #[error(transparent)] - Deser(#[from] DeserError), -} - -impl HotKvReadError for MdbxError { - fn into_hot_kv_error(self) -> HotKvError { - match self { - MdbxError::Mdbx(e) => HotKvError::from_err(e), - MdbxError::Deser(e) => HotKvError::Deser(e), - MdbxError::Reth(e) => HotKvError::from_err(e), - } - } -} - -impl From for DatabaseError { - fn from(value: DeserError) -> Self { - DatabaseError::Other(value.to_string()) - } -} - -impl HotKv for DatabaseEnv { - type RoTx = Tx; - type RwTx = Tx; - - fn reader(&self) -> Result { - self.tx().map_err(HotKvError::from_err) - } - - fn writer(&self) -> Result { - self.tx_mut().map_err(HotKvError::from_err) - } -} - -impl HotKvRead for Tx -where - K: TransactionKind, -{ - type Error = MdbxError; - - fn raw_get<'a>( - &'a self, - table: &str, - key: &[u8], - ) -> Result>, Self::Error> { - let dbi = self.inner.open_db(Some(table)).map(|db| db.dbi())?; - - self.inner.get(dbi, key.as_ref()).map_err(MdbxError::Mdbx) - } - - fn raw_get_dual<'a>( - &'a self, - _table: &str, - _key1: &[u8], - _key2: &[u8], - ) -> Result>, Self::Error> { - unimplemented!("Not implemented: raw_get_dual. Use get_dual instead."); - } - - fn get_dual( - &self, - key1: &T::Key, - key2: &T::Key2, - ) -> Result, Self::Error> { - let mut key1_buf = [0u8; MAX_KEY_SIZE]; - let key1_bytes = key1.encode_key(&mut key1_buf); - - // K2 slice must be EXACTLY the size of the fixed value size, if the - // table has one. This is a bit ugly, and results in an extra - // allocation for fixed-size values. This could be avoided using - // max value size. - let value_bytes = if T::IS_FIXED_VAL { - let buf = [0u8; MAX_KEY_SIZE + MAX_FIXED_VAL_SIZE]; - let _ = key2.encode_key(&mut buf[..MAX_KEY_SIZE].try_into().unwrap()); - - let kv_size = ::SIZE + T::FIXED_VAL_SIZE.unwrap(); - - let db = self.inner.open_db(Some(T::NAME))?; - let mut cursor = self.inner.cursor(&db).map_err(MdbxError::Mdbx)?; - cursor.get_both_range(key1_bytes, &buf[..kv_size]).map_err(MdbxError::Mdbx) - } else { - let mut buf = [0u8; MAX_KEY_SIZE]; - let encoded = key2.encode_key(&mut buf); - - let db = self.inner.open_db(Some(T::NAME))?; - let mut cursor = self.inner.cursor(&db).map_err(MdbxError::Mdbx)?; - cursor.get_both_range::>(key1_bytes, encoded).map_err(MdbxError::Mdbx) - }; - - let Some(value_bytes) = value_bytes? else { - return Ok(None); - }; - // we need to strip the key2 prefix from the value bytes before decoding - let value_bytes = &value_bytes[<::Key2 as KeySer>::SIZE..]; - - T::Value::decode_value(value_bytes).map(Some).map_err(Into::into) - } -} - -impl HotKvWrite for Tx { - fn queue_raw_put(&mut self, table: &str, key: &[u8], value: &[u8]) -> Result<(), Self::Error> { - let dbi = self.inner.open_db(Some(table)).map(|db| db.dbi())?; - - self.inner.put(dbi, key, value, WriteFlags::UPSERT).map(|_| ()).map_err(MdbxError::Mdbx) - } - - fn queue_raw_put_dual( - &mut self, - _table: &str, - _key1: &[u8], - _key2: &[u8], - _value: &[u8], - ) -> Result<(), Self::Error> { - unimplemented!("Not implemented: queue_raw_put_dual. Use queue_put_dual instead."); - } - - // Specialized put for dual-keyed tables. - fn queue_put_dual( - &mut self, - key1: &T::Key, - key2: &T::Key2, - value: &T::Value, - ) -> Result<(), Self::Error> { - let k2_size = ::SIZE; - let mut scratch = [0u8; MAX_KEY_SIZE]; - - // This will be the total length of key2 + value, reserved in mdbx - let encoded_len = k2_size + value.encoded_size(); - - // Prepend the value with k2. - let mut buf = BytesMut::with_capacity(encoded_len); - let encoded_k2 = key2.encode_key(&mut scratch); - buf.put_slice(encoded_k2); - value.encode_value_to(&mut buf); - - let encoded_k1 = key1.encode_key(&mut scratch); - // NB: DUPSORT and RESERVE are incompatible :( - let db = self.inner.open_db(Some(T::NAME))?; - self.inner.put(db.dbi(), encoded_k1, &buf, Default::default())?; - - Ok(()) - } - - fn queue_raw_delete(&mut self, table: &str, key: &[u8]) -> Result<(), Self::Error> { - let dbi = self.inner.open_db(Some(table)).map(|db| db.dbi())?; - - self.inner.del(dbi, key, None).map(|_| ()).map_err(MdbxError::Mdbx) - } - - fn queue_raw_clear(&mut self, table: &str) -> Result<(), Self::Error> { - // Future: port more of reth's db env with dbi caching to avoid - // repeated open_db calls - let dbi = self.inner.open_db(Some(table)).map(|db| db.dbi())?; - self.inner.clear_db(dbi).map(|_| ()).map_err(MdbxError::Mdbx) - } - - fn queue_raw_create( - &mut self, - table: &str, - dual_key: bool, - fixed_val: bool, - ) -> Result<(), Self::Error> { - let mut flags = Default::default(); - - if dual_key { - flags |= reth_libmdbx::DatabaseFlags::DUP_SORT; - if fixed_val { - flags |= reth_libmdbx::DatabaseFlags::DUP_FIXED; - } - } - - self.inner.create_db(Some(table), flags).map(|_| ()).map_err(MdbxError::Mdbx) - } - - fn raw_commit(self) -> Result<(), Self::Error> { - // when committing, mdbx returns true on failure - let res = self.inner.commit()?; - - if res.0 { Err(reth_libmdbx::Error::Other(1).into()) } else { Ok(()) } - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::{ - hot::{HotDbWrite, HotKv, HotKvRead, HotKvWrite}, - tables::hot, - }; - use alloy::{ - consensus::Sealed, - primitives::{Address, B256, BlockNumber, U256}, - }; - use reth::primitives::{Account, Bytecode, Header}; - use reth_db::DatabaseEnv; - - /// A test database wrapper that automatically cleans up on drop - struct TestDb { - db: DatabaseEnv, - #[allow(dead_code)] - temp_dir: tempfile::TempDir, - } - - impl std::ops::Deref for TestDb { - type Target = DatabaseEnv; - - fn deref(&self) -> &Self::Target { - &self.db - } - } - - /// Create a temporary MDBX database for testing that will be automatically cleaned up - fn create_test_db() -> TestDb { - let temp_dir = tempfile::tempdir().expect("Failed to create temp directory"); - - // Create the database - let db = reth_db::create_db(&temp_dir, Default::default()).unwrap(); - - // Create tables from the `crate::tables::hot` module - let mut writer = db.writer().unwrap(); - - writer.queue_create::().unwrap(); - writer.queue_create::().unwrap(); - writer.queue_create::().unwrap(); - writer.queue_create::().unwrap(); - writer.queue_create::().unwrap(); - writer.queue_create::().unwrap(); - writer.queue_create::().unwrap(); - writer.queue_create::().unwrap(); - writer.queue_create::().unwrap(); - - writer.commit().expect("Failed to commit table creation"); - - TestDb { db, temp_dir } - } - - /// Create test data - fn create_test_account() -> (Address, Account) { - let address = Address::from_slice(&[0x1; 20]); - let account = Account { - nonce: 42, - balance: U256::from(1000u64), - bytecode_hash: Some(B256::from_slice(&[0x2; 32])), - }; - (address, account) - } - - fn create_test_bytecode() -> (B256, Bytecode) { - let hash = B256::from_slice(&[0x2; 32]); - let code = reth::primitives::Bytecode::new_raw(vec![0x60, 0x80, 0x60, 0x40].into()); - (hash, code) - } - - fn create_test_header() -> (BlockNumber, Header) { - let block_number = 12345; - let header = Header { - number: block_number, - gas_limit: 8000000, - gas_used: 100000, - timestamp: 1640995200, - parent_hash: B256::from_slice(&[0x3; 32]), - state_root: B256::from_slice(&[0x4; 32]), - ..Default::default() - }; - (block_number, header) - } - - #[test] - fn test_hotkv_basic_operations() { - let db = create_test_db(); - let (address, account) = create_test_account(); - let (hash, bytecode) = create_test_bytecode(); - - // Test HotKv::writer() and basic write operations - { - let mut writer: Tx = db.writer().unwrap(); - - // Create tables first - writer.queue_create::().unwrap(); - - // Write account data - writer.queue_put::(&address, &account).unwrap(); - writer.queue_put::(&hash, &bytecode).unwrap(); - - // Commit the transaction - writer.raw_commit().unwrap(); - } - - // Test HotKv::reader() and basic read operations - { - let reader: Tx = db.reader().unwrap(); - - // Read account data - let read_account: Option = - reader.get::(&address).unwrap(); - assert_eq!(read_account, Some(account)); - - // Read bytecode - let read_bytecode: Option = reader.get::(&hash).unwrap(); - assert_eq!(read_bytecode, Some(bytecode)); - - // Test non-existent data - let nonexistent_addr = Address::from_slice(&[0xff; 20]); - let nonexistent_account: Option = - reader.get::(&nonexistent_addr).unwrap(); - assert_eq!(nonexistent_account, None); - } - } - - #[test] - fn test_raw_operations() { - let db = create_test_db(); - - let table_name = "test_table"; - let key = b"test_key"; - let value = b"test_value"; - - // Test raw write operations - { - let mut writer: Tx = db.writer().unwrap(); - - // Create table - writer.queue_raw_create(table_name, false, false).unwrap(); - - // Put raw data - writer.queue_raw_put(table_name, key, value).unwrap(); - - writer.raw_commit().unwrap(); - } - - // Test raw read operations - { - let reader: Tx = db.reader().unwrap(); - - let read_value = reader.raw_get(table_name, key).unwrap(); - assert_eq!(read_value.as_deref(), Some(value.as_slice())); - - // Test non-existent key - let nonexistent = reader.raw_get(table_name, b"nonexistent").unwrap(); - assert_eq!(nonexistent, None); - } - - // Test raw delete - { - let mut writer: Tx = db.writer().unwrap(); - - writer.queue_raw_delete(table_name, key).unwrap(); - writer.raw_commit().unwrap(); - } - - // Verify deletion - { - let reader: Tx = db.reader().unwrap(); - let deleted_value = reader.raw_get(table_name, key).unwrap(); - assert_eq!(deleted_value, None); - } - } - - #[test] - fn test_dual_keyed_operations() { - let db = create_test_db(); - - let address = Address::from_slice(&[0x1; 20]); - let storage_key = B256::from_slice(&[0x5; 32]); - let storage_value = U256::from(999u64); - - // Test dual-keyed table operations - { - let mut writer: Tx = db.writer().unwrap(); - - // Put storage data using dual keys - writer - .queue_put_dual::(&address, &storage_key, &storage_value) - .unwrap(); - - writer.raw_commit().unwrap(); - } - - // Test reading dual-keyed data - { - let reader: Tx = db.reader().unwrap(); - - // Read storage using dual key lookup - let read_value = - reader.get_dual::(&address, &storage_key).unwrap().unwrap(); - - assert_eq!(read_value, storage_value); - } - } - - #[test] - fn test_table_management() { - let db = create_test_db(); - - // Add some data - let (block_number, header) = create_test_header(); - { - let mut writer: Tx = db.writer().unwrap(); - writer.queue_put::(&block_number, &header).unwrap(); - writer.raw_commit().unwrap(); - } - - // Verify data exists - { - let reader: Tx = db.reader().unwrap(); - let read_header: Option
= reader.get::(&block_number).unwrap(); - assert_eq!(read_header, Some(header.clone())); - } - - // Clear the table - { - let mut writer: Tx = db.writer().unwrap(); - writer.queue_clear::().unwrap(); - writer.raw_commit().unwrap(); - } - - // Verify table is empty - { - let reader: Tx = db.reader().unwrap(); - let read_header: Option
= reader.get::(&block_number).unwrap(); - assert_eq!(read_header, None); - } - } - - #[test] - fn test_batch_operations() { - let db = create_test_db(); - - // Create test data - let accounts: Vec<(Address, Account)> = (0..10) - .map(|i| { - let mut addr_bytes = [0u8; 20]; - addr_bytes[19] = i; - let address = Address::from_slice(&addr_bytes); - let account = Account { - nonce: i.into(), - balance: U256::from((i as u64) * 100), - bytecode_hash: None, - }; - (address, account) - }) - .collect(); - - // Test batch writes - { - let mut writer: Tx = db.writer().unwrap(); - - // Write multiple accounts - for (address, account) in &accounts { - writer.queue_put::(address, account).unwrap(); - } - - writer.raw_commit().unwrap(); - } - - // Test batch reads - { - let reader: Tx = db.reader().unwrap(); - - for (address, expected_account) in &accounts { - let read_account: Option = - reader.get::(address).unwrap(); - assert_eq!(read_account.as_ref(), Some(expected_account)); - } - } - - // Test batch get_many - { - let reader: Tx = db.reader().unwrap(); - let addresses: Vec
= accounts.iter().map(|(addr, _)| *addr).collect(); - let read_accounts: Vec<(_, Option)> = - reader.get_many::(addresses.iter()).unwrap(); - - for (i, (_, expected_account)) in accounts.iter().enumerate() { - assert_eq!(read_accounts[i].1.as_ref(), Some(expected_account)); - } - } - } - - #[test] - fn test_transaction_isolation() { - let db = create_test_db(); - let (address, account) = create_test_account(); - - // Setup initial data - { - let mut writer: Tx = db.writer().unwrap(); - writer.queue_put::(&address, &account).unwrap(); - writer.raw_commit().unwrap(); - } - - // Start a reader transaction - let reader: Tx = db.reader().unwrap(); - - // Modify data in a writer transaction - { - let mut writer: Tx = db.writer().unwrap(); - let modified_account = - Account { nonce: 999, balance: U256::from(9999u64), bytecode_hash: None }; - writer.queue_put::(&address, &modified_account).unwrap(); - writer.raw_commit().unwrap(); - } - - // Reader should still see original data (snapshot isolation) - { - let read_account: Option = - reader.get::(&address).unwrap(); - assert_eq!(read_account, Some(account)); - } - - // New reader should see modified data - { - let new_reader: Tx = db.reader().unwrap(); - let read_account: Option = - new_reader.get::(&address).unwrap(); - assert_eq!(read_account.unwrap().nonce, 999); - } - } - - #[test] - fn test_multiple_readers() { - let db = create_test_db(); - let (address, account) = create_test_account(); - - // Setup data - { - let mut writer: Tx = db.writer().unwrap(); - writer.queue_put::(&address, &account).unwrap(); - writer.raw_commit().unwrap(); - } - - // Create multiple readers - let reader1: Tx = db.reader().unwrap(); - let reader2: Tx = db.reader().unwrap(); - let reader3: Tx = db.reader().unwrap(); - - // All readers should see the same data - let account1: Option = reader1.get::(&address).unwrap(); - let account2: Option = reader2.get::(&address).unwrap(); - let account3: Option = reader3.get::(&address).unwrap(); - - assert_eq!(account1, Some(account)); - assert_eq!(account2, Some(account)); - assert_eq!(account3, Some(account)); - } - - #[test] - fn test_error_handling() { - let db = create_test_db(); - - // Test reading from non-existent table - { - let reader: Tx = db.reader().unwrap(); - let result = reader.raw_get("nonexistent_table", b"key"); - - // Should handle gracefully (may return None or error depending on MDBX behavior) - match result { - Ok(None) => {} // This is fine - Err(_) => {} // This is also acceptable for non-existent table - Ok(Some(_)) => panic!("Should not return data for non-existent table"), - } - } - - // Test writing to a table without creating it first - { - let mut writer: Tx = db.writer().unwrap(); - let (address, account) = create_test_account(); - - // This should handle the case where table doesn't exist - let result = writer.queue_put::(&address, &account); - match result { - Ok(_) => { - // If it succeeds, commit should work - writer.raw_commit().unwrap(); - } - Err(_) => { - // If it fails, that's expected behavior - } - } - } - } - - #[test] - fn test_serialization_roundtrip() { - let db = create_test_db(); - - // Test various data types - let (block_number, header) = create_test_header(); - let header = Sealed::new(header); - - { - let mut writer: Tx = db.writer().unwrap(); - - // Write different types - writer.queue_put::(&block_number, &header).unwrap(); - writer.queue_put::(&header.hash(), &block_number).unwrap(); - - writer.raw_commit().unwrap(); - } - - { - let reader: Tx = db.reader().unwrap(); - - // Read and verify - let read_header: Option
= reader.get::(&block_number).unwrap(); - assert_eq!(read_header.as_ref(), Some(header.inner())); - - let read_hash: Option = reader.get::(&header.hash()).unwrap(); - assert_eq!(read_hash, Some(header.number)); - } - } - - #[test] - fn test_large_data() { - let db = create_test_db(); - - // Create a large bytecode - let hash = B256::from_slice(&[0x8; 32]); - let large_code_vec: Vec = (0..10000).map(|i| (i % 256) as u8).collect(); - let large_bytecode = Bytecode::new_raw(large_code_vec.clone().into()); - - { - let mut writer: Tx = db.writer().unwrap(); - writer.queue_create::().unwrap(); - writer.queue_put::(&hash, &large_bytecode).unwrap(); - writer.raw_commit().unwrap(); - } - - { - let reader: Tx = db.reader().unwrap(); - let read_bytecode: Option = reader.get::(&hash).unwrap(); - assert_eq!(read_bytecode, Some(large_bytecode)); - } - } -} diff --git a/crates/storage/src/hot/mem.rs b/crates/storage/src/hot/mem.rs deleted file mode 100644 index b7e84d2..0000000 --- a/crates/storage/src/hot/mem.rs +++ /dev/null @@ -1,828 +0,0 @@ -use crate::{ - hot::{HotKv, HotKvError, HotKvRead, HotKvWrite}, - ser::MAX_KEY_SIZE, -}; -use std::{ - borrow::Cow, - collections::BTreeMap, - sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard}, -}; - -type Table = BTreeMap<[u8; MAX_KEY_SIZE], bytes::Bytes>; -type Store = BTreeMap; - -type TableOp = BTreeMap<[u8; MAX_KEY_SIZE], QueuedKvOp>; -type OpStore = BTreeMap; - -/// A simple in-memory key-value store using [`BTreeMap`]s. -/// -/// The store is backed by an [`RwLock`]. As a result, this implementation -/// supports concurrent multiple concurrent read transactions, but write -/// transactions are exclusive, and cannot overlap with other read or write -/// transactions. -/// -/// This implementation is primarily intended for testing and -/// development purposes. -#[derive(Clone)] -pub struct MemKv { - map: Arc>, -} - -impl core::fmt::Debug for MemKv { - fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { - f.debug_struct("MemKv").finish() - } -} - -impl MemKv { - /// Create a new empty in-memory KV store. - pub fn new() -> Self { - Self { map: Arc::new(RwLock::new(BTreeMap::new())) } - } - - #[track_caller] - fn key(k: &[u8]) -> [u8; MAX_KEY_SIZE] { - assert!(k.len() <= MAX_KEY_SIZE, "Key length exceeds MAX_KEY_SIZE"); - let mut buf = [0u8; MAX_KEY_SIZE]; - buf[..k.len()].copy_from_slice(k); - buf - } - - #[track_caller] - fn dual_key(k1: &[u8], k2: &[u8]) -> [u8; MAX_KEY_SIZE] { - assert!(k1.len() + k2.len() <= MAX_KEY_SIZE, "Combined key length exceeds MAX_KEY_SIZE"); - let mut buf = [0u8; MAX_KEY_SIZE]; - buf[..k1.len()].copy_from_slice(k1); - buf[k1.len()..k1.len() + k2.len()].copy_from_slice(k2); - buf - } -} - -impl Default for MemKv { - fn default() -> Self { - Self::new() - } -} - -/// Read-only transaction for MemKv. -pub struct MemKvRoTx { - guard: RwLockReadGuard<'static, Store>, - - // Keep the store alive while the transaction exists - _store: Arc>, -} - -impl core::fmt::Debug for MemKvRoTx { - fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { - f.debug_struct("MemKvRoTx").finish() - } -} - -// SAFETY: MemKvRoTx holds a read guard which ensures the data remains valid -unsafe impl Send for MemKvRoTx {} -unsafe impl Sync for MemKvRoTx {} - -/// Read-write transaction for MemKv. -pub struct MemKvRwTx { - guard: RwLockWriteGuard<'static, Store>, - queued_ops: OpStore, - - // Keep the store alive while the transaction exists - _store: Arc>, -} - -impl MemKvRwTx { - fn commit_inner(&mut self) { - let ops = std::mem::take(&mut self.queued_ops); - - for (table, table_op) in ops.into_iter() { - table_op.apply(&table, &mut self.guard); - } - } - - /// Downgrade the transaction to a read-only transaction without - /// committing, discarding queued changes. - pub fn downgrade(self) -> MemKvRoTx { - let guard = RwLockWriteGuard::downgrade(self.guard); - - MemKvRoTx { guard, _store: self._store } - } - - /// Commit the transaction and downgrade to a read-only transaction. - pub fn commit_downgrade(mut self) -> MemKvRoTx { - self.commit_inner(); - - let guard = RwLockWriteGuard::downgrade(self.guard); - - MemKvRoTx { guard, _store: self._store } - } -} - -impl core::fmt::Debug for MemKvRwTx { - fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { - f.debug_struct("MemKvRwTx").finish() - } -} - -/// Queued key-value operation -#[derive(Debug, Clone)] -enum QueuedKvOp { - Delete, - Put { value: bytes::Bytes }, -} - -impl QueuedKvOp { - /// Apply the op to a table - fn apply(self, table: &mut Table, key: [u8; MAX_KEY_SIZE]) { - match self { - QueuedKvOp::Put { value } => { - table.insert(key, value); - } - QueuedKvOp::Delete => { - table.remove(&key); - } - } - } -} - -/// Queued table operation -#[derive(Debug)] -enum QueuedTableOp { - Modify { ops: TableOp }, - Clear { new_table: TableOp }, -} - -impl Default for QueuedTableOp { - fn default() -> Self { - QueuedTableOp::Modify { ops: TableOp::new() } - } -} - -impl QueuedTableOp { - const fn is_clear(&self) -> bool { - matches!(self, QueuedTableOp::Clear { .. }) - } - - fn get(&self, key: &[u8; MAX_KEY_SIZE]) -> Option<&QueuedKvOp> { - match self { - QueuedTableOp::Modify { ops } => ops.get(key), - QueuedTableOp::Clear { new_table } => new_table.get(key), - } - } - - fn put(&mut self, key: [u8; MAX_KEY_SIZE], op: QueuedKvOp) { - match self { - QueuedTableOp::Modify { ops } | QueuedTableOp::Clear { new_table: ops } => { - ops.insert(key, op); - } - } - } - - fn delete(&mut self, key: [u8; MAX_KEY_SIZE]) { - match self { - QueuedTableOp::Modify { ops } | QueuedTableOp::Clear { new_table: ops } => { - ops.insert(key, QueuedKvOp::Delete); - } - } - } - - /// Get mutable reference to the inner ops if applicable - fn apply(self, key: &str, store: &mut Store) { - match self { - QueuedTableOp::Modify { ops } => { - let table = store.entry(key.to_owned()).or_default(); - for (key, op) in ops { - op.apply(table, key); - } - } - QueuedTableOp::Clear { new_table } => { - let mut table = Table::new(); - for (k, op) in new_table { - op.apply(&mut table, k); - } - - // replace the table entirely - store.insert(key.to_owned(), table); - } - } - } -} - -// SAFETY: MemKvRwTx holds a write guard which ensures exclusive access -unsafe impl Send for MemKvRwTx {} - -impl HotKv for MemKv { - type RoTx = MemKvRoTx; - type RwTx = MemKvRwTx; - - fn reader(&self) -> Result { - let guard = self - .map - .try_read() - .map_err(|_| HotKvError::Inner("Failed to acquire read lock".into()))?; - - // SAFETY: This is safe-ish, as we ensure the map is not dropped until - // the guard is also dropped. - let guard: RwLockReadGuard<'static, Store> = unsafe { std::mem::transmute(guard) }; - - Ok(MemKvRoTx { guard, _store: self.map.clone() }) - } - - fn writer(&self) -> Result { - let guard = self.map.try_write().map_err(|_| HotKvError::WriteLocked)?; - - // SAFETY: This is safe-ish, as we ensure the map is not dropped until - // the guard is also dropped. - let guard: RwLockWriteGuard<'static, Store> = unsafe { std::mem::transmute(guard) }; - - Ok(MemKvRwTx { guard, _store: self.map.clone(), queued_ops: OpStore::new() }) - } -} - -impl HotKvRead for MemKvRoTx { - type Error = HotKvError; - - fn raw_get<'a>( - &'a self, - table: &str, - key: &[u8], - ) -> Result>, Self::Error> { - // Check queued operations first (read-your-writes consistency) - let key = MemKv::key(key); - - // SAFETY: The guard ensures the map remains valid - - Ok(self - .guard - .get(table) - .and_then(|t| t.get(&key)) - .map(|bytes| Cow::Borrowed(bytes.as_ref()))) - } - - fn raw_get_dual<'a>( - &'a self, - table: &str, - key1: &[u8], - key2: &[u8], - ) -> Result>, Self::Error> { - let key = MemKv::dual_key(key1, key2); - self.raw_get(table, &key) - } -} - -impl HotKvRead for MemKvRwTx { - type Error = HotKvError; - - fn raw_get<'a>( - &'a self, - table: &str, - key: &[u8], - ) -> Result>, Self::Error> { - // Check queued operations first (read-your-writes consistency) - let key = MemKv::key(key); - - if let Some(table) = self.queued_ops.get(table) { - if table.is_clear() { - return Ok(None); - } - - match table.get(&key) { - Some(QueuedKvOp::Put { value }) => { - return Ok(Some(Cow::Borrowed(value.as_ref()))); - } - Some(QueuedKvOp::Delete) => { - return Ok(None); - } - None => {} - } - } - - // If not found in queued ops, check the underlying map - Ok(self - .guard - .get(table) - .and_then(|t| t.get(&key)) - .map(|bytes| Cow::Borrowed(bytes.as_ref()))) - } - - fn raw_get_dual<'a>( - &'a self, - table: &str, - key1: &[u8], - key2: &[u8], - ) -> Result>, Self::Error> { - let key = MemKv::dual_key(key1, key2); - self.raw_get(table, &key) - } -} - -impl HotKvWrite for MemKvRwTx { - fn queue_raw_put(&mut self, table: &str, key: &[u8], value: &[u8]) -> Result<(), Self::Error> { - let key = MemKv::key(key); - - let value_bytes = bytes::Bytes::copy_from_slice(value); - - self.queued_ops - .entry(table.to_owned()) - .or_default() - .put(key, QueuedKvOp::Put { value: value_bytes }); - Ok(()) - } - - fn queue_raw_put_dual( - &mut self, - table: &str, - key1: &[u8], - key2: &[u8], - value: &[u8], - ) -> Result<(), Self::Error> { - let key = MemKv::dual_key(key1, key2); - self.queue_raw_put(table, &key, value) - } - - fn queue_raw_delete(&mut self, table: &str, key: &[u8]) -> Result<(), Self::Error> { - let key = MemKv::key(key); - - self.queued_ops.entry(table.to_owned()).or_default().delete(key); - Ok(()) - } - - fn queue_raw_clear(&mut self, table: &str) -> Result<(), Self::Error> { - self.queued_ops - .insert(table.to_owned(), QueuedTableOp::Clear { new_table: TableOp::new() }); - Ok(()) - } - - fn queue_raw_create( - &mut self, - _table: &str, - _dual_key: bool, - _dual_fixed: bool, - ) -> Result<(), Self::Error> { - Ok(()) - } - - fn raw_commit(mut self) -> Result<(), Self::Error> { - // Apply all queued operations to the map - self.commit_inner(); - - // The write guard is automatically dropped here, releasing the lock - Ok(()) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::tables::Table; - use alloy::primitives::{Address, U256}; - use bytes::Bytes; - - // Test table definitions - #[derive(Debug)] - struct TestTable; - - impl Table for TestTable { - const NAME: &'static str = "test_table"; - - type Key = u64; - type Value = Bytes; - } - - #[derive(Debug)] - struct AddressTable; - - impl Table for AddressTable { - const NAME: &'static str = "addresses"; - type Key = Address; - type Value = U256; - } - - #[test] - fn test_new_store() { - let store = MemKv::new(); - let reader = store.reader().unwrap(); - - // Empty store should return None for any key - assert!(reader.raw_get("test", &[1, 2, 3]).unwrap().is_none()); - } - - #[test] - fn test_basic_put_get() { - let store = MemKv::new(); - - // Write some data - { - let mut writer = store.writer().unwrap(); - writer.queue_raw_put("table1", &[1, 2, 3], b"value1").unwrap(); - writer.queue_raw_put("table1", &[4, 5, 6], b"value2").unwrap(); - writer.raw_commit().unwrap(); - } - - // Read the data back - { - let reader = store.reader().unwrap(); - let value1 = reader.raw_get("table1", &[1, 2, 3]).unwrap(); - let value2 = reader.raw_get("table1", &[4, 5, 6]).unwrap(); - let missing = reader.raw_get("table1", &[7, 8, 9]).unwrap(); - - assert_eq!(value1.as_deref(), Some(b"value1" as &[u8])); - assert_eq!(value2.as_deref(), Some(b"value2" as &[u8])); - assert!(missing.is_none()); - } - } - - #[test] - fn test_multiple_tables() { - let store = MemKv::new(); - - // Write to different tables - { - let mut writer = store.writer().unwrap(); - writer.queue_raw_put("table1", &[1], b"table1_value").unwrap(); - writer.queue_raw_put("table2", &[1], b"table2_value").unwrap(); - writer.raw_commit().unwrap(); - } - - // Read from different tables - { - let reader = store.reader().unwrap(); - let value1 = reader.raw_get("table1", &[1]).unwrap(); - let value2 = reader.raw_get("table2", &[1]).unwrap(); - - assert_eq!(value1.as_deref(), Some(b"table1_value" as &[u8])); - assert_eq!(value2.as_deref(), Some(b"table2_value" as &[u8])); - } - } - - #[test] - fn test_overwrite_value() { - let store = MemKv::new(); - - // Write initial value - { - let mut writer = store.writer().unwrap(); - writer.queue_raw_put("table1", &[1], b"original").unwrap(); - writer.raw_commit().unwrap(); - } - - // Overwrite with new value - { - let mut writer = store.writer().unwrap(); - writer.queue_raw_put("table1", &[1], b"updated").unwrap(); - writer.raw_commit().unwrap(); - } - - // Check the value was updated - { - let reader = store.reader().unwrap(); - let value = reader.raw_get("table1", &[1]).unwrap(); - assert_eq!(value.as_deref(), Some(b"updated" as &[u8])); - } - } - - #[test] - fn test_read_your_writes() { - let store = MemKv::new(); - let mut writer = store.writer().unwrap(); - - // Queue some operations but don't commit yet - writer.queue_raw_put("table1", &[1], b"queued_value").unwrap(); - - // Should be able to read the queued value - let value = writer.raw_get("table1", &[1]).unwrap(); - assert_eq!(value.as_deref(), Some(b"queued_value" as &[u8])); - - writer.raw_commit().unwrap(); - - // After commit, other readers should see it - { - let reader = store.reader().unwrap(); - let value = reader.raw_get("table1", &[1]).unwrap(); - assert_eq!(value.as_deref(), Some(b"queued_value" as &[u8])); - } - } - - #[test] - fn test_typed_operations() { - let store = MemKv::new(); - - // Write using typed interface - { - let mut writer = store.writer().unwrap(); - writer.queue_put::(&42u64, &Bytes::from_static(b"hello world")).unwrap(); - writer.queue_put::(&100u64, &Bytes::from_static(b"another value")).unwrap(); - writer.raw_commit().unwrap(); - } - - // Read using typed interface - { - let reader = store.reader().unwrap(); - let value1 = reader.get::(&42u64).unwrap(); - let value2 = reader.get::(&100u64).unwrap(); - let missing = reader.get::(&999u64).unwrap(); - - assert_eq!(value1, Some(Bytes::from_static(b"hello world"))); - assert_eq!(value2, Some(Bytes::from_static(b"another value"))); - assert!(missing.is_none()); - } - } - - #[test] - fn test_address_table() { - let store = MemKv::new(); - - let addr1 = Address::from([0x11; 20]); - let addr2 = Address::from([0x22; 20]); - let balance1 = U256::from(1000u64); - let balance2 = U256::from(2000u64); - - // Write address data - { - let mut writer = store.writer().unwrap(); - writer.queue_put::(&addr1, &balance1).unwrap(); - writer.queue_put::(&addr2, &balance2).unwrap(); - writer.raw_commit().unwrap(); - } - - // Read address data - { - let reader = store.reader().unwrap(); - let bal1 = reader.get::(&addr1).unwrap(); - let bal2 = reader.get::(&addr2).unwrap(); - - assert_eq!(bal1, Some(balance1)); - assert_eq!(bal2, Some(balance2)); - } - } - - #[test] - fn test_batch_operations() { - let store = MemKv::new(); - - let entries = [ - (1u64, Bytes::from_static(b"first")), - (2u64, Bytes::from_static(b"second")), - (3u64, Bytes::from_static(b"third")), - ]; - - // Write batch - { - let mut writer = store.writer().unwrap(); - let entry_refs: Vec<_> = entries.iter().map(|(k, v)| (k, v)).collect(); - writer.queue_put_many::(entry_refs).unwrap(); - writer.raw_commit().unwrap(); - } - - // Read batch - { - let reader = store.reader().unwrap(); - let keys: Vec<_> = entries.iter().map(|(k, _)| k).collect(); - let values = reader.get_many::(keys).unwrap(); - - assert_eq!(values.len(), 3); - assert_eq!(values[0], (&1u64, Some(Bytes::from_static(b"first")))); - assert_eq!(values[1], (&2u64, Some(Bytes::from_static(b"second")))); - assert_eq!(values[2], (&3u64, Some(Bytes::from_static(b"third")))); - } - } - - #[test] - fn test_concurrent_readers() { - let store = MemKv::new(); - - // Write some initial data - { - let mut writer = store.writer().unwrap(); - writer.queue_raw_put("table1", &[1], b"value1").unwrap(); - writer.raw_commit().unwrap(); - } - - // Multiple readers should be able to read concurrently - let reader1 = store.reader().unwrap(); - let reader2 = store.reader().unwrap(); - - let value1 = reader1.raw_get("table1", &[1]).unwrap(); - let value2 = reader2.raw_get("table1", &[1]).unwrap(); - - assert_eq!(value1.as_deref(), Some(b"value1" as &[u8])); - assert_eq!(value2.as_deref(), Some(b"value1" as &[u8])); - } - - #[test] - fn test_write_lock_exclusivity() { - let store = MemKv::new(); - - // Get a writer - let _writer1 = store.writer().unwrap(); - - // Second writer should fail - match store.writer() { - Err(HotKvError::WriteLocked) => {} // Expected - Ok(_) => panic!("Should not be able to get second writer"), - Err(e) => panic!("Unexpected error: {:?}", e), - } - } - - #[test] - fn test_empty_values() { - let store = MemKv::new(); - - { - let mut writer = store.writer().unwrap(); - writer.queue_raw_put("table1", &[1], b"").unwrap(); - writer.raw_commit().unwrap(); - } - - { - let reader = store.reader().unwrap(); - let value = reader.raw_get("table1", &[1]).unwrap(); - assert_eq!(value.as_deref(), Some(b"" as &[u8])); - } - } - - #[test] - fn test_multiple_operations_same_transaction() { - let store = MemKv::new(); - - { - let mut writer = store.writer().unwrap(); - - // Multiple operations on same key - last one should win - writer.queue_raw_put("table1", &[1], b"first").unwrap(); - writer.queue_raw_put("table1", &[1], b"second").unwrap(); - writer.queue_raw_put("table1", &[1], b"third").unwrap(); - - // Read-your-writes should return the latest value - let value = writer.raw_get("table1", &[1]).unwrap(); - assert_eq!(value.as_deref(), Some(b"third" as &[u8])); - - writer.raw_commit().unwrap(); - } - - { - let reader = store.reader().unwrap(); - let value = reader.raw_get("table1", &[1]).unwrap(); - assert_eq!(value.as_deref(), Some(b"third" as &[u8])); - } - } - - #[test] - fn test_isolation() { - let store = MemKv::new(); - - // Write initial value - { - let mut writer = store.writer().unwrap(); - writer.queue_raw_put("table1", &[1], b"original").unwrap(); - writer.raw_commit().unwrap(); - } - - // Start a read transaction - { - let reader = store.reader().unwrap(); - let original_value = reader.raw_get("table1", &[1]).unwrap(); - assert_eq!(original_value.as_deref(), Some(b"original" as &[u8])); - } - - // Update the value in a separate transaction - { - let mut writer = store.writer().unwrap(); - writer.queue_raw_put("table1", &[1], b"updated").unwrap(); - writer.raw_commit().unwrap(); - } - - // The value should now be latest for new readers - { - // New reader should see the updated value - let new_reader = store.reader().unwrap(); - let updated_value = new_reader.raw_get("table1", &[1]).unwrap(); - assert_eq!(updated_value.as_deref(), Some(b"updated" as &[u8])); - } - } - - #[test] - fn test_rollback_on_drop() { - let store = MemKv::new(); - - { - let mut writer = store.writer().unwrap(); - writer.queue_raw_put("table1", &[1], b"should_not_persist").unwrap(); - // Drop without committing - } - - // Value should not be persisted - { - let reader = store.reader().unwrap(); - let value = reader.raw_get("table1", &[1]).unwrap(); - assert!(value.is_none()); - } - } - - #[test] - fn write_two_tables() { - let store = MemKv::new(); - - { - let mut writer = store.writer().unwrap(); - writer.queue_raw_put("table1", &[1], b"value1").unwrap(); - writer.queue_raw_put("table2", &[2], b"value2").unwrap(); - writer.raw_commit().unwrap(); - } - - { - let reader = store.reader().unwrap(); - let value1 = reader.raw_get("table1", &[1]).unwrap(); - let value2 = reader.raw_get("table2", &[2]).unwrap(); - - assert_eq!(value1.as_deref(), Some(b"value1" as &[u8])); - assert_eq!(value2.as_deref(), Some(b"value2" as &[u8])); - } - } - - #[test] - fn test_downgrades() { - let store = MemKv::new(); - { - // Write some data - // Start a read-write transaction - let mut rw_tx = store.writer().unwrap(); - rw_tx.queue_raw_put("table1", &[1, 2, 3], b"value1").unwrap(); - rw_tx.queue_raw_put("table1", &[4, 5, 6], b"value2").unwrap(); - - let ro_tx = rw_tx.commit_downgrade(); - - // Read the data back - let value1 = ro_tx.raw_get("table1", &[1, 2, 3]).unwrap(); - let value2 = ro_tx.raw_get("table1", &[4, 5, 6]).unwrap(); - - assert_eq!(value1.as_deref(), Some(b"value1" as &[u8])); - assert_eq!(value2.as_deref(), Some(b"value2" as &[u8])); - } - - { - // Start another read-write transaction - let mut rw_tx = store.writer().unwrap(); - rw_tx.queue_raw_put("table2", &[7, 8, 9], b"value3").unwrap(); - - // Value should not be set - let ro_tx = rw_tx.downgrade(); - - // Read the data back - let value3 = ro_tx.raw_get("table2", &[7, 8, 9]).unwrap(); - - assert!(value3.is_none()); - } - } - - #[test] - fn test_clear_table() { - let store = MemKv::new(); - - { - let mut writer = store.writer().unwrap(); - writer.queue_raw_put("table1", &[1], b"value1").unwrap(); - writer.queue_raw_put("table1", &[2], b"value2").unwrap(); - writer.raw_commit().unwrap(); - } - - { - let reader = store.reader().unwrap(); - - let value1 = reader.raw_get("table1", &[1]).unwrap(); - let value2 = reader.raw_get("table1", &[2]).unwrap(); - - assert_eq!(value1.as_deref(), Some(b"value1" as &[u8])); - assert_eq!(value2.as_deref(), Some(b"value2" as &[u8])); - } - - { - let mut writer = store.writer().unwrap(); - - let value1 = writer.raw_get("table1", &[1]).unwrap(); - let value2 = writer.raw_get("table1", &[2]).unwrap(); - - assert_eq!(value1.as_deref(), Some(b"value1" as &[u8])); - assert_eq!(value2.as_deref(), Some(b"value2" as &[u8])); - - writer.queue_raw_clear("table1").unwrap(); - - let value1 = writer.raw_get("table1", &[1]).unwrap(); - let value2 = writer.raw_get("table1", &[2]).unwrap(); - - assert!(value1.is_none()); - assert!(value2.is_none()); - - writer.raw_commit().unwrap(); - } - - { - let reader = store.reader().unwrap(); - let value1 = reader.raw_get("table1", &[1]).unwrap(); - let value2 = reader.raw_get("table1", &[2]).unwrap(); - - assert!(value1.is_none()); - assert!(value2.is_none()); - } - } -} diff --git a/crates/storage/src/hot/mod.rs b/crates/storage/src/hot/mod.rs index 6f4786c..4e73f9c 100644 --- a/crates/storage/src/hot/mod.rs +++ b/crates/storage/src/hot/mod.rs @@ -1,16 +1,5 @@ -mod db_traits; -pub use db_traits::{HotDbRead, HotDbWrite, HotHistoryRead, HotHistoryWrite}; +/// Hot storage models and traits. +pub mod model; -mod error; -pub use error::{HotKvError, HotKvReadError, HotKvResult}; - -mod mem; -pub use mem::{MemKv, MemKvRoTx, MemKvRwTx}; - -mod mdbx; - -mod revm; -pub use revm::{RevmRead, RevmWrite}; - -mod traits; -pub use traits::{HotKv, HotKvRead, HotKvWrite, KeyValue}; +mod impls; +pub use impls::{mdbx, mem}; diff --git a/crates/storage/src/hot/db_traits.rs b/crates/storage/src/hot/model/db_traits.rs similarity index 96% rename from crates/storage/src/hot/db_traits.rs rename to crates/storage/src/hot/model/db_traits.rs index 25ef573..9b08f49 100644 --- a/crates/storage/src/hot/db_traits.rs +++ b/crates/storage/src/hot/model/db_traits.rs @@ -1,5 +1,5 @@ use crate::{ - hot::{HotKvRead, HotKvWrite}, + hot::model::{HotKvRead, HotKvWrite}, tables::hot::{self as tables}, }; use alloy::primitives::{Address, B256, U256}; @@ -55,7 +55,10 @@ pub trait HotDbRead: HotKvRead + sealed::Sealed { impl HotDbRead for T where T: HotKvRead {} -/// Trait for database write operations. +/// Trait for database write operations. This trait is low-level, and usage may +/// leave the database in an inconsistent state if not used carefully. Users +/// should prefer [`HotHistoryWrite`] or higher-level abstractions when +/// possible. pub trait HotDbWrite: HotKvWrite + sealed::Sealed { /// Write a block header. This will leave the DB in an inconsistent state /// until the corresponding header number is also written. Users should @@ -182,22 +185,17 @@ pub trait HotHistoryWrite: HotDbWrite { self.queue_put_dual::(address, &latest_height, touched) } - /// Write a storage change (before state) for an account at a specific + /// Write an account change (pre-state) for an account at a specific /// block. - fn write_storage_change( + fn write_account_change( &mut self, block_number: u64, address: Address, - slot: &B256, - value: &U256, + pre_state: &Account, ) -> Result<(), Self::Error> { - let block_number_address = BlockNumberAddress((block_number, address)); - self.queue_put_dual::(&block_number_address, slot, value) + self.queue_put_dual::(&block_number, &address, pre_state) } - /// Write an account change (pre-state) for an account at a specific - /// block. - /// Write storage history, by highest block number and touched block /// numbers. fn write_storage_history( @@ -213,20 +211,22 @@ pub trait HotHistoryWrite: HotDbWrite { /// Write a storage change (before state) for an account at a specific /// block. - fn write_account_change( + fn write_storage_change( &mut self, block_number: u64, address: Address, - pre_state: &Account, + slot: &B256, + value: &U256, ) -> Result<(), Self::Error> { - self.queue_put_dual::(&block_number, &address, pre_state) + let block_number_address = BlockNumberAddress((block_number, address)); + self.queue_put_dual::(&block_number_address, slot, value) } } impl HotHistoryWrite for T where T: HotDbWrite + HotKvWrite {} mod sealed { - use crate::hot::HotKvRead; + use crate::hot::model::HotKvRead; /// Sealed trait to prevent external implementations of HotDbReader and HotDbWriter. #[allow(dead_code, unreachable_pub)] diff --git a/crates/storage/src/hot/error.rs b/crates/storage/src/hot/model/error.rs similarity index 100% rename from crates/storage/src/hot/error.rs rename to crates/storage/src/hot/model/error.rs diff --git a/crates/storage/src/hot/model/mod.rs b/crates/storage/src/hot/model/mod.rs new file mode 100644 index 0000000..f499db7 --- /dev/null +++ b/crates/storage/src/hot/model/mod.rs @@ -0,0 +1,38 @@ +mod db_traits; +pub use db_traits::{HotDbRead, HotDbWrite, HotHistoryRead, HotHistoryWrite}; + +mod error; +pub use error::{HotKvError, HotKvReadError, HotKvResult}; + +mod revm; +pub use revm::{RevmRead, RevmWrite}; + +mod traits; +pub use traits::{HotKv, HotKvRead, HotKvWrite}; + +mod traverse; +pub use traverse::{ + DualKeyedTraverse, DualTableCursor, DualTableTraverse, KvTraverse, KvTraverseMut, TableCursor, + TableTraverse, TableTraverseMut, +}; + +use crate::tables::{DualKeyed, Table}; +use std::borrow::Cow; + +/// A key-value pair from a table. +pub type GetManyItem<'a, T> = (&'a ::Key, Option<::Value>); + +/// A key-value tuple from a table. +pub type KeyValue = (::Key, ::Value); + +/// A raw key-value pair. +pub type RawKeyValue<'a> = (Cow<'a, [u8]>, RawValue<'a>); + +/// A raw value. +pub type RawValue<'a> = Cow<'a, [u8]>; + +/// A raw dual key-value tuple. +pub type RawDualKeyValue<'a> = (Cow<'a, [u8]>, RawValue<'a>, RawValue<'a>); + +/// A dual key-value tuple from a table. +pub type DualKeyValue = (::Key, ::Key2, ::Value); diff --git a/crates/storage/src/hot/revm.rs b/crates/storage/src/hot/model/revm.rs similarity index 93% rename from crates/storage/src/hot/revm.rs rename to crates/storage/src/hot/model/revm.rs index 7b7bcbb..598da12 100644 --- a/crates/storage/src/hot/revm.rs +++ b/crates/storage/src/hot/model/revm.rs @@ -1,7 +1,7 @@ use crate::{ - hot::{HotKvError, HotKvRead, HotKvWrite, KeyValue}, + hot::model::{GetManyItem, HotKvError, HotKvRead, HotKvWrite}, tables::{ - DualKeyed, Table, + DualKeyed, SingleKey, Table, hot::{self, Bytecodes, PlainAccountState}, }, }; @@ -40,6 +40,15 @@ impl fmt::Debug for RevmRead { impl HotKvRead for RevmRead { type Error = U::Error; + type Traverse<'a> + = U::Traverse<'a> + where + U: 'a; + + fn raw_traverse<'a>(&'a self, table: &str) -> Result, Self::Error> { + self.reader.raw_traverse(table) + } + fn raw_get<'a>( &'a self, table: &str, @@ -57,7 +66,7 @@ impl HotKvRead for RevmRead { self.reader.raw_get_dual(table, key1, key2) } - fn get(&self, key: &T::Key) -> Result, Self::Error> { + fn get(&self, key: &T::Key) -> Result, Self::Error> { self.reader.get::(key) } @@ -69,10 +78,10 @@ impl HotKvRead for RevmRead { self.reader.get_dual::(key1, key2) } - fn get_many<'a, T, I>(&self, keys: I) -> Result>, Self::Error> + fn get_many<'a, T, I>(&self, keys: I) -> Result>, Self::Error> where T::Key: 'a, - T: Table, + T: SingleKey, I: IntoIterator, { self.reader.get_many::(keys) @@ -109,6 +118,15 @@ impl fmt::Debug for RevmWrite { impl HotKvRead for RevmWrite { type Error = U::Error; + type Traverse<'a> + = U::Traverse<'a> + where + U: 'a; + + fn raw_traverse<'a>(&'a self, table: &str) -> Result, Self::Error> { + self.writer.raw_traverse(table) + } + fn raw_get<'a>( &'a self, table: &str, @@ -126,7 +144,7 @@ impl HotKvRead for RevmWrite { self.writer.raw_get_dual(table, key1, key2) } - fn get(&self, key: &T::Key) -> Result, Self::Error> { + fn get(&self, key: &T::Key) -> Result, Self::Error> { self.writer.get::(key) } @@ -138,10 +156,10 @@ impl HotKvRead for RevmWrite { self.writer.get_dual::(key1, key2) } - fn get_many<'a, T, I>(&self, keys: I) -> Result>, Self::Error> + fn get_many<'a, T, I>(&self, keys: I) -> Result>, Self::Error> where T::Key: 'a, - T: Table, + T: SingleKey, I: IntoIterator, { self.writer.get_many::(keys) @@ -149,6 +167,18 @@ impl HotKvRead for RevmWrite { } impl HotKvWrite for RevmWrite { + type TraverseMut<'a> + = U::TraverseMut<'a> + where + U: 'a; + + fn raw_traverse_mut<'a>( + &'a mut self, + table: &str, + ) -> Result, Self::Error> { + self.writer.raw_traverse_mut(table) + } + fn queue_raw_put(&mut self, table: &str, key: &[u8], value: &[u8]) -> Result<(), Self::Error> { self.writer.queue_raw_put(table, key, value) } @@ -184,7 +214,11 @@ impl HotKvWrite for RevmWrite { self.writer.raw_commit() } - fn queue_put(&mut self, key: &T::Key, value: &T::Value) -> Result<(), Self::Error> { + fn queue_put( + &mut self, + key: &T::Key, + value: &T::Value, + ) -> Result<(), Self::Error> { self.writer.queue_put::(key, value) } @@ -197,13 +231,13 @@ impl HotKvWrite for RevmWrite { self.writer.queue_put_dual::(key1, key2, value) } - fn queue_delete(&mut self, key: &T::Key) -> Result<(), Self::Error> { + fn queue_delete(&mut self, key: &T::Key) -> Result<(), Self::Error> { self.writer.queue_delete::(key) } fn queue_put_many<'a, 'b, T, I>(&mut self, entries: I) -> Result<(), Self::Error> where - T: Table, + T: SingleKey, T::Key: 'a, T::Value: 'b, I: IntoIterator, @@ -408,7 +442,10 @@ where mod tests { use super::*; use crate::{ - hot::{HotKv, HotKvRead, HotKvWrite, MemKv}, + hot::{ + mem::MemKv, + model::{HotKv, HotKvRead, HotKvWrite}, + }, tables::hot::{Bytecodes, PlainAccountState}, }; use alloy::primitives::{Address, B256, U256}; diff --git a/crates/storage/src/hot/traits.rs b/crates/storage/src/hot/model/traits.rs similarity index 78% rename from crates/storage/src/hot/traits.rs rename to crates/storage/src/hot/model/traits.rs index 567ea5d..07942b8 100644 --- a/crates/storage/src/hot/traits.rs +++ b/crates/storage/src/hot/model/traits.rs @@ -1,16 +1,13 @@ -use std::borrow::Cow; - use crate::{ - hot::{ - HotKvError, HotKvReadError, + hot::model::{ + DualTableCursor, GetManyItem, HotKvError, HotKvReadError, KvTraverse, KvTraverseMut, + TableCursor, revm::{RevmRead, RevmWrite}, }, ser::{KeySer, MAX_KEY_SIZE, ValSer}, - tables::{DualKeyed, Table}, + tables::{DualKeyed, SingleKey, Table}, }; - -/// A key-value pair from a table. -pub type KeyValue<'a, T> = (&'a ::Key, Option<::Value>); +use std::borrow::Cow; /// Trait for hot storage. This is a KV store with read/write transactions. #[auto_impl::auto_impl(&, Arc, Box)] @@ -67,6 +64,14 @@ pub trait HotKvRead { /// Error type for read operations. type Error: HotKvReadError; + /// The cursor type for traversing key-value pairs. + type Traverse<'a>: KvTraverse + where + Self: 'a; + + /// Get a raw cursor to traverse the database. + fn raw_traverse<'a>(&'a self, table: &str) -> Result, Self::Error>; + /// Get a raw value from a specific table. /// /// The `key` buf must be <= [`MAX_KEY_SIZE`] bytes. Implementations are @@ -92,8 +97,25 @@ pub trait HotKvRead { key2: &[u8], ) -> Result>, Self::Error>; + /// Traverse a specific table. Returns a typed cursor wrapper. + fn traverse<'a, T: SingleKey>( + &'a self, + ) -> Result, T, Self::Error>, Self::Error> { + let cursor = self.raw_traverse(T::NAME)?; + Ok(TableCursor::new(cursor)) + } + + /// Traverse a specific dual-keyed table. Returns a typed dual-keyed + /// cursor wrapper. + fn traverse_dual<'a, T: DualKeyed>( + &'a self, + ) -> Result, T, Self::Error>, Self::Error> { + let cursor = self.raw_traverse(T::NAME)?; + Ok(DualTableCursor::new(cursor)) + } + /// Get a value from a specific table. - fn get(&self, key: &T::Key) -> Result, Self::Error> { + fn get(&self, key: &T::Key) -> Result, Self::Error> { let mut key_buf = [0u8; MAX_KEY_SIZE]; let key_bytes = key.encode_key(&mut key_buf); debug_assert!( @@ -152,10 +174,10 @@ pub trait HotKvRead { /// /// If any error occurs during retrieval or deserialization, the entire /// operation will return an error. - fn get_many<'a, T, I>(&self, keys: I) -> Result>, Self::Error> + fn get_many<'a, T, I>(&self, keys: I) -> Result>, Self::Error> where T::Key: 'a, - T: Table, + T: SingleKey, I: IntoIterator, { let mut key_buf = [0u8; MAX_KEY_SIZE]; @@ -173,6 +195,17 @@ pub trait HotKvRead { /// Trait for hot storage write transactions. pub trait HotKvWrite: HotKvRead { + /// The mutable cursor type for traversing key-value pairs. + type TraverseMut<'a>: KvTraverseMut + where + Self: 'a; + + /// Get a raw mutable cursor to traverse the database. + fn raw_traverse_mut<'a>( + &'a mut self, + table: &str, + ) -> Result, Self::Error>; + /// Queue a raw put operation. /// /// The `key` buf must be <= [`MAX_KEY_SIZE`] bytes. Implementations are @@ -215,8 +248,31 @@ pub trait HotKvWrite: HotKvRead { fixed_val: bool, ) -> Result<(), Self::Error>; + /// Traverse a specific table. Returns a mutable typed cursor wrapper. + /// If invoked for a dual-keyed table, it will traverse the primary keys + /// only, and the return value may be implementation-defined. + fn traverse_mut<'a, T: SingleKey>( + &'a mut self, + ) -> Result, T, Self::Error>, Self::Error> { + let cursor = self.raw_traverse_mut(T::NAME)?; + Ok(TableCursor::new(cursor)) + } + + /// Traverse a specific dual-keyed table. Returns a mutable typed + /// dual-keyed cursor wrapper. + fn traverse_dual_mut<'a, T: DualKeyed>( + &'a mut self, + ) -> Result, T, Self::Error>, Self::Error> { + let cursor = self.raw_traverse_mut(T::NAME)?; + Ok(DualTableCursor::new(cursor)) + } + /// Queue a put operation for a specific table. - fn queue_put(&mut self, key: &T::Key, value: &T::Value) -> Result<(), Self::Error> { + fn queue_put( + &mut self, + key: &T::Key, + value: &T::Value, + ) -> Result<(), Self::Error> { let mut key_buf = [0u8; MAX_KEY_SIZE]; let key_bytes = key.encode_key(&mut key_buf); let value_bytes = value.encoded(); @@ -241,7 +297,7 @@ pub trait HotKvWrite: HotKvRead { } /// Queue a delete operation for a specific table. - fn queue_delete(&mut self, key: &T::Key) -> Result<(), Self::Error> { + fn queue_delete(&mut self, key: &T::Key) -> Result<(), Self::Error> { let mut key_buf = [0u8; MAX_KEY_SIZE]; let key_bytes = key.encode_key(&mut key_buf); @@ -251,7 +307,7 @@ pub trait HotKvWrite: HotKvRead { /// Queue many put operations for a specific table. fn queue_put_many<'a, 'b, T, I>(&mut self, entries: I) -> Result<(), Self::Error> where - T: Table, + T: SingleKey, T::Key: 'a, T::Value: 'b, I: IntoIterator, diff --git a/crates/storage/src/hot/model/traverse.rs b/crates/storage/src/hot/model/traverse.rs new file mode 100644 index 0000000..12cefd4 --- /dev/null +++ b/crates/storage/src/hot/model/traverse.rs @@ -0,0 +1,399 @@ +//! Cursor traversal traits and typed wrappers for database navigation. + +use crate::{ + hot::model::{DualKeyValue, HotKvReadError, KeyValue, RawDualKeyValue, RawKeyValue, RawValue}, + ser::{KeySer, MAX_KEY_SIZE}, + tables::{DualKeyed, Table}, +}; +use std::ops::Range; + +/// Trait for traversing key-value pairs in the database. +pub trait KvTraverse { + /// Set position to the first key-value pair in the database, and return + /// the KV pair. + fn first<'a>(&'a mut self) -> Result>, E>; + + /// Set position to the last key-value pair in the database, and return the + /// KV pair. + fn last<'a>(&'a mut self) -> Result>, E>; + + /// Set the cursor to specific key in the database, and return the EXACT KV + /// pair if it exists. + fn exact<'a>(&'a mut self, key: &[u8]) -> Result>, E>; + + /// Seek to the next key-value pair AT OR ABOVE the specified key in the + /// database, and return that KV pair. + fn lower_bound<'a>(&'a mut self, key: &[u8]) -> Result>, E>; + + /// Get the next key-value pair in the database, and advance the cursor. + /// + /// Returning `Ok(None)` indicates the cursor is past the end of the + /// database. + fn read_next<'a>(&'a mut self) -> Result>, E>; + + /// Get the previous key-value pair in the database, and move the cursor. + /// + /// Returning `Ok(None)` indicates the cursor is before the start of the + /// database. + fn read_prev<'a>(&'a mut self) -> Result>, E>; +} + +/// Trait for traversing key-value pairs in the database with mutation +/// capabilities. +pub trait KvTraverseMut: KvTraverse { + /// Delete the current key-value pair in the database. + fn delete_current(&mut self) -> Result<(), E>; + + /// Delete a range of key-value pairs in the database, from `start_key` + fn delete_range(&mut self, range: Range<&[u8]>) -> Result<(), E> { + let _ = self.exact(range.start)?; + while let Some((key, _value)) = self.read_next()? { + if key.as_ref() >= range.end { + break; + } + self.delete_current()?; + } + Ok(()) + } +} + +/// Trait for traversing dual-keyed key-value pairs in the database. +pub trait DualKeyedTraverse: KvTraverse { + /// Set the cursor to specific dual key in the database, and return the + /// EXACT KV pair if it exists. + /// + /// Returning `Ok(None)` indicates the exact dual key does not exist. + fn exact_dual<'a>(&'a mut self, key1: &[u8], key2: &[u8]) -> Result>, E>; + + /// Seek to the next key-value pair AT or ABOVE the specified dual key in + /// the database, and return that KV pair. + /// + /// Returning `Ok(None)` indicates there are no more key-value pairs above + /// the specified dual key. + fn next_dual_above<'a>( + &'a mut self, + key1: &[u8], + key2: &[u8], + ) -> Result>, E>; + + /// Move the cursor to the next distinct key1, and return the first + /// key-value pair with that key1. + /// + /// Returning `Ok(None)` indicates there are no more distinct key1 values. + fn next_k1<'a>(&'a mut self) -> Result>, E>; + + /// Move the cursor to the next distinct key2 for the current key1, and + /// return the first key-value pair with that key2. + fn next_k2<'a>(&'a mut self) -> Result>, E>; +} + +// ============================================================================ +// Typed Extension Traits +// ============================================================================ + +/// Extension trait for typed table traversal. +/// +/// This trait provides type-safe access to table entries by encoding keys +/// and decoding values according to the table's schema. +pub trait TableTraverse: KvTraverse { + /// Get the first key-value pair in the table. + fn first(&mut self) -> Result>, E> { + KvTraverse::first(self)?.map(T::decode_kv_tuple).transpose().map_err(Into::into) + } + + /// Get the last key-value pair in the table. + fn last(&mut self) -> Result>, E> { + KvTraverse::last(self)?.map(T::decode_kv_tuple).transpose().map_err(Into::into) + } + + /// Set the cursor to a specific key and return the EXACT value if it exists. + fn exact(&mut self, key: &T::Key) -> Result, E> { + let mut key_buf = [0u8; MAX_KEY_SIZE]; + let key_bytes = key.encode_key(&mut key_buf); + + KvTraverse::exact(self, key_bytes)?.map(T::decode_value).transpose().map_err(Into::into) + } + + /// Seek to the next key-value pair AT OR ABOVE the specified key. + fn lower_bound(&mut self, key: &T::Key) -> Result>, E> { + let mut key_buf = [0u8; MAX_KEY_SIZE]; + let key_bytes = key.encode_key(&mut key_buf); + + KvTraverse::lower_bound(self, key_bytes)? + .map(T::decode_kv_tuple) + .transpose() + .map_err(Into::into) + } + + /// Get the next key-value pair and advance the cursor. + fn read_next(&mut self) -> Result>, E> { + KvTraverse::read_next(self)?.map(T::decode_kv_tuple).transpose().map_err(Into::into) + } + + /// Get the previous key-value pair and move the cursor backward. + fn read_prev(&mut self) -> Result>, E> { + KvTraverse::read_prev(self)?.map(T::decode_kv_tuple).transpose().map_err(Into::into) + } +} + +/// Blanket implementation of `TableTraverse` for any cursor that implements `KvTraverse`. +impl TableTraverse for C +where + C: KvTraverse, + T: Table, + E: HotKvReadError, +{ +} + +/// Extension trait for typed table traversal with mutation capabilities. +pub trait TableTraverseMut: KvTraverseMut { + /// Delete the current key-value pair. + fn delete_current(&mut self) -> Result<(), E> { + KvTraverseMut::delete_current(self) + } + + /// Delete a range of key-value pairs. + fn delete_range(&mut self, range: Range) -> Result<(), E> { + let mut start_key_buf = [0u8; MAX_KEY_SIZE]; + let mut end_key_buf = [0u8; MAX_KEY_SIZE]; + let start_key_bytes = range.start.encode_key(&mut start_key_buf); + let end_key_bytes = range.end.encode_key(&mut end_key_buf); + + KvTraverseMut::delete_range(self, start_key_bytes..end_key_bytes) + } +} + +/// Blanket implementation of `TableTraverseMut` for any cursor that implements `KvTraverseMut`. +impl TableTraverseMut for C +where + C: KvTraverseMut, + T: Table, + E: HotKvReadError, +{ +} + +/// A typed cursor wrapper for traversing dual-keyed tables. +/// +/// This is an extension trait rather than a wrapper struct because MDBX +/// requires specialized implementations for DUPSORT tables that need access +/// to the table type `T` to handle fixed-size values correctly. +pub trait DualTableTraverse { + /// Return the EXACT value for the specified dual key if it exists. + fn exact_dual(&mut self, key1: &T::Key, key2: &T::Key2) -> Result, E> { + let Some((k1, k2, v)) = self.next_dual_above(key1, key2)? else { + return Ok(None); + }; + + if k1 == *key1 && k2 == *key2 { Ok(Some(v)) } else { Ok(None) } + } + + /// Seek to the next key-value pair AT or ABOVE the specified dual key. + fn next_dual_above( + &mut self, + key1: &T::Key, + key2: &T::Key2, + ) -> Result>, E>; + + /// Seek to the next distinct key1, and return the first key-value pair with that key1. + fn next_k1(&mut self) -> Result>, E>; + + /// Seek to the next distinct key2 for the current key1. + fn next_k2(&mut self) -> Result>, E>; +} + +// ============================================================================ +// Wrapper Structs +// ============================================================================ + +use core::marker::PhantomData; + +/// A wrapper struct for typed table traversal. +/// +/// This struct wraps a raw cursor and provides type-safe access to table +/// entries. It implements `TableTraverse` by delegating to the inner +/// cursor. +#[derive(Debug)] +pub struct TableCursor { + inner: C, + _marker: PhantomData (T, E)>, +} + +impl TableCursor { + /// Create a new typed table cursor wrapper. + pub const fn new(cursor: C) -> Self { + Self { inner: cursor, _marker: PhantomData } + } + + /// Get a reference to the inner cursor. + pub const fn inner(&self) -> &C { + &self.inner + } + + /// Get a mutable reference to the inner cursor. + pub fn inner_mut(&mut self) -> &mut C { + &mut self.inner + } + + /// Consume the wrapper and return the inner cursor. + pub fn into_inner(self) -> C { + self.inner + } +} + +impl TableCursor +where + C: KvTraverse, + T: Table, + E: HotKvReadError, +{ + /// Get the first key-value pair in the table. + pub fn first(&mut self) -> Result>, E> { + TableTraverse::::first(&mut self.inner) + } + + /// Get the last key-value pair in the table. + pub fn last(&mut self) -> Result>, E> { + TableTraverse::::last(&mut self.inner) + } + + /// Set the cursor to a specific key and return the EXACT value if it exists. + pub fn exact(&mut self, key: &T::Key) -> Result, E> { + TableTraverse::::exact(&mut self.inner, key) + } + + /// Seek to the next key-value pair AT OR ABOVE the specified key. + pub fn lower_bound(&mut self, key: &T::Key) -> Result>, E> { + TableTraverse::::lower_bound(&mut self.inner, key) + } + + /// Get the next key-value pair and advance the cursor. + pub fn read_next(&mut self) -> Result>, E> { + TableTraverse::::read_next(&mut self.inner) + } + + /// Get the previous key-value pair and move the cursor backward. + pub fn read_prev(&mut self) -> Result>, E> { + TableTraverse::::read_prev(&mut self.inner) + } +} + +impl TableCursor +where + C: KvTraverseMut, + T: Table, + E: HotKvReadError, +{ + /// Delete the current key-value pair. + pub fn delete_current(&mut self) -> Result<(), E> { + TableTraverseMut::::delete_current(&mut self.inner) + } + + /// Delete a range of key-value pairs. + pub fn delete_range(&mut self, range: Range) -> Result<(), E> { + TableTraverseMut::::delete_range(&mut self.inner, range) + } +} + +/// A wrapper struct for typed dual-keyed table traversal. +/// +/// This struct wraps a raw cursor and provides type-safe access to dual-keyed +/// table entries. It delegates to the `DualTableTraverse` trait +/// implementation on the inner cursor. +#[derive(Debug)] +pub struct DualTableCursor { + inner: C, + _marker: PhantomData (T, E)>, +} + +impl DualTableCursor { + /// Create a new typed dual-keyed table cursor wrapper. + pub const fn new(cursor: C) -> Self { + Self { inner: cursor, _marker: PhantomData } + } + + /// Get a reference to the inner cursor. + pub const fn inner(&self) -> &C { + &self.inner + } + + /// Get a mutable reference to the inner cursor. + pub fn inner_mut(&mut self) -> &mut C { + &mut self.inner + } + + /// Consume the wrapper and return the inner cursor. + pub fn into_inner(self) -> C { + self.inner + } +} + +impl DualTableCursor +where + C: DualTableTraverse, + T: DualKeyed, + E: HotKvReadError, +{ + /// Return the EXACT value for the specified dual key if it exists. + pub fn exact_dual(&mut self, key1: &T::Key, key2: &T::Key2) -> Result, E> { + DualTableTraverse::::exact_dual(&mut self.inner, key1, key2) + } + + /// Seek to the next key-value pair AT or ABOVE the specified dual key. + pub fn next_dual_above( + &mut self, + key1: &T::Key, + key2: &T::Key2, + ) -> Result>, E> { + DualTableTraverse::::next_dual_above(&mut self.inner, key1, key2) + } + + /// Seek to the next distinct key1, and return the first key-value pair with that key1. + pub fn next_k1(&mut self) -> Result>, E> { + DualTableTraverse::::next_k1(&mut self.inner) + } + + /// Seek to the next distinct key2 for the current key1. + pub fn next_k2(&mut self) -> Result>, E> { + DualTableTraverse::::next_k2(&mut self.inner) + } +} + +// Also provide access to single-key traversal methods for dual-keyed cursors +impl DualTableCursor +where + C: KvTraverse, + T: DualKeyed, + E: HotKvReadError, +{ + /// Get the first key-value pair in the table (raw traversal). + pub fn first(&mut self) -> Result>, E> { + TableTraverse::::first(&mut self.inner) + } + + /// Get the last key-value pair in the table (raw traversal). + pub fn last(&mut self) -> Result>, E> { + TableTraverse::::last(&mut self.inner) + } + + /// Get the next key-value pair and advance the cursor. + pub fn read_next(&mut self) -> Result>, E> { + TableTraverse::::read_next(&mut self.inner) + } + + /// Get the previous key-value pair and move the cursor backward. + pub fn read_prev(&mut self) -> Result>, E> { + TableTraverse::::read_prev(&mut self.inner) + } +} + +impl DualTableCursor +where + C: KvTraverseMut, + T: DualKeyed, + E: HotKvReadError, +{ + /// Delete the current key-value pair. + pub fn delete_current(&mut self) -> Result<(), E> { + TableTraverseMut::::delete_current(&mut self.inner) + } +} diff --git a/crates/storage/src/ser/impls.rs b/crates/storage/src/ser/impls.rs index bcfb819..6c4f3d1 100644 --- a/crates/storage/src/ser/impls.rs +++ b/crates/storage/src/ser/impls.rs @@ -301,7 +301,7 @@ impl KeySer for BlockNumberAddress { const SIZE: usize = u64::SIZE + Address::SIZE; fn encode_key<'a: 'c, 'b: 'c, 'c>(&'a self, buf: &'b mut [u8; MAX_KEY_SIZE]) -> &'c [u8] { - buf.copy_from_slice(&self.0.0.to_be_bytes()); + buf[0..8].copy_from_slice(&self.0.0.to_be_bytes()); buf[8..28].copy_from_slice(self.0.1.as_ref()); &buf[..Self::SIZE] } diff --git a/crates/storage/src/ser/reth_impls.rs b/crates/storage/src/ser/reth_impls.rs index 66ce986..24d5dd6 100644 --- a/crates/storage/src/ser/reth_impls.rs +++ b/crates/storage/src/ser/reth_impls.rs @@ -13,9 +13,7 @@ use reth::{ }; use reth_db_api::{ BlockNumberList, - models::{ - AccountBeforeTx, ShardedKey, StoredBlockBodyIndices, storage_sharded_key::StorageShardedKey, - }, + models::{AccountBeforeTx, ShardedKey, StoredBlockBodyIndices}, }; // Specialized impls for the sharded key types. This was implemented @@ -61,36 +59,6 @@ macro_rules! sharded_key { sharded_key!(B256); sharded_key!(Address); -impl KeySer for StorageShardedKey { - const SIZE: usize = Address::SIZE + B256::SIZE + u64::SIZE; - - fn encode_key<'a: 'c, 'b: 'c, 'c>(&'a self, buf: &'b mut [u8; MAX_KEY_SIZE]) -> &'c [u8] { - buf[0..Address::SIZE].copy_from_slice(self.address.as_slice()); - buf[Address::SIZE..Address::SIZE + B256::SIZE] - .copy_from_slice(self.sharded_key.key.as_slice()); - buf[Address::SIZE + B256::SIZE..Self::SIZE] - .copy_from_slice(&self.sharded_key.highest_block_number.to_be_bytes()); - - &buf[0..Self::SIZE] - } - - fn decode_key(mut data: &[u8]) -> Result { - if data.len() < Self::SIZE { - return Err(DeserError::InsufficientData { needed: Self::SIZE, available: data.len() }); - } - - let address = Address::from_slice(&data[0..Address::SIZE]); - data = &data[Address::SIZE..]; - - let storage_key = B256::from_slice(&data[0..B256::SIZE]); - data = &data[B256::SIZE..]; - - let highest_block_number = u64::from_be_bytes(data[0..8].try_into().unwrap()); - - Ok(Self { address, sharded_key: ShardedKey { key: storage_key, highest_block_number } }) - } -} - macro_rules! by_props { (@size $($prop:ident),* $(,)?) => { { @@ -1931,39 +1899,6 @@ mod tests { test_key_roundtrip(&key3); } - #[test] - fn test_storage_sharded_key_roundtrips() { - // Test basic cases - let key1 = StorageShardedKey { - address: Address::ZERO, - sharded_key: ShardedKey { key: B256::ZERO, highest_block_number: 0 }, - }; - test_key_roundtrip(&key1); - - // Test with maximum values - let key2 = StorageShardedKey { - address: Address::repeat_byte(0xFF), - sharded_key: ShardedKey { - key: B256::repeat_byte(0xFF), - highest_block_number: u64::MAX, - }, - }; - test_key_roundtrip(&key2); - - // Test with realistic values - let key3 = StorageShardedKey { - address: Address::from([ - 0xd8, 0xda, 0x6b, 0xf2, 0x69, 0x64, 0xaf, 0x9d, 0x7e, 0xed, 0x9e, 0x03, 0xe5, 0x34, - 0x15, 0xd3, 0x7a, 0xa9, 0x60, 0x45, - ]), // Example address - sharded_key: ShardedKey { - key: keccak256(b"storage_slot"), - highest_block_number: 1000000, - }, - }; - test_key_roundtrip(&key3); - } - #[test] fn test_sharded_key_b256_ordering() { let keys = vec![ @@ -1992,36 +1927,6 @@ mod tests { test_key_ordering(&keys); } - #[test] - fn test_storage_sharded_key_ordering() { - let keys = vec![ - StorageShardedKey { - address: Address::ZERO, - sharded_key: ShardedKey { key: B256::ZERO, highest_block_number: 0 }, - }, - StorageShardedKey { - address: Address::ZERO, - sharded_key: ShardedKey { key: B256::ZERO, highest_block_number: 1 }, - }, - StorageShardedKey { - address: Address::ZERO, - sharded_key: ShardedKey { key: B256::from([0x01; 32]), highest_block_number: 0 }, - }, - StorageShardedKey { - address: Address::from([0x01; 20]), - sharded_key: ShardedKey { key: B256::ZERO, highest_block_number: 0 }, - }, - StorageShardedKey { - address: Address::repeat_byte(0xFF), - sharded_key: ShardedKey { - key: B256::repeat_byte(0xFF), - highest_block_number: u64::MAX, - }, - }, - ]; - test_key_ordering(&keys); - } - #[test] fn test_key_decode_insufficient_data() { // Test ShardedKey with insufficient data @@ -2043,15 +1948,6 @@ mod tests { } other => panic!("Expected InsufficientData error, got: {:?}", other), } - - // Test StorageShardedKey with insufficient data - match StorageShardedKey::decode_key(&short_data) { - Err(DeserError::InsufficientData { needed, available }) => { - assert_eq!(needed, StorageShardedKey::SIZE); - assert_eq!(available, 10); - } - other => panic!("Expected InsufficientData error, got: {:?}", other), - } } #[test] @@ -2071,36 +1967,4 @@ mod tests { // Test that ordering is preserved across boundaries test_key_ordering(&boundary_keys); } - - #[test] - fn test_storage_sharded_key_component_ordering() { - // Test that address takes precedence in ordering - let addr1 = Address::from([0x01; 20]); - let addr2 = Address::from([0x02; 20]); - - let key1 = StorageShardedKey { - address: addr1, - sharded_key: ShardedKey { - key: B256::repeat_byte(0xFF), - highest_block_number: u64::MAX, - }, - }; - - let key2 = StorageShardedKey { - address: addr2, - sharded_key: ShardedKey { key: B256::ZERO, highest_block_number: 0 }, - }; - - // key1 should be less than key2 because addr1 < addr2 - assert!(key1 < key2); - - // This should be reflected in the encoding - let mut buf1 = [0u8; MAX_KEY_SIZE]; - let mut buf2 = [0u8; MAX_KEY_SIZE]; - - let encoded1 = key1.encode_key(&mut buf1); - let encoded2 = key2.encode_key(&mut buf2); - - assert!(encoded1 < encoded2, "Encoding should preserve ordering"); - } } diff --git a/crates/storage/src/ser/traits.rs b/crates/storage/src/ser/traits.rs index 88e1fef..7d3fcb4 100644 --- a/crates/storage/src/ser/traits.rs +++ b/crates/storage/src/ser/traits.rs @@ -14,7 +14,7 @@ pub const MAX_KEY_SIZE: usize = 64; /// /// In practice, keys are often hashes, addresses, numbers, or composites /// of these. -pub trait KeySer: Ord + Sized { +pub trait KeySer: PartialOrd + Ord + Sized + Clone + core::fmt::Debug { /// The fixed size of the serialized key in bytes. /// Must satisfy `SIZE <= MAX_KEY_SIZE`. const SIZE: usize; diff --git a/crates/storage/src/tables/macros.rs b/crates/storage/src/tables/macros.rs index 85e479f..77cf8ce 100644 --- a/crates/storage/src/tables/macros.rs +++ b/crates/storage/src/tables/macros.rs @@ -16,6 +16,7 @@ macro_rules! table { type Key = $key; type Value = $value; } + }; ( @@ -30,6 +31,8 @@ macro_rules! table { false, None ); + + impl crate::tables::SingleKey for $name {} }; diff --git a/crates/storage/src/tables/mod.rs b/crates/storage/src/tables/mod.rs index 572daac..f93dad5 100644 --- a/crates/storage/src/tables/mod.rs +++ b/crates/storage/src/tables/mod.rs @@ -7,7 +7,10 @@ pub mod cold; /// Tables that are hot, or conditionally hot. pub mod hot; -use crate::ser::{KeySer, ValSer}; +use crate::{ + hot::model::{DualKeyValue, KeyValue}, + ser::{DeserError, KeySer, ValSer}, +}; /// The maximum size of a dual key (in bytes). pub const MAX_FIXED_VAL_SIZE: usize = 64; @@ -40,6 +43,42 @@ pub trait Table { type Key: KeySer; /// The value type. type Value: ValSer; + + /// Shortcut to decode a key. + fn decode_key(data: impl AsRef<[u8]>) -> Result { + ::decode_key(data.as_ref()) + } + + /// Shortcut to decode a value. + fn decode_value(data: impl AsRef<[u8]>) -> Result { + ::decode_value(data.as_ref()) + } + + /// Shortcut to decode a key-value pair. + fn decode_kv( + key_data: impl AsRef<[u8]>, + value_data: impl AsRef<[u8]>, + ) -> Result, DeserError> { + let key = Self::decode_key(key_data)?; + let value = Self::decode_value(value_data)?; + Ok((key, value)) + } + + /// Shortcut to decode a key-value tuple. + fn decode_kv_tuple( + data: (impl AsRef<[u8]>, impl AsRef<[u8]>), + ) -> Result, DeserError> { + Self::decode_kv(data.0, data.1) + } +} + +/// Trait for tables with a single key. +pub trait SingleKey: Table { + /// Compile-time assertions for the single-keyed table. + #[doc(hidden)] + const ASSERT: () = { + assert!(!Self::DUAL_KEY, "SingleKey tables must have DUAL_KEY = false"); + }; } /// Trait for tables with two keys. @@ -56,4 +95,39 @@ pub trait DualKeyed: Table { const ASSERT: () = { assert!(Self::DUAL_KEY, "DualKeyed tables must have DUAL_KEY = true"); }; + + /// Shortcut to decode the second key. + fn decode_key2(data: impl AsRef<[u8]>) -> Result { + ::decode_key(data.as_ref()) + } + + /// Shortcut to decode a prepended value. This is useful for some table + /// implementations. + fn decode_prepended_value( + data: impl AsRef<[u8]>, + ) -> Result<(Self::Key2, Self::Value), DeserError> { + let data = data.as_ref(); + let key = Self::decode_key2(&data[..Self::Key2::SIZE])?; + let value = Self::decode_value(&data[Self::Key2::SIZE..])?; + Ok((key, value)) + } + + /// Shortcut to decode a dual key-value triplet. + fn decode_kkv( + key1_data: impl AsRef<[u8]>, + key2_data: impl AsRef<[u8]>, + value_data: impl AsRef<[u8]>, + ) -> Result, DeserError> { + let key1 = Self::decode_key(key1_data)?; + let key2 = Self::decode_key2(key2_data)?; + let value = Self::decode_value(value_data)?; + Ok((key1, key2, value)) + } + + /// Shortcut to decode a dual key-value tuple. + fn decode_kkv_tuple( + data: (impl AsRef<[u8]>, impl AsRef<[u8]>, impl AsRef<[u8]>), + ) -> Result, DeserError> { + Self::decode_kkv(data.0, data.1, data.2) + } }