[aptos-vm] Refactor type tag processing

aptos-move/aptos-gas-schedule/src/ver.rs

@@ -8,6 +8,9 @@
 ///   - Changing how gas is calculated in any way
 ///
 /// Change log:
+/// - V31:
+///    - Gas charging for modules used in type tags
+///
 /// - V22
 ///    - Gas parameters for enums
 ///    - Gas parameters for new native function `bcs::serialized_size`
@@ -90,4 +93,8 @@ pub mod gas_feature_versions {
     pub const RELEASE_V1_23: u64 = 27;
     pub const RELEASE_V1_24: u64 = 28;
     pub const RELEASE_V1_26: u64 = 30;
+    pub const RELEASE_V1_27: u64 = 31;
+    pub const RELEASE_V1_28: u64 = 32;
+    pub const RELEASE_V1_29: u64 = 33;
+    pub const RELEASE_V1_30: u64 = 34;
 }

aptos-move/aptos-vm/src/aptos_vm.rs

@@ -39,7 +39,10 @@ use aptos_framework::{
 };
 use aptos_gas_algebra::{Gas, GasQuantity, NumBytes, Octa};
 use aptos_gas_meter::{AptosGasMeter, GasAlgebra};
-use aptos_gas_schedule::{AptosGasParameters, VMGasParameters};
+use aptos_gas_schedule::{
+    gas_feature_versions::{RELEASE_V1_10, RELEASE_V1_27},
+    AptosGasParameters, VMGasParameters,
+};
 use aptos_logger::{enabled, prelude::*, Level};
 use aptos_metrics_core::TimerHelper;
 #[cfg(any(test, feature = "testing"))]
@@ -776,16 +779,16 @@ impl AptosVM {
         Ok((VMStatus::Executed, output))
     }
 
-    fn validate_and_execute_script(
+    fn validate_and_execute_script<'a>(
         &self,
         session: &mut SessionExt,
         serialized_signers: &SerializedSigners,
         code_storage: &impl AptosCodeStorage,
         // Note: cannot use AptosGasMeter because it is not implemented for
         //       UnmeteredGasMeter.
         gas_meter: &mut impl GasMeter,
-        traversal_context: &mut TraversalContext,
-        script: &Script,
+        traversal_context: &mut TraversalContext<'a>,
+        script: &'a Script,
     ) -> Result<(), VMStatus> {
         if !self
             .features()
@@ -803,14 +806,22 @@ impl AptosVM {
         // Note: Feature gating is needed here because the traversal of the dependencies could
         //       result in shallow-loading of the modules and therefore subtle changes in
         //       the error semantics.
-        if self.gas_feature_version() >= 15 {
+        if self.gas_feature_version() >= RELEASE_V1_10 {
             session.check_script_dependencies_and_check_gas(
                 code_storage,
                 gas_meter,
                 traversal_context,
                 script.code(),
             )?;
         }
+        if self.gas_feature_version() >= RELEASE_V1_27 {
+            session.check_type_tag_dependencies_and_charge_gas(
+                code_storage,
+                gas_meter,
+                traversal_context,
+                script.ty_args(),
+            )?;
+        }
 
         let func = session.load_script(code_storage, script.code(), script.ty_args())?;
 
@@ -868,12 +879,11 @@ impl AptosVM {
         gas_meter: &mut impl AptosGasMeter,
         traversal_context: &mut TraversalContext,
         entry_fn: &EntryFunction,
-        _txn_data: &TransactionMetadata,
     ) -> Result<(), VMStatus> {
         // Note: Feature gating is needed here because the traversal of the dependencies could
         //       result in shallow-loading of the modules and therefore subtle changes in
         //       the error semantics.
-        if self.gas_feature_version() >= 15 {
+        if self.gas_feature_version() >= RELEASE_V1_10 {
             let module_id = traversal_context
                 .referenced_module_ids
                 .alloc(entry_fn.module().clone());
@@ -885,6 +895,15 @@ impl AptosVM {
             )?;
         }
 
+        if self.gas_feature_version() >= RELEASE_V1_27 {
+            session.check_type_tag_dependencies_and_charge_gas(
+                module_storage,
+                gas_meter,
+                traversal_context,
+                entry_fn.ty_args(),
+            )?;
+        }
+
         let function = session.load_function(
             module_storage,
             entry_fn.module(),
@@ -988,7 +1007,6 @@ impl AptosVM {
                         gas_meter,
                         traversal_context,
                         entry_fn,
-                        txn_data,
                     )
                 })?;
             },
@@ -1118,7 +1136,6 @@ impl AptosVM {
                                 payload.multisig_address,
                                 entry_function,
                                 new_published_modules_loaded,
-                                txn_data,
                                 change_set_configs,
                             )?;
                             let has_modules_published_to_special_address =
@@ -1263,7 +1280,6 @@ impl AptosVM {
                     txn_payload.multisig_address,
                     &entry_function,
                     new_published_modules_loaded,
-                    txn_data,
                     change_set_configs,
                 ),
         };
@@ -1408,7 +1424,6 @@ impl AptosVM {
         multisig_address: AccountAddress,
         payload: &EntryFunction,
         new_published_modules_loaded: &mut bool,
-        txn_data: &TransactionMetadata,
         change_set_configs: &ChangeSetConfigs,
     ) -> Result<UserSessionChangeSet, VMStatus> {
         // If txn args are not valid, we'd still consider the transaction as executed but
@@ -1422,7 +1437,6 @@ impl AptosVM {
                 gas_meter,
                 traversal_context,
                 payload,
-                txn_data,
             )
         })?;
 
@@ -1596,7 +1610,7 @@ impl AptosVM {
         // Note: Feature gating is needed here because the traversal of the dependencies could
         //       result in shallow-loading of the modules and therefore subtle changes in
         //       the error semantics.
-        if self.gas_feature_version() >= 15 {
+        if self.gas_feature_version() >= RELEASE_V1_10 {
             // Charge old versions of existing modules, in case of upgrades.
             for module in modules.iter() {
                 let addr = module.self_addr();

third_party/move/move-core/types/src/language_storage.rs

@@ -96,6 +96,42 @@ impl TypeTag {
             Struct(s) => s.to_canonical_string(),
         }
     }
+
+    pub fn struct_tag(&self) -> Option<&StructTag> {
+        use TypeTag::*;
+        match self {
+            Struct(struct_tag) => Some(struct_tag.as_ref()),
+            Bool | U8 | U16 | U32 | U64 | U128 | U256 | Address | Signer | Vector(_) => None,
+        }
+    }
+
+    pub fn preorder_traversal_iter(&self) -> impl Iterator<Item = &TypeTag> {
+        TypeTagPreorderTraversalIter { stack: vec![self] }
+    }
+}
+
+struct TypeTagPreorderTraversalIter<'a> {
+    stack: Vec<&'a TypeTag>,
+}
+
+impl<'a> Iterator for TypeTagPreorderTraversalIter<'a> {
+    type Item = &'a TypeTag;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        use TypeTag::*;
+
+        match self.stack.pop() {
+            Some(ty) => {
+                match ty {
+                    Signer | Bool | Address | U8 | U16 | U32 | U64 | U128 | U256 => (),
+                    Vector(ty) => self.stack.push(ty),
+                    Struct(struct_tag) => self.stack.extend(struct_tag.type_args.iter().rev()),
+                }
+                Some(ty)
+            },
+            None => None,
+        }
+    }
 }
 
 impl FromStr for TypeTag {
@@ -361,8 +397,28 @@ mod tests {
         collections::hash_map::DefaultHasher,
         hash::{Hash, Hasher},
         mem,
+        str::FromStr,
     };
 
+    #[test]
+    fn test_tag_iter() {
+        let tag = TypeTag::from_str("vector<0x1::a::A<u8, 0x2::b::B, vector<vector<0x3::c::C>>>>")
+            .unwrap();
+        let actual_tags = tag.preorder_traversal_iter().collect::<Vec<_>>();
+        let expected_tags = [
+            tag.clone(),
+            TypeTag::from_str("0x1::a::A<u8, 0x2::b::B, vector<vector<0x3::c::C>>>").unwrap(),
+            TypeTag::from_str("u8").unwrap(),
+            TypeTag::from_str("0x2::b::B").unwrap(),
+            TypeTag::from_str("vector<vector<0x3::c::C>>").unwrap(),
+            TypeTag::from_str("vector<0x3::c::C>").unwrap(),
+            TypeTag::from_str("0x3::c::C").unwrap(),
+        ];
+        for (actual_tag, expected_tag) in actual_tags.into_iter().zip(expected_tags) {
+            assert_eq!(actual_tag, &expected_tag);
+        }
+    }
+
     #[test]
     fn test_type_tag_serde() {
         let a = TypeTag::Struct(Box::new(StructTag {

third_party/move/move-vm/runtime/src/session.rs

@@ -28,7 +28,7 @@ use move_vm_types::{
     loaded_data::runtime_types::{StructNameIndex, StructType, Type, TypeBuilder},
     values::{GlobalValue, Value},
 };
-use std::{borrow::Borrow, sync::Arc};
+use std::{borrow::Borrow, collections::BTreeSet, sync::Arc};
 
 pub struct Session<'r, 'l> {
     pub(crate) move_vm: &'l MoveVM,
@@ -535,6 +535,34 @@ impl<'r, 'l> Session<'r, 'l> {
             .ok()
     }
 
+    pub fn check_type_tag_dependencies_and_charge_gas(
+        &mut self,
+        module_storage: &impl ModuleStorage,
+        gas_meter: &mut impl GasMeter,
+        traversal_context: &mut TraversalContext,
+        ty_tags: &[TypeTag],
+    ) -> VMResult<()> {
+        // Charge gas based on the distinct ordered module ids.
+        let ordered_ty_tags = ty_tags
+            .iter()
+            .flat_map(|ty_tag| ty_tag.preorder_traversal_iter())
+            .filter_map(TypeTag::struct_tag)
+            .map(|struct_tag| {
+                let module_id = traversal_context
+                    .referenced_module_ids
+                    .alloc(struct_tag.module_id());
+                (module_id.address(), module_id.name())
+            })
+            .collect::<BTreeSet<_>>();
+
+        self.check_dependencies_and_charge_gas(
+            module_storage,
+            gas_meter,
+            traversal_context,
+            ordered_ty_tags,
+        )
+    }
+
     pub fn check_dependencies_and_charge_gas<'a, I>(
         &mut self,
         module_storage: &impl ModuleStorage,