[gas] allow natives to read the gas balance

language/move-vm/runtime/src/interpreter.rs

@@ -22,6 +22,7 @@ use move_vm_types::{
     data_store::DataStore,
     gas::{GasMeter, SimpleInstruction},
     loaded_data::runtime_types::Type,
+    natives::function::NativeResult,
     values::{
         self, GlobalValue, IntegerValue, Locals, Reference, Struct, StructRef, VMValueCast, Value,
         Vector, VectorRef,
@@ -382,7 +383,13 @@ impl Interpreter {
             }
         }
 
-        let mut native_context = NativeContext::new(self, data_store, resolver, extensions);
+        let mut native_context = NativeContext::new(
+            self,
+            data_store,
+            resolver,
+            extensions,
+            gas_meter.balance_internal(),
+        );
         let native_function = function.get_native()?;
 
         gas_meter.charge_native_function_before_execution(
@@ -397,17 +404,27 @@ impl Interpreter {
 
         // Note(Gas): The order by which gas is charged / error gets returned MUST NOT be modified
         //            here or otherwise it becomes an incompatible change!!!
-        let return_values = match result.result {
-            Ok(vals) => {
-                gas_meter.charge_native_function(result.cost, Some(vals.iter()))?;
-                vals
-            }
-            Err(code) => {
-                gas_meter.charge_native_function(
-                    result.cost,
+        let return_values = match result {
+            NativeResult::Success { cost, ret_vals } => {
+                gas_meter.charge_native_function(cost, Some(ret_vals.iter()))?;
+                ret_vals
+            }
+            NativeResult::Abort { cost, abort_code } => {
+                gas_meter.charge_native_function(cost, Option::<std::iter::Empty<&Value>>::None)?;
+                return Err(PartialVMError::new(StatusCode::ABORTED).with_sub_status(abort_code));
+            }
+            NativeResult::OutOfGas { partial_cost } => {
+                let err = match gas_meter.charge_native_function(
+                    partial_cost,
                     Option::<std::iter::Empty<&Value>>::None,
-                )?;
-                return Err(PartialVMError::new(StatusCode::ABORTED).with_sub_status(code));
+                ) {
+                    Err(err) if err.major_status() == StatusCode::OUT_OF_GAS => err,
+                    Ok(_) | Err(_) => PartialVMError::new(StatusCode::UNKNOWN_INVARIANT_VIOLATION_ERROR).with_message(
+                        "The partial cost returned by the native function did not cause the gas meter to trigger an OutOfGas error, at least one of them is violating the contract".to_string()
+                    ),
+                };
+
+                return Err(err);
             }
         };
 

language/move-vm/runtime/src/native_functions.rs

@@ -8,6 +8,7 @@ use crate::{
 use move_binary_format::errors::{ExecutionState, PartialVMError, PartialVMResult};
 use move_core_types::{
     account_address::AccountAddress,
+    gas_algebra::InternalGas,
     identifier::Identifier,
     language_storage::TypeTag,
     value::MoveTypeLayout,
@@ -94,6 +95,7 @@ pub struct NativeContext<'a, 'b> {
     data_store: &'a mut dyn DataStore,
     resolver: &'a Resolver<'a>,
     extensions: &'a mut NativeContextExtensions<'b>,
+    gas_balance: InternalGas,
 }
 
 impl<'a, 'b> NativeContext<'a, 'b> {
@@ -102,12 +104,14 @@ impl<'a, 'b> NativeContext<'a, 'b> {
         data_store: &'a mut dyn DataStore,
         resolver: &'a Resolver<'a>,
         extensions: &'a mut NativeContextExtensions<'b>,
+        gas_balance: InternalGas,
     ) -> Self {
         Self {
             interpreter,
             data_store,
             resolver,
             extensions,
+            gas_balance,
         }
     }
 }
@@ -172,4 +176,8 @@ impl<'a, 'b> NativeContext<'a, 'b> {
     pub fn stack_frames(&self, count: usize) -> ExecutionState {
         self.interpreter.get_stack_frames(count)
     }
+
+    pub fn gas_balance(&self) -> InternalGas {
+        self.gas_balance
+    }
 }

language/move-vm/test-utils/src/gas_schedule.rs

@@ -261,6 +261,10 @@ fn get_simple_instruction_opcode(instr: SimpleInstruction) -> Opcodes {
 }
 
 impl<'b> GasMeter for GasStatus<'b> {
+    fn balance_internal(&self) -> InternalGas {
+        self.gas_left
+    }
+
     /// Charge an instruction and fail if not enough gas units are left.
     fn charge_simple_instr(&mut self, instr: SimpleInstruction) -> PartialVMResult<()> {
         self.charge_instr(get_simple_instruction_opcode(instr))

language/move-vm/types/src/gas.rs

@@ -70,6 +70,8 @@ pub enum SimpleInstruction {
 /// Trait that defines a generic gas meter interface, allowing clients of the Move VM to implement
 /// their own metering scheme.
 pub trait GasMeter {
+    fn balance_internal(&self) -> InternalGas;
+
     /// Charge an instruction and fail if not enough gas units are left.
     fn charge_simple_instr(&mut self, instr: SimpleInstruction) -> PartialVMResult<()>;
 
@@ -237,6 +239,10 @@ pub trait GasMeter {
 pub struct UnmeteredGasMeter;
 
 impl GasMeter for UnmeteredGasMeter {
+    fn balance_internal(&self) -> InternalGas {
+        u64::MAX.into()
+    }
+
     fn charge_simple_instr(&mut self, _instr: SimpleInstruction) -> PartialVMResult<()> {
         Ok(())
     }

language/move-vm/types/src/natives/function.rs

@@ -32,18 +32,26 @@ pub use move_core_types::{gas_algebra::InternalGas, vm_status::StatusCode};
 /// is a VM invariant violation which should have been forbidden by the verifier.
 /// Errors (typically user errors and aborts) that are logically part of the function execution
 /// must be expressed in a `NativeResult` with a cost and a VMStatus.
-pub struct NativeResult {
-    /// Result of execution. This is either the return values or the error to report.
-    pub cost: InternalGas,
-    pub result: Result<SmallVec<[Value; 1]>, u64>,
+pub enum NativeResult {
+    Success {
+        cost: InternalGas,
+        ret_vals: SmallVec<[Value; 1]>,
+    },
+    Abort {
+        cost: InternalGas,
+        abort_code: u64,
+    },
+    OutOfGas {
+        partial_cost: InternalGas,
+    },
 }
 
 impl NativeResult {
     /// Return values of a successful execution.
     pub fn ok(cost: InternalGas, values: SmallVec<[Value; 1]>) -> Self {
-        NativeResult {
+        NativeResult::Success {
             cost,
-            result: Ok(values),
+            ret_vals: values,
         }
     }
 
@@ -52,10 +60,20 @@ impl NativeResult {
     /// The only thing the funciton can specify is its abort code, as if it had invoked the `Abort`
     /// bytecode instruction
     pub fn err(cost: InternalGas, abort_code: u64) -> Self {
-        NativeResult {
-            cost,
-            result: Err(abort_code),
-        }
+        NativeResult::Abort { cost, abort_code }
+    }
+
+    /// A special variant indicating that the native has determined there is not enough
+    /// balance to cover the full cost to get all the work done.
+    ///
+    /// Along with the ability to get the gas balance from the native context, this offers
+    /// natives a way to emulate incremental gas metering, avoiding doing expensive operations
+    /// before charging for gas.
+    ///
+    /// The natives are still required to return a partial cost, which the VM will pass
+    /// to the gas meter for proper bookkeeping.
+    pub fn out_of_gas(partial_cost: InternalGas) -> Self {
+        NativeResult::OutOfGas { partial_cost }
     }
 
     /// Convert a PartialVMResult<()> into a PartialVMResult<NativeResult>