[vm-rewrite] Rootless linkage

external-crates/move/crates/move-cli/src/sandbox/commands/publish.rs

@@ -14,7 +14,7 @@ use move_vm_runtime::{
     runtime::MoveRuntime,
     shared::{linkage_context::LinkageContext, types::PackageStorageId},
 };
-use std::collections::{BTreeSet, HashMap};
+use std::collections::{BTreeMap, BTreeSet};
 
 pub fn publish(
     natives: impl IntoIterator<Item = NativeFunctionRecord>,
@@ -55,7 +55,7 @@ pub fn publish(
         .collect::<Vec<_>>();
 
     // Build the dependency map from the package
-    let mut dependency_map = HashMap::new();
+    let mut dependency_map = BTreeMap::new();
     for (name, unit) in package.deps_compiled_units.iter() {
         let unit_address = *unit.unit.module.self_id().address();
         if let Some(other) = dependency_map.insert(unit_address, unit_address) {
@@ -77,7 +77,7 @@ pub fn publish(
     let mut gas_status = get_gas_status(cost_table, None)?;
 
     // Create a `LinkageContext`
-    let linkage_context = LinkageContext::new(package_storage_id, dependency_map);
+    let linkage_context = LinkageContext::new(dependency_map);
 
     // Serialize the modules into a package to prepare them for publishing
     let pkg = StoredPackage::from_module_for_testing_with_linkage(

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

@@ -411,7 +411,6 @@ impl OnDiskStateView {
             "Found circular dependencies during dependency generation."
         );
         let linkage_context = LinkageContext::new(
-            *package_address,
             all_dependencies
                 .into_iter()
                 .map(|id| (id, id))

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

@@ -41,11 +41,7 @@ use move_vm_runtime::{
 };
 use once_cell::sync::Lazy;
 
-use std::{
-    collections::{BTreeMap, HashMap},
-    path::Path,
-    sync::Arc,
-};
+use std::{collections::BTreeMap, path::Path, sync::Arc};
 
 const STD_ADDR: AccountAddress = AccountAddress::ONE;
 
@@ -87,24 +83,29 @@ impl Linkage {
                 .add_entry(*runtime_id, *storage_id)
                 .map_err(|err| err.finish(Location::Undefined))?;
         }
-        // Remap the root address if it is in the linkage table since it may be overriden
-        existing_linkage.root_package = existing_linkage
-            .linkage_table
-            .get(&existing_linkage.root_package)
-            .unwrap_or(&existing_linkage.root_package)
-            .clone();
         Ok(existing_linkage)
     }
 }
 
 impl PublishLinkageArgs {
     pub fn overlay(&self, existing_linkage: LinkageContext) -> VMResult<LinkageContext> {
-        let mut linkage = self.linkage.overlay(existing_linkage)?;
-        if let Some(location) = self.location {
-            linkage.root_package = location;
-        }
+        self.linkage.overlay(existing_linkage)
+    }
 
-        Ok(linkage)
+    pub fn resolve_publication_location(
+        &self,
+        sender: AccountAddress,
+        linkage: &LinkageContext,
+    ) -> AccountAddress {
+        // 1. Use location if specified; otherwise
+        // 2. Use the sender address (remapped if it exists in the linkage table).
+        self.location.unwrap_or(
+            linkage
+                .linkage_table
+                .get(&sender)
+                .copied()
+                .expect("Sender address not found in linkage table"),
+        )
     }
 }
 
@@ -218,8 +219,7 @@ impl<'a> MoveTestAdapter<'a> for SimpleRuntimeTestAdapter {
                 let move_stdlib = MOVE_STDLIB_COMPILED.to_vec();
                 let sender = *move_stdlib.first().unwrap().self_id().address();
                 println!("generating stdlib linkage");
-                let linkage_context =
-                    LinkageContext::new(sender, HashMap::from([(sender, sender)]));
+                let linkage_context = LinkageContext::new(BTreeMap::from([(sender, sender)]));
                 println!("calling stdlib publish with address {sender:?}");
                 let pkg = StoredPackage::from_module_for_testing_with_linkage(
                     sender,
@@ -271,16 +271,19 @@ impl<'a> MoveTestAdapter<'a> for SimpleRuntimeTestAdapter {
             sender,
             &pub_modules,
         )?)?;
+        let storage_id = extra_args.resolve_publication_location(sender, &linkage_context);
         println!("linkage: {linkage_context:#?}");
+        println!("publication location: {storage_id}");
         println!("doing publish");
         let pkg = StoredPackage::from_module_for_testing_with_linkage(
-            sender,
+            storage_id,
             linkage_context,
             pub_modules.clone(),
         )?;
+        let runtime_id = pkg.runtime_id();
         let publish_result = self.perform_action(gas_budget, |inner_adapter, _gas_status| {
             let pkg = pkg.into_serialized_package();
-            inner_adapter.publish_package(sender, pkg)?;
+            inner_adapter.publish_package(runtime_id, pkg)?;
             println!("done");
             Ok(())
         });
@@ -322,18 +325,21 @@ impl<'a> MoveTestAdapter<'a> for SimpleRuntimeTestAdapter {
             sender,
             &pub_modules,
         )?)?;
+        let storage_id = extra_args.resolve_publication_location(sender, &linkage_context);
         println!("linkage: {linkage_context:#?}");
+        println!("publication location: {storage_id}");
         let mut gas_meter = Self::make_gas_status(gas_budget);
         println!("doing verification");
         let pkg = StoredPackage::from_module_for_testing_with_linkage(
-            sender,
+            storage_id,
             linkage_context,
             pub_modules.clone(),
         )?;
+        let runtime_id = pkg.runtime_id();
         let pkg = pkg.into_serialized_package();
         let (verif_pkg, mut publish_vm) = self
             .perform_action_with_gas(&mut gas_meter, |inner_adapter, _gas_status| {
-                inner_adapter.verify_package(sender, pkg)
+                inner_adapter.verify_package(runtime_id, pkg)
             })?;
         println!("doing calls");
         let signers: Vec<_> = signers

external-crates/move/crates/move-vm-runtime/src/dev_utils/compilation_utils.rs

@@ -114,6 +114,6 @@ pub fn compile_packages(
 ) -> BTreeMap<RuntimePackageId, StoredPackage> {
     compile_packages_in_file(filename, dependencies)
         .into_iter()
-        .map(|p| (p.linkage_context.root_package(), p))
+        .map(|p| (p.package_id, p))
         .collect()
 }

external-crates/move/crates/move-vm-runtime/src/dev_utils/in_memory_test_adapter.rs

@@ -123,7 +123,7 @@ impl VMTestAdapter<InMemoryStorage> for InMemoryTestAdapter {
             // TODO: VM error instead?
             panic!("Did not find runtime ID {runtime_id} in linkage context.");
         };
-        assert!(storage_id == package.storage_id);
+        assert_eq!(storage_id, package.storage_id);
         let mut gas_meter = GasStatus::new_unmetered();
         self.runtime.validate_package(
             &self.storage,
@@ -204,7 +204,6 @@ impl VMTestAdapter<InMemoryStorage> for InMemoryTestAdapter {
             "Found circular dependencies during linkage generation."
         );
         let linkage_context = LinkageContext::new(
-            storage_id,
             all_dependencies
                 .into_iter()
                 .map(|id| (id, id))

external-crates/move/crates/move-vm-runtime/src/dev_utils/storage.rs

@@ -3,7 +3,10 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use crate::{
-    shared::{linkage_context::LinkageContext, types::PackageStorageId},
+    shared::{
+        linkage_context::LinkageContext,
+        types::{PackageStorageId, RuntimePackageId},
+    },
     validation::verification::ast as verif_ast,
 };
 use anyhow::Result;
@@ -15,7 +18,7 @@ use move_core_types::{
     language_storage::ModuleId,
     resolver::{ModuleResolver, MoveResolver, SerializedPackage, TypeOrigin},
 };
-use std::collections::{BTreeMap, HashMap};
+use std::collections::BTreeMap;
 
 // -------------------------------------------------------------------------------------------------
 // Types
@@ -36,6 +39,8 @@ pub struct DeltaStorage<'a, 'b, S> {
 /// Simple in-memory representation of packages
 #[derive(Debug, Clone)]
 pub struct StoredPackage {
+    /// The package ID (address) for this package.
+    pub package_id: PackageStorageId,
     pub modules: BTreeMap<Identifier, Vec<u8>>,
     /// For each dependency (including transitive dependencies), maps runtime package ID to the
     /// storage ID of the package that is to be used for the linkage rooted at this package.
@@ -70,20 +75,21 @@ impl<'a, 'b, S: MoveResolver> DeltaStorage<'a, 'b, S> {
 impl StoredPackage {
     fn empty(storage_id: AccountAddress) -> Self {
         Self {
+            package_id: storage_id,
             modules: BTreeMap::new(),
-            linkage_context: LinkageContext::new(storage_id, HashMap::new()),
+            linkage_context: LinkageContext::new(BTreeMap::new()),
             type_origin_table: vec![],
         }
     }
 
     pub fn from_modules_for_testing(
-        storage_id: AccountAddress,
+        storage_id: PackageStorageId,
         modules: Vec<CompiledModule>,
     ) -> Result<Self> {
         assert!(!modules.is_empty());
         // Map the modules in this package to `storage_id` and generate the identity linkage for
         // all deps.
-        let mut linkage_table = HashMap::new();
+        let mut linkage_table = BTreeMap::new();
         let type_origin_table = generate_type_origins(storage_id, &modules);
         let modules: BTreeMap<_, _> = modules
             .into_iter()
@@ -104,14 +110,15 @@ impl StoredPackage {
             .collect::<Result<_>>()?;
 
         Ok(Self {
+            package_id: storage_id,
             modules,
-            linkage_context: LinkageContext::new(storage_id, linkage_table),
+            linkage_context: LinkageContext::new(linkage_table),
             type_origin_table,
         })
     }
 
     pub fn from_module_for_testing_with_linkage(
-        storage_id: AccountAddress,
+        storage_id: PackageStorageId,
         linkage_context: LinkageContext,
         modules: Vec<CompiledModule>,
     ) -> Result<Self> {
@@ -126,6 +133,7 @@ impl StoredPackage {
             .collect::<Result<_>>()?;
 
         Ok(Self {
+            package_id: storage_id,
             modules,
             linkage_context,
             type_origin_table,
@@ -134,6 +142,7 @@ impl StoredPackage {
 
     pub fn from_verified_package(verified_package: verif_ast::Package) -> Self {
         Self {
+            package_id: verified_package.storage_id,
             modules: verified_package
                 .as_modules()
                 .into_iter()
@@ -147,7 +156,6 @@ impl StoredPackage {
                 })
                 .collect(),
             linkage_context: LinkageContext::new(
-                verified_package.storage_id,
                 verified_package.linkage_table.into_iter().collect(),
             ),
             type_origin_table: verified_package.type_origin_table,
@@ -156,12 +164,26 @@ impl StoredPackage {
 
     pub fn into_serialized_package(self) -> SerializedPackage {
         SerializedPackage {
-            storage_id: self.linkage_context.root_package(),
+            storage_id: self.package_id,
             modules: self.modules.into_values().collect(),
             linkage_table: self.linkage_context.linkage_table.into_iter().collect(),
             type_origin_table: self.type_origin_table,
         }
     }
+
+    pub fn runtime_id(&self) -> RuntimePackageId {
+        self.linkage_context
+            .linkage_table
+            .iter()
+            .find_map(|(k, v)| {
+                if *v == self.package_id {
+                    Some(*k)
+                } else {
+                    None
+                }
+            })
+            .expect("address not found in linkage table")
+    }
 }
 
 pub fn generate_type_origins(
@@ -206,10 +228,8 @@ impl InMemoryStorage {
     }
 
     pub fn publish_package(&mut self, stored_package: StoredPackage) {
-        self.accounts.insert(
-            stored_package.linkage_context.root_package(),
-            stored_package,
-        );
+        self.accounts
+            .insert(stored_package.package_id, stored_package);
     }
 
     pub fn publish_or_overwrite_module(

external-crates/move/crates/move-vm-runtime/src/dev_utils/vm_test_adapter.rs

@@ -13,7 +13,7 @@ use crate::{
 };
 use move_binary_format::{errors::VMResult, file_format::CompiledModule};
 use move_core_types::resolver::{MoveResolver, SerializedPackage};
-use std::collections::HashMap;
+use std::collections::BTreeMap;
 
 // FIXME(cswords): support gas
 
@@ -76,10 +76,7 @@ pub trait VMTestAdapter<Storage: MoveResolver + Sync + Send> {
     /// Retrieve the linkage context for the given package in `Storage`.
     fn get_linkage_context(&self, package_id: PackageStorageId) -> VMResult<LinkageContext> {
         let pkg = self.get_package_from_store(&package_id)?;
-        Ok(LinkageContext::new(
-            package_id,
-            HashMap::from_iter(pkg.linkage_table),
-        ))
+        Ok(LinkageContext::new(BTreeMap::from_iter(pkg.linkage_table)))
     }
 
     fn get_package_from_store(&self, package_id: &PackageStorageId) -> VMResult<SerializedPackage>;