[Storage] Add value index for JMT leaf nodes

Closes: aptos-labs#466

storage/aptosdb/src/lib.rs

@@ -244,6 +244,7 @@ impl AptosDB {
             JELLYFISH_MERKLE_NODE_CF_NAME,
             LEDGER_COUNTERS_CF_NAME,
             STALE_NODE_INDEX_CF_NAME,
+            STATE_VALUE_INDEX_CF_NAME,
             TRANSACTION_CF_NAME,
             TRANSACTION_ACCUMULATOR_CF_NAME,
             TRANSACTION_BY_ACCOUNT_CF_NAME,

storage/aptosdb/src/schema/mod.rs

@@ -15,13 +15,13 @@ 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 transaction;
 pub(crate) mod transaction_accumulator;
 pub(crate) mod transaction_by_account;
 pub(crate) mod transaction_by_hash;
 pub(crate) mod transaction_info;
 pub(crate) mod write_set;
-pub(crate) mod state_value_index;
 
 use anyhow::{ensure, Result};
 use schemadb::ColumnFamilyName;

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

@@ -18,8 +18,7 @@
 
 use crate::schema::{ensure_slice_len_eq, ensure_slice_len_gt, STATE_VALUE_INDEX_CF_NAME};
 use anyhow::Result;
-use aptos_types::state_store::state_key::StateKey;
-use aptos_types::transaction::Version;
+use aptos_types::{state_store::state_key::StateKey, transaction::Version};
 use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt};
 use schemadb::{
     define_schema,

storage/aptosdb/src/state_store/mod.rs

@@ -6,13 +6,13 @@
 #[cfg(test)]
 mod state_store_test;
 
-use crate::state_value_index::StateValueIndexSchema;
 use crate::{
     change_set::ChangeSet,
     ledger_counters::LedgerCounter,
     schema::{
         jellyfish_merkle_node::JellyfishMerkleNodeSchema, stale_node_index::StaleNodeIndexSchema,
     },
+    state_value_index::StateValueIndexSchema,
     AptosDbError,
 };
 use anyhow::{ensure, Result};
@@ -49,7 +49,7 @@ impl StateStore {
         Self { db }
     }
 
-    /// Get the account state blob given account address and root hash of state Merkle tree
+    /// Get the state value with proof given the state key and root hash of state Merkle tree
     pub fn get_value_with_proof_by_version(
         &self,
         state_key: &StateKey,
@@ -63,6 +63,61 @@ impl StateStore {
         ))
     }
 
+    /// Get the state value given the state key and root hash of state Merkle tree by using the
+    /// state value index. Only used for testing for now but should replace the
+    /// `get_value_with_proof_by_version` call for VM execution to fetch the value without proof.
+    #[cfg(test)]
+    pub fn get_value_by_version(
+        &self,
+        state_key: &StateKey,
+        version: Version,
+    ) -> Result<Option<StateValue>> {
+        match self.get_jmt_leaf_node_key(state_key, version)? {
+            Some(node_key) => {
+                let state_value =
+                    if let Some(node) = self.db.get::<JellyfishMerkleNodeSchema>(&node_key)? {
+                        match node {
+                            // Only if the node is a leaf node then the value
+                            // is present in the DB.
+                            Node::Leaf(leaf) => Some(leaf.value().value.clone()),
+                            _ => None,
+                        }
+                    } else {
+                        None
+                    };
+                Ok(state_value)
+            }
+            None => Ok(None),
+        }
+    }
+
+    /// Returns the value index in the form of number of nibbles for given pair of state key and version
+    /// which can be used to index into the JMT leaf.
+    #[cfg(test)]
+    fn get_jmt_leaf_node_key(
+        &self,
+        state_key: &StateKey,
+        version: Version,
+    ) -> Result<Option<NodeKey>> {
+        let mut iter = self.db.iter::<StateValueIndexSchema>(Default::default())?;
+        iter.seek_for_prev(&(state_key.clone(), version))?;
+        Ok(iter
+            .next()
+            .transpose()?
+            .and_then(|((db_state_key, db_version), num_nibbles)| {
+                if *state_key == db_state_key {
+                    Some(NodeKey::new(
+                        // It is possible that the db_version is not equal to the version passed,
+                        // but it should be strictly less than or equal to the version.
+                        db_version,
+                        NibblePath::new_from_state_key(state_key, num_nibbles as usize),
+                    ))
+                } else {
+                    None
+                }
+            }))
+    }
+
     /// Gets the proof that proves a range of accounts.
     pub fn get_value_range_proof(
         &self,

storage/aptosdb/src/state_store/state_store_test.rs

@@ -73,7 +73,18 @@ fn prune_stale_indices(
     .unwrap();
 }
 
-fn verify_state_in_store(
+fn verify_value_and_proof(
+    store: &StateStore,
+    address: AccountAddress,
+    expected_value: Option<&AccountStateBlob>,
+    version: Version,
+    root: HashValue,
+) {
+    verify_value_and_proof_in_store(store, address, expected_value, version, root);
+    verify_value_index_in_store(store, address, expected_value, version);
+}
+
+fn verify_value_and_proof_in_store(
     store: &StateStore,
     address: AccountAddress,
     expected_value: Option<&AccountStateBlob>,
@@ -99,6 +110,23 @@ fn verify_state_in_store(
         .unwrap();
 }
 
+fn verify_value_index_in_store(
+    store: &StateStore,
+    address: AccountAddress,
+    expected_value: Option<&AccountStateBlob>,
+    version: Version,
+) {
+    let value = store
+        .get_value_by_version(&StateKey::AccountAddressKey(address), version)
+        .unwrap();
+    assert_eq!(
+        value
+            .map(|x| AccountStateBlob::try_from(x).unwrap())
+            .as_ref(),
+        expected_value
+    );
+}
+
 #[test]
 fn test_empty_store() {
     let tmp_dir = TempPath::new();
@@ -132,12 +160,14 @@ fn test_state_store_reader_writer() {
         0, /* expected_nodes_retired */
         0, /* expected_blobs_retired */
     );
-    verify_state_in_store(store, address1, Some(&value1), 0, root);
-    verify_state_in_store(store, address2, None, 0, root);
-    verify_state_in_store(store, address3, None, 0, root);
+    verify_value_and_proof(store, address1, Some(&value1), 0, root);
+
+    verify_value_and_proof(store, address2, None, 0, root);
+    verify_value_and_proof(store, address3, None, 0, root);
 
     // Insert address 1 with updated value1, address2 with value 2 and address3 with value3 and
     // verify new states.
+
     root = put_account_state_set(
         store,
         vec![
@@ -150,9 +180,10 @@ fn test_state_store_reader_writer() {
         1, /* expected_nodes_retired */
         1, /* expected_blobs_retired */
     );
-    verify_state_in_store(store, address1, Some(&value1_update), 1, root);
-    verify_state_in_store(store, address2, Some(&value2), 1, root);
-    verify_state_in_store(store, address3, Some(&value3), 1, root);
+
+    verify_value_and_proof(store, address1, Some(&value1_update), 1, root);
+    verify_value_and_proof(store, address2, Some(&value2), 1, root);
+    verify_value_and_proof(store, address3, Some(&value3), 1, root);
 }
 
 #[test]
@@ -213,7 +244,7 @@ fn test_retired_records() {
             1, /* target_least_readable_version */
             0, /* limit */
         );
-        verify_state_in_store(store, address1, Some(&value1), 0, root0);
+        verify_value_and_proof(store, address1, Some(&value1), 0, root0);
     }
     // Prune till version=1.
     {
@@ -227,9 +258,9 @@ fn test_retired_records() {
             .get_value_with_proof_by_version(&StateKey::AccountAddressKey(address2), 0)
             .is_err());
         // root1 is still there.
-        verify_state_in_store(store, address1, Some(&value1), 1, root1);
-        verify_state_in_store(store, address2, Some(&value2_update), 1, root1);
-        verify_state_in_store(store, address3, Some(&value3), 1, root1);
+        verify_value_and_proof(store, address1, Some(&value1), 1, root1);
+        verify_value_and_proof(store, address2, Some(&value2_update), 1, root1);
+        verify_value_and_proof(store, address3, Some(&value3), 1, root1);
     }
     // Prune till version=2.
     {
@@ -243,9 +274,9 @@ fn test_retired_records() {
             .get_value_with_proof_by_version(&StateKey::AccountAddressKey(address2), 1)
             .is_err());
         // root2 is still there.
-        verify_state_in_store(store, address1, Some(&value1), 2, root2);
-        verify_state_in_store(store, address2, Some(&value2_update), 2, root2);
-        verify_state_in_store(store, address3, Some(&value3_update), 2, root2);
+        verify_value_and_proof(store, address1, Some(&value1), 2, root2);
+        verify_value_and_proof(store, address2, Some(&value2_update), 2, root2);
+        verify_value_and_proof(store, address3, Some(&value3_update), 2, root2);
     }
 }
 

storage/jellyfish-merkle/src/iterator/mod.rs

@@ -125,7 +125,7 @@ where
         let mut done = false;
 
         let mut current_node_key = NodeKey::new_empty_path(version);
-        let nibble_path = NibblePath::new(starting_key.to_vec());
+        let nibble_path = NibblePath::new_even(starting_key.to_vec());
         let mut nibble_iter = nibble_path.nibbles();
 
         while let Node::Internal(internal_node) = reader.get_node(&current_node_key)? {

storage/jellyfish-merkle/src/lib.rs

@@ -619,7 +619,7 @@ where
         version: Version,
         tree_cache: &mut TreeCache<R, V>,
     ) -> Result<()> {
-        let nibble_path = NibblePath::new(key.to_vec());
+        let nibble_path = NibblePath::new_even(key.to_vec());
 
         // Get the root node. If this is the first operation, it would get the root node from the
         // underlying db. Otherwise it most likely would come from `cache`.
@@ -758,7 +758,8 @@ where
         // visited part of the nibble iter of the incoming key and advances the existing leaf
         // nibble iterator by the length of that prefix.
         let mut visited_nibble_iter = nibble_iter.visited_nibbles();
-        let existing_leaf_nibble_path = NibblePath::new(existing_leaf_node.account_key().to_vec());
+        let existing_leaf_nibble_path =
+            NibblePath::new_even(existing_leaf_node.account_key().to_vec());
         let mut existing_leaf_nibble_iter = existing_leaf_nibble_path.nibbles();
         skip_common_prefix(&mut visited_nibble_iter, &mut existing_leaf_nibble_iter);
 
@@ -876,7 +877,7 @@ where
         // Empty tree just returns proof with no sibling hash.
         let mut next_node_key = NodeKey::new_empty_path(version);
         let mut siblings = vec![];
-        let nibble_path = NibblePath::new(key.to_vec());
+        let nibble_path = NibblePath::new_even(key.to_vec());
         let mut nibble_iter = nibble_path.nibbles();
 
         // We limit the number of loops here deliberately to avoid potential cyclic graph bugs

storage/jellyfish-merkle/src/node_type/mod.rs

@@ -60,7 +60,7 @@ impl NodeKey {
 
     /// A shortcut to generate a node key consisting of a version and an empty nibble path.
     pub fn new_empty_path(version: Version) -> Self {
-        Self::new(version, NibblePath::new(vec![]))
+        Self::new(version, NibblePath::new_even(vec![]))
     }
 
     /// Gets the version.
@@ -123,7 +123,7 @@ impl NodeKey {
             nibble_bytes
         );
         let nibble_path = if num_nibbles % 2 == 0 {
-            NibblePath::new(nibble_bytes)
+            NibblePath::new_even(nibble_bytes)
         } else {
             let padding = nibble_bytes.last().unwrap() & 0x0f;
             ensure!(

storage/jellyfish-merkle/src/restore/mod.rs

@@ -322,7 +322,7 @@ where
 
     /// Restores one account.
     fn add_one(&mut self, new_key: HashValue, new_value: V) {
-        let nibble_path = NibblePath::new(new_key.to_vec());
+        let nibble_path = NibblePath::new_even(new_key.to_vec());
         let mut nibbles = nibble_path.nibbles();
 
         for i in 0..ROOT_NIBBLE_HEIGHT {

storage/jellyfish-merkle/src/tree_cache/tree_cache_test.rs

@@ -9,7 +9,7 @@ use aptos_types::nibble::nibble_path::NibblePath;
 fn random_leaf_with_key(next_version: Version) -> (Node<ValueBlob>, NodeKey) {
     let address = HashValue::random();
     let node = Node::new_leaf(address, ValueBlob::from(HashValue::random().to_vec()));
-    let node_key = NodeKey::new(next_version, NibblePath::new(address.to_vec()));
+    let node_key = NodeKey::new(next_version, NibblePath::new_even(address.to_vec()));
     (node, node_key)
 }
 

storage/scratchpad/src/sparse_merkle/mod.rs

@@ -486,7 +486,7 @@ where
             .with_label_values(&["generate_node_hashes"])
             .start_timer();
         let mut node_hashes = HashMap::new();
-        let mut nibble_path = NibblePath::new(vec![]);
+        let mut nibble_path = NibblePath::new_even(vec![]);
         self.collect_new_hashes(
             touched_accounts.as_slice(),
             self.smt.root_weak(),
@@ -516,7 +516,7 @@ where
         if level_within_nibble == 0 {
             if depth_in_nibble != 0 {
                 cur_nibble_path
-                    .push(NibblePath::new(keys[0].to_vec()).get_nibble(depth_in_nibble - 1));
+                    .push(NibblePath::new_even(keys[0].to_vec()).get_nibble(depth_in_nibble - 1));
             }
             node_hashes.insert(cur_nibble_path.clone(), subtree.hash());
         }
@@ -559,7 +559,7 @@ where
                 if level_within_nibble != 0 {
                     let mut leaf_nibble_path = cur_nibble_path.clone();
                     leaf_nibble_path
-                        .push(NibblePath::new(keys[0].to_vec()).get_nibble(depth_in_nibble));
+                        .push(NibblePath::new_even(keys[0].to_vec()).get_nibble(depth_in_nibble));
                     node_hashes.insert(leaf_nibble_path, subtree.hash());
                 }
             }