[Storage] Add State Key to the JMT leaf node

Closes: aptos-labs#455

execution/executor/tests/db_bootstrapper_test.rs

@@ -22,7 +22,7 @@ use aptos_types::{
         OnChainConfig, ValidatorSet,
     },
     proof::SparseMerkleRangeProof,
-    state_store::{state_key::StateKey, state_value::StateValue},
+    state_store::{state_key::StateKey, state_value::StateKeyAndValue},
     transaction::{
         authenticator::AuthenticationKey, ChangeSet, Transaction, Version, WriteSetPayload,
         PRE_GENESIS_VERSION,
@@ -205,7 +205,7 @@ fn get_configuration(db: &DbReaderWriter) -> ConfigurationResource {
 fn get_state_backup(
     db: &Arc<AptosDB>,
 ) -> (
-    Vec<(HashValue, StateValue)>,
+    Vec<(HashValue, StateKeyAndValue)>,
     SparseMerkleRangeProof,
     HashValue,
 ) {
@@ -231,7 +231,7 @@ fn get_state_backup(
 
 fn restore_state_to_db(
     db: &Arc<AptosDB>,
-    accounts: Vec<(HashValue, StateValue)>,
+    accounts: Vec<(HashValue, StateKeyAndValue)>,
     proof: SparseMerkleRangeProof,
     root_hash: HashValue,
     version: Version,

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

@@ -15,7 +15,10 @@ use aptos_types::{
     epoch_state::EpochState,
     ledger_info::{LedgerInfo, LedgerInfoWithSignatures},
     proof::SparseMerkleRangeProof,
-    state_store::state_value::StateValueChunkWithProof,
+    state_store::{
+        state_key::StateKey,
+        state_value::{StateKeyAndValue, StateValueChunkWithProof},
+    },
     transaction::{
         RawTransaction, Script, SignedTransaction, Transaction, TransactionListWithProof,
         TransactionOutput, TransactionOutputListWithProof, TransactionPayload, TransactionStatus,
@@ -127,7 +130,10 @@ impl AptosDataClient for MockAptosDataClient {
         // Create epoch ending ledger infos according to the requested epochs
         let mut account_blobs = vec![];
         for _ in start_index..=end_index {
-            account_blobs.push((HashValue::random(), vec![].into()));
+            account_blobs.push((
+                HashValue::random(),
+                StateKeyAndValue::new(StateKey::Raw(vec![]), vec![].into()),
+            ));
         }
 
         // Create an account states chunk with proof

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

@@ -18,7 +18,10 @@ use aptos_types::{
     event::EventKey,
     ledger_info::{LedgerInfo, LedgerInfoWithSignatures},
     proof::{SparseMerkleRangeProof, TransactionInfoListWithProof},
-    state_store::state_value::StateValueChunkWithProof,
+    state_store::{
+        state_key::StateKey,
+        state_value::{StateKeyAndValue, StateValueChunkWithProof},
+    },
     transaction::{
         RawTransaction, Script, SignedTransaction, Transaction, TransactionListWithProof,
         TransactionOutput, TransactionOutputListWithProof, TransactionPayload, TransactionStatus,
@@ -96,7 +99,10 @@ async fn test_get_account_states_with_proof() {
     let chunk_size = end_account_index - start_account_index + 1;
     let mut account_blobs = vec![];
     for _ in 0..chunk_size {
-        account_blobs.push((HashValue::zero(), vec![].into()));
+        account_blobs.push((
+            HashValue::zero(),
+            StateKeyAndValue::new(StateKey::Raw(vec![]), vec![].into()),
+        ));
     }
     let expected_response =
         StorageServiceResponse::AccountStatesChunkWithProof(StateValueChunkWithProof {
@@ -742,7 +748,10 @@ impl DbReader for MockDbReader {
         // Create empty account blobs
         let mut account_blobs = vec![];
         for _ in 0..chunk_size {
-            account_blobs.push((HashValue::zero(), vec![].into()));
+            account_blobs.push((
+                HashValue::zero(),
+                StateKeyAndValue::new(StateKey::Raw(vec![]), vec![].into()),
+            ));
         }
 
         // Create an account states chunk with proof

storage/aptosdb/src/aptosdb_test.rs

@@ -641,7 +641,10 @@ fn test_get_latest_tree_state() {
             &NodeKey::new_empty_path(PRE_GENESIS_VERSION),
             &Node::new_leaf(
                 StateKey::AccountAddressKey(address).hash(),
-                StateValue::from(blob.clone()),
+                StateKeyAndValue::new(
+                    StateKey::AccountAddressKey(address),
+                    StateValue::from(blob.clone()),
+                ),
             ),
         )
         .unwrap();

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::StateValue,
+    state_store::state_value::StateKeyAndValue,
     transaction::{Transaction, TransactionInfo, Version},
 };
 use itertools::zip_eq;
@@ -103,7 +103,7 @@ impl BackupHandler {
     pub fn get_account_iter(
         &self,
         version: Version,
-    ) -> Result<Box<dyn Iterator<Item = Result<(HashValue, StateValue)>> + Send + Sync>> {
+    ) -> Result<Box<dyn Iterator<Item = Result<(HashValue, StateKeyAndValue)>> + Send + Sync>> {
         let iterator = JellyfishMerkleIterator::new(
             Arc::clone(&self.state_store),
             version,

storage/aptosdb/src/backup/restore_handler.rs

@@ -12,7 +12,7 @@ use aptos_types::{
     contract_event::ContractEvent,
     ledger_info::LedgerInfoWithSignatures,
     proof::definition::LeafCount,
-    state_store::state_value::StateValue,
+    state_store::state_value::StateKeyAndValue,
     transaction::{Transaction, TransactionInfo, Version, PRE_GENESIS_VERSION},
 };
 use schemadb::DB;
@@ -53,7 +53,7 @@ impl RestoreHandler {
         &self,
         version: Version,
         expected_root_hash: HashValue,
-    ) -> Result<JellyfishMerkleRestore<StateValue>> {
+    ) -> Result<JellyfishMerkleRestore<StateKeyAndValue>> {
         JellyfishMerkleRestore::new_overwrite(
             Arc::clone(&self.state_store),
             version,

storage/aptosdb/src/lib.rs

@@ -73,7 +73,9 @@ use aptos_types::{
     state_proof::StateProof,
     state_store::{
         state_key::StateKey,
-        state_value::{StateValue, StateValueChunkWithProof, StateValueWithProof},
+        state_value::{
+            StateKeyAndValue, StateValue, StateValueChunkWithProof, StateValueWithProof,
+        },
     },
     transaction::{
         AccountTransactionsWithProof, Transaction, TransactionInfo, TransactionListWithProof,
@@ -1301,7 +1303,7 @@ impl DbWriter for AptosDB {
         &self,
         version: Version,
         expected_root_hash: HashValue,
-    ) -> Result<Box<dyn StateSnapshotReceiver<StateValue>>> {
+    ) -> Result<Box<dyn StateSnapshotReceiver<StateKeyAndValue>>> {
         gauged_api("get_state_snapshot_receiver", || {
             self.state_store
                 .get_snapshot_receiver(version, expected_root_hash)

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::StateValue, transaction::Version};
+use aptos_types::{state_store::state_value::StateKeyAndValue, 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<StateValue>;
+type Node = aptos_jellyfish_merkle::node_type::Node<StateKeyAndValue>;
 
 define_schema!(
     JellyfishMerkleNodeSchema,

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

@@ -12,7 +12,7 @@ proptest! {
     fn test_jellyfish_merkle_node_schema(
         node_key in any::<NodeKey>(),
         account_key in any::<HashValue>(),
-        value in any::<StateValue>(),
+        value in any::<StateKeyAndValue>(),
     ) {
         assert_encode_decode::<JellyfishMerkleNodeSchema>(
             &node_key,

storage/aptosdb/src/state_store/mod.rs

@@ -25,7 +25,7 @@ use aptos_types::{
     proof::{SparseMerkleProof, SparseMerkleRangeProof},
     state_store::{
         state_key::StateKey,
-        state_value::{StateValue, StateValueChunkWithProof},
+        state_value::{StateKeyAndValue, StateValue, StateValueChunkWithProof},
     },
     transaction::Version,
 };
@@ -34,9 +34,9 @@ use schemadb::{SchemaBatch, DB};
 use std::{collections::HashMap, sync::Arc};
 use storage_interface::StateSnapshotReceiver;
 
-type LeafNode = aptos_jellyfish_merkle::node_type::LeafNode<StateValue>;
-type Node = aptos_jellyfish_merkle::node_type::Node<StateValue>;
-type NodeBatch = aptos_jellyfish_merkle::NodeBatch<StateValue>;
+type LeafNode = aptos_jellyfish_merkle::node_type::LeafNode<StateKeyAndValue>;
+type Node = aptos_jellyfish_merkle::node_type::Node<StateKeyAndValue>;
+type NodeBatch = aptos_jellyfish_merkle::NodeBatch<StateKeyAndValue>;
 
 #[derive(Debug)]
 pub(crate) struct StateStore {
@@ -54,7 +54,12 @@ impl StateStore {
         state_key: &StateKey,
         version: Version,
     ) -> Result<(Option<StateValue>, SparseMerkleProof<StateValue>)> {
-        JellyfishMerkleTree::new(self).get_with_proof(state_key.hash(), version)
+        let (state_key_value_option, proof) =
+            JellyfishMerkleTree::new(self).get_with_proof(state_key.hash(), version)?;
+        Ok((
+            state_key_value_option.map(|x| x.value),
+            SparseMerkleProof::from(proof),
+        ))
     }
 
     /// Gets the proof that proves a range of accounts.
@@ -77,16 +82,26 @@ impl StateStore {
     ) -> Result<Vec<HashValue>> {
         let value_sets = value_state_sets
             .into_iter()
-            .map(|account_states| {
-                account_states
+            .map(|value_set| {
+                value_set
                     .iter()
-                    .map(|(addr, blob)| (addr.hash(), blob))
+                    .map(|(key, value)| {
+                        (
+                            key.hash(),
+                            StateKeyAndValue::new(key.clone(), value.clone()),
+                        )
+                    })
                     .collect::<Vec<_>>()
             })
             .collect::<Vec<_>>();
 
+        let value_sets_ref = value_sets
+            .iter()
+            .map(|value_set| value_set.iter().map(|(x, y)| (*x, y)).collect::<Vec<_>>())
+            .collect::<Vec<_>>();
+
         let (new_root_hash_vec, tree_update_batch) = JellyfishMerkleTree::new(self)
-            .batch_put_value_sets(value_sets, node_hashes, first_version)?;
+            .batch_put_value_sets(value_sets_ref, node_hashes, first_version)?;
 
         let num_versions = new_root_hash_vec.len();
         assert_eq!(num_versions, tree_update_batch.node_stats.len());
@@ -160,15 +175,15 @@ impl StateStore {
         let result_iter =
             JellyfishMerkleIterator::new_by_index(Arc::clone(self), version, first_index)?
                 .take(chunk_size);
-        let account_blobs: Vec<(HashValue, StateValue)> =
+        let state_key_values: Vec<(HashValue, StateKeyAndValue)> =
             process_results(result_iter, |iter| iter.collect())?;
         ensure!(
-            !account_blobs.is_empty(),
+            !state_key_values.is_empty(),
             AptosDbError::NotFound(format!("State chunk starting at {}", first_index)),
         );
-        let last_index = (account_blobs.len() - 1 + first_index) as u64;
-        let first_key = account_blobs.first().expect("checked to exist").0;
-        let last_key = account_blobs.last().expect("checked to exist").0;
+        let last_index = (state_key_values.len() - 1 + first_index) as u64;
+        let first_key = state_key_values.first().expect("checked to exist").0;
+        let last_key = state_key_values.last().expect("checked to exist").0;
         let proof = self.get_value_range_proof(last_key, version)?;
         let root_hash = self.get_root_hash(version)?;
 
@@ -177,7 +192,7 @@ impl StateStore {
             last_index,
             first_key,
             last_key,
-            raw_values: account_blobs,
+            raw_values: state_key_values,
             proof,
             root_hash,
         })
@@ -187,7 +202,7 @@ impl StateStore {
         self: &Arc<Self>,
         version: Version,
         expected_root_hash: HashValue,
-    ) -> Result<Box<dyn StateSnapshotReceiver<StateValue>>> {
+    ) -> Result<Box<dyn StateSnapshotReceiver<StateKeyAndValue>>> {
         Ok(Box::new(JellyfishMerkleRestore::new_overwrite(
             Arc::clone(self),
             version,
@@ -196,7 +211,7 @@ impl StateStore {
     }
 }
 
-impl TreeReader<StateValue> for StateStore {
+impl TreeReader<StateKeyAndValue> for StateStore {
     fn get_node_option(&self, node_key: &NodeKey) -> Result<Option<Node>> {
         self.db.get::<JellyfishMerkleNodeSchema>(node_key)
     }
@@ -263,7 +278,7 @@ impl TreeReader<StateValue> for StateStore {
     }
 }
 
-impl TreeWriter<StateValue> for StateStore {
+impl TreeWriter<StateKeyAndValue> for StateStore {
     fn write_node_batch(&self, node_batch: &NodeBatch) -> Result<()> {
         let mut batch = SchemaBatch::new();
         add_node_batch(&mut batch, node_batch)?;

storage/aptosdb/src/state_store/state_store_test.rs

@@ -274,7 +274,7 @@ proptest! {
                 .unwrap();
             let mut expected_values: Vec<_> = kvs[..=i]
                 .iter()
-                .map(|(key, value)| (key.hash(), value.clone()))
+                .map(|(key, value)| (key.hash(), StateKeyAndValue::new(key.clone(), value.clone())))
                 .collect();
             expected_values.sort_unstable_by_key(|item| item.0);
             prop_assert_eq!(actual_values, expected_values);
@@ -314,6 +314,7 @@ proptest! {
             .clone()
             .into_iter()
             .take(batch1_size)
+            .map(|(key, value)| (key, StateKeyAndValue::new(StateKey::Raw(vec![]), value)))
             .collect();
         let rightmost_of_batch1 = batch1.last().map(|(key, _value)| *key).unwrap();
         let proof_of_batch1 = store1
@@ -325,6 +326,7 @@ proptest! {
         let batch2: Vec<_> = ordered_input
             .into_iter()
             .skip(batch1_size)
+            .map(|(key, value)| (key, StateKeyAndValue::new(StateKey::Raw(vec![]), value)))
             .collect();
         let rightmost_of_batch2 = batch2.last().map(|(key, _value)| *key).unwrap();
         let proof_of_batch2 = store1
@@ -412,6 +414,7 @@ proptest! {
         let batch1: Vec<_> = ordered_input
             .into_iter()
             .take(batch1_size)
+            .map(|(key, value)| (key, StateKeyAndValue::new(StateKey::Raw(vec![]), value)))
             .collect();
         let rightmost_of_batch1 = batch1.last().map(|(key, _value)| *key).unwrap();
         let proof_of_batch1 = store1

storage/backup/backup-cli/src/backup_types/state_snapshot/backup.rs

@@ -15,7 +15,7 @@ use aptos_crypto::HashValue;
 use aptos_logger::prelude::*;
 use aptos_types::{
     ledger_info::LedgerInfoWithSignatures, proof::TransactionInfoWithProof,
-    state_store::state_value::StateValue, transaction::Version,
+    state_store::state_value::StateKeyAndValue, transaction::Version,
 };
 use bytes::Bytes;
 use once_cell::sync::Lazy;
@@ -155,7 +155,7 @@ impl StateSnapshotBackupController {
     }
 
     fn parse_key(record: &Bytes) -> Result<HashValue> {
-        let (key, _): (HashValue, StateValue) = bcs::from_bytes(record)?;
+        let (key, _): (HashValue, StateKeyAndValue) = bcs::from_bytes(record)?;
         Ok(key)
     }
 

storage/backup/backup-cli/src/backup_types/state_snapshot/restore.rs

@@ -25,7 +25,7 @@ use aptos_crypto::HashValue;
 use aptos_logger::prelude::*;
 use aptos_types::{
     ledger_info::LedgerInfoWithSignatures, proof::TransactionInfoWithProof,
-    state_store::state_value::StateValue, transaction::Version,
+    state_store::state_value::StateKeyAndValue, transaction::Version,
 };
 use std::sync::Arc;
 use storage_interface::StateSnapshotReceiver;
@@ -145,7 +145,7 @@ impl StateSnapshotRestoreController {
     async fn read_state_value(
         &self,
         file_handle: FileHandle,
-    ) -> Result<Vec<(HashValue, StateValue)>> {
+    ) -> Result<Vec<(HashValue, StateKeyAndValue)>> {
         let mut file = self.storage.open_for_read(&file_handle).await?;
 
         let mut chunk = vec![];