Newtype the definition maps to prevent iteration when it is not stable

external-crates/move/crates/move-cli/src/sandbox/utils/on_disk_state_view.rs

@@ -368,7 +368,7 @@ impl OnDiskStateView {
             };
 
             // Process each module and add its dependencies to the to_process list
-            for (_, module) in &pkg.modules {
+            for module in pkg.modules.values() {
                 let module = CompiledModule::deserialize_with_defaults(module).unwrap();
                 let deps = module
                     .immediate_dependencies()

external-crates/move/crates/move-transactional-test-runner/src/tasks.rs

@@ -297,7 +297,7 @@ impl<ExtraValueArgs: ParsableValue> PublishRunCommand<ExtraValueArgs> {
             .map(|call| {
                 let name = parse_qualified_module_access(&call[0])?;
                 let args: Vec<ParsedValue<ExtraValueArgs>> = call[1..]
-                    .into_iter()
+                    .iter()
                     .map(|arg| {
                         ParsedValue::<ExtraValueArgs>::parse(arg)
                             .map_err(|e| anyhow!("Failed to parse argument: {}", e))

external-crates/move/crates/move-vm-runtime/src/execution/dispatch_tables.rs

@@ -54,35 +54,39 @@ use std::{
 ///
 /// TODO(tzakian): The representation can be optimized to use a more efficient data structure for
 /// vtable/cross-package function resolution but we will keep it simple for now.
-#[derive(Debug, Clone)]
+#[derive(Debug)]
 pub struct VMDispatchTables {
     pub(crate) vm_config: Arc<VMConfig>,
     pub(crate) loaded_packages: BTreeMap<RuntimePackageId, Arc<Package>>,
     /// Representation of runtime type depths. This is separate from the underlying packages to
     /// avoid grabbing write-locks and toward the possibility these may change based on linkage
     /// (e.g., type ugrades or similar).
     /// [SAFETY] Ordering of inner maps is not guaranteed
-    pub(crate) type_depths: BTreeMap<RuntimePackageId, HashMap<IntraPackageKey, DepthFormula>>,
+    pub(crate) type_depths: BTreeMap<RuntimePackageId, DefinitionMap<DepthFormula>>,
     /// Defining ID Set -- a set of all defining IDs on any types mentioned in the package.
     /// [SAFETY] Ordering is not guaranteed
     #[allow(dead_code)]
-    pub(crate) defining_id_origins: HashMap<PackageStorageId, RuntimePackageId>,
+    pub(crate) defining_id_origins: BTreeMap<PackageStorageId, RuntimePackageId>,
 }
 
 /// A `PackageVTable` is a collection of pointers indexed by the module and name
 /// within the package.
 #[derive(Debug)]
 pub struct PackageVirtualTable {
     /// Representation of runtime functions.
-    /// [SAFETY] Ordering is not guaranteed
-    pub functions: HashMap<IntraPackageKey, VMPointer<Function>>,
+    pub functions: DefinitionMap<VMPointer<Function>>,
     /// Representation of runtime types.
-    /// [SAFETY] Ordering is not guaranteed
-    pub types: HashMap<IntraPackageKey, VMPointer<DatatypeDescriptor>>,
+    pub types: DefinitionMap<VMPointer<DatatypeDescriptor>>,
     /// Defining ID Set -- a set of all defining IDs on any types mentioned in the package.
     pub defining_ids: BTreeSet<PackageStorageId>,
 }
 
+/// This is a lookup-only map for recording information about module members in loaded package
+/// modules. It exposes an intentionally spartan interface to prevent any unexpected behavior
+/// (e.g., unstable iteration ordering) that Rust's standard collections run afoul of.
+#[derive(Debug)]
+pub struct DefinitionMap<Value>(HashMap<IntraPackageKey, Value>);
+
 /// runtime_address::module_name::function_name
 /// NB: This relies on no boxing -- if this introduces boxes, the arena allocation in the execution
 /// AST will leak memory.
@@ -128,12 +132,12 @@ pub struct DepthFormula {
 impl VMDispatchTables {
     /// Create a new RuntimeVTables instance.
     /// NOTE: This assumes linkage has already occured.
-    pub fn new(
+    pub(crate) fn new(
         vm_config: Arc<VMConfig>,
         loaded_packages: BTreeMap<RuntimePackageId, Arc<Package>>,
     ) -> VMResult<Self> {
         let defining_id_origins = {
-            let mut defining_id_map = HashMap::new();
+            let mut defining_id_map = BTreeMap::new();
             for (addr, pkg) in &loaded_packages {
                 for defining_id in &pkg.vtable.defining_ids {
                     if let Some(prev) = defining_id_map.insert(*defining_id, *addr) {
@@ -394,16 +398,10 @@ impl VMDispatchTables {
         let depth_formula =
             self.calculate_depth_of_datatype_and_cache(datatype_name, &mut depth_cache)?;
         for (datatype_key, depth) in depth_cache {
-            let _prev_depth = self
-                .type_depths
+            self.type_depths
                 .entry(datatype_key.package_key)
                 .or_default()
-                .insert(datatype_key.inner_pkg_key, depth);
-            debug_assert!(
-                _prev_depth.is_none(),
-                "recomputed type depth formula for {}",
-                datatype_key.to_string()?
-            );
+                .insert(datatype_key.inner_pkg_key, depth)?;
         }
         Ok(depth_formula)
     }
@@ -953,12 +951,73 @@ impl DepthFormula {
 // ------------------------------------------------------------------------
 
 impl PackageVirtualTable {
-    pub fn new() -> Self {
+    pub fn new(
+        functions: DefinitionMap<VMPointer<Function>>,
+        types: DefinitionMap<VMPointer<DatatypeDescriptor>>,
+    ) -> Self {
+        // [SAFETY] This is unordered, but okay because we are making a set anyway.
+        let defining_ids = types
+            .0
+            .values()
+            .map(|ty| ty.defining_id.address())
+            .collect::<BTreeSet<_>>()
+            .into_iter()
+            .copied()
+            .collect();
         Self {
-            functions: HashMap::new(),
-            types: HashMap::new(),
-            defining_ids: BTreeSet::new(),
+            functions,
+            types,
+            defining_ids,
+        }
+    }
+}
+
+impl<T> DefinitionMap<T> {
+    /// Create a new, empty DefinitionMap
+    pub fn empty() -> Self {
+        Self(HashMap::new())
+    }
+
+    /// Indicates if the DefinitionMap is empty or not.
+    pub fn is_empty(&self) -> bool {
+        self.0.is_empty()
+    }
+
+    /// Number of elements in the DefinitionMap
+    pub fn len(&self) -> usize {
+        self.0.len()
+    }
+
+    /// Creates a new DefintionMap, producing an error if a duplicate key is found.
+    pub fn from_unique(
+        items: impl IntoIterator<Item = (IntraPackageKey, T)>,
+    ) -> PartialVMResult<Self> {
+        let mut map = HashMap::new();
+        for (name, value) in items {
+            if map.insert(name, value).is_some() {
+                return Err(
+                    PartialVMError::new(StatusCode::UNKNOWN_INVARIANT_VIOLATION_ERROR)
+                        .with_message(format!("Duplicate key {}", name.to_string()?)),
+                );
+            }
+        }
+        Ok(Self(map))
+    }
+
+    /// Retrieve a key from the definition map.
+    pub fn get(&self, key: &IntraPackageKey) -> Option<&T> {
+        self.0.get(key)
+    }
+
+    /// Private to this module, as it is only used in depth formula calculations.
+    fn insert(&mut self, key: IntraPackageKey, value: T) -> PartialVMResult<()> {
+        if self.0.insert(key, value).is_some() {
+            return Err(
+                PartialVMError::new(StatusCode::UNKNOWN_INVARIANT_VIOLATION_ERROR)
+                    .with_message(format!("Duplicate key {}", key.to_string()?)),
+            );
         }
+        Ok(())
     }
 }
 
@@ -1020,9 +1079,9 @@ impl IntraPackageKey {
 // Default
 // -------------------------------------------------------------------------------------------------
 
-impl Default for PackageVirtualTable {
+impl<T> Default for DefinitionMap<T> {
     fn default() -> Self {
-        Self::new()
+        Self(HashMap::new())
     }
 }
 

external-crates/move/crates/move-vm-runtime/src/jit/execution/translate.rs

@@ -12,7 +12,7 @@ use crate::{
     },
     dbg_println,
     execution::{
-        dispatch_tables::{IntraPackageKey, PackageVirtualTable, VirtualTableKey},
+        dispatch_tables::{DefinitionMap, IntraPackageKey, PackageVirtualTable, VirtualTableKey},
         values::Value,
     },
     jit::{execution::ast::*, optimization::ast as input},
@@ -50,7 +50,10 @@ struct PackageContext<'natives> {
 
     // NB: All things except for types are allocated into this arena.
     pub package_arena: Arena,
-    pub vtable: PackageVirtualTable,
+
+    pub vtable_funs: DefinitionMap<VMPointer<Function>>,
+    pub vtable_types: DefinitionMap<VMPointer<DatatypeDescriptor>>,
+    // pub vtable: PackageVirtualTable,
 }
 
 struct FunctionContext<'pkg_ctxt, 'natives> {
@@ -77,47 +80,41 @@ struct Definitions {
 impl PackageContext<'_> {
     fn insert_vtable_functions(
         &mut self,
-        vtable: impl IntoIterator<Item = VMPointer<Function>>,
+        functions: impl IntoIterator<Item = VMPointer<Function>>,
     ) -> PartialVMResult<()> {
-        for func in vtable {
-            let fun_ref = func.ptr_clone().to_ref();
-            if self
-                .vtable
-                .functions
-                .insert(func.name.inner_pkg_key, func)
-                .is_some()
-            {
-                return Err(
-                    PartialVMError::new(StatusCode::UNKNOWN_INVARIANT_VIOLATION_ERROR)
-                        .with_message(format!(
-                            "Duplicate key {}::{}",
-                            self.storage_id,
-                            fun_ref.name.inner_pkg_key.to_string()?
-                        )),
-                );
-            }
-        }
+        let funs = functions
+            .into_iter()
+            .map(|ptr| {
+                let name = ptr.name.inner_pkg_key;
+                (name, ptr)
+            })
+            .collect::<Vec<_>>();
+        let vtable_funs = DefinitionMap::from_unique(funs)?;
+        debug_assert!(
+            self.vtable_funs.is_empty(),
+            "Non-empty function map during loading."
+        );
+        self.vtable_funs = vtable_funs;
         Ok(())
     }
 
     fn insert_vtable_datatypes(
         &mut self,
         datatype_descriptors: Vec<VMPointer<DatatypeDescriptor>>,
     ) -> PartialVMResult<()> {
-        for ptr in datatype_descriptors.into_iter() {
-            self.vtable.defining_ids.insert(*ptr.defining_id.address());
-            let name = ptr.intra_package_name();
-            if self.vtable.types.insert(name, ptr).is_some() {
-                return Err(
-                    PartialVMError::new(StatusCode::UNKNOWN_INVARIANT_VIOLATION_ERROR)
-                        .with_message(format!(
-                            "Duplicate key {}::{}",
-                            self.storage_id,
-                            name.to_string()?
-                        )),
-                );
-            }
-        }
+        let datatypes = datatype_descriptors
+            .into_iter()
+            .map(|ptr| {
+                let name = ptr.intra_package_name();
+                (name, ptr)
+            })
+            .collect::<Vec<_>>();
+        let vtable_types = DefinitionMap::from_unique(datatypes)?;
+        debug_assert!(
+            self.vtable_types.is_empty(),
+            "Non-empty type map during loading."
+        );
+        self.vtable_types = vtable_types;
         Ok(())
     }
 
@@ -133,8 +130,7 @@ impl PackageContext<'_> {
             return None;
         }
         // TODO(vm-rewrite): Have this return an error if the function was not found.
-        self.vtable
-            .functions
+        self.vtable_funs
             .get(&vtable_entry.inner_pkg_key)
             .map(|f| f.ptr_clone())
     }
@@ -213,7 +209,8 @@ pub fn package(
         runtime_id,
         loaded_modules: BTreeMap::new(),
         package_arena: Arena::new(),
-        vtable: PackageVirtualTable::new(),
+        vtable_funs: DefinitionMap::empty(),
+        vtable_types: DefinitionMap::empty(),
         type_origin_table,
     };
 
@@ -258,10 +255,13 @@ pub fn package(
         runtime_id,
         loaded_modules,
         package_arena,
-        vtable,
+        vtable_funs,
+        vtable_types,
         type_origin_table: _,
     } = package_context;
 
+    let vtable = PackageVirtualTable::new(vtable_funs, vtable_types);
+
     Ok(Package {
         storage_id,
         runtime_id,

external-crates/move/crates/move-vm-runtime/src/unit_tests/package_cache_tests.rs

@@ -158,10 +158,6 @@ fn load_package_external_package_calls_no_types() {
     assert_eq!(l_pkg.runtime.loaded_modules.len(), 1);
     assert_eq!(l_pkg.runtime.storage_id, package2_address);
     assert_eq!(l_pkg.runtime.vtable.functions.len(), 1);
-
-    for fptr in l_pkg.runtime.vtable.functions.values() {
-        println!("{:#?}", fptr.to_ref().code());
-    }
 }
 
 #[test]