refactor: Receipt handling to be done in VM logic

runtime/near-vm-logic/src/gas_counter.rs

@@ -275,10 +275,17 @@ impl GasCounter {
     pub fn burnt_gas(&self) -> Gas {
         self.fast_counter.burnt_gas
     }
+
+    /// Amount of gas used through promises and amount burned.
     pub fn used_gas(&self) -> Gas {
         self.promises_gas + self.fast_counter.burnt_gas
     }
 
+    /// Remaining gas based on the amount of prepaid gas not yet used.
+    pub fn unused_gas(&self) -> Gas {
+        self.prepaid_gas - self.used_gas()
+    }
+
     pub fn profile_data(&self) -> ProfileData {
         self.profile.clone()
     }

runtime/near-vm-logic/src/lib.rs

@@ -7,6 +7,7 @@ mod dependencies;
 pub mod gas_counter;
 mod logic;
 pub mod mocks;
+pub(crate) mod receipt_manager;
 pub mod serde_with;
 #[cfg(test)]
 mod tests;
@@ -20,6 +21,7 @@ pub use near_primitives_core::config::*;
 pub use near_primitives_core::profile;
 pub use near_primitives_core::types::ProtocolVersion;
 pub use near_vm_errors::{HostError, VMLogicError};
+pub use receipt_manager::ReceiptMetadata;
 pub use types::ReturnData;
 
 pub use gas_counter::with_ext_cost_counter;

runtime/near-vm-logic/src/logic.rs

@@ -1,9 +1,10 @@
 use crate::context::VMContext;
 use crate::dependencies::{External, MemoryLike};
 use crate::gas_counter::{FastGasCounter, GasCounter};
+use crate::receipt_manager::ReceiptManager;
 use crate::types::{PromiseIndex, PromiseResult, ReceiptIndex, ReturnData};
 use crate::utils::split_method_names;
-use crate::ValuePtr;
+use crate::{ReceiptMetadata, ValuePtr};
 use byteorder::ByteOrder;
 use near_crypto::Secp256K1Signature;
 use near_primitives::version::is_implicit_account_creation_enabled;
@@ -60,14 +61,14 @@ pub struct VMLogic<'a> {
 
     /// The DAG of promises, indexed by promise id.
     promises: Vec<Promise>,
-    /// Record the accounts towards which the receipts are directed.
-    receipt_to_account: HashMap<ReceiptIndex, AccountId>,
-
     /// Tracks the total log length. The sum of length of all logs.
     total_log_length: u64,
 
     /// Current protocol version that is used for the function call.
     current_protocol_version: ProtocolVersion,
+
+    /// Handles the receipts generated through execution.
+    receipt_manager: ReceiptManager,
 }
 
 /// Promises API allows to create a DAG-structure that defines dependencies between smart contract
@@ -117,6 +118,7 @@ impl<'a> VMLogic<'a> {
             Some(ViewConfig { max_gas_burnt: max_gas_burnt_view }) => max_gas_burnt_view,
             None => config.limit_config.max_gas_burnt,
         };
+
         let current_account_locked_balance = context.account_locked_balance;
         let gas_counter = GasCounter::new(
             config.ext_costs.clone(),
@@ -140,12 +142,24 @@ impl<'a> VMLogic<'a> {
             logs: vec![],
             registers: HashMap::new(),
             promises: vec![],
-            receipt_to_account: HashMap::new(),
             total_log_length: 0,
             current_protocol_version,
+            receipt_manager: ReceiptManager::default(),
         }
     }
 
+    #[allow(dead_code)]
+    #[cfg(test)]
+    pub(crate) fn receipt_manager(&self) -> &ReceiptManager {
+        &self.receipt_manager
+    }
+
+    #[allow(dead_code)]
+    #[cfg(test)]
+    pub(crate) fn gas_counter(&self) -> &GasCounter {
+        &self.gas_counter
+    }
+
     // ###########################
     // # Memory helper functions #
     // ###########################
@@ -1314,8 +1328,8 @@ impl<'a> VMLogic<'a> {
         let account_id = self.read_and_parse_account_id(account_id_ptr, account_id_len)?;
         let sir = account_id == self.context.current_account_id;
         self.pay_gas_for_new_receipt(sir, &[])?;
-        let new_receipt_idx = self.ext.create_receipt(vec![], account_id.clone())?;
-        self.receipt_to_account.insert(new_receipt_idx, account_id);
+        let new_receipt_idx =
+            self.receipt_manager.create_receipt(self.ext, vec![], account_id.clone())?;
 
         self.checked_push_promise(Promise::Receipt(new_receipt_idx))
     }
@@ -1368,21 +1382,22 @@ impl<'a> VMLogic<'a> {
         let sir = account_id == self.context.current_account_id;
         let deps: Vec<_> = receipt_dependencies
             .iter()
-            .map(|receipt_idx| self.get_account_by_receipt(receipt_idx) == &account_id)
+            .map(|&receipt_idx| self.get_account_by_receipt(receipt_idx) == &account_id)
             .collect();
         self.pay_gas_for_new_receipt(sir, &deps)?;
 
-        let new_receipt_idx = self.ext.create_receipt(receipt_dependencies, account_id.clone())?;
-        self.receipt_to_account.insert(new_receipt_idx, account_id);
+        let new_receipt_idx = self.receipt_manager.create_receipt(
+            self.ext,
+            receipt_dependencies,
+            account_id.clone(),
+        )?;
 
         self.checked_push_promise(Promise::Receipt(new_receipt_idx))
     }
 
     /// Helper function to return the account id towards which the receipt is directed.
-    fn get_account_by_receipt(&self, receipt_idx: &ReceiptIndex) -> &AccountId {
-        self.receipt_to_account
-            .get(receipt_idx)
-            .expect("promises and receipt_to_account should be consistent.")
+    fn get_account_by_receipt(&self, receipt_idx: ReceiptIndex) -> &AccountId {
+        self.receipt_manager.get_receipt_receiver(receipt_idx)
     }
 
     /// Helper function to return the receipt index corresponding to the given promise index.
@@ -1401,7 +1416,7 @@ impl<'a> VMLogic<'a> {
             Promise::NotReceipt(_) => Err(HostError::CannotAppendActionToJointPromise),
         }?;
 
-        let account_id = self.get_account_by_receipt(&receipt_idx);
+        let account_id = self.get_account_by_receipt(receipt_idx);
         let sir = account_id == &self.context.current_account_id;
         Ok((receipt_idx, sir))
     }
@@ -1436,7 +1451,7 @@ impl<'a> VMLogic<'a> {
             ActionCosts::create_account,
         )?;
 
-        self.ext.append_action_create_account(receipt_idx)?;
+        self.receipt_manager.append_action_create_account(receipt_idx)?;
         Ok(())
     }
 
@@ -1494,7 +1509,7 @@ impl<'a> VMLogic<'a> {
             ActionCosts::deploy_contract,
         )?;
 
-        self.ext.append_action_deploy_contract(receipt_idx, code)?;
+        self.receipt_manager.append_action_deploy_contract(receipt_idx, code)?;
         Ok(())
     }
 
@@ -1527,7 +1542,13 @@ impl<'a> VMLogic<'a> {
         gas: Gas,
     ) -> Result<()> {
         let append_action_fn = |vm: &mut Self, receipt_idx, method_name, arguments, amount, gas| {
-            vm.ext.append_action_function_call(receipt_idx, method_name, arguments, amount, gas)
+            vm.receipt_manager.append_action_function_call(
+                receipt_idx,
+                method_name,
+                arguments,
+                amount,
+                gas,
+            )
         };
         self.internal_promise_batch_action_function_call(
             promise_idx,
@@ -1590,7 +1611,7 @@ impl<'a> VMLogic<'a> {
         gas_weight: GasWeight,
     ) -> Result<()> {
         let append_action_fn = |vm: &mut Self, receipt_idx, method_name, arguments, amount, gas| {
-            vm.ext.append_action_function_call_weight(
+            vm.receipt_manager.append_action_function_call_weight(
                 receipt_idx,
                 method_name,
                 arguments,
@@ -1690,7 +1711,7 @@ impl<'a> VMLogic<'a> {
         let amount = self.memory_get_u128(amount_ptr)?;
 
         let (receipt_idx, sir) = self.promise_idx_to_receipt_idx_with_sir(promise_idx)?;
-        let receiver_id = self.get_account_by_receipt(&receipt_idx);
+        let receiver_id = self.get_account_by_receipt(receipt_idx);
         let is_receiver_implicit =
             is_implicit_account_creation_enabled(self.current_protocol_version)
                 && receiver_id.is_implicit();
@@ -1705,7 +1726,7 @@ impl<'a> VMLogic<'a> {
 
         self.deduct_balance(amount)?;
 
-        self.ext.append_action_transfer(receipt_idx, amount)?;
+        self.receipt_manager.append_action_transfer(receipt_idx, amount)?;
         Ok(())
     }
 
@@ -1751,7 +1772,7 @@ impl<'a> VMLogic<'a> {
             ActionCosts::stake,
         )?;
 
-        self.ext.append_action_stake(receipt_idx, amount, public_key)?;
+        self.receipt_manager.append_action_stake(receipt_idx, amount, public_key)?;
         Ok(())
     }
 
@@ -1796,7 +1817,11 @@ impl<'a> VMLogic<'a> {
             ActionCosts::add_key,
         )?;
 
-        self.ext.append_action_add_key_with_full_access(receipt_idx, public_key, nonce)?;
+        self.receipt_manager.append_action_add_key_with_full_access(
+            receipt_idx,
+            public_key,
+            nonce,
+        )?;
         Ok(())
     }
 
@@ -1862,7 +1887,7 @@ impl<'a> VMLogic<'a> {
             ActionCosts::function_call,
         )?;
 
-        self.ext.append_action_add_key_with_function_call(
+        self.receipt_manager.append_action_add_key_with_function_call(
             receipt_idx,
             public_key,
             nonce,
@@ -1913,7 +1938,7 @@ impl<'a> VMLogic<'a> {
             ActionCosts::delete_key,
         )?;
 
-        self.ext.append_action_delete_key(receipt_idx, public_key)?;
+        self.receipt_manager.append_action_delete_key(receipt_idx, public_key)?;
         Ok(())
     }
 
@@ -1956,7 +1981,7 @@ impl<'a> VMLogic<'a> {
             ActionCosts::delete_account,
         )?;
 
-        self.ext.append_action_delete_account(receipt_idx, beneficiary_id)?;
+        self.receipt_manager.append_action_delete_account(receipt_idx, beneficiary_id)?;
         Ok(())
     }
 
@@ -2621,10 +2646,11 @@ impl<'a> VMLogic<'a> {
         #[cfg(feature = "protocol_feature_function_call_weight")]
         if !self.context.is_view() {
             // Distribute unused gas to scheduled function calls
-            let unused_gas = self.context.prepaid_gas - self.gas_counter.used_gas();
+            let unused_gas = self.gas_counter.unused_gas();
 
-            // Distribute the unused gas and prepay for the gas.
-            if matches!(self.ext.distribute_unused_gas(unused_gas), GasDistribution::All) {
+            // Spend all remaining gas by distributing it among function calls that specify
+            // a gas weight
+            if let GasDistribution::All = self.receipt_manager.distribute_unused_gas(unused_gas) {
                 self.gas_counter.prepay_gas(unused_gas).unwrap();
             }
         }
@@ -2643,6 +2669,7 @@ impl<'a> VMLogic<'a> {
             used_gas,
             logs: self.logs,
             profile,
+            action_receipts: self.receipt_manager.action_receipts,
         }
     }
 
@@ -2675,6 +2702,7 @@ pub struct VMOutcome {
     pub logs: Vec<String>,
     /// Data collected from making a contract call
     pub profile: ProfileData,
+    pub action_receipts: Vec<(AccountId, ReceiptMetadata)>,
 }
 
 impl std::fmt::Debug for VMOutcome {