[static ptb] Fix PTB gas stack height (#24440)

## Description 

- The new PTB gas metering was a bit sloppy, and did not properly pop
from the stack
- Added an invariant violation for when it is not zero

## Test plan 

- Ran tests

---

## Release notes

Check each box that your changes affect. If none of the boxes relate to
your changes, release notes aren't required.

For each box you select, include information after the relevant heading
that describes the impact of your changes that a user might notice and
any actions they must take to implement updates.

- [ ] Protocol: 
- [ ] Nodes (Validators and Full nodes): 
- [ ] gRPC:
- [ ] JSON-RPC: 
- [ ] GraphQL: 
- [ ] CLI: 
- [ ] Rust SDK:
- [ ] Indexing Framework:

Author: tnowacki

crates/sui-adapter-transactional-tests/tests/size_limits/event_limits_tests_out_of_gas.snap

@@ -22,9 +22,9 @@ Execution Error: ExecutionError: ExecutionError { inner: ExecutionErrorInner { k
 task 4, lines 73-75:
 //# run Test::M1::emit_event_with_size --args 259000 --gas-budget 1000000000
 Error: Transaction Effects Status: Insufficient Gas.
-Execution Error: ExecutionError: ExecutionError { inner: ExecutionErrorInner { kind: InsufficientGas, source: Some(VMError { major_status: OUT_OF_GAS, sub_status: None, message: None, exec_state: None, location: Module(ModuleId { address: Test, name: Identifier("M1") }), indices: [], offsets: [(FunctionDefinitionIndex(0), 26)] }), command: Some(0) } }
+Execution Error: ExecutionError: ExecutionError { inner: ExecutionErrorInner { kind: InsufficientGas, source: Some(VMError { major_status: OUT_OF_GAS, sub_status: None, message: None, exec_state: None, location: Module(ModuleId { address: Test, name: Identifier("M1") }), indices: [], offsets: [(FunctionDefinitionIndex(0), 37)] }), command: Some(0) } }
 
 task 5, line 76:
 //# run Test::M1::emit_n_events_with_size --args 3 256000 --gas-budget 1000000000 --summarize
 Error: Transaction Effects Status: Insufficient Gas.
-Execution Error: ExecutionError: ExecutionError { inner: ExecutionErrorInner { kind: InsufficientGas, source: Some(VMError { major_status: OUT_OF_GAS, sub_status: None, message: None, exec_state: None, location: Module(ModuleId { address: Test, name: Identifier("M1") }), indices: [], offsets: [(FunctionDefinitionIndex(0), 35)] }), command: Some(0) } }
+Execution Error: ExecutionError: ExecutionError { inner: ExecutionErrorInner { kind: InsufficientGas, source: Some(VMError { major_status: OUT_OF_GAS, sub_status: None, message: None, exec_state: None, location: Module(ModuleId { address: Test, name: Identifier("M1") }), indices: [], offsets: [(FunctionDefinitionIndex(0), 27)] }), command: Some(0) } }

crates/sui-adapter-transactional-tests/tests/upgrade/type_name_unit_tests.snap

@@ -32,4 +32,4 @@ gas summary: computation_cost: 3000000, storage_cost: 988000,  storage_rebate: 0
 task 6, line 177:
 //# run A2::m::generic_cases --args @A0 @A1
 mutated: object(0,1)
-gas summary: computation_cost: 196000000, storage_cost: 988000,  storage_rebate: 978120, non_refundable_storage_fee: 9880
+gas summary: computation_cost: 195000000, storage_cost: 988000,  storage_rebate: 978120, non_refundable_storage_fee: 9880

crates/sui-types/src/gas_model/tables.rs

@@ -115,7 +115,7 @@ impl GasStatus {
     /// code that does not have to charge the user.
     pub fn new_unmetered() -> Self {
         Self {
-            gas_model_version: 4,
+            gas_model_version: 11,
             gas_left: InternalGas::new(0),
             gas_price: 1,
             initial_budget: InternalGas::new(0),
@@ -337,6 +337,14 @@ impl GasStatus {
     pub fn instructions_executed(&self) -> u64 {
         self.instructions_executed
     }
+
+    pub fn stack_height_current(&self) -> u64 {
+        self.stack_height_current
+    }
+
+    pub fn stack_size_current(&self) -> u64 {
+        self.stack_size_current
+    }
 }
 
 pub fn zero_cost_schedule() -> CostTable {

sui-execution/latest/sui-adapter/src/gas_meter.rs

@@ -259,10 +259,15 @@ impl GasMeter for SuiGasMeter<'_> {
         // TODO(tzakian): We should account for this elsewhere as the owner of data the
         // reference points to won't be on the stack. For now though, we treat it as adding to the
         // stack size.
+        let (pushes, pops) = if reduce_stack_size(self.0.gas_model_version) {
+            (0, 2)
+        } else {
+            (1, 2)
+        };
         self.0.charge(
             1,
-            1,
-            2,
+            pushes,
+            pops,
             abstract_memory_size(self.0, new_val).into(),
             abstract_memory_size(self.0, old_val).into(),
         )
@@ -355,7 +360,12 @@ impl GasMeter for SuiGasMeter<'_> {
 
     fn charge_vec_swap(&mut self, _ty: impl TypeView) -> PartialVMResult<()> {
         let size_decrease = REFERENCE_SIZE + Type::U64.size() + Type::U64.size();
-        self.0.charge(1, 1, 1, 0, size_decrease.into())
+        let (pushes, pops) = if reduce_stack_size(self.0.gas_model_version) {
+            (0, 3)
+        } else {
+            (1, 1)
+        };
+        self.0.charge(1, pushes, pops, 0, size_decrease.into())
     }
 
     fn charge_drop_frame(
@@ -400,6 +410,11 @@ fn reweight_move_loc(gas_model_version: u64) -> bool {
     gas_model_version > 10
 }
 
+fn reduce_stack_size(gas_model_version: u64) -> bool {
+    // Reducing stack size is only done in gas model versions 10 and above.
+    gas_model_version > 10
+}
+
 fn size_config_for_gas_model_version(
     gas_model_version: u64,
     should_traverse_refs: bool,

sui-execution/latest/sui-adapter/src/static_programmable_transactions/execution/context.rs

@@ -32,7 +32,7 @@ use move_core_types::{
 use move_trace_format::format::MoveTraceBuilder;
 use move_vm_runtime::native_extensions::NativeContextExtensions;
 use move_vm_types::{
-    gas::GasMeter,
+    gas::{GasMeter, SimpleInstruction},
     values::{VMValueCast, Value as VMValue},
 };
 use std::{
@@ -78,11 +78,14 @@ macro_rules! object_runtime_mut {
 }
 
 macro_rules! charge_gas_ {
-    ($gas_charger:expr, $env:expr, $case:ident, $value_view:expr) => {{
+    ($gas_charger:expr, $env:expr, $call:ident($($args:expr),*)) => {{
         SuiGasMeter($gas_charger.move_gas_status_mut())
-            .$case($value_view)
+            .$call($($args),*)
             .map_err(|e| $env.convert_vm_error(e.finish(Location::Undefined)))
     }};
+    ($gas_charger:expr, $env:expr, $case:ident, $value_view:expr) => {
+        charge_gas_!($gas_charger, $env, $case($value_view))
+    };
 }
 
 macro_rules! charge_gas {
@@ -269,6 +272,7 @@ impl<'env, 'pc, 'vm, 'state, 'linkage, 'gas> Context<'env, 'pc, 'vm, 'state, 'li
             tx_context.clone(),
         );
 
+        debug_assert_eq!(gas_charger.move_gas_status().stack_height_current(), 0);
         let tx_context_value = Locals::new(vec![Some(Value::new_tx_context(
             tx_context.borrow().digest(),
         )?)])?;
@@ -515,13 +519,24 @@ impl<'env, 'pc, 'vm, 'state, 'linkage, 'gas> Context<'env, 'pc, 'vm, 'state, 'li
             }
         };
         Ok(match usage {
-            UsageKind::Move => local.move_()?,
+            UsageKind::Move => {
+                let value = local.move_()?;
+                charge_gas_!(self.gas_charger, self.env, charge_move_loc, &value)?;
+                value
+            }
             UsageKind::Copy => {
                 let value = local.copy()?;
                 charge_gas_!(self.gas_charger, self.env, charge_copy_loc, &value)?;
                 value
             }
-            UsageKind::Borrow => local.borrow()?,
+            UsageKind::Borrow => {
+                charge_gas_!(
+                    self.gas_charger,
+                    self.env,
+                    charge_simple_instr(SimpleInstruction::MutBorrowLoc)
+                )?;
+                local.borrow()?
+            }
         })
     }
 
@@ -550,7 +565,10 @@ impl<'env, 'pc, 'vm, 'state, 'linkage, 'gas> Context<'env, 'pc, 'vm, 'state, 'li
     where
         VMValue: VMValueCast<V>,
     {
+        let before_height = self.gas_charger.move_gas_status().stack_height_current();
         let value = self.argument_value(arg)?;
+        let after_height = self.gas_charger.move_gas_status().stack_height_current();
+        debug_assert_eq!(before_height.saturating_add(1), after_height);
         let value: V = value.cast()?;
         Ok(value)
     }
@@ -569,6 +587,36 @@ impl<'env, 'pc, 'vm, 'state, 'linkage, 'gas> Context<'env, 'pc, 'vm, 'state, 'li
         Ok(())
     }
 
+    pub fn charge_command(
+        &mut self,
+        is_move_call: bool,
+        num_args: usize,
+        num_return: usize,
+    ) -> Result<(), ExecutionError> {
+        let move_gas_status = self.gas_charger.move_gas_status_mut();
+        let before_size = move_gas_status.stack_size_current();
+        // Pop all of the arguments
+        // If the return values came from the Move VM directly (via a Move call), pop those
+        // as well
+        let num_popped = if is_move_call {
+            num_args.checked_add(num_return).ok_or_else(|| {
+                make_invariant_violation!("usize overflow when charging gas for command",)
+            })?
+        } else {
+            num_args
+        };
+        move_gas_status
+            .charge(1, 0, num_popped as u64, 0, /* unused */ 1)
+            .map_err(|e| self.env.convert_vm_error(e.finish(Location::Undefined)))?;
+        let after_size = move_gas_status.stack_size_current();
+        assert_invariant!(
+            before_size == after_size,
+            "We assume currently that the stack size is not decremented. \
+            If this changes, we need to actually account for it here"
+        );
+        Ok(())
+    }
+
     pub fn copy_value(&mut self, value: &CtxValue) -> Result<CtxValue, ExecutionError> {
         Ok(CtxValue(copy_value(self.gas_charger, self.env, &value.0)?))
     }
@@ -838,11 +886,15 @@ impl<'env, 'pc, 'vm, 'state, 'linkage, 'gas> Context<'env, 'pc, 'vm, 'state, 'li
                 .map_err(|e| {
                     make_invariant_violation!("Failed to get tx context for init function: {}", e)
                 })?;
+            // balance the stack after borrowing the tx context
+            charge_gas!(self, charge_store_loc, &tx_context)?;
+
             let args = if has_otw {
                 vec![CtxValue(Value::one_time_witness()?), CtxValue(tx_context)]
             } else {
                 vec![CtxValue(tx_context)]
             };
+            debug_assert_eq!(self.gas_charger.move_gas_status().stack_height_current(), 0);
             let return_values = self.execute_function_bypass_visibility(
                 &module.self_id(),
                 INIT_FN_NAME,
@@ -862,7 +914,8 @@ impl<'env, 'pc, 'vm, 'state, 'linkage, 'gas> Context<'env, 'pc, 'vm, 'state, 'li
             assert_invariant!(
                 return_values.is_empty(),
                 "init should not have return values"
-            )
+            );
+            debug_assert_eq!(self.gas_charger.move_gas_status().stack_height_current(), 0);
         }
 
         Ok(())
@@ -1247,6 +1300,7 @@ fn load_object_arg_impl(
 
     let v = Value::deserialize(env, move_obj.contents(), ty)?;
     charge_gas_!(meter, env, charge_copy_loc, &v)?;
+    charge_gas_!(meter, env, charge_store_loc, &v)?;
     Ok((object_metadata, v))
 }
 
@@ -1263,6 +1317,7 @@ fn load_withdrawal_arg(
     } = withdrawal;
     let loaded = Value::funds_accumulator_withdrawal(*owner, *amount);
     charge_gas_!(meter, env, charge_copy_loc, &loaded)?;
+    charge_gas_!(meter, env, charge_store_loc, &loaded)?;
     Ok(loaded)
 }
 
@@ -1275,6 +1330,7 @@ fn load_pure_value(
     let loaded = Value::deserialize(env, bytes, metadata.ty.clone())?;
     // ByteValue::Receiving { id, version } => Value::receiving(*id, *version),
     charge_gas_!(meter, env, charge_copy_loc, &loaded)?;
+    charge_gas_!(meter, env, charge_store_loc, &loaded)?;
     Ok(loaded)
 }
 
@@ -1286,11 +1342,13 @@ fn load_receiving_value(
     let (id, version, _) = metadata.object_ref;
     let loaded = Value::receiving(id, version);
     charge_gas_!(meter, env, charge_copy_loc, &loaded)?;
+    charge_gas_!(meter, env, charge_store_loc, &loaded)?;
     Ok(loaded)
 }
 
 fn copy_value(meter: &mut GasCharger, env: &Env, value: &Value) -> Result<Value, ExecutionError> {
     charge_gas_!(meter, env, charge_copy_loc, value)?;
+    charge_gas_!(meter, env, charge_pop, value)?;
     value.copy()
 }
 

sui-execution/latest/sui-adapter/src/static_programmable_transactions/execution/interpreter.rs

@@ -66,6 +66,7 @@ pub fn execute_inner<'env, 'pc, 'vm, 'state, 'linkage, Mode: ExecutionMode>(
 where
     'pc: 'state,
 {
+    debug_assert_eq!(gas_charger.move_gas_status().stack_height_current(), 0);
     let T::Transaction {
         bytes,
         objects,
@@ -147,6 +148,12 @@ fn execute_command<Mode: ExecutionMode>(
         drop_values,
         consumed_shared_objects: _,
     } = c;
+    assert_invariant!(
+        context.gas_charger.move_gas_status().stack_height_current() == 0,
+        "stack height did not start at 0"
+    );
+    let is_move_call = matches!(command, T::Command__::MoveCall(_));
+    let num_args = command.arguments_len();
     let mut args_to_update = vec![];
     let result = match command {
         T::Command__::MoveCall(move_call) => {
@@ -323,11 +330,16 @@ fn execute_command<Mode: ExecutionMode>(
         result.len() == drop_values.len(),
         "result values and drop values mismatch"
     );
+    context.charge_command(is_move_call, num_args, result.len())?;
     let result = result
         .into_iter()
         .zip(drop_values)
         .map(|(value, drop)| if !drop { Some(value) } else { None })
         .collect::<Vec<_>>();
     context.result(result)?;
+    assert_invariant!(
+        context.gas_charger.move_gas_status().stack_height_current() == 0,
+        "stack height did not end at 0"
+    );
     Ok(())
 }

sui-execution/latest/sui-adapter/src/static_programmable_transactions/typing/ast.rs

@@ -7,7 +7,7 @@ use crate::static_programmable_transactions::{
 use indexmap::IndexSet;
 use move_core_types::{account_address::AccountAddress, u256::U256};
 use move_vm_types::values::VectorSpecialization;
-use std::{cell::OnceCell, vec};
+use std::cell::OnceCell;
 use sui_types::base_types::{ObjectID, ObjectRef};
 
 //**************************************************************************************************
@@ -229,21 +229,21 @@ impl Argument__ {
 }
 
 impl Command__ {
-    pub fn arguments(&self) -> Vec<&Argument> {
+    pub fn arguments(&self) -> Box<dyn Iterator<Item = &Argument> + '_> {
         match self {
-            Command__::MoveCall(mc) => mc.arguments.iter().collect(),
+            Command__::MoveCall(mc) => Box::new(mc.arguments.iter()),
             Command__::TransferObjects(objs, addr) => {
-                objs.iter().chain(std::iter::once(addr)).collect()
+                Box::new(objs.iter().chain(std::iter::once(addr)))
             }
             Command__::SplitCoins(_, coin, amounts) => {
-                std::iter::once(coin).chain(amounts).collect()
+                Box::new(std::iter::once(coin).chain(amounts))
             }
             Command__::MergeCoins(_, target, sources) => {
-                std::iter::once(target).chain(sources).collect()
+                Box::new(std::iter::once(target).chain(sources))
             }
-            Command__::MakeMoveVec(_, elems) => elems.iter().collect(),
-            Command__::Publish(_, _, _) => vec![],
-            Command__::Upgrade(_, _, _, arg, _) => vec![arg],
+            Command__::MakeMoveVec(_, elems) => Box::new(elems.iter()),
+            Command__::Publish(_, _, _) => Box::new(std::iter::empty()),
+            Command__::Upgrade(_, _, _, arg, _) => Box::new(std::iter::once(arg)),
         }
     }
 
@@ -277,6 +277,20 @@ impl Command__ {
             Command__::Publish(_, _, _) => Box::new(std::iter::empty()),
         }
     }
+
+    pub fn arguments_len(&self) -> usize {
+        let n = match self {
+            Command__::MoveCall(mc) => mc.arguments.len(),
+            Command__::TransferObjects(objs, _) => objs.len().saturating_add(1),
+            Command__::SplitCoins(_, _, amounts) => amounts.len().saturating_add(1),
+            Command__::MergeCoins(_, _, sources) => sources.len().saturating_add(1),
+            Command__::MakeMoveVec(_, elems) => elems.len(),
+            Command__::Publish(_, _, _) => 0,
+            Command__::Upgrade(_, _, _, _, _) => 1,
+        };
+        debug_assert_eq!(self.arguments().count(), n);
+        n
+    }
 }
 
 //**************************************************************************************************