refactor: propagate iter changes from libmdbx

Author: prestwich

Cargo.toml

@@ -42,7 +42,7 @@ signet-storage = { version = "0.0.1", path = "./crates/storage" }
 signet-storage-types = { version = "0.0.1", path = "./crates/types" }
 
 # External, in-house
-signet-libmdbx = { version = "0.7.0" }
+signet-libmdbx = { version = "0.8.0" }
 
 signet-zenith = "0.16.0-rc.5"
 

crates/hot-mdbx/src/cursor.rs

@@ -4,11 +4,11 @@ use crate::{FixedSizeInfo, MdbxError};
 use signet_hot::{
     MAX_FIXED_VAL_SIZE, MAX_KEY_SIZE,
     model::{
-        DualKeyTraverse, DualKeyTraverseMut, KvTraverse, KvTraverseMut, RawDualKeyValue,
-        RawKeyValue, RawValue,
+        DualKeyItem, DualKeyTraverse, DualKeyTraverseMut, KvTraverse, KvTraverseMut,
+        RawDualKeyItem, RawDualKeyValue, RawKeyValue, RawValue,
     },
 };
-use signet_libmdbx::{Ro, Rw, RwSync, TransactionKind, tx::WriteMarker};
+use signet_libmdbx::{DupItem, Ro, Rw, RwSync, TransactionKind, tx::WriteMarker};
 use std::{
     borrow::Cow,
     ops::{Deref, DerefMut},
@@ -504,6 +504,61 @@ where
             None => Ok(None),
         }
     }
+
+    fn iter_items(
+        &mut self,
+    ) -> Result<impl Iterator<Item = Result<RawDualKeyItem<'_>, MdbxError>> + '_, MdbxError> {
+        if !self.fsi.is_dupsort() {
+            return Err(MdbxError::NotDupSort);
+        }
+
+        let key2_size = self.fsi.key2_size().ok_or(MdbxError::UnknownFixedSize)?;
+
+        // Use iter_dup_start which yields DupItem::NewKey/SameKey
+        let iter_dup = self.inner.iter_dup_start::<Cow<'_, [u8]>, Cow<'_, [u8]>>()?;
+        Ok(MdbxDualKeyItemIter { iter_dup, key2_size })
+    }
+}
+
+/// Iterator adapter that converts mdbx's [`DupItem`] to [`DualKeyItem`].
+///
+/// This iterator returns owned data to avoid lifetime complexities between
+/// the transaction lifetime and the iterator borrow lifetime.
+struct MdbxDualKeyItemIter<'tx, 'cur, K: TransactionKind> {
+    iter_dup: signet_libmdbx::tx::iter::IterDup<'tx, 'cur, K, Cow<'tx, [u8]>, Cow<'tx, [u8]>>,
+    key2_size: usize,
+}
+
+impl<'a, K: TransactionKind> Iterator for MdbxDualKeyItemIter<'_, 'a, K> {
+    type Item = Result<RawDualKeyItem<'a>, MdbxError>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        match self.iter_dup.borrow_next() {
+            Ok(Some(item)) => {
+                let result = match item {
+                    DupItem::NewKey(k1, v) => {
+                        let (k2, val) = split_cow_at_owned(v, self.key2_size);
+                        DualKeyItem::NewK1(Cow::Owned(k1.into_owned()), k2, val)
+                    }
+                    DupItem::SameKey(v) => {
+                        let (k2, val) = split_cow_at_owned(v, self.key2_size);
+                        DualKeyItem::SameK1(k2, val)
+                    }
+                };
+                Some(Ok(result))
+            }
+            Ok(None) => None,
+            Err(e) => Some(Err(MdbxError::from(e))),
+        }
+    }
+}
+
+/// Splits a [`Cow`] at the given index and returns owned [`Cow`]s.
+#[inline]
+fn split_cow_at_owned(cow: Cow<'_, [u8]>, at: usize) -> (Cow<'static, [u8]>, Cow<'static, [u8]>) {
+    let vec = cow.into_owned();
+    let (left, right) = vec.split_at(at);
+    (Cow::Owned(left.to_vec()), Cow::Owned(right.to_vec()))
 }
 
 impl<K: TransactionKind + WriteMarker> DualKeyTraverseMut<MdbxError> for Cursor<'_, K> {

crates/hot/src/mem.rs

@@ -5,8 +5,9 @@
 
 use crate::{
     model::{
-        DualKeyTraverse, DualKeyTraverseMut, HotKv, HotKvError, HotKvRead, HotKvReadError,
-        HotKvWrite, KvTraverse, KvTraverseMut, RawDualKeyValue, RawKeyValue, RawValue,
+        DualKeyItem, DualKeyTraverse, DualKeyTraverseMut, HotKv, HotKvError, HotKvRead,
+        HotKvReadError, HotKvWrite, KvTraverse, KvTraverseMut, RawDualKeyItem, RawDualKeyValue,
+        RawKeyValue, RawValue,
     },
     ser::{DeserError, MAX_KEY_SIZE},
 };
@@ -514,6 +515,47 @@ impl<'a> DualKeyTraverse<MemKvError> for MemKvCursor<'a> {
         self.set_current_key(*found_key);
         Ok(Some((found_k1, found_k2, Cow::Borrowed(value.as_ref()))))
     }
+
+    fn iter_items(
+        &mut self,
+    ) -> Result<impl Iterator<Item = Result<RawDualKeyItem<'_>, MemKvError>> + '_, MemKvError> {
+        DualKeyTraverse::first(self)?;
+        Ok(MemDualKeyItemIter { cursor: self, prev_k1: None })
+    }
+}
+
+/// Iterator that yields [`DualKeyItem`] for read-only memory cursors.
+struct MemDualKeyItemIter<'a, 'b> {
+    cursor: &'a mut MemKvCursor<'b>,
+    prev_k1: Option<Vec<u8>>,
+}
+
+impl<'a> Iterator for MemDualKeyItemIter<'a, '_> {
+    type Item = Result<RawDualKeyItem<'a>, MemKvError>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let result = DualKeyTraverse::read_next(self.cursor);
+        match result {
+            Ok(Some((k1, k2, v))) => {
+                let k1_vec = k1.as_ref().to_vec();
+                let is_new_k1 = self.prev_k1.as_ref() != Some(&k1_vec);
+                self.prev_k1 = Some(k1_vec);
+
+                let item = if is_new_k1 {
+                    DualKeyItem::NewK1(
+                        Cow::Owned(k1.into_owned()),
+                        Cow::Owned(k2.into_owned()),
+                        Cow::Owned(v.into_owned()),
+                    )
+                } else {
+                    DualKeyItem::SameK1(Cow::Owned(k2.into_owned()), Cow::Owned(v.into_owned()))
+                };
+                Some(Ok(item))
+            }
+            Ok(None) => None,
+            Err(e) => Some(Err(e)),
+        }
+    }
 }
 
 /// Memory cursor for read-write operations
@@ -1037,6 +1079,47 @@ impl<'a> DualKeyTraverse<MemKvError> for MemKvCursorMut<'a> {
             Cow::Owned(value.to_vec()),
         )))
     }
+
+    fn iter_items(
+        &mut self,
+    ) -> Result<impl Iterator<Item = Result<RawDualKeyItem<'_>, MemKvError>> + '_, MemKvError> {
+        DualKeyTraverse::first(self)?;
+        Ok(MemDualKeyItemIterMut { cursor: self, prev_k1: None })
+    }
+}
+
+/// Iterator that yields [`DualKeyItem`] for read-write memory cursors.
+struct MemDualKeyItemIterMut<'a, 'b> {
+    cursor: &'a mut MemKvCursorMut<'b>,
+    prev_k1: Option<Vec<u8>>,
+}
+
+impl<'a> Iterator for MemDualKeyItemIterMut<'a, '_> {
+    type Item = Result<RawDualKeyItem<'a>, MemKvError>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let result = DualKeyTraverse::read_next(self.cursor);
+        match result {
+            Ok(Some((k1, k2, v))) => {
+                let k1_vec = k1.as_ref().to_vec();
+                let is_new_k1 = self.prev_k1.as_ref() != Some(&k1_vec);
+                self.prev_k1 = Some(k1_vec);
+
+                let item = if is_new_k1 {
+                    DualKeyItem::NewK1(
+                        Cow::Owned(k1.into_owned()),
+                        Cow::Owned(k2.into_owned()),
+                        Cow::Owned(v.into_owned()),
+                    )
+                } else {
+                    DualKeyItem::SameK1(Cow::Owned(k2.into_owned()), Cow::Owned(v.into_owned()))
+                };
+                Some(Ok(item))
+            }
+            Ok(None) => None,
+            Err(e) => Some(Err(e)),
+        }
+    }
 }
 
 impl DualKeyTraverseMut<MemKvError> for MemKvCursorMut<'_> {

crates/hot/src/model/mod.rs

@@ -45,8 +45,9 @@ pub use traits::{HotKv, HotKvRead, HotKvWrite};
 
 mod traverse;
 pub use traverse::{
-    DualKeyTraverse, DualKeyTraverseMut, DualTableCursor, DualTableTraverse, DualTableTraverseMut,
-    K2Value, KvTraverse, KvTraverseMut, RawK2Value, TableCursor, TableTraverse, TableTraverseMut,
+    DualKeyItem, DualKeyTraverse, DualKeyTraverseMut, DualTableCursor, DualTableTraverse,
+    DualTableTraverseMut, K2Value, KvTraverse, KvTraverseMut, RawDualKeyItem, RawK2Value,
+    TableCursor, TableTraverse, TableTraverseMut,
 };
 
 use crate::tables::{DualKey, Table};

crates/hot/src/model/traverse.rs

@@ -18,6 +18,72 @@ pub type RawK2Value<'a> = (Cow<'a, [u8]>, RawValue<'a>);
 /// Typed k2-value pair for a dual-keyed table.
 pub type K2Value<T> = (<T as DualKey>::Key2, <T as crate::tables::Table>::Value);
 
+/// An item from a dual-key iterator.
+///
+/// This enum avoids cloning k1 for every value when iterating
+/// over dual-keyed tables. K1 is only provided when it changes.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum DualKeyItem<K1, K2, V> {
+    /// First entry for a new k1.
+    NewK1(K1, K2, V),
+    /// Additional k2/value for the current k1.
+    SameK1(K2, V),
+}
+
+impl<K1, K2, V> DualKeyItem<K1, K2, V> {
+    /// Returns the value, consuming self.
+    pub fn into_value(self) -> V {
+        match self {
+            Self::NewK1(_, _, v) | Self::SameK1(_, v) => v,
+        }
+    }
+
+    /// Returns a reference to the value.
+    pub const fn value(&self) -> &V {
+        match self {
+            Self::NewK1(_, _, v) | Self::SameK1(_, v) => v,
+        }
+    }
+
+    /// Returns the k2, consuming self.
+    pub fn into_k2(self) -> K2 {
+        match self {
+            Self::NewK1(_, k2, _) | Self::SameK1(k2, _) => k2,
+        }
+    }
+
+    /// Returns a reference to k2.
+    pub const fn k2(&self) -> &K2 {
+        match self {
+            Self::NewK1(_, k2, _) | Self::SameK1(k2, _) => k2,
+        }
+    }
+
+    /// Returns k1 if this is a NewK1 entry.
+    pub const fn k1(&self) -> Option<&K1> {
+        match self {
+            Self::NewK1(k1, _, _) => Some(k1),
+            Self::SameK1(_, _) => None,
+        }
+    }
+
+    /// Returns true if this item represents a new k1.
+    pub const fn is_new_k1(&self) -> bool {
+        matches!(self, Self::NewK1(..))
+    }
+
+    /// Convert to a tuple, requiring k1 to be provided for SameK1 variants.
+    pub fn into_tuple(self, current_k1: K1) -> (K1, K2, V) {
+        match self {
+            Self::NewK1(k1, k2, v) => (k1, k2, v),
+            Self::SameK1(k2, v) => (current_k1, k2, v),
+        }
+    }
+}
+
+/// Raw dual-key item: (k1?, k2, value) where k1 is only present on NewK1.
+pub type RawDualKeyItem<'a> = DualKeyItem<Cow<'a, [u8]>, Cow<'a, [u8]>, Cow<'a, [u8]>>;
+
 /// Trait for traversing key-value pairs in the database.
 pub trait KvTraverse<E: HotKvReadError> {
     /// Set position to the first key-value pair in the database, and return
@@ -246,6 +312,19 @@ pub trait DualKeyTraverse<E: HotKvReadError> {
         let done = found_k1.as_ref() != k1;
         Ok(RawDualKeyK2Iter { cursor: self, done, _marker: PhantomData })
     }
+
+    /// Position at first entry and return iterator yielding [`DualKeyItem`].
+    ///
+    /// This is more efficient than [`Self::iter()`] as it avoids cloning k1 for
+    /// every entry within the same k1 group. The iterator yields
+    /// [`DualKeyItem::NewK1`] when k1 changes and [`DualKeyItem::SameK1`] for
+    /// subsequent entries with the same k1.
+    ///
+    /// Backends implement this natively to leverage efficient "new key"
+    /// notifications from the underlying database.
+    fn iter_items(&mut self) -> Result<impl Iterator<Item = Result<RawDualKeyItem<'_>, E>> + '_, E>
+    where
+        Self: Sized;
 }
 
 /// Trait for traversing dual-keyed key-value pairs with mutation capabilities.

crates/hot/src/ser/traits.rs

@@ -17,7 +17,7 @@ pub const MAX_FIXED_VAL_SIZE: usize = 64;
 ///
 /// In practice, keys are often hashes, addresses, numbers, or composites
 /// of these.
-pub trait KeySer: PartialOrd + Ord + Sized + Clone + core::fmt::Debug {
+pub trait KeySer: PartialOrd + Ord + Sized + Clone + Copy + core::fmt::Debug {
     /// The fixed size of the serialized key in bytes.
     /// Must satisfy `SIZE <= MAX_KEY_SIZE`.
     const SIZE: usize;

crates/types/src/sharded.rs

@@ -4,7 +4,7 @@
 /// `Address | 200` -> data is from block 0 to 200.
 ///
 /// `Address | 300` -> data is from block 201 to 300.
-#[derive(Debug, Default, Clone, Eq, PartialEq, Ord, PartialOrd, Hash)]
+#[derive(Debug, Default, Clone, Copy, Eq, PartialEq, Ord, PartialOrd, Hash)]
 pub struct ShardedKey<T> {
     /// The key for this type.
     pub key: T,