From b28b84f6d805d909527d5cccc521dd62f6e3a8da Mon Sep 17 00:00:00 2001 From: stefan-mysten <135084671+stefan-mysten@users.noreply.github.com> Date: Fri, 28 Feb 2025 13:50:44 -0800 Subject: [PATCH] Fix tree shaking alg --- crates/sui/src/client_commands.rs | 60 +++++++++++++------ crates/sui/tests/cli_tests.rs | 45 ++++++++++++++ .../sui/tests/data/tree_shaking/K/.gitignore | 1 + .../sui/tests/data/tree_shaking/K/Move.toml | 7 +++ .../tests/data/tree_shaking/K/sources/k.move | 8 +++ .../tests/data/tree_shaking/K_v2/.gitignore | 1 + .../tests/data/tree_shaking/K_v2/Move.toml | 7 +++ .../data/tree_shaking/K_v2/sources/k.move | 12 ++++ .../tree_shaking/L_depends_on_K/.gitignore | 1 + .../tree_shaking/L_depends_on_K/Move.toml | 10 ++++ .../sources/l_depends_on_k.move | 8 +++ .../.gitignore | 1 + .../Move.toml | 10 ++++ ...on_l_and_k_v2_no_code_references_k_v2.move | 9 +++ 14 files changed, 162 insertions(+), 18 deletions(-) create mode 100644 crates/sui/tests/data/tree_shaking/K/.gitignore create mode 100644 crates/sui/tests/data/tree_shaking/K/Move.toml create mode 100644 crates/sui/tests/data/tree_shaking/K/sources/k.move create mode 100644 crates/sui/tests/data/tree_shaking/K_v2/.gitignore create mode 100644 crates/sui/tests/data/tree_shaking/K_v2/Move.toml create mode 100644 crates/sui/tests/data/tree_shaking/K_v2/sources/k.move create mode 100644 crates/sui/tests/data/tree_shaking/L_depends_on_K/.gitignore create mode 100644 crates/sui/tests/data/tree_shaking/L_depends_on_K/Move.toml create mode 100644 crates/sui/tests/data/tree_shaking/L_depends_on_K/sources/l_depends_on_k.move create mode 100644 crates/sui/tests/data/tree_shaking/M_depends_on_L_and_K_v2_no_code_references_K_v2/.gitignore create mode 100644 crates/sui/tests/data/tree_shaking/M_depends_on_L_and_K_v2_no_code_references_K_v2/Move.toml create mode 100644 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 diff --git a/crates/sui/src/client_commands.rs b/crates/sui/src/client_commands.rs index 7a883de745618..df4af77a7e3e7 100644 --- a/crates/sui/src/client_commands.rs +++ b/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::>(); - let pkg_id_to_name = pkgs - .iter() - .map(|(name, id)| (id, name)) - .collect::>(); + // 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::>(); + // 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(()) } diff --git a/crates/sui/tests/cli_tests.rs b/crates/sui/tests/cli_tests.rs index 1266054dc45a2..75114dd406c62 100644 --- a/crates/sui/tests/cli_tests.rs +++ b/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?; diff --git a/crates/sui/tests/data/tree_shaking/K/.gitignore b/crates/sui/tests/data/tree_shaking/K/.gitignore new file mode 100644 index 0000000000000..a007feab071f4 --- /dev/null +++ b/crates/sui/tests/data/tree_shaking/K/.gitignore @@ -0,0 +1 @@ +build/* diff --git a/crates/sui/tests/data/tree_shaking/K/Move.toml b/crates/sui/tests/data/tree_shaking/K/Move.toml new file mode 100644 index 0000000000000..677b468c23ff5 --- /dev/null +++ b/crates/sui/tests/data/tree_shaking/K/Move.toml @@ -0,0 +1,7 @@ +[package] +name = "K" +edition = "2024.beta" + +[addresses] +k = "0x0" + diff --git a/crates/sui/tests/data/tree_shaking/K/sources/k.move b/crates/sui/tests/data/tree_shaking/K/sources/k.move new file mode 100644 index 0000000000000..8a93cbd04b2e7 --- /dev/null +++ b/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; + } +} diff --git a/crates/sui/tests/data/tree_shaking/K_v2/.gitignore b/crates/sui/tests/data/tree_shaking/K_v2/.gitignore new file mode 100644 index 0000000000000..a007feab071f4 --- /dev/null +++ b/crates/sui/tests/data/tree_shaking/K_v2/.gitignore @@ -0,0 +1 @@ +build/* diff --git a/crates/sui/tests/data/tree_shaking/K_v2/Move.toml b/crates/sui/tests/data/tree_shaking/K_v2/Move.toml new file mode 100644 index 0000000000000..677b468c23ff5 --- /dev/null +++ b/crates/sui/tests/data/tree_shaking/K_v2/Move.toml @@ -0,0 +1,7 @@ +[package] +name = "K" +edition = "2024.beta" + +[addresses] +k = "0x0" + diff --git a/crates/sui/tests/data/tree_shaking/K_v2/sources/k.move b/crates/sui/tests/data/tree_shaking/K_v2/sources/k.move new file mode 100644 index 0000000000000..7cdda4c8eef43 --- /dev/null +++ b/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; + } +} diff --git a/crates/sui/tests/data/tree_shaking/L_depends_on_K/.gitignore b/crates/sui/tests/data/tree_shaking/L_depends_on_K/.gitignore new file mode 100644 index 0000000000000..a007feab071f4 --- /dev/null +++ b/crates/sui/tests/data/tree_shaking/L_depends_on_K/.gitignore @@ -0,0 +1 @@ +build/* diff --git a/crates/sui/tests/data/tree_shaking/L_depends_on_K/Move.toml b/crates/sui/tests/data/tree_shaking/L_depends_on_K/Move.toml new file mode 100644 index 0000000000000..74364ae48942e --- /dev/null +++ b/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" + diff --git a/crates/sui/tests/data/tree_shaking/L_depends_on_K/sources/l_depends_on_k.move b/crates/sui/tests/data/tree_shaking/L_depends_on_K/sources/l_depends_on_k.move new file mode 100644 index 0000000000000..af2bc3ae703ed --- /dev/null +++ b/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; + } +} diff --git a/crates/sui/tests/data/tree_shaking/M_depends_on_L_and_K_v2_no_code_references_K_v2/.gitignore b/crates/sui/tests/data/tree_shaking/M_depends_on_L_and_K_v2_no_code_references_K_v2/.gitignore new file mode 100644 index 0000000000000..a007feab071f4 --- /dev/null +++ b/crates/sui/tests/data/tree_shaking/M_depends_on_L_and_K_v2_no_code_references_K_v2/.gitignore @@ -0,0 +1 @@ +build/* diff --git a/crates/sui/tests/data/tree_shaking/M_depends_on_L_and_K_v2_no_code_references_K_v2/Move.toml b/crates/sui/tests/data/tree_shaking/M_depends_on_L_and_K_v2_no_code_references_K_v2/Move.toml new file mode 100644 index 0000000000000..8125bfb8557da --- /dev/null +++ b/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" diff --git a/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 b/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 new file mode 100644 index 0000000000000..98d2725a434ee --- /dev/null +++ b/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(); + } +}