Minor sov-db improvements

full-node/db/sov-db/src/ledger_db/mod.rs

@@ -166,7 +166,7 @@ impl LedgerDB {
         let iter = raw_iter.take(max_items);
         let mut out = Vec::with_capacity(max_items);
         for res in iter {
-            let (_, batch) = res?;
+            let batch = res?.value;
             out.push(batch)
         }
         Ok(out)
@@ -307,7 +307,7 @@ impl LedgerDB {
         iter.seek_to_last();
 
         match iter.next() {
-            Some(Ok((version, _))) => Ok(Some(version.into())),
+            Some(Ok(item)) => Ok(Some(item.key.into())),
             Some(Err(e)) => Err(e),
             _ => Ok(None),
         }
@@ -319,7 +319,7 @@ impl LedgerDB {
         iter.seek_to_last();
 
         match iter.next() {
-            Some(Ok((slot_number, slot))) => Ok(Some((slot_number, slot))),
+            Some(Ok(item)) => Ok(Some(item.into_tuple())),
             Some(Err(e)) => Err(e),
             _ => Ok(None),
         }

full-node/db/sov-db/src/native_db.rs

@@ -50,7 +50,7 @@ impl NativeDB {
         let found = iter.next();
         match found {
             Some(result) => {
-                let ((found_key, found_version), value) = result?;
+                let ((found_key, found_version), value) = result?.into_tuple();
                 if &found_key == key {
                     anyhow::ensure!(found_version <= version, "Bug! iterator isn't returning expected values. expected a version <= {version:} but found {found_version:}");
                     Ok(value)

full-node/db/sov-db/src/state_db.rs

@@ -1,7 +1,7 @@
 use std::path::Path;
 use std::sync::{Arc, Mutex};
 
-use jmt::storage::{TreeReader, TreeWriter};
+use jmt::storage::{HasPreimage, TreeReader, TreeWriter};
 use jmt::{KeyHash, Version};
 use sov_schema_db::{SchemaBatch, DB};
 
@@ -69,7 +69,7 @@ impl StateDB {
         let found = iter.next();
         match found {
             Some(result) => {
-                let ((found_key, found_version), value) = result?;
+                let ((found_key, found_version), value) = result?.into_tuple();
                 if &found_key == key {
                     anyhow::ensure!(found_version <= version, "Bug! iterator isn't returning expected values. expected a version <= {version:} but found {found_version:}");
                     Ok(value)
@@ -98,7 +98,7 @@ impl StateDB {
         iter.seek_to_last();
 
         let version = match iter.next() {
-            Some(Ok((key, _))) => Some(key.version()),
+            Some(Ok(item)) => Some(item.key.version()),
             _ => None,
         };
         Ok(version)
@@ -153,6 +153,12 @@ impl TreeWriter for StateDB {
     }
 }
 
+impl HasPreimage for StateDB {
+    fn preimage(&self, key_hash: KeyHash) -> anyhow::Result<Option<Vec<u8>>> {
+        self.db.get::<KeyHashToKey>(&key_hash.0)
+    }
+}
+
 #[cfg(feature = "arbitrary")]
 pub mod arbitrary {
     //! Arbitrary definitions for the [`StateDB`].

full-node/db/sov-schema-db/src/iterator.rs

@@ -95,7 +95,7 @@ where
         }
     }
 
-    fn next_impl(&mut self) -> Result<Option<(S::Key, S::Value)>> {
+    fn next_impl(&mut self) -> Result<Option<IteratorOutput<S::Key, S::Value>>> {
         let _timer = SCHEMADB_ITER_LATENCY_SECONDS
             .with_label_values(&[S::COLUMN_FAMILY_NAME])
             .start_timer();
@@ -107,6 +107,7 @@ where
 
         let raw_key = self.db_iter.key().expect("db_iter.key() failed.");
         let raw_value = self.db_iter.value().expect("db_iter.value() failed.");
+        let value_size_bytes = raw_value.len();
         SCHEMADB_ITER_BYTES
             .with_label_values(&[S::COLUMN_FAMILY_NAME])
             .observe((raw_key.len() + raw_value.len()) as f64);
@@ -119,15 +120,32 @@ where
             ScanDirection::Backward => self.db_iter.prev(),
         }
 
-        Ok(Some((key, value)))
+        Ok(Some(IteratorOutput {
+            key,
+            value,
+            value_size_bytes,
+        }))
+    }
+}
+
+/// The output of [`SchemaIterator`]'s next_impl
+pub struct IteratorOutput<K, V> {
+    pub key: K,
+    pub value: V,
+    pub value_size_bytes: usize,
+}
+
+impl<K, V> IteratorOutput<K, V> {
+    pub fn into_tuple(self) -> (K, V) {
+        (self.key, self.value)
     }
 }
 
 impl<'a, S> Iterator for SchemaIterator<'a, S>
 where
     S: Schema,
 {
-    type Item = Result<(S::Key, S::Value)>;
+    type Item = Result<IteratorOutput<S::Key, S::Value>>;
 
     fn next(&mut self) -> Option<Self::Item> {
         self.next_impl().transpose()

full-node/db/sov-schema-db/src/lib.rs

@@ -31,6 +31,7 @@ use metrics::{
     SCHEMADB_BATCH_COMMIT_BYTES, SCHEMADB_BATCH_COMMIT_LATENCY_SECONDS, SCHEMADB_DELETES,
     SCHEMADB_GET_BYTES, SCHEMADB_GET_LATENCY_SECONDS, SCHEMADB_PUT_BYTES,
 };
+pub use rocksdb;
 use rocksdb::ReadOptions;
 pub use rocksdb::DEFAULT_COLUMN_FAMILY_NAME;
 use thiserror::Error;
@@ -127,7 +128,7 @@ impl DB {
         let k = schema_key.encode_key()?;
         let cf_handle = self.get_cf_handle(S::COLUMN_FAMILY_NAME)?;
 
-        let result = self.inner.get_cf(cf_handle, k)?;
+        let result = self.inner.get_pinned_cf(cf_handle, k)?;
         SCHEMADB_GET_BYTES
             .with_label_values(&[S::COLUMN_FAMILY_NAME])
             .observe(result.as_ref().map_or(0.0, |v| v.len() as f64));

full-node/db/sov-schema-db/tests/db_test.rs

@@ -132,7 +132,9 @@ fn test_schema_put_get() {
 fn collect_values<S: Schema>(db: &TestDB) -> Vec<(S::Key, S::Value)> {
     let mut iter = db.iter::<S>().expect("Failed to create iterator.");
     iter.seek_to_first();
-    iter.collect::<Result<Vec<_>, anyhow::Error>>().unwrap()
+    iter.map(|res| res.map(|item| item.into_tuple()))
+        .collect::<Result<Vec<_>, anyhow::Error>>()
+        .unwrap()
 }
 
 fn gen_expected_values(values: &[(u32, u32)]) -> Vec<(TestField, TestField)> {

full-node/db/sov-schema-db/tests/iterator_test.rs

@@ -19,7 +19,7 @@ define_schema!(TestSchema, TestCompositeField, TestField, "TestCF");
 type S = TestSchema;
 
 fn collect_values(iter: SchemaIterator<S>) -> Vec<u32> {
-    iter.map(|row| row.unwrap().1 .0).collect()
+    iter.map(|row| row.unwrap().value.0).collect()
 }
 
 fn decode_key(key: &[u8]) -> TestCompositeField {