[Storage] Follow up from adding index for JMT leaf

Closes: aptos-labs#487

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

@@ -33,7 +33,7 @@ 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_all(&self.0.encode()?)?;
         encoded.write_u64::<BigEndian>(self.1)?;
         Ok(encoded)
     }
@@ -43,7 +43,7 @@ impl KeyCodec<StateValueIndexSchema> for Key {
 
         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 state_key: StateKey = StateKey::decode(&data[..state_key_len])?;
         let version = (&data[state_key_len..]).read_u64::<BigEndian>()?;
         Ok((state_key, version))
     }

storage/aptosdb/src/state_store/mod.rs

@@ -74,18 +74,16 @@ impl StateStore {
     ) -> 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)
+                if let Some(Node::Leaf(leaf)) =
+                    self.db.get::<JellyfishMerkleNodeSchema>(&node_key)?
+                {
+                    Ok(Some(leaf.value().value.clone()))
+                } else {
+                    Err(anyhow::anyhow!(
+                        "Can't find value in JMT for state key {:?}",
+                        state_key
+                    ))
+                }
             }
             None => Ok(None),
         }

types/src/nibble/nibble_path/mod.rs

@@ -117,21 +117,15 @@ impl NibblePath {
         if num_nibbles % 2 == 1 {
             // Rounded up number of bytes to be considered
             let num_bytes = (num_nibbles + 1) / 2;
-            let mut nibble_bytes = bytes[..num_bytes].to_vec();
             checked_precondition!(bytes.len() >= num_bytes);
+            let mut nibble_bytes = bytes[..num_bytes].to_vec();
             // If number of nibbles is odd, make sure to pad the last nibble with 0s.
             let last_byte_padded = bytes[num_bytes - 1] & 0xF0;
             nibble_bytes[num_bytes - 1] = last_byte_padded;
-            NibblePath {
-                num_nibbles,
-                bytes: nibble_bytes,
-            }
+            NibblePath::new_odd(nibble_bytes)
         } else {
             checked_precondition!(bytes.len() >= num_nibbles / 2);
-            NibblePath {
-                num_nibbles,
-                bytes: bytes[..num_nibbles / 2].to_vec(),
-            }
+            NibblePath::new_even(bytes[..num_nibbles / 2].to_vec())
         }
     }
 

types/src/state_store/state_key.rs

@@ -8,9 +8,10 @@ use aptos_crypto::{
 };
 use aptos_crypto_derive::CryptoHasher;
 use move_core_types::account_address::AccountAddress;
-use num_derive::ToPrimitive;
-use num_traits::ToPrimitive;
+use num_derive::{FromPrimitive, ToPrimitive};
+use num_traits::FromPrimitive;
 use serde::{Deserialize, Serialize};
+use thiserror::Error;
 
 #[derive(
     Clone, Debug, CryptoHasher, Eq, PartialEq, Serialize, Deserialize, Ord, PartialOrd, Hash,
@@ -24,43 +25,49 @@ pub enum StateKey {
     Raw(Vec<u8>),
 }
 
-#[derive(ToPrimitive)]
+#[repr(u8)]
+#[derive(FromPrimitive, ToPrimitive)]
 enum StateKeyPrefix {
     AccountAddress,
     AccessPath,
     Raw = 255,
 }
 
-impl StateKeyPrefix {
-    fn to_bytes(&self) -> Vec<u8> {
-        let byte = self
-            .to_u8()
-            .expect("Failed to convert StateKeyPrefix to u8");
-        vec![byte]
-    }
-}
+impl StateKey {
+    /// Serializes to bytes for physical storage.
+    pub fn encode(&self) -> anyhow::Result<Vec<u8>> {
+        let mut out = vec![];
 
-struct RawStateKey {
-    bytes: Vec<u8>,
-}
-
-impl From<&StateKey> for RawStateKey {
-    fn from(key: &StateKey) -> Self {
-        let (prefix, raw_key) = match key {
-            StateKey::AccountAddressKey(account_address) => {
-                (StateKeyPrefix::AccountAddress, account_address.to_vec())
-            }
+        let (prefix, raw_key) = match self {
+            StateKey::AccountAddressKey(account_address) => (
+                StateKeyPrefix::AccountAddress,
+                bcs::to_bytes(account_address)?,
+            ),
             StateKey::AccessPath(access_path) => {
-                let mut raw_key = access_path.address.to_vec();
-                raw_key.extend(access_path.path.clone());
-                (StateKeyPrefix::AccessPath, raw_key)
+                (StateKeyPrefix::AccessPath, bcs::to_bytes(access_path)?)
             }
             StateKey::Raw(raw_bytes) => (StateKeyPrefix::Raw, raw_bytes.to_vec()),
         };
-        let mut bytes = prefix.to_bytes();
-        bytes.extend(raw_key);
+        out.push(prefix as u8);
+        out.extend(raw_key);
+        Ok(out)
+    }
 
-        Self { bytes }
+    /// Recovers from serialized bytes in physical storage.
+    pub fn decode(val: &[u8]) -> anyhow::Result<StateKey> {
+        if val.is_empty() {
+            return Err(StateKeyDecodeErr::EmptyInput.into());
+        }
+        let tag = val[0];
+        let state_key_tag = StateKeyPrefix::from_u8(tag)
+            .ok_or(StateKeyDecodeErr::UnknownTag { unknown_tag: tag })?;
+        match state_key_tag {
+            StateKeyPrefix::AccountAddress => {
+                Ok(StateKey::AccountAddressKey(bcs::from_bytes(&val[1..])?))
+            }
+            StateKeyPrefix::AccessPath => Ok(StateKey::AccessPath(bcs::from_bytes(&val[1..])?)),
+            StateKeyPrefix::Raw => Ok(StateKey::Raw(val[1..].to_vec())),
+        }
     }
 }
 
@@ -69,7 +76,24 @@ impl CryptoHash for StateKey {
 
     fn hash(&self) -> HashValue {
         let mut state = Self::Hasher::default();
-        state.update(RawStateKey::from(self).bytes.as_ref());
+        state.update(
+            self.encode()
+                .expect("Failed to serialize the state key")
+                .as_ref(),
+        );
         state.finish()
     }
 }
+
+/// Error thrown when a [`StateKey`] fails to be deserialized out of a byte sequence stored in physical
+/// storage, via [`StateKey::decode`].
+#[derive(Debug, Error, Eq, PartialEq)]
+pub enum StateKeyDecodeErr {
+    /// Input is empty.
+    #[error("Missing tag due to empty input")]
+    EmptyInput,
+
+    /// The first byte of the input is not a known tag representing one of the variants.
+    #[error("lead tag byte is unknown: {}", unknown_tag)]
+    UnknownTag { unknown_tag: u8 },
+}