feat: generic static file masks (#12785)

crates/blockchain-tree/src/externals.rs

@@ -2,7 +2,7 @@
 
 use alloy_primitives::{BlockHash, BlockNumber};
 use reth_consensus::Consensus;
-use reth_db::{static_file::HeaderMask, tables};
+use reth_db::{static_file::BlockHashMask, tables};
 use reth_db_api::{cursor::DbCursorRO, transaction::DbTx};
 use reth_node_types::{FullNodePrimitives, NodeTypesWithDB};
 use reth_primitives::{BlockBody, StaticFileSegment};
@@ -85,7 +85,7 @@ impl<N: ProviderNodeTypes, E> TreeExternals<N, E> {
             hashes.extend(range.clone().zip(static_file_provider.fetch_range_with_predicate(
                 StaticFileSegment::Headers,
                 range,
-                |cursor, number| cursor.get_one::<HeaderMask<BlockHash>>(number.into()),
+                |cursor, number| cursor.get_one::<BlockHashMask>(number.into()),
                 |_| true,
             )?));
         }

crates/cli/commands/src/db/get.rs

@@ -2,7 +2,9 @@ use alloy_consensus::Header;
 use alloy_primitives::{hex, BlockHash};
 use clap::Parser;
 use reth_db::{
-    static_file::{ColumnSelectorOne, ColumnSelectorTwo, HeaderMask, ReceiptMask, TransactionMask},
+    static_file::{
+        ColumnSelectorOne, ColumnSelectorTwo, HeaderWithHashMask, ReceiptMask, TransactionMask,
+    },
     tables, RawKey, RawTable, Receipts, TableViewer, Transactions,
 };
 use reth_db_api::table::{Decompress, DupSort, Table};
@@ -61,7 +63,7 @@ impl Command {
             Subcommand::StaticFile { segment, key, raw } => {
                 let (key, mask): (u64, _) = match segment {
                     StaticFileSegment::Headers => {
-                        (table_key::<tables::Headers>(&key)?, <HeaderMask<Header, BlockHash>>::MASK)
+                        (table_key::<tables::Headers>(&key)?, <HeaderWithHashMask<Header>>::MASK)
                     }
                     StaticFileSegment::Transactions => (
                         table_key::<tables::Transactions>(&key)?,

crates/stages/stages/src/stages/bodies.rs

@@ -523,9 +523,9 @@ mod tests {
             },
         };
         use alloy_consensus::Header;
-        use alloy_primitives::{BlockHash, BlockNumber, TxNumber, B256};
+        use alloy_primitives::{BlockNumber, TxNumber, B256};
         use futures_util::Stream;
-        use reth_db::{static_file::HeaderMask, tables};
+        use reth_db::{static_file::HeaderWithHashMask, tables};
         use reth_db_api::{
             cursor::DbCursorRO,
             models::{StoredBlockBodyIndices, StoredBlockOmmers},
@@ -813,7 +813,7 @@ mod tests {
                 for header in static_file_provider.fetch_range_iter(
                     StaticFileSegment::Headers,
                     *range.start()..*range.end() + 1,
-                    |cursor, number| cursor.get_two::<HeaderMask<Header, BlockHash>>(number.into()),
+                    |cursor, number| cursor.get_two::<HeaderWithHashMask<Header>>(number.into()),
                 )? {
                     let (header, hash) = header?;
                     self.headers.push_back(SealedHeader::new(header, hash));

crates/storage/db/Cargo.toml

@@ -48,7 +48,6 @@ page_size = { version = "0.6.0", optional = true }
 thiserror.workspace = true
 tempfile = { workspace = true, optional = true }
 derive_more.workspace = true
-paste.workspace = true
 rustc-hash = { workspace = true, optional = true }
 sysinfo = { version = "0.31", default-features = false, features = ["system"] }
 parking_lot = { workspace = true, optional = true }

crates/storage/db/src/static_file/mask.rs

@@ -1,38 +1,5 @@
 use reth_db_api::table::Decompress;
 
-/// Generic Mask helper struct for selecting specific column values to read and decompress.
-///
-/// #### Explanation:
-///
-/// A `NippyJar` static file row can contain multiple column values. To specify the column values
-/// to be read, a mask is utilized.
-///
-/// For example, a static file with three columns, if the first and last columns are queried, the
-/// mask `0b101` would be passed. To select only the second column, the mask `0b010` would be used.
-///
-/// Since each static file has its own column distribution, different wrapper types are necessary.
-/// For instance, `B256` might be the third column in the `Header` segment, while being the second
-/// column in another segment. Hence, `Mask<B256>` would only be applicable to one of these
-/// scenarios.
-///
-/// Alongside, the column selector traits (eg. [`ColumnSelectorOne`]) this provides a structured way
-/// to tie the types to be decoded to the mask necessary to query them.
-#[derive(Debug)]
-pub struct Mask<FIRST, SECOND = (), THIRD = ()>(std::marker::PhantomData<(FIRST, SECOND, THIRD)>);
-
-macro_rules! add_segments {
-    ($($segment:tt),+) => {
-        paste::paste! {
-            $(
-                #[doc = concat!("Mask for ", stringify!($segment), " static file segment. See [`Mask`] for more.")]
-                #[derive(Debug)]
-                pub struct [<$segment Mask>]<FIRST, SECOND = (), THIRD = ()>(Mask<FIRST, SECOND, THIRD>);
-            )+
-        }
-    };
-}
-add_segments!(Header, Receipt, Transaction);
-
 ///  Trait for specifying a mask to select one column value.
 pub trait ColumnSelectorOne {
     /// First desired column value
@@ -66,21 +33,45 @@ pub trait ColumnSelectorThree {
 #[macro_export]
 /// Add mask to select `N` column values from a specific static file segment row.
 macro_rules! add_static_file_mask {
-    ($mask_struct:tt, $type1:ty, $mask:expr) => {
+    ($(#[$attr:meta])* $mask_struct:ident $(<$generic:ident>)?, $type1:ty, $mask:expr) => {
-        impl ColumnSelectorOne for $mask_struct<$type1> {
+        $(#[$attr])*
+        #[derive(Debug)]
+        pub struct $mask_struct$(<$generic>)?$((std::marker::PhantomData<$generic>))?;
+
+        impl$(<$generic>)? ColumnSelectorOne for $mask_struct$(<$generic>)?
+        where
+            $type1: Send + Sync + std::fmt::Debug + reth_db_api::table::Decompress,
+        {
             type FIRST = $type1;
             const MASK: usize = $mask;
         }
     };
-    ($mask_struct:tt, $type1:ty, $type2:ty, $mask:expr) => {
+    ($(#[$attr:meta])* $mask_struct:ident $(<$generic:ident>)?, $type1:ty, $type2:ty, $mask:expr) => {
-        impl ColumnSelectorTwo for $mask_struct<$type1, $type2> {
+        $(#[$attr])*
+        #[derive(Debug)]
+        pub struct $mask_struct$(<$generic>)?$((std::marker::PhantomData<$generic>))?;
+
+        impl$(<$generic>)? ColumnSelectorTwo for $mask_struct$(<$generic>)?
+        where
+            $type1: Send + Sync + std::fmt::Debug + reth_db_api::table::Decompress,
+            $type2: Send + Sync + std::fmt::Debug + reth_db_api::table::Decompress,
+        {
             type FIRST = $type1;
             type SECOND = $type2;
             const MASK: usize = $mask;
         }
     };
-    ($mask_struct:tt, $type1:ty, $type2:ty, $type3:ty, $mask:expr) => {
+    ($(#[$attr:meta])* $mask_struct:ident $(<$generic:ident>)?, $type1:ty, $type2:ty, $type3:ty, $mask:expr) => {
-        impl ColumnSelectorThree for $mask_struct<$type1, $type2, $type3> {
+        $(#[$attr])*
+        #[derive(Debug)]
+        pub struct $mask_struct$(<$generic>)?$((std::marker::PhantomData<$generic>))?;
+
+        impl$(<$generic>)? ColumnSelectorThree for $mask_struct$(<$generic>)?
+        where
+            $type1: Send + Sync + std::fmt::Debug + reth_db_api::table::Decompress,
+            $type2: Send + Sync + std::fmt::Debug + reth_db_api::table::Decompress,
+            $type3: Send + Sync + std::fmt::Debug + reth_db_api::table::Decompress,
+        {
             type FIRST = $type1;
             type SECOND = $type2;
             type THIRD = $type3;

crates/storage/db/src/static_file/masks.rs

@@ -1,23 +1,44 @@
-use super::{ReceiptMask, TransactionMask};
 use crate::{
     add_static_file_mask,
-    static_file::mask::{ColumnSelectorOne, ColumnSelectorTwo, HeaderMask},
+    static_file::mask::{ColumnSelectorOne, ColumnSelectorTwo},
-    HeaderTerminalDifficulties, RawValue, Receipts, Transactions,
+    HeaderTerminalDifficulties,
 };
-use alloy_consensus::Header;
 use alloy_primitives::BlockHash;
 use reth_db_api::table::Table;
 
 // HEADER MASKS
-add_static_file_mask!(HeaderMask, Header, 0b001);
+add_static_file_mask! {
-add_static_file_mask!(HeaderMask, <HeaderTerminalDifficulties as Table>::Value, 0b010);
+    #[doc = "Mask for selecting a single header from Headers static file segment"]
-add_static_file_mask!(HeaderMask, BlockHash, 0b100);
+    HeaderMask<H>, H, 0b001
-add_static_file_mask!(HeaderMask, Header, BlockHash, 0b101);
+}
-add_static_file_mask!(HeaderMask, <HeaderTerminalDifficulties as Table>::Value, BlockHash, 0b110);
+add_static_file_mask! {
+    #[doc = "Mask for selecting a total difficulty value from Headers static file segment"]
+    TotalDifficultyMask, <HeaderTerminalDifficulties as Table>::Value, 0b010
+}
+add_static_file_mask! {
+    #[doc = "Mask for selecting a block hash value from Headers static file segment"]
+    BlockHashMask, BlockHash, 0b100
+}
+add_static_file_mask! {
+    #[doc = "Mask for selecting a header along with block hash from Headers static file segment"]
+    HeaderWithHashMask<H>, H, BlockHash, 0b101
+}
+add_static_file_mask! {
+    #[doc = "Mask for selecting a total difficulty along with block hash from Headers static file segment"]
+    TDWithHashMask,
+    <HeaderTerminalDifficulties as Table>::Value,
+    BlockHash,
+    0b110
+}
 
 // RECEIPT MASKS
-add_static_file_mask!(ReceiptMask, <Receipts as Table>::Value, 0b1);
+add_static_file_mask! {
+    #[doc = "Mask for selecting a single receipt from Receipts static file segment"]
+    ReceiptMask<R>, R, 0b1
+}
 
 // TRANSACTION MASKS
-add_static_file_mask!(TransactionMask, <Transactions as Table>::Value, 0b1);
+add_static_file_mask! {
-add_static_file_mask!(TransactionMask, RawValue<<Transactions as Table>::Value>, 0b1);
+    #[doc = "Mask for selecting a single transaction from Transactions static file segment"]
+    TransactionMask<T>, T, 0b1
+}

crates/storage/db/src/static_file/mod.rs

@@ -17,6 +17,7 @@ use reth_primitives::{
 };
 
 mod masks;
+pub use masks::*;
 
 /// Alias type for a map of [`StaticFileSegment`] and sorted lists of existing static file ranges.
 type SortedStaticFiles =

crates/storage/provider/src/providers/static_file/jar.rs

@@ -10,8 +10,10 @@ use alloy_consensus::Header;
 use alloy_eips::BlockHashOrNumber;
 use alloy_primitives::{Address, BlockHash, BlockNumber, TxHash, TxNumber, B256, U256};
 use reth_chainspec::ChainInfo;
-use reth_db::static_file::{HeaderMask, ReceiptMask, StaticFileCursor, TransactionMask};
+use reth_db::static_file::{
-use reth_db_api::models::CompactU256;
+    BlockHashMask, HeaderMask, HeaderWithHashMask, ReceiptMask, StaticFileCursor, TDWithHashMask,
+    TotalDifficultyMask, TransactionMask,
+};
 use reth_node_types::NodePrimitives;
 use reth_primitives::{
     Receipt, SealedHeader, TransactionMeta, TransactionSigned, TransactionSignedNoHash,
@@ -90,7 +92,7 @@ impl<N: NodePrimitives> HeaderProvider for StaticFileJarProvider<'_, N> {
     fn header(&self, block_hash: &BlockHash) -> ProviderResult<Option<Header>> {
         Ok(self
             .cursor()?
-            .get_two::<HeaderMask<Header, BlockHash>>(block_hash.into())?
+            .get_two::<HeaderWithHashMask<Header>>(block_hash.into())?
             .filter(|(_, hash)| hash == block_hash)
             .map(|(header, _)| header))
     }
@@ -102,13 +104,13 @@ impl<N: NodePrimitives> HeaderProvider for StaticFileJarProvider<'_, N> {
     fn header_td(&self, block_hash: &BlockHash) -> ProviderResult<Option<U256>> {
         Ok(self
             .cursor()?
-            .get_two::<HeaderMask<CompactU256, BlockHash>>(block_hash.into())?
+            .get_two::<TDWithHashMask>(block_hash.into())?
             .filter(|(_, hash)| hash == block_hash)
             .map(|(td, _)| td.into()))
     }
 
     fn header_td_by_number(&self, num: BlockNumber) -> ProviderResult<Option<U256>> {
-        Ok(self.cursor()?.get_one::<HeaderMask<CompactU256>>(num.into())?.map(Into::into))
+        Ok(self.cursor()?.get_one::<TotalDifficultyMask>(num.into())?.map(Into::into))
     }
 
     fn headers_range(&self, range: impl RangeBounds<BlockNumber>) -> ProviderResult<Vec<Header>> {
@@ -129,7 +131,7 @@ impl<N: NodePrimitives> HeaderProvider for StaticFileJarProvider<'_, N> {
     fn sealed_header(&self, number: BlockNumber) -> ProviderResult<Option<SealedHeader>> {
         Ok(self
             .cursor()?
-            .get_two::<HeaderMask<Header, BlockHash>>(number.into())?
+            .get_two::<HeaderWithHashMask<Header>>(number.into())?
             .map(|(header, hash)| SealedHeader::new(header, hash)))
     }
 
@@ -145,7 +147,7 @@ impl<N: NodePrimitives> HeaderProvider for StaticFileJarProvider<'_, N> {
 
         for number in range {
             if let Some((header, hash)) =
-                cursor.get_two::<HeaderMask<Header, BlockHash>>(number.into())?
+                cursor.get_two::<HeaderWithHashMask<Header>>(number.into())?
             {
                 let sealed = SealedHeader::new(header, hash);
                 if !predicate(&sealed) {
@@ -160,7 +162,7 @@ impl<N: NodePrimitives> HeaderProvider for StaticFileJarProvider<'_, N> {
 
 impl<N: NodePrimitives> BlockHashReader for StaticFileJarProvider<'_, N> {
     fn block_hash(&self, number: u64) -> ProviderResult<Option<B256>> {
-        self.cursor()?.get_one::<HeaderMask<BlockHash>>(number.into())
+        self.cursor()?.get_one::<BlockHashMask>(number.into())
     }
 
     fn canonical_hashes_range(
@@ -172,7 +174,7 @@ impl<N: NodePrimitives> BlockHashReader for StaticFileJarProvider<'_, N> {
         let mut hashes = Vec::with_capacity((end - start) as usize);
 
         for number in start..end {
-            if let Some(hash) = cursor.get_one::<HeaderMask<BlockHash>>(number.into())? {
+            if let Some(hash) = cursor.get_one::<BlockHashMask>(number.into())? {
                 hashes.push(hash)
             }
         }
@@ -200,7 +202,7 @@ impl<N: NodePrimitives> BlockNumReader for StaticFileJarProvider<'_, N> {
         let mut cursor = self.cursor()?;
 
         Ok(cursor
-            .get_one::<HeaderMask<BlockHash>>((&hash).into())?
+            .get_one::<BlockHashMask>((&hash).into())?
             .and_then(|res| (res == hash).then(|| cursor.number()).flatten()))
     }
 }

crates/storage/provider/src/providers/static_file/manager.rs

@@ -19,14 +19,14 @@ use parking_lot::RwLock;
 use reth_chainspec::{ChainInfo, ChainSpecProvider};
 use reth_db::{
     lockfile::StorageLock,
-    static_file::{iter_static_files, HeaderMask, ReceiptMask, StaticFileCursor, TransactionMask},
+    static_file::{
+        iter_static_files, BlockHashMask, HeaderMask, HeaderWithHashMask, ReceiptMask,
+        StaticFileCursor, TDWithHashMask, TransactionMask,
+    },
     tables,
 };
 use reth_db_api::{
-    cursor::DbCursorRO,
+    cursor::DbCursorRO, models::StoredBlockBodyIndices, table::Table, transaction::DbTx,
-    models::{CompactU256, StoredBlockBodyIndices},
-    table::Table,
-    transaction::DbTx,
 };
 use reth_nippy_jar::{NippyJar, NippyJarChecker, CONFIG_FILE_EXTENSION};
 use reth_node_types::NodePrimitives;
@@ -1236,7 +1236,7 @@ impl<N: NodePrimitives> HeaderProvider for StaticFileProvider<N> {
         self.find_static_file(StaticFileSegment::Headers, |jar_provider| {
             Ok(jar_provider
                 .cursor()?
-                .get_two::<HeaderMask<Header, BlockHash>>(block_hash.into())?
+                .get_two::<HeaderWithHashMask<Header>>(block_hash.into())?
                 .and_then(|(header, hash)| {
                     if &hash == block_hash {
                         return Some(header)
@@ -1262,7 +1262,7 @@ impl<N: NodePrimitives> HeaderProvider for StaticFileProvider<N> {
         self.find_static_file(StaticFileSegment::Headers, |jar_provider| {
             Ok(jar_provider
                 .cursor()?
-                .get_two::<HeaderMask<CompactU256, BlockHash>>(block_hash.into())?
+                .get_two::<TDWithHashMask>(block_hash.into())?
                 .and_then(|(td, hash)| (&hash == block_hash).then_some(td.0)))
         })
     }
@@ -1310,7 +1310,7 @@ impl<N: NodePrimitives> HeaderProvider for StaticFileProvider<N> {
             to_range(range),
             |cursor, number| {
                 Ok(cursor
-                    .get_two::<HeaderMask<Header, BlockHash>>(number.into())?
+                    .get_two::<HeaderWithHashMask<Header>>(number.into())?
                     .map(|(header, hash)| SealedHeader::new(header, hash)))
             },
             predicate,
@@ -1331,7 +1331,7 @@ impl<N: NodePrimitives> BlockHashReader for StaticFileProvider<N> {
         self.fetch_range_with_predicate(
             StaticFileSegment::Headers,
             start..end,
-            |cursor, number| cursor.get_one::<HeaderMask<BlockHash>>(number.into()),
+            |cursor, number| cursor.get_one::<BlockHashMask>(number.into()),
             |_| true,
         )
     }