[storage] Splitting (state_key, state_key_and_value) from JMT

Closes: aptos-labs#1012

.gitignore

@@ -77,6 +77,7 @@ target-out-docker
 
 # Emacs backup files
 *~
+.\#*
 .*\#
 \#*\#
 

execution/executor-types/src/lib.rs

@@ -36,7 +36,7 @@ use storage_interface::DbReader;
 pub use executed_chunk::ExecutedChunk;
 use storage_interface::{in_memory_state::InMemoryState, verified_state_view::VerifiedStateView};
 
-type SparseMerkleProof = aptos_types::proof::SparseMerkleProof<StateValue>;
+type SparseMerkleProof = aptos_types::proof::SparseMerkleProof;
 
 pub trait ChunkExecutorTrait: Send + Sync {
     /// Verifies the transactions based on the provided proofs and ledger info. If the transactions
@@ -390,7 +390,7 @@ impl ProofReader {
     }
 }
 
-impl ProofRead<StateValue> for ProofReader {
+impl ProofRead for ProofReader {
     fn get_proof(&self, key: HashValue) -> Option<&SparseMerkleProof> {
         self.proofs.get(&key)
     }

execution/executor/tests/db_bootstrapper_test.rs

@@ -4,7 +4,7 @@
 #![forbid(unsafe_code)]
 
 use anyhow::Result;
-use aptos_crypto::{ed25519::Ed25519PrivateKey, HashValue, PrivateKey, Uniform};
+use aptos_crypto::{ed25519::Ed25519PrivateKey, hash::CryptoHash, HashValue, PrivateKey, Uniform};
 use aptos_state_view::account_with_state_view::AsAccountWithStateView;
 use aptos_temppath::TempPath;
 use aptos_transaction_builder::aptos_stdlib;
@@ -195,7 +195,7 @@ fn get_configuration(db: &DbReaderWriter) -> ConfigurationResource {
 fn get_state_backup(
     db: &Arc<AptosDB>,
 ) -> (
-    Vec<(HashValue, StateKeyAndValue)>,
+    Vec<(StateKey, StateKeyAndValue)>,
     SparseMerkleRangeProof,
     HashValue,
 ) {
@@ -206,7 +206,7 @@ fn get_state_backup(
         .collect::<Result<Vec<_>>>()
         .unwrap();
     let proof = backup_handler
-        .get_account_state_range_proof(accounts.last().unwrap().0, 1)
+        .get_account_state_range_proof(CryptoHash::hash(&accounts.last().unwrap().0), 1)
         .unwrap();
     let db_reader: Arc<dyn DbReader> = db.clone();
     let root_hash = db
@@ -222,7 +222,7 @@ fn get_state_backup(
 
 fn restore_state_to_db(
     db: &Arc<AptosDB>,
-    accounts: Vec<(HashValue, StateKeyAndValue)>,
+    accounts: Vec<(StateKey, StateKeyAndValue)>,
     proof: SparseMerkleRangeProof,
     root_hash: HashValue,
     version: Version,

state-sync/state-sync-v2/data-streaming-service/src/tests/utils.rs

@@ -163,7 +163,7 @@ impl AptosDataClient for MockAptosDataClient {
         let mut account_blobs = vec![];
         for _ in start_index..=end_index {
             account_blobs.push((
-                HashValue::random(),
+                StateKey::Raw(HashValue::random().to_vec()),
                 StateKeyAndValue::new(StateKey::Raw(vec![]), vec![].into()),
             ));
         }

state-sync/state-sync-v2/state-sync-driver/src/tests/mocks.rs

@@ -269,7 +269,7 @@ mock! {
             &self,
             state_key: &StateKey,
             version: Version,
-        ) -> Result<(Option<StateValue>, SparseMerkleProof<StateValue>)>;
+        ) -> Result<(Option<StateValue>, SparseMerkleProof)>;
 
         fn get_latest_tree_state(&self) -> Result<TreeState>;
 
@@ -311,7 +311,7 @@ mock! {
             &self,
             version: Version,
             expected_root_hash: HashValue,
-        ) -> Result<Box<dyn StateSnapshotReceiver<StateKeyAndValue>>>;
+        ) -> Result<Box<dyn StateSnapshotReceiver<StateKey, StateKeyAndValue>>>;
 
         fn finalize_state_snapshot(
             &self,
@@ -335,8 +335,8 @@ mock! {
 // This automatically creates a MockSnapshotReceiver.
 mock! {
     pub SnapshotReceiver {}
-    impl StateSnapshotReceiver<StateKeyAndValue> for SnapshotReceiver {
-        fn add_chunk(&mut self, chunk: Vec<(HashValue, StateKeyAndValue)>, proof: SparseMerkleRangeProof) -> Result<()>;
+    impl StateSnapshotReceiver<StateKey, StateKeyAndValue> for SnapshotReceiver {
+        fn add_chunk(&mut self, chunk: Vec<(StateKey, StateKeyAndValue)>, proof: SparseMerkleRangeProof) -> Result<()>;
 
         fn finish(self) -> Result<()>;
 

state-sync/storage-service/server/src/tests.rs

@@ -1545,7 +1545,7 @@ mock! {
             &self,
             state_key: &StateKey,
             version: Version,
-        ) -> Result<(Option<StateValue>, SparseMerkleProof<StateValue>)>;
+        ) -> Result<(Option<StateValue>, SparseMerkleProof)>;
 
         fn get_latest_tree_state(&self) -> Result<TreeState>;
 

storage/aptosdb/src/backup/backup_handler.rs

@@ -18,7 +18,7 @@ use aptos_types::{
     contract_event::ContractEvent,
     ledger_info::LedgerInfoWithSignatures,
     proof::{SparseMerkleRangeProof, TransactionAccumulatorRangeProof, TransactionInfoWithProof},
-    state_store::state_value::StateKeyAndValue,
+    state_store::{state_key::StateKey, state_value::StateKeyAndValue},
     transaction::{Transaction, TransactionInfo, Version},
 };
 use itertools::zip_eq;
@@ -103,18 +103,22 @@ impl BackupHandler {
     pub fn get_account_iter(
         &self,
         version: Version,
-    ) -> Result<Box<dyn Iterator<Item = Result<(HashValue, StateKeyAndValue)>> + Send + Sync>> {
-        let iterator = JellyfishMerkleIterator::new(
-            Arc::clone(&self.state_store),
-            version,
-            HashValue::zero(),
-        )?
-        .enumerate()
-        .map(move |(idx, res)| {
-            BACKUP_STATE_SNAPSHOT_VERSION.set(version as i64);
-            BACKUP_STATE_SNAPSHOT_LEAF_IDX.set(idx as i64);
-            res
-        });
+    ) -> Result<Box<dyn Iterator<Item = Result<(StateKey, StateKeyAndValue)>> + Send + Sync>> {
+        let store = Arc::clone(&self.state_store);
+        let iterator =
+            JellyfishMerkleIterator::new(Arc::clone(&store), version, HashValue::zero())?
+                .enumerate()
+                .map(move |(idx, res)| {
+                    BACKUP_STATE_SNAPSHOT_VERSION.set(version as i64);
+                    BACKUP_STATE_SNAPSHOT_LEAF_IDX.set(idx as i64);
+                    match res {
+                        Ok((_hashed_key, key_and_version)) => Ok((
+                            key_and_version.0.clone(),
+                            store.get_value_at_version(&key_and_version)?,
+                        )),
+                        Err(err) => Err(err),
+                    }
+                });
         Ok(Box::new(iterator))
     }
 

storage/aptosdb/src/backup/restore_handler.rs

@@ -7,12 +7,12 @@ use crate::{
 };
 use anyhow::Result;
 use aptos_crypto::{hash::SPARSE_MERKLE_PLACEHOLDER_HASH, HashValue};
-use aptos_jellyfish_merkle::restore::JellyfishMerkleRestore;
+use aptos_jellyfish_merkle::restore::StateSnapshotRestore;
 use aptos_types::{
     contract_event::ContractEvent,
     ledger_info::LedgerInfoWithSignatures,
     proof::definition::LeafCount,
-    state_store::state_value::StateKeyAndValue,
+    state_store::{state_key::StateKey, state_value::StateKeyAndValue},
     transaction::{Transaction, TransactionInfo, Version, PRE_GENESIS_VERSION},
 };
 use schemadb::DB;
@@ -53,8 +53,8 @@ impl RestoreHandler {
         &self,
         version: Version,
         expected_root_hash: HashValue,
-    ) -> Result<JellyfishMerkleRestore<StateKeyAndValue>> {
-        JellyfishMerkleRestore::new_overwrite(
+    ) -> Result<StateSnapshotRestore<StateKey, StateKeyAndValue>> {
+        StateSnapshotRestore::new_overwrite(
             Arc::clone(&self.state_store),
             version,
             expected_root_hash,

storage/aptosdb/src/lib.rs

@@ -235,7 +235,7 @@ impl AptosDB {
             JELLYFISH_MERKLE_NODE_CF_NAME,
             LEDGER_COUNTERS_CF_NAME,
             STALE_NODE_INDEX_CF_NAME,
-            STATE_VALUE_INDEX_CF_NAME,
+            STATE_VALUE_CF_NAME,
             TRANSACTION_CF_NAME,
             TRANSACTION_ACCUMULATOR_CF_NAME,
             TRANSACTION_BY_ACCOUNT_CF_NAME,
@@ -611,8 +611,14 @@ impl AptosDB {
                 .iter()
                 .map(|txn_to_commit| txn_to_commit.jf_node_hashes())
                 .collect::<Option<Vec<_>>>();
+            self.state_store.merklize_value_sets(
+                account_state_sets.clone(),
+                node_hashes,
+                first_version,
+                cs,
+            )?;
             self.state_store
-                .put_value_sets(account_state_sets, node_hashes, first_version, cs)?;
+                .put_value_sets(account_state_sets, first_version, cs)?;
         }
 
         // Event updates. Gather event accumulator root hashes.
@@ -1070,7 +1076,7 @@ impl DbReader for AptosDB {
         &self,
         state_store_key: &StateKey,
         version: Version,
-    ) -> Result<(Option<StateValue>, SparseMerkleProof<StateValue>)> {
+    ) -> Result<(Option<StateValue>, SparseMerkleProof)> {
         gauged_api("get_account_state_with_proof_by_version", || {
             self.state_store
                 .get_value_with_proof_by_version(state_store_key, version)
@@ -1358,7 +1364,7 @@ impl DbWriter for AptosDB {
         &self,
         version: Version,
         expected_root_hash: HashValue,
-    ) -> Result<Box<dyn StateSnapshotReceiver<StateKeyAndValue>>> {
+    ) -> Result<Box<dyn StateSnapshotReceiver<StateKey, StateKeyAndValue>>> {
         gauged_api("get_state_snapshot_receiver", || {
             self.state_store
                 .get_snapshot_receiver(version, expected_root_hash)

storage/aptosdb/src/pruner/state_store/test.rs

@@ -22,8 +22,11 @@ fn put_value_set(
         .collect();
 
     let root = state_store
-        .put_value_sets(vec![&value_set], None, version, &mut cs)
+        .merklize_value_sets(vec![&value_set], None, version, &mut cs)
         .unwrap()[0];
+    state_store
+        .put_value_sets(vec![&value_set], version, &mut cs)
+        .unwrap();
     db.write_schemas(cs.batch).unwrap();
     state_store.set_latest_version(version);
 

storage/aptosdb/src/schema/jellyfish_merkle_node/mod.rs

@@ -11,15 +11,15 @@
 use crate::schema::JELLYFISH_MERKLE_NODE_CF_NAME;
 use anyhow::Result;
 use aptos_jellyfish_merkle::node_type::NodeKey;
-use aptos_types::{state_store::state_value::StateKeyAndValue, transaction::Version};
+use aptos_types::{state_store::state_key::StateKey, transaction::Version};
 use byteorder::{BigEndian, WriteBytesExt};
 use schemadb::{
     define_schema,
     schema::{KeyCodec, SeekKeyCodec, ValueCodec},
 };
 use std::mem::size_of;
 
-type Node = aptos_jellyfish_merkle::node_type::Node<StateKeyAndValue>;
+type Node = aptos_jellyfish_merkle::node_type::Node<StateKey>;
 
 define_schema!(
     JellyfishMerkleNodeSchema,

storage/aptosdb/src/schema/jellyfish_merkle_node/test.rs

@@ -12,11 +12,13 @@ proptest! {
     fn test_jellyfish_merkle_node_schema(
         node_key in any::<NodeKey>(),
         account_key in any::<HashValue>(),
-        value in any::<StateKeyAndValue>(),
+        value_hash in any::<HashValue>(),
+        state_key in any::<StateKey>(),
+        version in any::<Version>()
     ) {
         assert_encode_decode::<JellyfishMerkleNodeSchema>(
             &node_key,
-            &Node::new_leaf(account_key, value),
+            &Node::new_leaf(account_key, value_hash, (state_key, version)),
         );
     }
 }

storage/aptosdb/src/schema/mod.rs

@@ -15,7 +15,7 @@ pub(crate) mod jellyfish_merkle_node;
 pub(crate) mod ledger_counters;
 pub(crate) mod ledger_info;
 pub(crate) mod stale_node_index;
-pub(crate) mod state_value_index;
+pub(crate) mod state_value;
 pub(crate) mod transaction;
 pub(crate) mod transaction_accumulator;
 pub(crate) mod transaction_by_account;
@@ -34,7 +34,7 @@ pub const EVENT_CF_NAME: ColumnFamilyName = "event";
 pub const JELLYFISH_MERKLE_NODE_CF_NAME: ColumnFamilyName = "jellyfish_merkle_node";
 pub const LEDGER_COUNTERS_CF_NAME: ColumnFamilyName = "ledger_counters";
 pub const STALE_NODE_INDEX_CF_NAME: ColumnFamilyName = "stale_node_index";
-pub const STATE_VALUE_INDEX_CF_NAME: ColumnFamilyName = "state_value_index";
+pub const STATE_VALUE_CF_NAME: ColumnFamilyName = "state_value";
 pub const TRANSACTION_CF_NAME: ColumnFamilyName = "transaction";
 pub const TRANSACTION_ACCUMULATOR_CF_NAME: ColumnFamilyName = "transaction_accumulator";
 pub const TRANSACTION_BY_ACCOUNT_CF_NAME: ColumnFamilyName = "transaction_by_account";
@@ -80,6 +80,7 @@ pub mod fuzzing {
             assert_no_panic_decoding::<super::ledger_counters::LedgerCountersSchema>(data);
             assert_no_panic_decoding::<super::ledger_info::LedgerInfoSchema>(data);
             assert_no_panic_decoding::<super::stale_node_index::StaleNodeIndexSchema>(data);
+            assert_no_panic_decoding::<super::state_value::StateValueSchema>(data);
             assert_no_panic_decoding::<super::transaction::TransactionSchema>(data);
             assert_no_panic_decoding::<super::transaction_accumulator::TransactionAccumulatorSchema>(
                 data,

storage/aptosdb/src/schema/state_value_index/mod.rs → storage/aptosdb/src/schema/state_value/mod.rs

@@ -1,25 +1,25 @@
 // Copyright (c) Aptos
 // SPDX-License-Identifier: Apache-2.0
 
-//! This module defines the physical storage schema for state value index, which is used
-//! to access the state value directly without needing to walk through the JMT.
+//! This module defines the physical storage schema for state value, which is used
+//! to access the state value directly.
 //!
 //! An Index Key in this data set has 2 pieces of information:
 //!     1. The state key
 //!     2. The version associated with the key
-//! The value associated with the key is nibble length, which combined with the state key and
-//! version can give us access to the JMT leaf associated with corresponding key and version
+//! The value associated with the key is the the serialized State Key and Value.
 //!
-//!
-//! ```text
-//! |<---------key--------> |<-----value----->|
-//! |  state_key, version  |   num of nibbles |
+//! //! ```text
+//! |<-------- key -------->|<----- value ----->|
+//! |  state key  | version |  statekeyandvalue  |
 //! ```
 
-use crate::schema::{ensure_slice_len_eq, ensure_slice_len_gt, STATE_VALUE_INDEX_CF_NAME};
+use crate::schema::{ensure_slice_len_gt, STATE_VALUE_CF_NAME};
 use anyhow::Result;
 use aptos_types::{
-    state_store::{state_key::StateKey, state_key_prefix::StateKeyPrefix},
+    state_store::{
+        state_key::StateKey, state_key_prefix::StateKeyPrefix, state_value::StateKeyAndValue,
+    },
     transaction::Version,
 };
 use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt};
@@ -31,9 +31,9 @@ use std::{io::Write, mem::size_of};
 
 type Key = (StateKey, Version);
 
-define_schema!(StateValueIndexSchema, Key, u8, STATE_VALUE_INDEX_CF_NAME);
+define_schema!(StateValueSchema, Key, StateKeyAndValue, STATE_VALUE_CF_NAME);
 
-impl KeyCodec<StateValueIndexSchema> for Key {
+impl KeyCodec<StateValueSchema> for Key {
     fn encode_key(&self) -> Result<Vec<u8>> {
         let mut encoded = vec![];
         encoded.write_all(&self.0.encode()?)?;
@@ -52,20 +52,17 @@ impl KeyCodec<StateValueIndexSchema> for Key {
     }
 }
 
-impl ValueCodec<StateValueIndexSchema> for u8 {
+impl ValueCodec<StateValueSchema> for StateKeyAndValue {
     fn encode_value(&self) -> Result<Vec<u8>> {
-        let mut encoded = vec![];
-        encoded.write_u8(*self)?;
-        Ok(encoded)
+        bcs::to_bytes(self).map_err(Into::into)
     }
 
     fn decode_value(data: &[u8]) -> Result<Self> {
-        ensure_slice_len_eq(data, 1)?;
-        Ok(data[0])
+        bcs::from_bytes(data).map_err(Into::into)
     }
 }
 
-impl SeekKeyCodec<StateValueIndexSchema> for &StateKeyPrefix {
+impl SeekKeyCodec<StateValueSchema> for &StateKeyPrefix {
     fn encode_seek_key(&self) -> Result<Vec<u8>> {
         self.encode()
     }

storage/aptosdb/src/schema/state_value_index/test.rs → storage/aptosdb/src/schema/state_value/test.rs

@@ -10,10 +10,10 @@ proptest! {
     fn test_encode_decode(
         state_key in any::<StateKey>(),
         version in any::<Version>(),
-        num_nibbles in any::<u8>(),
+        v in any::<StateKeyAndValue>(),
     ) {
-        assert_encode_decode::<StateValueIndexSchema>(&(state_key, version), &num_nibbles);
+        assert_encode_decode::<StateValueSchema>(&(state_key, version), &v);
     }
 }
 
-test_no_panic_decoding!(StateValueIndexSchema);
+test_no_panic_decoding!(StateValueSchema);