store: add methods to check for presence of a value without reading it out

core/store/src/genesis/state_applier.rs

@@ -1,6 +1,6 @@
 use crate::flat::FlatStateChanges;
 use crate::{
-    get_account, get_received_data, set, set_access_key, set_account, set_code,
+    get_account, has_received_data, set, set_access_key, set_account, set_code,
     set_delayed_receipt, set_postponed_receipt, set_received_data, ShardTries, TrieUpdate,
 };
 
@@ -277,9 +277,8 @@ impl GenesisStateApplier {
             let mut pending_data_count: u32 = 0;
             for data_id in &action_receipt.input_data_ids {
                 storage.modify(|state_update| {
-                    if get_received_data(state_update, account_id, *data_id)
+                    if !has_received_data(state_update, account_id, *data_id)
                         .expect("Genesis storage error")
-                        .is_none()
                     {
                         pending_data_count += 1;
                         set(

core/store/src/lib.rs

@@ -703,6 +703,14 @@ pub fn get_received_data(
     get(trie, &TrieKey::ReceivedData { receiver_id: receiver_id.clone(), data_id })
 }
 
+pub fn has_received_data(
+    trie: &dyn TrieAccess,
+    receiver_id: &AccountId,
+    data_id: CryptoHash,
+) -> Result<bool, StorageError> {
+    trie.contains_key(&TrieKey::ReceivedData { receiver_id: receiver_id.clone(), data_id })
+}
+
 pub fn set_postponed_receipt(state_update: &mut TrieUpdate, receipt: &Receipt) {
     let key = TrieKey::PostponedReceipt {
         receiver_id: receipt.receiver_id.clone(),

core/store/src/trie/mod.rs

@@ -365,6 +365,11 @@ pub trait TrieAccess {
     /// root are already known by the object rather than being passed as
     /// argument.
     fn get(&self, key: &TrieKey) -> Result<Option<Vec<u8>>, StorageError>;
+
+    /// Check if the key is present.
+    ///
+    /// Equivalent to `Self::get(k)?.is_some()`, but avoids reading out the value.
+    fn contains_key(&self, key: &TrieKey) -> Result<bool, StorageError>;
 }
 
 /// Stores reference count addition for some key-value pair in DB.
@@ -1360,6 +1365,47 @@ impl Trie {
         Ok(bytes.to_vec())
     }
 
+    /// Check if the column contains a value with the given `key`.
+    ///
+    /// This method is guaranteed to not inspect the value stored for this key, which would
+    /// otherwise have potential gas cost implications.
+    pub fn contains_key(&self, key: &[u8]) -> Result<bool, StorageError> {
+        self.contains_key_mode(key, KeyLookupMode::FlatStorage)
+    }
+
+    /// Check if the column contains a value with the given `key`.
+    ///
+    /// This method is guaranteed to not inspect the value stored for this key, which would
+    /// otherwise have potential gas cost implications.
+    pub fn contains_key_mode(&self, key: &[u8], mode: KeyLookupMode) -> Result<bool, StorageError> {
+        let charge_gas_for_trie_node_access =
+            mode == KeyLookupMode::Trie || self.charge_gas_for_trie_node_access;
+        if self.memtries.is_some() {
+            return Ok(self
+                .lookup_from_memory(key, charge_gas_for_trie_node_access, |_| ())?
+                .is_some());
+        }
+
+        'flat: {
+            let KeyLookupMode::FlatStorage = mode else { break 'flat };
+            let Some(flat_storage_chunk_view) = &self.flat_storage_chunk_view else { break 'flat };
+            let value = flat_storage_chunk_view.contains_key(key)?;
+            if self.recorder.is_some() {
+                // If recording, we need to look up in the trie as well to record the trie nodes,
+                // as they are needed to prove the value. Also, it's important that this lookup
+                // is done even if the key was not found, because intermediate trie nodes may be
+                // needed to prove the non-existence of the key.
+                let value_ref_from_trie =
+                    self.lookup_from_state_column(NibbleSlice::new(key), false)?;
+                debug_assert_eq!(&value_ref_from_trie.is_some(), &value);
+            }
+        }
+
+        Ok(self
+            .lookup_from_state_column(NibbleSlice::new(key), charge_gas_for_trie_node_access)?
+            .is_some())
+    }
+
     /// Retrieves an `OptimizedValueRef`` for the given key. See `OptimizedValueRef`.
     ///
     /// `mode`: whether we will try to perform the lookup through flat storage or trie.
@@ -1493,6 +1539,10 @@ impl TrieAccess for Trie {
     fn get(&self, key: &TrieKey) -> Result<Option<Vec<u8>>, StorageError> {
         Trie::get(self, &key.to_vec())
     }
+
+    fn contains_key(&self, key: &TrieKey) -> Result<bool, StorageError> {
+        Trie::contains_key(&self, &key.to_vec())
+    }
 }
 
 /// Counts trie nodes reads during tx/receipt execution for proper storage costs charging.
@@ -1719,6 +1769,41 @@ mod tests {
         }
     }
 
+    #[test]
+    fn test_contains_key() {
+        let sid = ShardUId::single_shard();
+        let bid = CryptoHash::default();
+        let tries = TestTriesBuilder::new().with_flat_storage().build();
+        let initial = vec![
+            (vec![99, 44, 100, 58, 58, 49], Some(vec![1])),
+            (vec![99, 44, 100, 58, 58, 50], Some(vec![1])),
+            (vec![99, 44, 100, 58, 58, 50, 51], Some(vec![1])),
+        ];
+        let root = test_populate_trie(&tries, &Trie::EMPTY_ROOT, sid, initial);
+        let trie = tries.get_trie_with_block_hash_for_shard(sid, root, &bid, false);
+        assert!(trie.has_flat_storage_chunk_view());
+        assert!(trie.contains_key_mode(&[99, 44, 100, 58, 58, 49], KeyLookupMode::Trie).unwrap());
+        assert!(trie
+            .contains_key_mode(&[99, 44, 100, 58, 58, 49], KeyLookupMode::FlatStorage)
+            .unwrap());
+        assert!(!trie.contains_key_mode(&[99, 44, 100, 58, 58, 48], KeyLookupMode::Trie).unwrap());
+        assert!(!trie
+            .contains_key_mode(&[99, 44, 100, 58, 58, 48], KeyLookupMode::FlatStorage)
+            .unwrap());
+        let changes = vec![(vec![99, 44, 100, 58, 58, 49], None)];
+        let root = test_populate_trie(&tries, &root, sid, changes);
+        let trie = tries.get_trie_with_block_hash_for_shard(sid, root, &bid, false);
+        assert!(trie.has_flat_storage_chunk_view());
+        assert!(trie.contains_key_mode(&[99, 44, 100, 58, 58, 50], KeyLookupMode::Trie).unwrap());
+        assert!(trie
+            .contains_key_mode(&[99, 44, 100, 58, 58, 50], KeyLookupMode::FlatStorage)
+            .unwrap());
+        assert!(!trie
+            .contains_key_mode(&[99, 44, 100, 58, 58, 49], KeyLookupMode::FlatStorage)
+            .unwrap());
+        assert!(!trie.contains_key_mode(&[99, 44, 100, 58, 58, 49], KeyLookupMode::Trie).unwrap());
+    }
+
     #[test]
     fn test_equal_leafs() {
         let initial = vec![

core/store/src/trie/update.rs

@@ -80,6 +80,18 @@ impl TrieUpdate {
         Ok(result)
     }
 
+    pub fn contains_key(&self, key: &TrieKey) -> Result<bool, StorageError> {
+        let key = key.to_vec();
+        if self.prospective.contains_key(&key) {
+            return Ok(true);
+        } else if let Some(changes_with_trie_key) = self.committed.get(&key) {
+            if let Some(RawStateChange { data, .. }) = changes_with_trie_key.changes.last() {
+                return Ok(data.is_some());
+            }
+        }
+        self.trie.contains_key(&key)
+    }
+
     pub fn get(&self, key: &TrieKey) -> Result<Option<Vec<u8>>, StorageError> {
         let key = key.to_vec();
         if let Some(key_value) = self.prospective.get(&key) {
@@ -163,6 +175,10 @@ impl crate::TrieAccess for TrieUpdate {
     fn get(&self, key: &TrieKey) -> Result<Option<Vec<u8>>, StorageError> {
         TrieUpdate::get(self, key)
     }
+
+    fn contains_key(&self, key: &TrieKey) -> Result<bool, StorageError> {
+        TrieUpdate::contains_key(&self, key)
+    }
 }
 
 #[cfg(test)]

runtime/runtime/src/lib.rs

@@ -37,9 +37,9 @@ use near_primitives::utils::{
 };
 use near_primitives::version::{ProtocolFeature, ProtocolVersion};
 use near_store::{
-    get, get_account, get_postponed_receipt, get_received_data, remove_postponed_receipt, set,
-    set_account, set_delayed_receipt, set_postponed_receipt, set_received_data, PartialStorage,
-    StorageError, Trie, TrieChanges, TrieUpdate,
+    get, get_account, get_postponed_receipt, get_received_data, has_received_data,
+    remove_postponed_receipt, set, set_account, set_delayed_receipt, set_postponed_receipt,
+    set_received_data, PartialStorage, StorageError, Trie, TrieChanges, TrieUpdate,
 };
 use near_store::{set_access_key, set_code};
 use near_vm_runner::logic::types::PromiseResult;
@@ -967,7 +967,7 @@ impl Runtime {
                 // If not, then we will postpone this receipt for later.
                 let mut pending_data_count: u32 = 0;
                 for data_id in &action_receipt.input_data_ids {
-                    if get_received_data(state_update, account_id, *data_id)?.is_none() {
+                    if !has_received_data(state_update, account_id, *data_id)? {
                         pending_data_count += 1;
                         // The data for a given data_id is not available, so we save a link to this
                         // receipt_id for the pending data_id into the state.