feat: Switch TrieUpdate API to use TrieKey trait (#2278)

Introducing `TrieKey` trait to work with keys within `TrieUpdate` instead of using raw `Vec<u8>`.

This allows to have additional logic for keys. E.g. if we want to hash AccountID for keys (see #2215) and later return all touched accounts for a block (see #2253) we need a way to map `sha256(account_id)` from a raw key back to the `account_id`.

Then `TrieKey` may implement `fn get_account_id_with_hash() -> Option<(AccountId, CryptoHash)>` logic that provides Account ID and hash of this Account ID. With this info we can remember pairs of touched Accounts in memory of `TrieUpdate` and now we can map back hashes here: https://github.com/nearprotocol/nearcore/blob/62f00b211915790d06e822da866fb01ceb6cdc89/core/store/src/trie/mod.rs#L578

Removes iterators from VMLogic (cause they somewhat block changes from this PR).
Fixes #2276 

##Test Plan:
- unit tests and CI

core/primitives/src/utils.rs

@@ -52,9 +52,14 @@ pub mod col {
     pub const DELAYED_RECEIPT: &[u8] = &[8];
 }
 
+pub trait TrieKey: AsRef<[u8]> + Into<Vec<u8>> {}
+
 #[derive(derive_more::AsRef, derive_more::Into)]
+#[as_ref(forward)]
 struct KeyForColumnAccountId(Vec<u8>);
 
+impl TrieKey for KeyForColumnAccountId {}
+
 impl KeyForColumnAccountId {
     pub fn estimate_len(column: &[u8], account_id: &AccountId) -> usize {
         column.len() + account_id.len()
@@ -87,6 +92,8 @@ impl KeyForColumnAccountId {
 #[as_ref(forward)]
 pub struct KeyForAccount(Vec<u8>);
 
+impl TrieKey for KeyForAccount {}
+
 impl KeyForAccount {
     pub fn estimate_len(account_id: &AccountId) -> usize {
         KeyForColumnAccountId::estimate_len(col::ACCOUNT, account_id)
@@ -118,11 +125,18 @@ impl KeyForAccount {
 #[as_ref(forward)]
 pub struct KeyForAccessKey(Vec<u8>);
 
+impl TrieKey for KeyForAccessKey {}
+
 impl KeyForAccessKey {
     fn estimate_prefix_len(account_id: &AccountId) -> usize {
         KeyForColumnAccountId::estimate_len(col::ACCESS_KEY, account_id) + col::ACCESS_KEY.len()
     }
 
+    /// This is not safe and should only be used internally for iterating over access keys for reading.
+    pub fn from_raw_key(key: &[u8]) -> Self {
+        Self(key.to_vec())
+    }
+
     pub fn estimate_len(account_id: &AccountId, public_key: &PublicKey) -> usize {
         let serialized_public_key =
             public_key.try_to_vec().expect("Failed to serialize public key");
@@ -194,6 +208,8 @@ impl KeyForAccessKey {
 #[as_ref(forward)]
 pub struct KeyForData(Vec<u8>);
 
+impl TrieKey for KeyForData {}
+
 impl KeyForData {
     pub fn estimate_len(account_id: &AccountId, data: &[u8]) -> usize {
         KeyForAccount::estimate_len(&account_id) + ACCOUNT_DATA_SEPARATOR.len() + data.len()
@@ -213,13 +229,26 @@ impl KeyForData {
         Self::get_prefix_with_capacity(account_id, 0)
     }
 
+    pub fn with_suffix(&self, suffix: &[u8]) -> Self {
+        let mut raw_key = Vec::with_capacity(self.0.len() + suffix.len());
+        raw_key.extend(&self.0);
+        raw_key.extend(suffix);
+        Self(raw_key)
+    }
+
     pub fn new(account_id: &AccountId, data: &[u8]) -> Self {
         let mut key = Self::get_prefix_with_capacity(&account_id, data.len());
         key.0.extend(data);
         debug_assert_eq!(key.0.len(), Self::estimate_len(&account_id, &data));
         key
     }
 
+    /// Not safe, use only for genesis reads.
+    /// TODO(#2215): Remove once AccountId hashing is implemented.
+    pub fn from_raw_key(key: Vec<u8>) -> Self {
+        Self(key)
+    }
+
     pub fn parse_account_id<K: AsRef<[u8]>>(raw_key: K) -> Result<AccountId, std::io::Error> {
         let account_id_prefix =
             KeyForColumnAccountId::parse_account_id_prefix(col::ACCOUNT, raw_key.as_ref())?;
@@ -266,6 +295,8 @@ impl KeyForData {
 #[as_ref(forward)]
 pub struct KeyForContractCode(Vec<u8>);
 
+impl TrieKey for KeyForContractCode {}
+
 impl KeyForContractCode {
     pub fn new(account_id: &AccountId) -> Self {
         Self(KeyForColumnAccountId::with_capacity(col::CODE, account_id, 0).into())
@@ -287,6 +318,8 @@ impl KeyForContractCode {
 #[as_ref(forward)]
 pub struct KeyForReceivedData(Vec<u8>);
 
+impl TrieKey for KeyForReceivedData {}
+
 impl KeyForReceivedData {
     pub fn new(account_id: &AccountId, data_id: &CryptoHash) -> Self {
         let mut key: Vec<u8> = KeyForColumnAccountId::with_capacity(
@@ -305,6 +338,8 @@ impl KeyForReceivedData {
 #[as_ref(forward)]
 pub struct KeyForPostponedReceiptId(Vec<u8>);
 
+impl TrieKey for KeyForPostponedReceiptId {}
+
 impl KeyForPostponedReceiptId {
     pub fn new(account_id: &AccountId, data_id: &CryptoHash) -> Self {
         let mut key: Vec<u8> = KeyForColumnAccountId::with_capacity(
@@ -323,6 +358,8 @@ impl KeyForPostponedReceiptId {
 #[as_ref(forward)]
 pub struct KeyForPendingDataCount(Vec<u8>);
 
+impl TrieKey for KeyForPendingDataCount {}
+
 impl KeyForPendingDataCount {
     pub fn new(account_id: &AccountId, receipt_id: &CryptoHash) -> Self {
         let mut key: Vec<u8> = KeyForColumnAccountId::with_capacity(
@@ -341,6 +378,8 @@ impl KeyForPendingDataCount {
 #[as_ref(forward)]
 pub struct KeyForPostponedReceipt(Vec<u8>);
 
+impl TrieKey for KeyForPostponedReceipt {}
+
 impl KeyForPostponedReceipt {
     pub fn new(account_id: &AccountId, receipt_id: &CryptoHash) -> Self {
         let mut key: Vec<u8> = KeyForColumnAccountId::with_capacity(
@@ -359,6 +398,8 @@ impl KeyForPostponedReceipt {
 #[as_ref(forward)]
 pub struct KeyForDelayedReceipt(Vec<u8>);
 
+impl TrieKey for KeyForDelayedReceipt {}
+
 impl KeyForDelayedReceipt {
     pub fn new(index: u64) -> Self {
         let index_bytes = index.to_le_bytes();
@@ -369,6 +410,18 @@ impl KeyForDelayedReceipt {
     }
 }
 
+#[derive(derive_more::AsRef, derive_more::Into)]
+#[as_ref(forward)]
+pub struct KeyForDelayedReceiptIndices(Vec<u8>);
+
+impl TrieKey for KeyForDelayedReceiptIndices {}
+
+impl KeyForDelayedReceiptIndices {
+    pub fn new() -> Self {
+        Self(col::DELAYED_RECEIPT_INDICES.to_vec())
+    }
+}
+
 pub fn create_nonce_with_nonce(base: &CryptoHash, salt: u64) -> CryptoHash {
     let mut nonce: Vec<u8> = base.as_ref().to_owned();
     nonce.extend(index_to_bytes(salt));

core/store/src/lib.rs

@@ -19,7 +19,7 @@ use near_primitives::serialize::to_base;
 use near_primitives::types::{AccountId, StorageUsage};
 use near_primitives::utils::{
     KeyForAccessKey, KeyForAccount, KeyForContractCode, KeyForData, KeyForPostponedReceipt,
-    KeyForReceivedData,
+    KeyForReceivedData, TrieKey,
 };
 
 use crate::db::{DBOp, DBTransaction, Database, RocksDB};
@@ -231,11 +231,11 @@ pub fn create_store(path: &str) -> Arc<Store> {
 /// Reads an object from Trie.
 /// # Errors
 /// see StorageError
-pub fn get<T: BorshDeserialize, K: AsRef<[u8]>>(
+pub fn get<T: BorshDeserialize, K: TrieKey>(
     state_update: &TrieUpdate,
-    key: K,
+    key: &K,
 ) -> Result<Option<T>, StorageError> {
-    state_update.get(key.as_ref()).and_then(|opt| {
+    state_update.get(key).and_then(|opt| {
         opt.map_or_else(
             || Ok(None),
             |data| {
@@ -250,9 +250,9 @@ pub fn get<T: BorshDeserialize, K: AsRef<[u8]>>(
 }
 
 /// Writes an object into Trie.
-pub fn set<T: BorshSerialize, K: Into<Vec<u8>>>(state_update: &mut TrieUpdate, key: K, value: &T) {
+pub fn set<T: BorshSerialize, K: TrieKey>(state_update: &mut TrieUpdate, key: K, value: &T) {
     let data = value.try_to_vec().expect("Borsh serializer is not expected to ever fail");
-    state_update.set(key.into(), data);
+    state_update.set(key, data);
 }
 
 /// Number of bytes account and all of it's other data occupies in the storage.
@@ -322,19 +322,19 @@ pub fn get_access_key_raw(
     state_update: &TrieUpdate,
     key: &[u8],
 ) -> Result<Option<AccessKey>, StorageError> {
-    get(state_update, key)
+    get(state_update, &KeyForAccessKey::from_raw_key(&key))
 }
 
 pub fn set_code(state_update: &mut TrieUpdate, account_id: &AccountId, code: &ContractCode) {
-    state_update.set(KeyForContractCode::new(account_id).into(), code.code.clone());
+    state_update.set(KeyForContractCode::new(account_id), code.code.clone());
 }
 
 pub fn get_code(
     state_update: &TrieUpdate,
     account_id: &AccountId,
 ) -> Result<Option<ContractCode>, StorageError> {
     state_update
-        .get(KeyForContractCode::new(account_id))
+        .get(&KeyForContractCode::new(account_id))
         .map(|opt| opt.map(|code| ContractCode::new(code.to_vec())))
 }
 
@@ -345,7 +345,7 @@ pub fn remove_account(
 ) -> Result<(), StorageError> {
     state_update.remove(KeyForAccount::new(account_id));
     state_update.remove(KeyForContractCode::new(account_id));
-    state_update.remove_starts_with(KeyForAccessKey::get_prefix(account_id))?;
-    state_update.remove_starts_with(KeyForData::get_prefix(account_id))?;
+    state_update.remove_starts_with(&KeyForAccessKey::get_prefix(account_id))?;
+    state_update.remove_starts_with(&KeyForData::get_prefix(account_id))?;
     Ok(())
 }

core/store/src/trie/trie_tests.rs

@@ -5,6 +5,7 @@ use crate::trie::POISONED_LOCK_ERR;
 use crate::{PartialStorage, Trie, TrieUpdate};
 use near_primitives::errors::StorageError;
 use near_primitives::hash::{hash, CryptoHash};
+use near_primitives::utils::TrieKey;
 use rand::seq::SliceRandom;
 use rand::Rng;
 use std::collections::{HashMap, HashSet};
@@ -83,6 +84,22 @@ where
     println!("Success");
 }
 
+struct TestKey(pub Vec<u8>);
+
+impl AsRef<[u8]> for TestKey {
+    fn as_ref(&self) -> &[u8] {
+        self.0.as_ref()
+    }
+}
+
+impl Into<Vec<u8>> for TestKey {
+    fn into(self) -> Vec<u8> {
+        self.0
+    }
+}
+
+impl TrieKey for TestKey {}
+
 #[test]
 fn test_reads_with_incomplete_storage() {
     let mut rng = rand::thread_rng();
@@ -125,7 +142,7 @@ fn test_reads_with_incomplete_storage() {
             let trie_update_keys = |trie: Arc<Trie>| -> Result<_, StorageError> {
                 let trie_update = TrieUpdate::new(trie, state_root);
                 let mut keys = Vec::new();
-                trie_update.for_keys_with_prefix(key_prefix, |key| {
+                trie_update.for_keys_with_prefix(&TestKey(key_prefix.to_vec()), |key| {
                     keys.push(key.to_vec());
                 })?;
                 Ok(keys)