Try to use key hash instead of key

crates/sovereign-sdk/full-node/db/sov-db/Cargo.toml

@@ -1,11 +1,11 @@
 [package]
 name = "sov-db"
-description = "A high-level DB interface for the Sovereign SDK"
-license = "Apache-2.0"                                          # This license is inherited from Aptos 
-edition = { workspace = true }
 authors = { workspace = true }
+edition = { workspace = true }
 homepage = { workspace = true }
+license = "Apache-2.0" # This license is inherited from Aptos
 repository = { workspace = true }
+description = "A high-level DB interface for the Sovereign SDK"
 
 version = { workspace = true }
 readme = "README.md"
@@ -16,44 +16,38 @@ resolver = "2"
 [dependencies]
 # Maintained by sovereign labs
 jmt = { workspace = true }
+sov-rollup-interface = { path = "../../../rollup-interface", features = ["native"] }
 sov-schema-db = { path = "../sov-schema-db" }
-sov-rollup-interface = { path = "../../../rollup-interface", features = [
-    "native",
-] }
 
 # External
 anyhow = { workspace = true, default-features = true }
 arbitrary = { workspace = true, optional = true }
+bincode = { workspace = true }
+borsh = { workspace = true, default-features = true, features = ["bytes", "rc"] }
 byteorder = { workspace = true, default-features = true }
-borsh = { workspace = true, default-features = true, features = [
-    "bytes",
-    "rc",
-] }
+hex = { workspace = true }
 proptest = { workspace = true, optional = true, default-features = true }
 proptest-derive = { workspace = true, optional = true }
+rocksdb = { workspace = true }
 serde = { workspace = true, default-features = true, features = ["rc"] }
+sha2 = { workspace = true }
 tempfile = { workspace = true, optional = true }
-rocksdb = { workspace = true }
-bincode = { workspace = true }
 tokio = { workspace = true }
 tracing = { workspace = true }
-hex = { workspace = true }
-
 
 [dev-dependencies]
-sov-mock-da = { path = "../../../adapters/mock-da", features = ["native"] }
-tempfile = { workspace = true }
 criterion = "0.5.1"
 rand = { workspace = true }
 sha2 = { workspace = true }
-
+sov-mock-da = { path = "../../../adapters/mock-da", features = ["native"] }
+tempfile = { workspace = true }
 
 [features]
 arbitrary = [
-    "dep:arbitrary",
-    "dep:proptest",
-    "dep:proptest-derive",
-    "dep:tempfile",
+  "dep:arbitrary",
+  "dep:proptest",
+  "dep:proptest-derive",
+  "dep:tempfile",
 ]
 
 [[bench]]

crates/sovereign-sdk/full-node/db/sov-db/benches/state_db_bench.rs

@@ -10,6 +10,7 @@ use jmt::storage::TreeWriter;
 use jmt::{JellyfishMerkleTree, KeyHash};
 use rand::rngs::StdRng;
 use rand::{Rng, SeedableRng};
+use sov_db::schema::types::JmtNodeHashKey;
 use sov_db::state_db::StateDB;
 use sov_schema_db::snapshot::{DbSnapshot, NoopQueryManager, ReadOnlyLock};
 
@@ -73,20 +74,24 @@ fn prepare_data(size: usize) -> TestData {
         .put_value_set_with_proof(batch, 1)
         .expect("JMT update must succeed");
 
-    db.put_preimages(key_preimages).unwrap();
-
     db.write_node_batch(&tree_update.node_batch).unwrap();
 
     // Sanity check:
     let version = db.get_next_version() - 1;
     for chunk in raw_data.chunks(2) {
         let key = &chunk[0];
+        let key_hash: JmtNodeHashKey = (key, version).into();
         let value = chunk[1].clone();
-        let res = db.get_value_option_by_key(version, key).unwrap();
+        let res = db
+            .get_value_option_by_key(version, &KeyHash(key_hash.hash()))
+            .unwrap();
         assert_eq!(Some(value), res);
     }
 
-    let random_value = db.get_value_option_by_key(version, &random_key).unwrap();
+    let random_key_hash: JmtNodeHashKey = (&random_key, version).into();
+    let random_value = db
+        .get_value_option_by_key(version, &KeyHash(random_key_hash.hash()))
+        .unwrap();
     assert!(random_value.is_some());
 
     TestData {
@@ -106,7 +111,11 @@ fn bench_random_read(g: &mut BenchmarkGroup<WallTime>, size: usize) {
         |b, i| {
             b.iter(|| {
                 let (db, key, version) = i;
-                let result = black_box(db.get_value_option_by_key(*version, key).unwrap());
+                let hash_key: JmtNodeHashKey = (key, *version).into();
+                let result = black_box(
+                    db.get_value_option_by_key(*version, &KeyHash(hash_key.hash()))
+                        .unwrap(),
+                );
                 assert!(result.is_some());
                 black_box(result);
             })
@@ -127,7 +136,11 @@ fn bench_largest_read(g: &mut BenchmarkGroup<WallTime>, size: usize) {
         |b, i| {
             b.iter(|| {
                 let (db, key, version) = i;
-                let result = black_box(db.get_value_option_by_key(*version, key).unwrap());
+                let hash_key: JmtNodeHashKey = (key, *version).into();
+                let result = black_box(
+                    db.get_value_option_by_key(*version, &KeyHash(hash_key.hash()))
+                        .unwrap(),
+                );
                 assert!(result.is_some());
                 black_box(result);
             })
@@ -148,7 +161,11 @@ fn bench_not_found_read(g: &mut BenchmarkGroup<WallTime>, size: usize) {
         |b, i| {
             b.iter(|| {
                 let (db, key, version) = i;
-                let result = black_box(db.get_value_option_by_key(*version, key).unwrap());
+                let hash_key: JmtNodeHashKey = (key, *version).into();
+                let result = black_box(
+                    db.get_value_option_by_key(*version, &KeyHash(hash_key.hash()))
+                        .unwrap(),
+                );
                 assert!(result.is_none());
                 black_box(result);
             })

crates/sovereign-sdk/full-node/db/sov-db/src/schema/tables.rs

@@ -25,28 +25,25 @@
 //! Module Accessory State Table:
 //! - `(ModuleAddress, Key) -> Value`
 
+use anyhow::anyhow;
 use borsh::{BorshDeserialize, BorshSerialize};
 use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt};
-use jmt::storage::{NibblePath, Node, NodeKey};
+use jmt::storage::{Node, NodeKey};
 use jmt::Version;
 use sov_rollup_interface::da::SequencerCommitment;
 use sov_rollup_interface::stf::{Event, EventKey, StateDiff};
 use sov_schema_db::schema::{KeyDecoder, KeyEncoder, ValueCodec};
 use sov_schema_db::{CodecError, SeekKeyEncoder};
 
 use super::types::{
-    AccessoryKey, AccessoryStateValue, BatchNumber, DbHash, EventNumber, JmtValue, L2HeightRange,
-    SlotNumber, StateKey, StoredBatch, StoredProof, StoredSlot, StoredSoftBatch, StoredTransaction,
-    StoredVerifiedProof, TxNumber,
+    AccessoryKey, AccessoryStateValue, BatchNumber, DbHash, EventNumber, JmtNodeHashKey, JmtValue,
+    L2HeightRange, SlotNumber, StoredBatch, StoredProof, StoredSlot, StoredSoftBatch,
+    StoredTransaction, StoredVerifiedProof, TxNumber,
 };
 
 /// A list of all tables used by the StateDB. These tables store rollup state - meaning
 /// account balances, nonces, etc.
-pub const STATE_TABLES: &[&str] = &[
-    KeyHashToKey::table_name(),
-    JmtValues::table_name(),
-    JmtNodes::table_name(),
-];
+pub const STATE_TABLES: &[&str] = &[JmtValues::table_name(), JmtNodes::table_name()];
 
 /// A list of all tables used by the LedgerDB. These tables store rollup "history" - meaning
 /// transaction, events, receipts, etc.
@@ -317,7 +314,7 @@ define_table_with_default_codec!(
 
 define_table_without_codec!(
     /// The source of truth for JMT nodes
-    (JmtNodes) NodeKey => Node
+    (JmtNodes) JmtNodeHashKey => Node
 );
 
 define_table_with_default_codec!(
@@ -341,19 +338,23 @@ impl KeyEncoder<JmtNodes> for NodeKey {
         Ok(output)
     }
 }
-impl KeyDecoder<JmtNodes> for NodeKey {
+
+impl KeyEncoder<JmtNodes> for JmtNodeHashKey {
+    fn encode_key(&self) -> sov_schema_db::schema::Result<Vec<u8>> {
+        let mut output = vec![];
+        BorshSerialize::serialize(self, &mut output)?;
+        Ok(output)
+    }
+}
+impl KeyDecoder<JmtNodes> for JmtNodeHashKey {
     fn decode_key(data: &[u8]) -> sov_schema_db::schema::Result<Self> {
-        if data.len() < 8 {
-            return Err(CodecError::InvalidKeyLength {
-                expected: 9,
-                got: data.len(),
-            });
-        }
-        let mut version = [0u8; 8];
-        version.copy_from_slice(&data[..8]);
-        let version = u64::from_be_bytes(version);
-        let nibble_path = NibblePath::deserialize_reader(&mut &data[8..])?;
-        Ok(Self::new(version, nibble_path))
+        let data = match BorshDeserialize::try_from_slice(data).map_err(CodecError::Io) {
+            Ok(data) => data,
+            Err(_) => {
+                return Err(CodecError::Wrapped(anyhow!("Could not deserialize key")));
+            }
+        };
+        Ok(data)
     }
 }
 
@@ -369,36 +370,24 @@ impl ValueCodec<JmtNodes> for Node {
 
 define_table_without_codec!(
     /// The source of truth for JMT values by version
-    (JmtValues) (StateKey, Version) => JmtValue
+    (JmtValues) JmtNodeHashKey => JmtValue
 );
 
-impl<T: AsRef<[u8]> + PartialEq + core::fmt::Debug> KeyEncoder<JmtValues> for (T, Version) {
+impl KeyEncoder<JmtValues> for JmtNodeHashKey {
     fn encode_key(&self) -> sov_schema_db::schema::Result<Vec<u8>> {
-        let mut out =
-            Vec::with_capacity(self.0.as_ref().len() + std::mem::size_of::<Version>() + 8);
-        self.0
-            .as_ref()
-            .serialize(&mut out)
-            .map_err(CodecError::from)?;
-        // Write the version in big-endian order so that sorting order is based on the most-significant bytes of the key
-        out.write_u64::<BigEndian>(self.1)
-            .expect("serialization to vec is infallible");
-        Ok(out)
+        borsh::to_vec(self).map_err(CodecError::from)
     }
 }
 
-impl<T: AsRef<[u8]> + PartialEq + core::fmt::Debug> SeekKeyEncoder<JmtValues> for (T, Version) {
+impl SeekKeyEncoder<JmtValues> for JmtNodeHashKey {
     fn encode_seek_key(&self) -> sov_schema_db::schema::Result<Vec<u8>> {
-        <(T, Version) as KeyEncoder<JmtValues>>::encode_key(self)
+        <JmtNodeHashKey as KeyEncoder<JmtValues>>::encode_key(self)
     }
 }
 
-impl KeyDecoder<JmtValues> for (StateKey, Version) {
+impl KeyDecoder<JmtValues> for JmtNodeHashKey {
     fn decode_key(data: &[u8]) -> sov_schema_db::schema::Result<Self> {
-        let mut cursor = std::io::Cursor::new(data);
-        let key: Vec<u8> = BorshDeserialize::deserialize_reader(&mut cursor)?;
-        let version = cursor.read_u64::<BigEndian>()?;
-        Ok((key, version))
+        Ok(BorshDeserialize::deserialize_reader(&mut &data[..])?)
     }
 }
 
@@ -412,14 +401,6 @@ impl ValueCodec<JmtValues> for JmtValue {
     }
 }
 
-define_table_with_default_codec!(
-    /// A mapping from key-hashes to their preimages and latest version. Since we store raw
-    /// key-value pairs instead of keyHash->value pairs,
-    /// this table is required to implement the `jmt::TreeReader` trait,
-    /// which requires the ability to fetch values by hash.
-    (KeyHashToKey) [u8;32] => StateKey
-);
-
 define_table_without_codec!(
     /// Non-JMT state stored by a module for JSON-RPC use.
     (ModuleAccessoryState) (AccessoryKey, Version) => AccessoryStateValue

crates/sovereign-sdk/full-node/db/sov-db/src/schema/types.rs

@@ -2,6 +2,8 @@ use std::marker::PhantomData;
 use std::sync::Arc;
 
 use borsh::{BorshDeserialize, BorshSerialize};
+use jmt::storage::NodeKey;
+use jmt::{KeyHash, Version};
 use serde::de::DeserializeOwned;
 use serde::{Deserialize, Serialize};
 use sov_rollup_interface::rpc::{
@@ -57,7 +59,6 @@ pub type AccessoryStateValue = Option<Vec<u8>>;
 pub type DbHash = [u8; 32];
 /// The "value" half of a key/value pair from the JMT
 pub type JmtValue = Option<Vec<u8>>;
-pub(crate) type StateKey = Vec<u8>;
 
 /// The on-disk format of a slot. Specifies the batches contained in the slot
 /// and the hash of the da block. TODO(@preston-evans98): add any additional data
@@ -412,3 +413,53 @@ pub mod arbitrary {
         }
     }
 }
+
+/// Defines a wrapper around the key hash and version
+#[derive(Debug, PartialEq, BorshDeserialize, BorshSerialize)]
+pub struct JmtNodeHashKey {
+    key: [u8; 32],
+    version: Version,
+}
+
+impl JmtNodeHashKey {
+    /// Instantiate a new hash key
+    pub fn new(key: [u8; 32], version: Version) -> Self {
+        Self { key, version }
+    }
+
+    /// Return the version
+    pub fn version(&self) -> Version {
+        self.version
+    }
+
+    /// Return the hash
+    pub fn hash(&self) -> [u8; 32] {
+        self.key
+    }
+}
+
+impl From<&NodeKey> for JmtNodeHashKey {
+    fn from(node_key: &NodeKey) -> Self {
+        let mut out = vec![];
+        let _ = BorshSerialize::serialize(&node_key, &mut out);
+        let key = KeyHash::with::<sha2::Sha256>(out);
+        Self {
+            key: key.0,
+            version: node_key.version(),
+        }
+    }
+}
+
+impl From<(&Vec<u8>, Version)> for JmtNodeHashKey {
+    fn from(data: (&Vec<u8>, Version)) -> Self {
+        let node_key = data.0;
+        let version = data.1;
+        let mut out = vec![];
+        let _ = BorshSerialize::serialize(&node_key, &mut out);
+        let key = KeyHash::with::<sha2::Sha256>(out);
+        Self {
+            key: key.0,
+            version,
+        }
+    }
+}