refactor: expand table, contract dualkeyed

Author: prestwich

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<T: crate::tables::DualKeyed>(
         &self,
-        key1: &T::K1,
-        key2: &T::K2,
+        key1: &T::Key,
+        key2: &T::Key2,
     ) -> Result<Option<T::Value>, 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) = <T as Table>::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[<<T as DualKeyed>::K2 as KeySer>::SIZE..];
+        let value_bytes = &value_bytes[<<T as DualKeyed>::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<RW> {
     // Specialized put for dual-keyed tables.
     fn queue_put_dual<T: crate::tables::DualKeyed>(
         &mut self,
-        key1: &T::K1,
-        key2: &T::K2,
+        key1: &T::Key,
+        key2: &T::Key2,
         value: &T::Value,
     ) -> Result<(), Self::Error> {
-        let k2_size = <T::K2 as KeySer>::SIZE;
+        let k2_size = <T::Key2 as KeySer>::SIZE;
         let mut scratch = [0u8; MAX_KEY_SIZE];
 
         // This will be the total length of key2 + value, reserved in mdbx

crates/storage/src/hot/revm.rs

@@ -63,8 +63,8 @@ impl<U: HotKvRead> HotKvRead for RevmRead<U> {
 
     fn get_dual<T: DualKeyed>(
         &self,
-        key1: &T::K1,
-        key2: &T::K2,
+        key1: &T::Key,
+        key2: &T::Key2,
     ) -> Result<Option<T::Value>, Self::Error> {
         self.reader.get_dual::<T>(key1, key2)
     }
@@ -132,8 +132,8 @@ impl<U: HotKvWrite> HotKvRead for RevmWrite<U> {
 
     fn get_dual<T: DualKeyed>(
         &self,
-        key1: &T::K1,
-        key2: &T::K2,
+        key1: &T::Key,
+        key2: &T::Key2,
     ) -> Result<Option<T::Value>, Self::Error> {
         self.writer.get_dual::<T>(key1, key2)
     }
@@ -190,8 +190,8 @@ impl<U: HotKvWrite> HotKvWrite for RevmWrite<U> {
 
     fn queue_put_dual<T: DualKeyed>(
         &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::<T>(key1, key2, value)

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<T: DualKeyed>(
         &self,
-        key1: &T::K1,
-        key2: &T::K2,
+        key1: &T::Key,
+        key2: &T::Key2,
     ) -> Result<Option<T::Value>, 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<T: DualKeyed>(
         &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.

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<BlockNumber => Header>,
+    Headers<BlockNumber => Header>
+}
 
+table! {
     /// Records block numbers by hash.
-    HeaderNumbers<B256 => BlockNumber>,
+    HeaderNumbers<B256 => BlockNumber>
+}
 
-    /// Records the canonical chain header hashes, by height.
-    CanonicalHeaders<BlockNumber => B256>,
+table! {
+        /// Records the canonical chain header hashes, by height.
+        CanonicalHeaders<BlockNumber => B256>
+}
 
+table! {
     /// Records contract Bytecode, by its hash.
-    Bytecodes<B256 => Bytecode>,
-
-    /// Records plain account states, keyed by address.
-    PlainAccountState<Address => Account>,
-
+    Bytecodes<B256 => Bytecode>
+}
+table! {
+     /// Records plain account states, keyed by address.
+    PlainAccountState<Address => Account>
+}
+table! {
     /// Records account state change history, keyed by address.
-    AccountsHistory<ShardedKey<Address> => BlockNumberList>,
-
+    AccountsHistory<ShardedKey<Address> => BlockNumberList>
+}
+table! {
     /// Records storage state change history, keyed by address and storage key.
-    StorageHistory<StorageShardedKey => BlockNumberList>,
-
+    StorageHistory<StorageShardedKey => BlockNumberList>
 }
 
-tables! {
+table! {
     /// Records plain storage states, keyed by address and storage key.
-    PlainStorageState<Address => B256 => U256> size: Some(32 + 32),
+    PlainStorageState<Address => B256 => U256> is 32 + 32
+}
 
+table! {
     /// Records account states before transactions, keyed by (address, block number).
-    StorageChangeSets<BlockNumberAddress => B256 => StorageEntry> size: Some(32 + 32 + 32),
+    StorageChangeSets<BlockNumberAddress => B256 => StorageEntry> is 32 + 32 + 32
+}
 
+table! {
     /// Records account states before transactions, keyed by (address, block number).
-    AccountChangeSets<BlockNumberAddress => Address => AccountBeforeTx> size: None,
+    AccountChangeSets<BlockNumberAddress => Address => AccountBeforeTx>
 }

crates/storage/src/tables/macros.rs

@@ -1,14 +1,17 @@
-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)]
         pub struct $name;
 
         impl crate::tables::Table for $name {
             const NAME: &'static str = stringify!($name);
+            const DUAL_KEY: bool = $dual;
+            const FIXED_VAL_SIZE: Option<usize> = $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<usize> = { $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;
+        }
     };
 }

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<usize> = 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<usize> = 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<T> 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");
+    };
 }