[move] Improve & test struct index map

third_party/move/move-vm/runtime/Cargo.toml

@@ -14,6 +14,7 @@ edition = "2021"
 ambassador = { workspace = true }
 better_any = { workspace = true }
 bytes = { workspace = true }
+claims = { workspace = true }
 fail = { workspace = true }
 hashbrown = { workspace = true }
 lazy_static = { workspace = true }
@@ -33,8 +34,8 @@ move-vm-types = { path = "../types" }
 
 [dev-dependencies]
 anyhow = { workspace = true }
-claims = { workspace = true }
 hex = { workspace = true }
+move-binary-format = { path = "../../move-binary-format", features = ["fuzzing"] }
 move-compiler = { path = "../../move-compiler" }
 move-ir-compiler = { path = "../../move-ir-compiler" }
 proptest = { workspace = true }

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

@@ -841,8 +841,11 @@ impl Interpreter {
                 )
             },
         };
-        let struct_name = loader.get_struct_name(struct_idx);
-        if let Some(access) = AccessInstance::new(kind, &struct_name, instance, addr) {
+        let struct_name = loader
+            .runtime_environment()
+            .struct_name_index_map()
+            .idx_to_struct_name(struct_idx);
+        if let Some(access) = AccessInstance::new(kind, struct_name, instance, addr) {
             self.access_control.check_access(access)?
         }
         Ok(())

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

@@ -310,24 +310,16 @@ impl Loader {
     // Internal helpers
     //
 
-    pub(crate) fn get_struct_name(&self, struct_idx: StructNameIndex) -> StructIdentifier {
-        match self {
-            Self::V1(loader) => loader
-                .struct_name_index_map()
-                .idx_to_struct_name(struct_idx),
-            Self::V2(loader) => loader
-                .struct_name_index_map()
-                .idx_to_struct_name(struct_idx),
-        }
-    }
-
     pub fn fetch_struct_ty_by_idx(
         &self,
         idx: StructNameIndex,
         module_store: &ModuleStorageAdapter,
         module_storage: &dyn ModuleStorageV2,
     ) -> PartialVMResult<Arc<StructType>> {
-        let struct_name = self.struct_name_index_map().idx_to_struct_name(idx);
+        // Ensure that the struct name index map is unlocked immediately by getting an arced
+        // struct name.
+        let struct_name = self.struct_name_index_map().idx_to_struct_name_ref(idx);
+
         match self {
             Loader::V1(_) => {
                 module_store.get_struct_type_by_identifier(&struct_name.name, &struct_name.module)
@@ -1727,16 +1719,9 @@ impl Loader {
             .iter()
             .map(|ty| self.type_to_type_tag_impl(ty, gas_context))
             .collect::<PartialVMResult<Vec<_>>>()?;
-
-        let name = self
+        let struct_tag = self
             .struct_name_index_map()
-            .idx_to_struct_name(struct_name_idx);
-        let struct_tag = StructTag {
-            address: *name.module.address(),
-            module: name.module.name().to_owned(),
-            name: name.name.clone(),
-            type_args,
-        };
+            .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;
@@ -1823,8 +1808,8 @@ impl Loader {
                 // times. Right now these are Aggregator and AggregatorSnapshot.
                 let struct_name = self
                     .struct_name_index_map()
-                    .idx_to_struct_name(struct_name_idx);
-                let maybe_mapping = self.get_identifier_mapping_kind(&struct_name);
+                    .idx_to_struct_name_ref(struct_name_idx);
+                let maybe_mapping = self.get_identifier_mapping_kind(struct_name.as_ref());
 
                 let field_tys = fields
                     .iter()
@@ -2243,12 +2228,12 @@ impl Loader {
             //       behavior, e.g., make the cache available per thread.
             let struct_name = self
                 .struct_name_index_map()
-                .idx_to_struct_name(struct_name_idx);
+                .idx_to_struct_name_ref(struct_name_idx);
 
             // TODO(loader_v2): Revisit this, because now we do share the VM...
             println!(
                 "ERROR: Depth formula for struct '{}' and formula {:?} (struct type: {:?}) is already cached: {:?}",
-                &struct_name,
+                struct_name.as_ref(),
                 formula,
                 struct_type.as_ref(),
                 f

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

@@ -38,12 +38,6 @@ impl PartiallyVerifiedModule {
     ) -> impl DoubleEndedIterator<Item = (&AccountAddress, &IdentStr)> {
         self.0.immediate_dependencies_iter()
     }
-
-    pub fn immediate_friends_iter(
-        &self,
-    ) -> impl DoubleEndedIterator<Item = (&AccountAddress, &IdentStr)> {
-        self.0.immediate_friends_iter()
-    }
 }
 
 /// Wrapper around partially verified compiled script, i.e., one that passed

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

@@ -1,43 +1,96 @@
 // Copyright (c) The Move Contributors
 // SPDX-License-Identifier: Apache-2.0
 
+use claims::assert_none;
+use move_core_types::language_storage::{StructTag, TypeTag};
 use move_vm_types::loaded_data::runtime_types::{StructIdentifier, StructNameIndex};
 use parking_lot::RwLock;
-use std::collections::BTreeMap;
+use std::{collections::BTreeMap, sync::Arc};
 
 #[derive(Clone)]
 struct IndexMap<T: Clone + Ord> {
     forward_map: BTreeMap<T, usize>,
-    backward_map: Vec<T>,
+    backward_map: Vec<Arc<T>>,
 }
 
-/// A data structure to cache struct identifiers (address, module name, struct name) and
-/// use indices instead, to save on the memory consumption and avoid unnecessary cloning.
+/// A data structure to cache struct identifiers (address, module name, struct name) and use
+/// indices instead, to save on the memory consumption and avoid unnecessary cloning. It
+/// guarantees that the same struct name identifier always corresponds to a unique index.
 pub(crate) struct StructNameIndexMap(RwLock<IndexMap<StructIdentifier>>);
 
 impl StructNameIndexMap {
+    /// Returns an empty map with no entries.
     pub(crate) fn empty() -> Self {
         Self(RwLock::new(IndexMap {
             forward_map: BTreeMap::new(),
             backward_map: vec![],
         }))
     }
 
+    /// Maps the struct identifier into an index. If the identifier already exists returns the
+    /// corresponding index. This function guarantees that for any struct identifiers A and B,
+    /// if A == B, they have the same indices.
     pub(crate) fn struct_name_to_idx(&self, struct_name: StructIdentifier) -> StructNameIndex {
+        // Note that we take a write lock here once, instead of (*): taking a read lock, checking
+        // if the index is cached, re-acquiring the (write) lock, and checking again, as it makes
+        // things faster.
+        // Note that even if we do (*), we MUST check if another thread has cached the index before
+        // we reached this point for correctness. If we do not do this, we can end up evicting the
+        // previous index, and end up with multiple indices corresponding to the same struct. As
+        // indices are stored inside types, type comparison breaks!
         let mut index_map = self.0.write();
+
+        // Index is cached, return early.
         if let Some(idx) = index_map.forward_map.get(&struct_name) {
             return StructNameIndex(*idx);
         }
+
+        // Otherwise, the cache is locked and the struct name is not present. We simply add it
+        // to the cache and return the corresponding index.
         let idx = index_map.backward_map.len();
-        index_map.forward_map.insert(struct_name.clone(), idx);
-        index_map.backward_map.push(struct_name);
+        let prev_idx = index_map.forward_map.insert(struct_name.clone(), idx);
+        assert_none!(prev_idx, "Indexing map should never evict cached entries");
+
+        index_map.backward_map.push(Arc::new(struct_name));
         StructNameIndex(idx)
     }
 
-    pub(crate) fn idx_to_struct_name(&self, idx: StructNameIndex) -> StructIdentifier {
-        // TODO(loader_v2): Avoid cloning here, changed for now to avoid deadlocks.
+    /// Returns the reference of the struct name corresponding to the index. Here, we wrap the
+    /// name into an [Arc] to ensure that the lock is released.
+    pub(crate) fn idx_to_struct_name_ref(&self, idx: StructNameIndex) -> Arc<StructIdentifier> {
         self.0.read().backward_map[idx.0].clone()
     }
+
+    /// Returns the clone of the struct name corresponding to the index. The clone ensures that
+    /// the lock is released before the control returns to the caller.
+    pub(crate) fn idx_to_struct_name(&self, idx: StructNameIndex) -> StructIdentifier {
+        self.0.read().backward_map[idx.0].as_ref().clone()
+    }
+
+    /// Returns the struct tag corresponding to the struct name and the provided type arguments.
+    pub(crate) fn idx_to_struct_tag(
+        &self,
+        idx: StructNameIndex,
+        ty_args: Vec<TypeTag>,
+    ) -> StructTag {
+        let index_map = self.0.read();
+        let struct_name = index_map.backward_map[idx.0].as_ref();
+        StructTag {
+            address: *struct_name.module.address(),
+            module: struct_name.module.name().to_owned(),
+            name: struct_name.name.clone(),
+            type_args: ty_args,
+        }
+    }
+
+    #[cfg(test)]
+    fn len(self) -> usize {
+        let index_map = self.0.into_inner();
+        let forward_map_len = index_map.forward_map.len();
+        let backward_map_len = index_map.backward_map.len();
+        assert_eq!(forward_map_len, backward_map_len);
+        forward_map_len
+    }
 }
 
 impl Clone for StructNameIndexMap {
@@ -52,6 +105,8 @@ mod test {
     use move_core_types::{
         account_address::AccountAddress, identifier::Identifier, language_storage::ModuleId,
     };
+    use proptest::{arbitrary::any, collection::vec, proptest, strategy::Strategy};
+    use std::sync::Arc;
 
     fn make_struct_name(module_name: &str, struct_name: &str) -> StructIdentifier {
         let module = ModuleId::new(AccountAddress::ONE, Identifier::new(module_name).unwrap());
@@ -63,7 +118,7 @@ mod test {
     #[should_panic]
     fn test_index_map_must_contain_idx() {
         let struct_name_idx_map = StructNameIndexMap::empty();
-        let _ = struct_name_idx_map.idx_to_struct_name(StructNameIndex(0));
+        let _ = struct_name_idx_map.idx_to_struct_name_ref(StructNameIndex(0));
     }
 
     #[test]
@@ -81,15 +136,73 @@ mod test {
         assert_eq!(bar_idx.0, 1);
 
         // Check that struct names actually correspond to indices.
-        let returned_foo = struct_name_idx_map.idx_to_struct_name(foo_idx);
-        assert_eq!(&returned_foo, &foo);
-        let returned_bar = struct_name_idx_map.idx_to_struct_name(bar_idx);
-        assert_eq!(&returned_bar, &bar);
+        let returned_foo = struct_name_idx_map.idx_to_struct_name_ref(foo_idx);
+        assert_eq!(returned_foo.as_ref(), &foo);
+        let returned_bar = struct_name_idx_map.idx_to_struct_name_ref(bar_idx);
+        assert_eq!(returned_bar.as_ref(), &bar);
 
         // Re-check indices on second access.
         let foo_idx = struct_name_idx_map.struct_name_to_idx(foo);
         assert_eq!(foo_idx.0, 0);
         let bar_idx = struct_name_idx_map.struct_name_to_idx(bar);
         assert_eq!(bar_idx.0, 1);
+
+        assert_eq!(struct_name_idx_map.len(), 2);
+    }
+
+    fn struct_name_strategy() -> impl Strategy<Value = StructIdentifier> {
+        let address = any::<AccountAddress>();
+        let module_identifier = any::<Identifier>();
+        let struct_identifier = any::<Identifier>();
+        (address, module_identifier, struct_identifier).prop_map(|(a, m, name)| StructIdentifier {
+            module: ModuleId::new(a, m),
+            name,
+        })
+    }
+
+    proptest! {
+        #[test]
+        fn test_index_map_concurrent_arvitrary_struct_names(struct_names in vec(struct_name_strategy(), 30..100),
+        ) {
+            let struct_name_idx_map = Arc::new(StructNameIndexMap::empty());
+
+            // Each thread caches a struct name, and reads it to check if the cached result is
+            // still the same.
+            std::thread::scope(|s| {
+                for struct_name in &struct_names {
+                    s.spawn({
+                        let struct_name_idx_map = struct_name_idx_map.clone();
+                        move || {
+                            let idx = struct_name_idx_map.struct_name_to_idx(struct_name.clone());
+                            let actual_struct_name = struct_name_idx_map.idx_to_struct_name_ref(idx);
+                            assert_eq!(actual_struct_name.as_ref(), struct_name);
+                        }
+                    });
+                }
+            });
+        }
+    }
+
+    #[test]
+    fn test_index_map_concurrent_single_struct_name() {
+        let struct_name_idx_map = Arc::new(StructNameIndexMap::empty());
+        let struct_name = make_struct_name("foo", "Foo");
+
+        // Each threads tries to cache the same struct name.
+        std::thread::scope(|s| {
+            for _ in 0..50 {
+                s.spawn({
+                    let struct_name_idx_map = struct_name_idx_map.clone();
+                    let struct_name = struct_name.clone();
+                    move || {
+                        struct_name_idx_map.struct_name_to_idx(struct_name);
+                    }
+                });
+            }
+        });
+
+        // Only a single struct name mast be cached!
+        let struct_name_idx_map = Arc::into_inner(struct_name_idx_map).unwrap();
+        assert_eq!(struct_name_idx_map.len(), 1);
     }
 }

third_party/move/move-vm/types/src/loaded_data/runtime_access_specifier.rs

@@ -494,13 +494,13 @@ impl AddressSpecifierFunction {
 impl AccessInstance {
     pub fn new(
         kind: AccessKind,
-        resource: &StructIdentifier,
+        resource: StructIdentifier,
         instance: &[Type],
         address: AccountAddress,
     ) -> Option<Self> {
         Some(AccessInstance {
             kind,
-            resource: resource.clone(),
+            resource,
             instance: instance.to_vec(),
             address,
         })
@@ -511,15 +511,15 @@ impl AccessInstance {
         instance: &[Type],
         address: AccountAddress,
     ) -> Option<Self> {
-        Self::new(AccessKind::Reads, resource, instance, address)
+        Self::new(AccessKind::Reads, resource.clone(), instance, address)
     }
 
     pub fn write(
         resource: &StructIdentifier,
         instance: &[Type],
         address: AccountAddress,
     ) -> Option<Self> {
-        Self::new(AccessKind::Writes, resource, instance, address)
+        Self::new(AccessKind::Writes, resource.clone(), instance, address)
     }
 }