Fix tree shaking alg

crates/sui/src/client_commands.rs

@@ -3169,37 +3169,61 @@ pub(crate) async fn pkg_tree_shake(
     with_unpublished_dependencies: bool,
     compiled_package: &mut CompiledPackage,
 ) -> Result<(), anyhow::Error> {
-    let pkgs = compiled_package.find_immediate_deps_pkgs_to_keep(with_unpublished_dependencies)?;
-    let pkg_ids = pkgs.values().cloned().collect::<Vec<_>>();
-    let pkg_id_to_name = pkgs
-        .iter()
-        .map(|(name, id)| (id, name))
-        .collect::<BTreeMap<_, _>>();
+    // these are packages that are immediate dependencies of the root package
+    let immediate_dep_packages =
+        compiled_package.find_immediate_deps_pkgs_to_keep(with_unpublished_dependencies)?;
+    let immediate_dep_pkgs_ids = immediate_dep_packages.values().cloned().collect::<Vec<_>>();
 
+    // if there's no lock file, the id will always be the one set in the [addresses] section
     let pkg_name_to_orig_id: BTreeMap<_, _> = compiled_package
         .package
         .deps_compiled_units
         .iter()
-        .map(|(pkg_name, module)| (pkg_name, ObjectID::from(module.unit.address.into_inner())))
+        .map(|(pkg_name, module)| (*pkg_name, ObjectID::from(module.unit.address.into_inner())))
         .collect();
 
-    let published_deps_packages = fetch_move_packages(read_api, &pkg_ids, &pkg_id_to_name).await?;
-
-    let linkage_table_ids: BTreeSet<_> = published_deps_packages
+    // need to find the trans dependencies of the immediate dep_pkgs
+    let pkg_id_to_name: BTreeMap<_, _> = immediate_dep_packages
         .iter()
-        .flat_map(|pkg| pkg.linkage_table().keys())
+        .map(|(name, id)| (id, name))
         .collect();
+    let immediate_dep_pkgs_linkage_tables: BTreeMap<_, _> =
+        fetch_move_packages(read_api, &immediate_dep_pkgs_ids, &pkg_id_to_name)
+            .await?
+            .iter()
+            .flat_map(|pkg| pkg.linkage_table())
+            .collect();
+
+    // for every immediate dep package, we need to use its linkage table to determine its
+    // transitive dependencies and ensure that we keep the required packages
+    let mut pkgs_to_keep: BTreeSet<&Symbol> = BTreeSet::new();
+    let published_pkgs = compiled_package.dependency_ids.published.clone();
+    for (linkage_orig_id, upgrade_info) in &immediate_dep_pkgs_linkage_tables {
+        // orig id
+        for (name, id) in &pkg_name_to_orig_id {
+            if *linkage_orig_id == id {
+                pkgs_to_keep.insert(name);
+            }
+        }
+
+        // for dependencies on pkgs that are upgrades, find the name of the published
+        // package by matching the upgraded package id with the published package id
+        for (name, id) in &published_pkgs {
+            if &upgrade_info.upgraded_id == id {
+                pkgs_to_keep.insert(name);
+            }
+        }
+    }
+
+    // need to also keep the immediate dep packages
+    for pkg in immediate_dep_packages.keys() {
+        pkgs_to_keep.insert(pkg);
+    }
 
     compiled_package
         .dependency_ids
         .published
-        .retain(|pkg_name, id| {
-            linkage_table_ids.contains(id)
-                || pkgs.contains_key(pkg_name)
-                || pkg_name_to_orig_id
-                    .get(pkg_name)
-                    .is_some_and(|orig_id| pkg_ids.contains(orig_id))
-        });
+        .retain(|pkg, _| pkgs_to_keep.contains(&pkg));
 
     Ok(())
 }

crates/sui/tests/cli_tests.rs

@@ -4522,6 +4522,51 @@ async fn test_tree_shaking_package_deps_on_pkg_upgrade_1() -> Result<(), anyhow:
     Ok(())
 }
 
+#[sim_test]
+async fn test_tree_shaking_package_deps_on_pkg_upgrade_2() -> Result<(), anyhow::Error> {
+    let mut test = TreeShakingTest::new().await?;
+
+    // Publish package K
+    let (package_k_id, cap) = test.publish_package("K", false).await?;
+    let package_path = test.package_path("K");
+    add_published_id_to_manifest(&package_path, &package_k_id, false)?;
+    // Upgrade package K (named K_v2)
+    std::fs::copy(
+        test.package_path("K").join("Move.lock"),
+        test.package_path("K_v2").join("Move.lock"),
+    )?;
+    let package_k_v2_id = test.upgrade_package("K_v2", cap).await?;
+
+    let package_path = test.package_path("K_v2");
+    add_published_id_to_manifest(&package_path, &package_k_v2_id, false)?;
+
+    let package_l_id = test
+        .publish_package_without_tree_shaking("L_depends_on_K")
+        .await;
+    let linkage_table_l = test.fetch_linkage_table(package_l_id).await;
+    assert!(
+        linkage_table_l.contains_key(&package_k_id),
+        "Package L should depend on K"
+    );
+
+    add_published_id_to_manifest(&test.package_path("L_depends_on_K"), &package_l_id, false)?;
+
+    let (package_m_id, _) = test
+        .publish_package("M_depends_on_L_and_K_v2_no_code_references_K_v2", false)
+        .await?;
+    let linkage_table_m = test.fetch_linkage_table(package_m_id).await;
+    assert!(
+        linkage_table_m.contains_key(&package_k_id),
+        "Package M should depend on K"
+    );
+
+    assert!(linkage_table_m
+        .get(&package_k_id)
+        .is_some_and(|x| x.upgraded_id == package_k_v2_id), "Package I should depend on A_v2 after upgrade, and the UpgradeInfo should have matching ids");
+
+    Ok(())
+}
+
 #[sim_test]
 async fn test_tree_shaking_package_system_deps() -> Result<(), anyhow::Error> {
     let mut test = TreeShakingTest::new().await?;

crates/sui/tests/data/tree_shaking/K/.gitignore

@@ -0,0 +1 @@
+build/*

crates/sui/tests/data/tree_shaking/K/Move.toml

@@ -0,0 +1,7 @@
+[package]
+name = "K"
+edition = "2024.beta"
+
+[addresses]
+k = "0x0"
+

crates/sui/tests/data/tree_shaking/K/sources/k.move

@@ -0,0 +1,8 @@
+// Copyright (c) Mysten Labs, Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+module k::k {
+    public fun k() {
+        let x = 1;
+    }
+}

crates/sui/tests/data/tree_shaking/K_v2/.gitignore

@@ -0,0 +1 @@
+build/*

crates/sui/tests/data/tree_shaking/K_v2/Move.toml

@@ -0,0 +1,7 @@
+[package]
+name = "K"
+edition = "2024.beta"
+
+[addresses]
+k = "0x0"
+

crates/sui/tests/data/tree_shaking/K_v2/sources/k.move

@@ -0,0 +1,12 @@
+// Copyright (c) Mysten Labs, Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+module k::k {
+    public fun k() {
+        let x = 1;
+    }
+
+    public fun k2() {
+        let x = 1;
+    }
+}

crates/sui/tests/data/tree_shaking/L_depends_on_K/.gitignore

@@ -0,0 +1 @@
+build/*

crates/sui/tests/data/tree_shaking/L_depends_on_K/Move.toml

@@ -0,0 +1,10 @@
+[package]
+name = "L_depends_on_K"
+edition = "2024.beta"
+
+[dependencies]
+K = { local = "../K" }
+
+[addresses]
+l_depends_on_k = "0x0"
+

crates/sui/tests/data/tree_shaking/L_depends_on_K/sources/l_depends_on_k.move

@@ -0,0 +1,8 @@
+// Copyright (c) Mysten Labs, Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+module l_depends_on_k::l_depends_on_k {
+    public fun l() {
+        let x = 1;
+    }
+}

crates/sui/tests/data/tree_shaking/M_depends_on_L_and_K_v2_no_code_references_K_v2/.gitignore

@@ -0,0 +1 @@
+build/*

crates/sui/tests/data/tree_shaking/M_depends_on_L_and_K_v2_no_code_references_K_v2/Move.toml

@@ -0,0 +1,10 @@
+[package]
+name = "M_depends_on_L_and_K_v2_no_code_references_K_v2"
+edition = "2024.beta"
+
+[dependencies]
+K = { local = "../K_v2", override = true }
+L_depends_on_K = { local = "../L_depends_on_K" }
+
+[addresses]
+m_depends_on_l_and_k_v2_no_code_references_k_v2 = "0x0"

crates/sui/tests/data/tree_shaking/M_depends_on_L_and_K_v2_no_code_references_K_v2/sources/m_depends_on_l_and_k_v2_no_code_references_k_v2.move

@@ -0,0 +1,9 @@
+// Copyright (c) Mysten Labs, Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+module m_depends_on_l_and_k_v2_no_code_references_k_v2::m_depends_on_l_and_k_v2_no_code_references_k_v2 {
+    public fun m() {
+        let k = 1;
+        l_depends_on_k::l_depends_on_k::l();
+    }
+}