[move] Move out tag cache (safe speculation)

third_party/move/move-vm/runtime/src/loader/mod.rs

@@ -56,6 +56,7 @@ use crate::{
     storage::{
         loader::LoaderV2, module_storage::FunctionValueExtensionAdapter,
         struct_name_index_map::StructNameIndexMap, ty_cache::StructInfoCache,
+        ty_tag_cache::TypeTagBuilder,
     },
 };
 pub use function::{Function, LoadedFunction};
@@ -1671,15 +1672,15 @@ pub const VALUE_DEPTH_MAX: u64 = 128;
 /// fields for struct types.
 const MAX_TYPE_TO_LAYOUT_NODES: u64 = 256;
 
-struct PseudoGasContext {
-    max_cost: u64,
-    cost: u64,
-    cost_base: u64,
-    cost_per_byte: u64,
+pub(crate) struct PseudoGasContext {
+    pub(crate) max_cost: u64,
+    pub(crate) cost: u64,
+    pub(crate) cost_base: u64,
+    pub(crate) cost_per_byte: u64,
 }
 
 impl PseudoGasContext {
-    fn charge(&mut self, amount: u64) -> PartialVMResult<()> {
+    pub(crate) fn charge(&mut self, amount: u64) -> PartialVMResult<()> {
         self.cost += amount;
         if self.cost > self.max_cost {
             Err(
@@ -1694,16 +1695,14 @@ impl PseudoGasContext {
     }
 }
 
-impl Loader {
+impl LoaderV1 {
     fn struct_name_to_type_tag(
         &self,
         struct_name_idx: StructNameIndex,
         ty_args: &[Type],
         gas_context: &mut PseudoGasContext,
-        module_storage: &dyn ModuleStorage,
     ) -> PartialVMResult<StructTag> {
-        let ty_cache = self.ty_cache(module_storage);
-        if let Some((struct_tag, gas)) = ty_cache.get_struct_tag(&struct_name_idx, ty_args) {
+        if let Some((struct_tag, gas)) = self.type_cache.get_struct_tag(&struct_name_idx, ty_args) {
             gas_context.charge(gas)?;
             return Ok(struct_tag.clone());
         }
@@ -1712,16 +1711,16 @@ impl Loader {
 
         let type_args = ty_args
             .iter()
-            .map(|ty| self.type_to_type_tag_impl(ty, gas_context, module_storage))
+            .map(|ty| self.type_to_type_tag_impl(ty, gas_context))
             .collect::<PartialVMResult<Vec<_>>>()?;
-
-        let struct_name_index_map = self.struct_name_index_map(module_storage);
-        let struct_tag = struct_name_index_map.idx_to_struct_tag(struct_name_idx, type_args)?;
+        let struct_tag = self
+            .name_cache
+            .idx_to_struct_tag(struct_name_idx, type_args)?;
 
         let size =
             (struct_tag.address.len() + struct_tag.module.len() + struct_tag.name.len()) as u64;
         gas_context.charge(size * gas_context.cost_per_byte)?;
-        ty_cache.store_struct_tag(
+        self.type_cache.store_struct_tag(
             struct_name_idx,
             ty_args.to_vec(),
             struct_tag.clone(),
@@ -1734,7 +1733,6 @@ impl Loader {
         &self,
         ty: &Type,
         gas_context: &mut PseudoGasContext,
-        module_storage: &dyn ModuleStorage,
     ) -> PartialVMResult<TypeTag> {
         gas_context.charge(gas_context.cost_base)?;
         Ok(match ty {
@@ -1748,17 +1746,16 @@ impl Loader {
             Type::Address => TypeTag::Address,
             Type::Signer => TypeTag::Signer,
             Type::Vector(ty) => {
-                let el_ty_tag = self.type_to_type_tag_impl(ty, gas_context, module_storage)?;
+                let el_ty_tag = self.type_to_type_tag_impl(ty, gas_context)?;
                 TypeTag::Vector(Box::new(el_ty_tag))
             },
             Type::Struct { idx, .. } => TypeTag::Struct(Box::new(self.struct_name_to_type_tag(
                 *idx,
                 &[],
                 gas_context,
-                module_storage,
             )?)),
             Type::StructInstantiation { idx, ty_args, .. } => TypeTag::Struct(Box::new(
-                self.struct_name_to_type_tag(*idx, ty_args, gas_context, module_storage)?,
+                self.struct_name_to_type_tag(*idx, ty_args, gas_context)?,
             )),
             Type::Reference(_) | Type::MutableReference(_) | Type::TyParam(_) => {
                 return Err(
@@ -1768,7 +1765,9 @@ impl Loader {
             },
         })
     }
+}
 
+impl Loader {
     fn struct_name_to_type_layout(
         &self,
         struct_name_idx: StructNameIndex,
@@ -2057,18 +2056,21 @@ impl Loader {
         }
 
         let count_before = *count;
-        let mut gas_context = PseudoGasContext {
-            cost: 0,
-            max_cost: self.vm_config().type_max_cost,
-            cost_base: self.vm_config().type_base_cost,
-            cost_per_byte: self.vm_config().type_byte_cost,
+        let struct_tag = match self {
+            Loader::V1(loader) => {
+                let mut gas_context = PseudoGasContext {
+                    cost: 0,
+                    max_cost: loader.vm_config().type_max_cost,
+                    cost_base: loader.vm_config().type_base_cost,
+                    cost_per_byte: loader.vm_config().type_byte_cost,
+                };
+                loader.struct_name_to_type_tag(struct_name_idx, ty_args, &mut gas_context)?
+            },
+            Loader::V2(_) => {
+                let ty_tag_builder = TypeTagBuilder::new(module_storage.runtime_environment());
+                ty_tag_builder.struct_ty_to_struct_tag(&struct_name_idx, ty_args)?
+            },
         };
-        let struct_tag = self.struct_name_to_type_tag(
-            struct_name_idx,
-            ty_args,
-            &mut gas_context,
-            module_storage,
-        )?;
         let fields = struct_type.fields(None)?;
 
         let field_layouts = fields
@@ -2263,13 +2265,21 @@ impl Loader {
         ty: &Type,
         module_storage: &dyn ModuleStorage,
     ) -> PartialVMResult<TypeTag> {
-        let mut gas_context = PseudoGasContext {
-            cost: 0,
-            max_cost: self.vm_config().type_max_cost,
-            cost_base: self.vm_config().type_base_cost,
-            cost_per_byte: self.vm_config().type_byte_cost,
-        };
-        self.type_to_type_tag_impl(ty, &mut gas_context, module_storage)
+        match self {
+            Loader::V1(loader) => {
+                let mut gas_context = PseudoGasContext {
+                    cost: 0,
+                    max_cost: self.vm_config().type_max_cost,
+                    cost_base: self.vm_config().type_base_cost,
+                    cost_per_byte: self.vm_config().type_byte_cost,
+                };
+                loader.type_to_type_tag_impl(ty, &mut gas_context)
+            },
+            Loader::V2(_) => {
+                let ty_tag_builder = TypeTagBuilder::new(module_storage.runtime_environment());
+                ty_tag_builder.ty_to_ty_tag(ty)
+            },
+        }
     }
 
     pub(crate) fn type_to_type_layout_with_identifier_mappings(

third_party/move/move-vm/runtime/src/storage/environment.rs

@@ -7,7 +7,7 @@ use crate::{
     native_functions::{NativeFunction, NativeFunctions},
     storage::{
         struct_name_index_map::StructNameIndexMap, ty_cache::StructInfoCache,
-        verified_module_cache::VERIFIED_MODULES_V2,
+        ty_tag_cache::TypeTagCache, verified_module_cache::VERIFIED_MODULES_V2,
     },
     Module, Script,
 };
@@ -52,6 +52,10 @@ pub struct RuntimeEnvironment {
     ///   We wrap the index map into an [Arc] so that on republishing these clones are cheap.
     struct_name_index_map: Arc<StructNameIndexMap>,
 
+    /// Caches struct tags for instantiated types. This cache can be used concurrently and
+    /// speculatively.
+    ty_tag_cache: Arc<TypeTagCache>,
+
     /// Type cache for struct layouts, tags and depths, shared across multiple threads.
     ///
     /// SAFETY:
@@ -100,6 +104,7 @@ impl RuntimeEnvironment {
             vm_config,
             natives,
             struct_name_index_map: Arc::new(StructNameIndexMap::empty()),
+            ty_tag_cache: Arc::new(TypeTagCache::empty()),
             ty_cache: StructInfoCache::empty(),
         }
     }
@@ -261,6 +266,12 @@ impl RuntimeEnvironment {
         &self.struct_name_index_map
     }
 
+    /// Returns the type tag cache used by this environment to store already constructed struct
+    /// tags.
+    pub(crate) fn ty_tag_cache(&self) -> &TypeTagCache {
+        &self.ty_tag_cache
+    }
+
     /// Returns the type cache owned by this runtime environment which stores information about
     /// struct layouts, tags and depth formulae.
     pub(crate) fn ty_cache(&self) -> &StructInfoCache {
@@ -276,6 +287,7 @@ impl RuntimeEnvironment {
     /// Flushes the struct information (type) cache. Flushing this cache does not invalidate struct
     /// name index map or module cache.
     pub fn flush_struct_info_cache(&self) {
+        self.ty_tag_cache.flush();
         self.ty_cache.flush();
     }
 
@@ -307,14 +319,15 @@ impl RuntimeEnvironment {
 }
 
 impl Clone for RuntimeEnvironment {
-    /// Returns the cloned environment. Struct re-indexing map and type caches are cloned and no
-    /// longer shared with the original environment.
+    /// Returns the cloned environment.
     fn clone(&self) -> Self {
         Self {
             vm_config: self.vm_config.clone(),
             natives: self.natives.clone(),
-            struct_name_index_map: self.struct_name_index_map.clone(),
             ty_cache: self.ty_cache.clone(),
+            // Cloning arcs below.
+            struct_name_index_map: self.struct_name_index_map.clone(),
+            ty_tag_cache: self.ty_tag_cache.clone(),
         }
     }
 }

third_party/move/move-vm/runtime/src/storage/mod.rs

@@ -4,6 +4,7 @@
 pub(crate) mod loader;
 pub(crate) mod struct_name_index_map;
 pub(crate) mod ty_cache;
+pub(crate) mod ty_tag_cache;
 mod verified_module_cache;
 
 pub mod code_storage;