Recording partial trie for storage proofs

Record every read of a trie node from storage to use
the recorded values for a proof of execution.

This includes all key paths for keys requested for reads, but also
in case of updates can include additional nodes which are
necessary to provide a canonical resulting trie (i.e. no branch nodes
with <= 1 branch and extension is always followed by a branch).

chain/chain/src/test_utils.rs

@@ -14,7 +14,7 @@ use near_primitives::transaction::{
 };
 use near_primitives::types::{AccountId, BlockIndex, MerkleHash, ShardId, ValidatorStake};
 use near_store::test_utils::create_test_store;
-use near_store::{Store, StoreUpdate, Trie, TrieChanges, WrappedTrieChanges};
+use near_store::{PartialStorage, Store, StoreUpdate, Trie, TrieChanges, WrappedTrieChanges};
 
 use crate::error::{Error, ErrorKind};
 use crate::types::{BlockHeader, ReceiptResult, RuntimeAdapter, Weight};
@@ -152,7 +152,7 @@ impl RuntimeAdapter for KeyValueRuntime {
         Ok((parent_hash, 0))
     }
 
-    fn apply_transactions(
+    fn apply_transactions_with_optional_storage_proof(
         &self,
         _shard_id: ShardId,
         state_root: &MerkleHash,
@@ -161,10 +161,19 @@ impl RuntimeAdapter for KeyValueRuntime {
         _block_hash: &CryptoHash,
         _receipts: &Vec<Vec<Receipt>>,
         transactions: &Vec<SignedTransaction>,
+        generate_storage_proof: bool,
     ) -> Result<
-        (WrappedTrieChanges, MerkleHash, Vec<TransactionLog>, ReceiptResult, Vec<ValidatorStake>),
+        (
+            WrappedTrieChanges,
+            MerkleHash,
+            Vec<TransactionLog>,
+            ReceiptResult,
+            Vec<ValidatorStake>,
+            Option<PartialStorage>,
+        ),
         Box<dyn std::error::Error>,
     > {
+        assert!(!generate_storage_proof);
         let mut tx_results = vec![];
         for tx in transactions {
             tx_results.push(TransactionLog {
@@ -183,6 +192,7 @@ impl RuntimeAdapter for KeyValueRuntime {
             tx_results,
             HashMap::default(),
             vec![],
+            None,
         ))
     }
 

chain/chain/src/types.rs

@@ -8,7 +8,7 @@ use near_primitives::receipt::Receipt;
 use near_primitives::rpc::QueryResponse;
 use near_primitives::transaction::{SignedTransaction, TransactionLog};
 use near_primitives::types::{AccountId, BlockIndex, MerkleHash, ShardId, ValidatorStake};
-use near_store::{StoreUpdate, WrappedTrieChanges};
+use near_store::{PartialStorage, StoreUpdate, WrappedTrieChanges};
 
 use crate::error::Error;
 
@@ -126,7 +126,7 @@ pub trait RuntimeAdapter: Send + Sync {
     fn apply_transactions(
         &self,
         shard_id: ShardId,
-        merkle_hash: &MerkleHash,
+        state_root: &MerkleHash,
         block_index: BlockIndex,
         prev_block_hash: &CryptoHash,
         block_hash: &CryptoHash,
@@ -135,6 +135,49 @@ pub trait RuntimeAdapter: Send + Sync {
     ) -> Result<
         (WrappedTrieChanges, MerkleHash, Vec<TransactionLog>, ReceiptResult, Vec<ValidatorStake>),
         Box<dyn std::error::Error>,
+    > {
+        self.apply_transactions_with_optional_storage_proof(
+            shard_id,
+            state_root,
+            block_index,
+            prev_block_hash,
+            block_hash,
+            receipts,
+            transactions,
+            false,
+        )
+        .map(
+            |(
+                trie_changes,
+                root,
+                tx_results,
+                receipt_result,
+                validator_proposals,
+                _partial_storage,
+            )| (trie_changes, root, tx_results, receipt_result, validator_proposals),
+        )
+    }
+
+    fn apply_transactions_with_optional_storage_proof(
+        &self,
+        shard_id: ShardId,
+        state_root: &MerkleHash,
+        block_index: BlockIndex,
+        prev_block_hash: &CryptoHash,
+        block_hash: &CryptoHash,
+        receipts: &Vec<Vec<Receipt>>,
+        transactions: &Vec<SignedTransaction>,
+        generate_storage_proof: bool,
+    ) -> Result<
+        (
+            WrappedTrieChanges,
+            MerkleHash,
+            Vec<TransactionLog>,
+            ReceiptResult,
+            Vec<ValidatorStake>,
+            Option<PartialStorage>,
+        ),
+        Box<dyn std::error::Error>,
     >;
 
     /// Query runtime with given `path` and `data`.

core/store/src/lib.rs

@@ -22,8 +22,8 @@ use near_protos::access_key as access_key_proto;
 use near_protos::account as account_proto;
 
 pub use crate::trie::{
-    update::TrieUpdate, update::TrieUpdateIterator, Trie, TrieChanges, TrieIterator,
-    WrappedTrieChanges,
+    update::TrieUpdate, update::TrieUpdateIterator, PartialStorage, Trie, TrieChanges,
+    TrieIterator, WrappedTrieChanges,
 };
 
 pub mod test_utils;