Add Value index

storage/aptosdb/src/schema/mod.rs

@@ -21,6 +21,7 @@ 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;
@@ -33,6 +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 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";

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

@@ -0,0 +1,67 @@
+// 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.
+//!
+//! 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
+//!
+//!
+//! //! ```text
+//! |<---------key--------> |<-----value----->|
+//! |  state_key, version  |   num of nibbles |
+//! ```
+
+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 byteorder::{BigEndian, ReadBytesExt, WriteBytesExt};
+use schemadb::{
+    define_schema,
+    schema::{KeyCodec, ValueCodec},
+};
+use std::{io::Write, mem::size_of};
+
+type Key = (StateKey, Version);
+
+define_schema!(StateValueIndexSchema, Key, u8, STATE_VALUE_INDEX_CF_NAME);
+
+impl KeyCodec<StateValueIndexSchema> for Key {
+    fn encode_key(&self) -> Result<Vec<u8>> {
+        let mut encoded = vec![];
+        encoded.write_all(&bcs::to_bytes(&self.0)?)?;
+        encoded.write_u64::<BigEndian>(self.1)?;
+        Ok(encoded)
+    }
+
+    fn decode_key(data: &[u8]) -> Result<Self> {
+        const VERSION_SIZE: usize = size_of::<Version>();
+
+        ensure_slice_len_gt(data, VERSION_SIZE)?;
+        let state_key_len = data.len() - VERSION_SIZE;
+        let state_key: StateKey = bcs::from_bytes(&data[..state_key_len])?;
+        let version = (&data[state_key_len..]).read_u64::<BigEndian>()?;
+        Ok((state_key, version))
+    }
+}
+
+impl ValueCodec<StateValueIndexSchema> for u8 {
+    fn encode_value(&self) -> Result<Vec<u8>> {
+        let mut encoded = vec![];
+        encoded.write_u8(*self)?;
+        Ok(encoded)
+    }
+
+    fn decode_value(data: &[u8]) -> Result<Self> {
+        ensure_slice_len_eq(data, 1)?;
+        Ok(data[0])
+    }
+}
+
+#[cfg(test)]
+mod test;

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

@@ -0,0 +1,19 @@
+// Copyright (c) Aptos
+// SPDX-License-Identifier: Apache-2.0
+
+use super::*;
+use proptest::prelude::*;
+use schemadb::{schema::fuzzing::assert_encode_decode, test_no_panic_decoding};
+
+proptest! {
+    #[test]
+    fn test_encode_decode(
+        state_key in any::<StateKey>(),
+        version in any::<Version>(),
+        num_nibbles in any::<u8>(),
+    ) {
+        assert_encode_decode::<StateValueIndexSchema>(&(state_key, version), &num_nibbles);
+    }
+}
+
+test_no_panic_decoding!(StateValueIndexSchema);

storage/aptosdb/src/state_store/mod.rs

@@ -6,6 +6,7 @@
 #[cfg(test)]
 mod state_store_test;
 
+use crate::state_value_index::StateValueIndexSchema;
 use crate::{
     change_set::ChangeSet,
     ledger_counters::LedgerCounter,
@@ -117,7 +118,7 @@ impl StateStore {
                 counter_bumps.bump(LedgerCounter::StaleStateNodes, stats.stale_nodes);
                 counter_bumps.bump(LedgerCounter::StaleStateLeaves, stats.stale_leaves);
             });
-        add_node_batch(&mut cs.batch, &tree_update_batch.node_batch)?;
+        add_node_batch_and_index(&mut cs.batch, &tree_update_batch.node_batch)?;
 
         tree_update_batch
             .stale_node_index_batch
@@ -281,15 +282,26 @@ impl TreeReader<StateKeyAndValue> 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)?;
+        add_node_batch_and_index(&mut batch, node_batch)?;
         self.db.write_schemas(batch)
     }
 }
 
-fn add_node_batch(batch: &mut SchemaBatch, node_batch: &NodeBatch) -> Result<()> {
+fn add_node_batch_and_index(batch: &mut SchemaBatch, node_batch: &NodeBatch) -> Result<()> {
     node_batch
         .iter()
-        .map(|(node_key, node)| batch.put::<JellyfishMerkleNodeSchema>(node_key, node))
+        .map(|(node_key, node)| {
+            batch.put::<JellyfishMerkleNodeSchema>(node_key, node)?;
+            // Add the value index for leaf nodes.
+            match node {
+                Node::Leaf(leaf) => batch.put::<StateValueIndexSchema>(
+                    &(leaf.value().key.clone(), node_key.version()),
+                    &(node_key.nibble_path().num_nibbles() as u8),
+                ),
+
+                _ => Ok(()),
+            }
+        })
         .collect::<Result<Vec<_>>>()?;
     Ok(())
 }

types/src/state_store/state_key.rs

@@ -40,8 +40,8 @@ impl StateKeyPrefix {
     }
 }
 
-struct RawStateKey {
-    bytes: Vec<u8>,
+pub struct RawStateKey {
+    pub bytes: Vec<u8>,
 }
 
 impl From<&StateKey> for RawStateKey {