feat: remove const generics from cursors and add segment masks (#5181)

Co-authored-by: Alexey Shekhirin <a.shekhirin@gmail.com>

crates/storage/db/src/snapshot/cursor.rs

@@ -0,0 +1,105 @@
+use super::mask::{ColumnSelectorOne, ColumnSelectorThree, ColumnSelectorTwo};
+use crate::table::Decompress;
+use derive_more::{Deref, DerefMut};
+use reth_interfaces::{RethError, RethResult};
+use reth_nippy_jar::{MmapHandle, NippyJar, NippyJarCursor};
+use reth_primitives::{snapshot::SegmentHeader, B256};
+
+/// Cursor of a snapshot segment.
+#[derive(Debug, Deref, DerefMut)]
+pub struct SnapshotCursor<'a>(NippyJarCursor<'a, SegmentHeader>);
+
+impl<'a> SnapshotCursor<'a> {
+    /// Returns a new [`SnapshotCursor`].
+    pub fn new(
+        jar: &'a NippyJar<SegmentHeader>,
+        mmap_handle: MmapHandle,
+    ) -> Result<Self, RethError> {
+        Ok(Self(NippyJarCursor::with_handle(jar, mmap_handle)?))
+    }
+
+    /// Gets a row of values.
+    pub fn get(
+        &mut self,
+        key_or_num: KeyOrNumber<'_>,
+        mask: usize,
+    ) -> RethResult<Option<Vec<&'_ [u8]>>> {
+        let row = match key_or_num {
+            KeyOrNumber::Key(k) => self.row_by_key_with_cols(k, mask),
+            KeyOrNumber::Number(n) => {
+                let offset = self.jar().user_header().start();
+                if offset > n {
+                    return Ok(None)
+                }
+                self.row_by_number_with_cols((n - offset) as usize, mask)
+            }
+        }?;
+
+        Ok(row)
+    }
+
+    /// Gets one column value from a row.
+    pub fn get_one<M: ColumnSelectorOne>(
+        &mut self,
+        key_or_num: KeyOrNumber<'_>,
+    ) -> RethResult<Option<M::FIRST>> {
+        let row = self.get(key_or_num, M::MASK)?;
+
+        match row {
+            Some(row) => Ok(Some(M::FIRST::decompress(row[0])?)),
+            None => Ok(None),
+        }
+    }
+
+    /// Gets two column values from a row.
+    pub fn get_two<M: ColumnSelectorTwo>(
+        &mut self,
+        key_or_num: KeyOrNumber<'_>,
+    ) -> RethResult<Option<(M::FIRST, M::SECOND)>> {
+        let row = self.get(key_or_num, M::MASK)?;
+
+        match row {
+            Some(row) => Ok(Some((M::FIRST::decompress(row[0])?, M::SECOND::decompress(row[1])?))),
+            None => Ok(None),
+        }
+    }
+
+    /// Gets three column values from a row.
+    #[allow(clippy::type_complexity)]
+    pub fn get_three<M: ColumnSelectorThree>(
+        &mut self,
+        key_or_num: KeyOrNumber<'_>,
+    ) -> RethResult<Option<(M::FIRST, M::SECOND, M::THIRD)>> {
+        let row = self.get(key_or_num, M::MASK)?;
+
+        match row {
+            Some(row) => Ok(Some((
+                M::FIRST::decompress(row[0])?,
+                M::SECOND::decompress(row[1])?,
+                M::THIRD::decompress(row[2])?,
+            ))),
+            None => Ok(None),
+        }
+    }
+}
+
+/// Either a key _or_ a block/tx number
+#[derive(Debug)]
+pub enum KeyOrNumber<'a> {
+    /// A slice used as a key. Usually a block/tx hash
+    Key(&'a [u8]),
+    /// A block/tx number
+    Number(u64),
+}
+
+impl<'a> From<&'a B256> for KeyOrNumber<'a> {
+    fn from(value: &'a B256) -> Self {
+        KeyOrNumber::Key(value.as_slice())
+    }
+}
+
+impl<'a> From<u64> for KeyOrNumber<'a> {
+    fn from(value: u64) -> Self {
+        KeyOrNumber::Number(value)
+    }
+}

crates/storage/db/src/snapshot/generation.rs

@@ -1,15 +1,13 @@
-//! reth's snapshot creation from database tables and access
-
 use crate::{
     abstraction::cursor::DbCursorRO,
-    table::{Decompress, Key, Table},
+    table::{Key, Table},
     transaction::DbTx,
     RawKey, RawTable,
 };
-use derive_more::{Deref, DerefMut};
-use reth_interfaces::{RethError, RethResult};
-use reth_nippy_jar::{ColumnResult, MmapHandle, NippyJar, NippyJarCursor, PHFKey};
-use reth_primitives::{snapshot::SegmentHeader, B256};
+
+use reth_interfaces::RethResult;
+use reth_nippy_jar::{ColumnResult, NippyJar, PHFKey};
+
 use reth_tracing::tracing::*;
 use serde::{Deserialize, Serialize};
 use std::{error::Error as StdError, ops::RangeInclusive};
@@ -104,117 +102,3 @@ macro_rules! generate_snapshot_func {
 }
 
 generate_snapshot_func!((T1), (T1, T2), (T1, T2, T3), (T1, T2, T3, T4), (T1, T2, T3, T4, T5),);
-
-/// Cursor of a snapshot segment.
-#[derive(Debug, Deref, DerefMut)]
-pub struct SnapshotCursor<'a>(NippyJarCursor<'a, SegmentHeader>);
-
-impl<'a> SnapshotCursor<'a> {
-    /// Returns a new [`SnapshotCursor`].
-    pub fn new(
-        jar: &'a NippyJar<SegmentHeader>,
-        mmap_handle: MmapHandle,
-    ) -> Result<Self, RethError> {
-        Ok(Self(NippyJarCursor::with_handle(jar, mmap_handle)?))
-    }
-
-    /// Gets a row of values.
-    pub fn get<const SELECTOR: usize, const COLUMNS: usize>(
-        &mut self,
-        key_or_num: KeyOrNumber<'_>,
-    ) -> RethResult<Option<Vec<&'_ [u8]>>> {
-        let row = match key_or_num {
-            KeyOrNumber::Hash(k) => self.row_by_key_with_cols::<SELECTOR, COLUMNS>(k),
-            KeyOrNumber::Number(n) => {
-                let offset = self.jar().user_header().start();
-                if offset > n {
-                    return Ok(None)
-                }
-                self.row_by_number_with_cols::<SELECTOR, COLUMNS>((n - offset) as usize)
-            }
-        }?;
-
-        Ok(row)
-    }
-
-    /// Gets one column value from a row.
-    pub fn get_one<T: Decompress, const SELECTOR: usize, const COLUMNS: usize>(
-        &mut self,
-        key_or_num: KeyOrNumber<'_>,
-    ) -> RethResult<Option<T>> {
-        let row = self.get::<SELECTOR, COLUMNS>(key_or_num)?;
-
-        match row {
-            Some(row) => Ok(Some(T::decompress(row[0])?)),
-            None => Ok(None),
-        }
-    }
-
-    /// Gets two column values from a row.
-    pub fn get_two<T: Decompress, K: Decompress, const SELECTOR: usize, const COLUMNS: usize>(
-        &mut self,
-        key_or_num: KeyOrNumber<'_>,
-    ) -> RethResult<Option<(T, K)>> {
-        let row = self.get::<SELECTOR, COLUMNS>(key_or_num)?;
-
-        match row {
-            Some(row) => Ok(Some((T::decompress(row[0])?, K::decompress(row[1])?))),
-            None => Ok(None),
-        }
-    }
-
-    /// Gets three column values from a row.
-    pub fn get_three<
-        T: Decompress,
-        K: Decompress,
-        J: Decompress,
-        const SELECTOR: usize,
-        const COLUMNS: usize,
-    >(
-        &mut self,
-        key_or_num: KeyOrNumber<'_>,
-    ) -> RethResult<Option<(T, K, J)>> {
-        let row = self.get::<SELECTOR, COLUMNS>(key_or_num)?;
-
-        match row {
-            Some(row) => {
-                Ok(Some((T::decompress(row[0])?, K::decompress(row[1])?, J::decompress(row[2])?)))
-            }
-            None => Ok(None),
-        }
-    }
-}
-
-/// Either a key _or_ a block number
-#[derive(Debug)]
-pub enum KeyOrNumber<'a> {
-    /// A slice used as a key. Usually a block hash
-    Hash(&'a [u8]),
-    /// A block number
-    Number(u64),
-}
-
-impl<'a> From<&'a B256> for KeyOrNumber<'a> {
-    fn from(value: &'a B256) -> Self {
-        KeyOrNumber::Hash(value.as_slice())
-    }
-}
-
-impl<'a> From<u64> for KeyOrNumber<'a> {
-    fn from(value: u64) -> Self {
-        KeyOrNumber::Number(value)
-    }
-}
-
-/// Snapshot segment total columns.
-pub const HEADER_COLUMNS: usize = 3;
-/// Selector for header.
-pub const S_HEADER: usize = 0b001;
-/// Selector for header td.
-pub const S_HEADER_TD: usize = 0b010;
-/// Selector for header hash.
-pub const S_HEADER_HASH: usize = 0b100;
-/// Selector for header td and header hash.
-pub const S_HEADER_TD_WITH_HASH: usize = 0b110;
-/// Selector for header and header hash.
-pub const S_HEADER_WITH_HASH: usize = 0b101;

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

@@ -0,0 +1,90 @@
+use crate::table::Decompress;
+
+/// Generic Mask helper struct for selecting specific column values to read and decompress.
+///
+/// #### Explanation:
+///
+/// A `NippyJar` snapshot row can contain multiple column values. To specify the column values
+/// to be read, a mask is utilized.
+///
+/// For example, a snapshot 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 snapshot 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), " snapshot 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
+    type FIRST: Decompress;
+    /// Mask to obtain desired values, should correspond to the order of columns in a snapshot.
+    const MASK: usize;
+}
+
+///  Trait for specifying a mask to select two column values.
+pub trait ColumnSelectorTwo {
+    /// First desired column value
+    type FIRST: Decompress;
+    /// Second desired column value
+    type SECOND: Decompress;
+    /// Mask to obtain desired values, should correspond to the order of columns in a snapshot.
+    const MASK: usize;
+}
+
+///  Trait for specifying a mask to select three column values.
+pub trait ColumnSelectorThree {
+    /// First desired column value
+    type FIRST: Decompress;
+    /// Second desired column value
+    type SECOND: Decompress;
+    /// Third desired column value
+    type THIRD: Decompress;
+    /// Mask to obtain desired values, should correspond to the order of columns in a snapshot.
+    const MASK: usize;
+}
+
+#[macro_export]
+/// Add mask to select `N` column values from a specific snapshot segment row.
+macro_rules! add_snapshot_mask {
+    ($mask_struct:tt, $type1:ty, $mask:expr) => {
+        impl ColumnSelectorOne for $mask_struct<$type1> {
+            type FIRST = $type1;
+            const MASK: usize = $mask;
+        }
+    };
+    ($mask_struct:tt, $type1:ty, $type2:ty, $mask:expr) => {
+        impl ColumnSelectorTwo for $mask_struct<$type1, $type2> {
+            type FIRST = $type1;
+            type SECOND = $type2;
+            const MASK: usize = $mask;
+        }
+    };
+    ($mask_struct:tt, $type1:ty, $type2:ty, $type3:ty, $mask:expr) => {
+        impl ColumnSelectorTwo for $mask_struct<$type1, $type2, $type3> {
+            type FIRST = $type1;
+            type SECOND = $type2;
+            type THIRD = $type3;
+            const MASK: usize = $mask;
+        }
+    };
+}

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

@@ -0,0 +1,28 @@
+use super::{ReceiptMask, TransactionMask};
+use crate::{
+    add_snapshot_mask,
+    snapshot::mask::{ColumnSelectorOne, ColumnSelectorTwo, HeaderMask},
+    table::Table,
+    CanonicalHeaders, HeaderTD, Receipts, Transactions,
+};
+use reth_primitives::{BlockHash, Header};
+
+// HEADER MASKS
+
+add_snapshot_mask!(HeaderMask, Header, 0b001);
+add_snapshot_mask!(HeaderMask, <HeaderTD as Table>::Value, 0b010);
+add_snapshot_mask!(HeaderMask, BlockHash, 0b100);
+
+add_snapshot_mask!(HeaderMask, Header, BlockHash, 0b101);
+add_snapshot_mask!(
+    HeaderMask,
+    <HeaderTD as Table>::Value,
+    <CanonicalHeaders as Table>::Value,
+    0b110
+);
+
+// RECEIPT MASKS
+add_snapshot_mask!(ReceiptMask, <Receipts as Table>::Value, 0b1);
+
+// TRANSACTION MASKS
+add_snapshot_mask!(TransactionMask, <Transactions as Table>::Value, 0b1);

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

@@ -0,0 +1,12 @@
+//! reth's snapshot database table import and access
+
+mod generation;
+pub use generation::*;
+
+mod cursor;
+pub use cursor::SnapshotCursor;
+
+mod mask;
+pub use mask::*;
+
+mod masks;