chore(blockifier): cap max_validate_gas and max_execute_gas

crates/blockifier/src/context.rs

@@ -5,13 +5,15 @@ use papyrus_config::{ParamPath, ParamPrivacyInput, SerializedParam};
 use serde::{Deserialize, Serialize};
 use starknet_api::block::{BlockInfo, FeeType, GasPriceVector};
 use starknet_api::core::{ChainId, ContractAddress};
+use starknet_api::execution_resources::GasAmount;
 use starknet_api::transaction::fields::{
     AllResourceBounds,
     GasVectorComputationMode,
     ValidResourceBounds,
 };
 
 use crate::bouncer::BouncerConfig;
+use crate::execution::call_info::CallInfo;
 use crate::transaction::objects::{
     CurrentTransactionInfo,
     HasRelatedFeeType,
@@ -49,7 +51,7 @@ impl TransactionContext {
             | TransactionInfo::Current(CurrentTransactionInfo {
                 resource_bounds: ValidResourceBounds::L1Gas(_),
                 ..
-            }) => self.block_context.versioned_constants.default_initial_gas_cost(),
+            }) => self.block_context.versioned_constants.default_initial_gas_amount(),
             TransactionInfo::Current(CurrentTransactionInfo {
                 resource_bounds: ValidResourceBounds::AllResources(AllResourceBounds { l2_gas, .. }),
                 ..
@@ -58,6 +60,33 @@ impl TransactionContext {
     }
 }
 
+pub(crate) struct GasCounter {
+    pub(crate) spent_gas: GasAmount,
+    pub(crate) remaining_gas: GasAmount,
+}
+
+impl GasCounter {
+    pub(crate) fn new(initial_gas: u64) -> Self {
+        GasCounter { spent_gas: GasAmount(0), remaining_gas: GasAmount(initial_gas) }
+    }
+
+    pub(crate) fn spend(&mut self, amount: GasAmount) {
+        self.spent_gas = self.spent_gas.checked_add(amount).expect("Gas overflow");
+        self.remaining_gas = self
+            .remaining_gas
+            .checked_sub(amount)
+            .expect("Overuse of gas; should have been caught earlier");
+    }
+
+    pub(crate) fn cap_usage(&self, amount: GasAmount) -> u64 {
+        self.remaining_gas.min(amount).0
+    }
+
+    pub(crate) fn subtract_used_gas(&mut self, call_info: &CallInfo) {
+        self.spend(GasAmount(call_info.execution.gas_consumed));
+    }
+}
+
 #[derive(Clone, Debug)]
 pub struct BlockContext {
     // TODO(Yoni, 1/10/2024): consider making these fields public.

crates/blockifier/src/transaction/account_transaction.rs

@@ -21,7 +21,7 @@ use starknet_api::transaction::fields::{
 use starknet_api::transaction::{constants, TransactionHash, TransactionVersion};
 use starknet_types_core::felt::Felt;
 
-use crate::context::{BlockContext, TransactionContext};
+use crate::context::{BlockContext, GasCounter, TransactionContext};
 use crate::execution::call_info::CallInfo;
 use crate::execution::contract_class::RunnableCompiledClass;
 use crate::execution::entry_point::{CallEntryPoint, CallType, EntryPointExecutionContext};
@@ -347,12 +347,28 @@ impl AccountTransaction {
         &self,
         state: &mut dyn State,
         tx_context: Arc<TransactionContext>,
-        remaining_gas: &mut u64,
+        remaining_gas: &mut GasCounter,
         validate: bool,
         limit_steps_by_resources: bool,
     ) -> TransactionExecutionResult<Option<CallInfo>> {
         if validate {
-            self.validate_tx(state, tx_context, remaining_gas, limit_steps_by_resources)
+            // TODO(Aner): cap the gas for validation.
+            let remaining_validation_gas = &mut remaining_gas
+                .cap_usage(tx_context.block_context.versioned_constants.validate_max_sierra_gas);
+            let call_info = self.validate_tx(
+                state,
+                tx_context,
+                remaining_validation_gas,
+                limit_steps_by_resources,
+            )?;
+            match call_info {
+                // TODO(Aner): Update the gas counter.
+                Some(call_info) => {
+                    remaining_gas.subtract_used_gas(&call_info);
+                    Ok(Some(call_info))
+                }
+                None => Ok(None),
+            }
         } else {
             Ok(None)
         }
@@ -479,20 +495,34 @@ impl AccountTransaction {
         &self,
         state: &mut S,
         context: &mut EntryPointExecutionContext,
-        remaining_gas: &mut u64,
+        remaining_gas: &mut GasCounter,
     ) -> TransactionExecutionResult<Option<CallInfo>> {
-        match &self.tx {
-            Transaction::Declare(tx) => tx.run_execute(state, context, remaining_gas),
-            Transaction::DeployAccount(tx) => tx.run_execute(state, context, remaining_gas),
-            Transaction::Invoke(tx) => tx.run_execute(state, context, remaining_gas),
+        // TODO(Aner): cap the gas usage for execution.
+        let remaining_execution_gas = &mut remaining_gas
+            .cap_usage(context.tx_context.block_context.versioned_constants.execute_max_sierra_gas);
+
+        let call_info = match &self.tx {
+            Transaction::Declare(tx) => tx.run_execute(state, context, remaining_execution_gas),
+            Transaction::DeployAccount(tx) => {
+                tx.run_execute(state, context, remaining_execution_gas)
+            }
+            Transaction::Invoke(tx) => tx.run_execute(state, context, remaining_execution_gas),
+        }?;
+        match call_info {
+            // TODO(Aner): Update the gas counter.
+            Some(call_info) => {
+                remaining_gas.subtract_used_gas(&call_info);
+                Ok(Some(call_info))
+            }
+            None => Ok(None),
         }
     }
 
     fn run_non_revertible<S: StateReader>(
         &self,
         state: &mut TransactionalState<'_, S>,
         tx_context: Arc<TransactionContext>,
-        remaining_gas: &mut u64,
+        remaining_gas: &mut GasCounter,
         validate: bool,
         charge_fee: bool,
     ) -> TransactionExecutionResult<ValidateExecuteCallInfo> {
@@ -552,7 +582,7 @@ impl AccountTransaction {
         &self,
         state: &mut TransactionalState<'_, S>,
         tx_context: Arc<TransactionContext>,
-        remaining_gas: &mut u64,
+        remaining_gas: &mut GasCounter,
         validate: bool,
         charge_fee: bool,
     ) -> TransactionExecutionResult<ValidateExecuteCallInfo> {
@@ -690,15 +720,14 @@ impl AccountTransaction {
     fn run_or_revert<S: StateReader>(
         &self,
         state: &mut TransactionalState<'_, S>,
-        remaining_gas: &mut u64,
+        remaining_gas: &mut GasCounter,
         tx_context: Arc<TransactionContext>,
         validate: bool,
         charge_fee: bool,
     ) -> TransactionExecutionResult<ValidateExecuteCallInfo> {
         if self.is_non_revertible(&tx_context.tx_info) {
             return self.run_non_revertible(state, tx_context, remaining_gas, validate, charge_fee);
         }
-
         self.run_revertible(state, tx_context, remaining_gas, validate, charge_fee)
     }
 }
@@ -723,7 +752,7 @@ impl<U: UpdatableState> ExecutableTransaction<U> for AccountTransaction {
         )?;
 
         // Run validation and execution.
-        let mut remaining_gas = tx_context.initial_sierra_gas();
+        let initial_gas = tx_context.initial_sierra_gas();
         let ValidateExecuteCallInfo {
             validate_call_info,
             execute_call_info,
@@ -737,7 +766,7 @@ impl<U: UpdatableState> ExecutableTransaction<U> for AccountTransaction {
                 },
         } = self.run_or_revert(
             state,
-            &mut remaining_gas,
+            &mut GasCounter::new(initial_gas),
             tx_context.clone(),
             execution_flags.validate,
             execution_flags.charge_fee,

crates/blockifier/src/versioned_constants.rs

@@ -247,6 +247,23 @@ impl VersionedConstants {
             * self.vm_resource_fee_cost().n_steps
     }
 
+    /// Default initial gas amount when L2 gas is not provided.
+    pub fn default_initial_gas_amount(&self) -> u64 {
+        (self.execute_max_sierra_gas.checked_add(self.validate_max_sierra_gas))
+            .expect("The default initial gas cost should be less than the maximum gas amount.")
+            .0
+    }
+
+    /// Returns the maximum gas amount for validation.
+    pub fn max_validation_sierra_gas(&self) -> GasAmount {
+        self.validate_max_sierra_gas
+    }
+
+    /// Returns the maximum gas amount for execute.
+    pub fn max_execution_sierra_gas(&self) -> GasAmount {
+        self.execute_max_sierra_gas
+    }
+
     /// Returns the default initial gas for VM mode transactions.
     pub fn default_initial_gas_cost(&self) -> u64 {
         self.os_constants.gas_costs.default_initial_gas_cost

crates/starknet_api/src/execution_resources.rs

@@ -9,7 +9,14 @@ use crate::transaction::fields::{Fee, Resource};
 
 #[cfg_attr(
     any(test, feature = "testing"),
-    derive(derive_more::Add, derive_more::Sum, derive_more::AddAssign, derive_more::Div)
+    derive(
+        derive_more::Add,
+        derive_more::Sum,
+        derive_more::AddAssign,
+        derive_more::Div,
+        derive_more::Sub,
+        derive_more::SubAssign
+    )
 )]
 #[derive(
     derive_more::Display,
@@ -55,6 +62,10 @@ impl GasAmount {
         self.0.checked_add(rhs.0).map(Self)
     }
 
+    pub fn checked_sub(self, rhs: Self) -> Option<Self> {
+        self.0.checked_sub(rhs.0).map(Self)
+    }
+
     pub const fn nonzero_saturating_mul(self, rhs: NonzeroGasPrice) -> Fee {
         rhs.saturating_mul(self)
     }