New alg for tree shaking

crates/sui-move-build/src/lib.rs

@@ -619,15 +619,12 @@ impl CompiledPackage {
         self.dependency_ids.published.values().cloned().collect()
     }
 
-    /// Tree-shake the package's dependencies to remove any that are not referenced in source code.
-    ///
-    /// This algorithm uses the set of root modules as the starting point to retrieve the
-    /// list of used packages that are immediate dependencies of these modules. Essentially, it
-    /// will remove any package that has no immediate module dependency to it.
-    ///
-    /// Then, it will recursively find all the transitive dependencies of the packages in the list
-    /// above and add them to the list of packages that need to be kept as dependencies.
-    pub fn tree_shake(&mut self, with_unpublished_deps: bool) -> Result<(), anyhow::Error> {
+    /// Find the map of packages that are immediate dependencies of the root modules, joined with
+    /// the set of bytecode dependencies.
+    pub fn find_immediate_deps_pkgs_to_keep(
+        &self,
+        with_unpublished_deps: bool,
+    ) -> Result<BTreeMap<Symbol, ObjectID>, anyhow::Error> {
         // Start from the root modules (or all modules if with_unpublished_deps is true as we
         // need to include modules with 0x0 address)
         let root_modules: Vec<_> = if with_unpublished_deps {
@@ -644,15 +641,14 @@ impl CompiledPackage {
         };
 
         // Find the immediate dependencies for each root module and store the package name
-        // in the used_immediate_packages set. This basically prunes the packages that are not used
+        // in the pkgs_to_keep set. This basically prunes the packages that are not used
         // based on the modules information.
         let mut pkgs_to_keep: BTreeSet<Symbol> = BTreeSet::new();
         let module_to_pkg_name: BTreeMap<_, _> = self
             .package
             .all_modules()
             .map(|m| (m.unit.module.self_id(), m.unit.package_name))
             .collect();
-        let mut used_immediate_packages: BTreeSet<Symbol> = BTreeSet::new();
 
         for module in &root_modules {
             let immediate_deps = module.module.immediate_dependencies();
@@ -661,31 +657,24 @@ impl CompiledPackage {
                     let Some(pkg_name) = pkg_name else {
                         bail!("Expected a package name but it's None")
                     };
-                    used_immediate_packages.insert(*pkg_name);
+                    pkgs_to_keep.insert(*pkg_name);
                 }
             }
         }
 
-        // Next, for each package from used_immediate_packages set, we need to find all the
-        // transitive dependencies. Those trans dependencies need to be included in the final list
-        // of package dependencies (note that the pkg itself will be added by the transitive deps
-        // function)
-        used_immediate_packages.iter().for_each(|pkg| {
-            self.dependency_graph
-                .add_transitive_dependencies(pkg, &mut pkgs_to_keep)
-        });
-
         // If a package depends on another published package that has only bytecode without source
         // code available, we need to include also that package as dep.
         pkgs_to_keep.extend(self.bytecode_deps.iter().map(|(name, _)| *name));
 
-        // Finally, filter out packages that are not in the union above from the
-        // dependency_ids.published field and return the package ids.
-        self.dependency_ids
+        // Finally, filter out packages that are published and exist in the manifest at the
+        // compilation time but are not referenced in the source code.
+        Ok(self
+            .dependency_ids
+            .clone()
             .published
-            .retain(|pkg_name, _| pkgs_to_keep.contains(pkg_name));
-
-        Ok(())
+            .into_iter()
+            .filter(|(pkg_name, _)| pkgs_to_keep.contains(&pkg_name))
+            .collect())
     }
 }
 

crates/sui-move/src/build.rs

@@ -67,7 +67,7 @@ impl Build {
         generate_struct_layouts: bool,
         chain_id: Option<String>,
     ) -> anyhow::Result<()> {
-        let mut pkg = BuildConfig {
+        let pkg = BuildConfig {
             config,
             run_bytecode_verifier: true,
             print_diags_to_stderr: true,
@@ -80,7 +80,6 @@ impl Build {
                 check_unpublished_dependencies(&pkg.dependency_ids.unpublished)?;
             }
 
-            pkg.tree_shake(with_unpublished_deps)?;
             println!(
                 "{}",
                 json!({

crates/sui/src/client_commands.rs

@@ -72,7 +72,7 @@ use sui_types::{
     gas_coin::GasCoin,
     message_envelope::Envelope,
     metrics::BytecodeVerifierMetrics,
-    move_package::UpgradeCap,
+    move_package::{MovePackage, UpgradeCap},
     object::Owner,
     parse_sui_type_tag,
     signature::GenericSignature,
@@ -954,9 +954,16 @@ impl SuiClientCommands {
                     )?;
                 }
 
-                let (upgrade_policy, compiled_package) =
+                let (upgrade_policy, mut compiled_package) =
                     upgrade_result.map_err(|e| anyhow!("{e}"))?;
 
+                filter_deps(
+                    &client,
+                    with_unpublished_dependencies,
+                    &mut compiled_package,
+                )
+                .await?;
+
                 let compiled_modules =
                     compiled_package.get_package_bytes(with_unpublished_dependencies);
                 let package_id = compiled_package.published_at.clone()?;
@@ -1080,7 +1087,13 @@ impl SuiClientCommands {
                     )?;
                 }
 
-                let compiled_package = compile_result?;
+                let mut compiled_package = compile_result?;
+                filter_deps(
+                    &client,
+                    with_unpublished_dependencies,
+                    &mut compiled_package,
+                )
+                .await?;
                 let compiled_modules =
                     compiled_package.get_package_bytes(with_unpublished_dependencies);
                 let dep_ids = compiled_package.get_published_dependencies_ids();
@@ -1774,9 +1787,7 @@ fn compile_package_simple(
         chain_id: chain_id.clone(),
     };
     let resolution_graph = config.resolution_graph(package_path, chain_id.clone())?;
-    let mut compiled_package =
-        build_from_resolution_graph(resolution_graph, false, false, chain_id)?;
-    compiled_package.tree_shake(false)?;
+    let compiled_package = build_from_resolution_graph(resolution_graph, false, false, chain_id)?;
 
     Ok(compiled_package)
 }
@@ -1790,15 +1801,14 @@ pub(crate) async fn upgrade_package(
     skip_dependency_verification: bool,
     env_alias: Option<String>,
 ) -> Result<(u8, CompiledPackage), anyhow::Error> {
-    let mut compiled_package = compile_package(
+    let compiled_package = compile_package(
         read_api,
         build_config,
         package_path,
         with_unpublished_dependencies,
         skip_dependency_verification,
     )
     .await?;
-    compiled_package.tree_shake(with_unpublished_dependencies)?;
 
     compiled_package.published_at.as_ref().map_err(|e| match e {
         PublishedAtError::NotPresent => {
@@ -1876,13 +1886,12 @@ pub(crate) async fn compile_package(
     if !with_unpublished_dependencies {
         check_unpublished_dependencies(&dependencies.unpublished)?;
     };
-    let mut compiled_package = build_from_resolution_graph(
+    let compiled_package = build_from_resolution_graph(
         resolution_graph,
         run_bytecode_verifier,
         print_diags_to_stderr,
         chain_id,
     )?;
-    compiled_package.tree_shake(with_unpublished_dependencies)?;
     let protocol_config = read_api.get_protocol_config(None).await?;
 
     // Check that the package's Move version is compatible with the chain's
@@ -3098,3 +3107,70 @@ async fn check_protocol_version_and_warn(client: &SuiClient) -> Result<(), anyho
 
     Ok(())
 }
+
+/// Fetch move packages based on the provided package IDs.
+pub async fn fetch_move_packages(
+    client: &SuiClient,
+    package_ids: &[ObjectID],
+) -> Result<Vec<MovePackage>, anyhow::Error> {
+    let objects = client
+        .read_api()
+        .multi_get_object_with_options(package_ids.to_vec(), SuiObjectDataOptions::bcs_lossless())
+        .await
+        .map_err(|e| anyhow!("{e}"))?;
+
+    objects
+        .into_iter()
+        .map(|o| {
+            let o = o.into_object().map_err(|e| anyhow!("{e}"))?;
+            let Some(SuiRawData::Package(p)) = o.bcs else {
+                bail!("Expected SuiRawData::Package but got something else");
+            };
+            p.to_move_package(u64::MAX /* safe as this pkg comes from the network */)
+                .map_err(|e| anyhow!("{e}"))
+        })
+        .collect()
+}
+
+/// Filter out a package's dependencies which are not referenced in the source code.
+async fn filter_deps(
+    client: &SuiClient,
+    with_unpublished_dependencies: bool,
+    compiled_package: &mut CompiledPackage,
+) -> Result<(), anyhow::Error> {
+    let pkgs = compiled_package.find_deps_pkgs_to_keep(with_unpublished_dependencies)?;
+    let pkg_ids = pkgs
+        .clone()
+        .into_iter()
+        .map(|(_, id)| id)
+        .collect::<Vec<_>>();
+
+    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())))
+        .collect();
+
+    let published_deps_packages = fetch_move_packages(client, &pkg_ids).await?;
+    let linkage_table_ids = published_deps_packages
+        .iter()
+        .flat_map(|pkg| {
+            let linkage_table = pkg.linkage_table();
+            linkage_table.keys().cloned().collect::<Vec<_>>()
+        })
+        .collect::<Vec<_>>();
+
+    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))
+        });
+
+    Ok(())
+}

crates/sui/tests/cli_tests.rs

@@ -104,7 +104,9 @@ impl TreeShakingTest {
         let temp_dir = tempfile::Builder::new().prefix("tree_shaking").tempdir()?;
         std::fs::create_dir_all(temp_dir.path()).unwrap();
         let tests_dir = PathBuf::from(TEST_DATA_DIR);
+        let framework_pkgs = PathBuf::from("../sui-framework/packages");
         copy_dir_all(tests_dir, temp_dir.path())?;
+        copy_dir_all(framework_pkgs, temp_dir.path())?;
 
         Ok(Self {
             test_cluster,
@@ -4330,7 +4332,7 @@ async fn test_tree_shaking_package_with_unused_dependency() -> Result<(), anyhow
 }
 
 #[sim_test]
-async fn test_tree_shaking_package_with_transitive_dependencies() -> Result<(), anyhow::Error> {
+async fn test_tree_shaking_package_with_transitive_dependencies1() -> Result<(), anyhow::Error> {
     let mut test = TreeShakingTest::new().await?;
 
     // Publish packages A and B
@@ -4519,3 +4521,19 @@ async fn test_tree_shaking_package_deps_on_pkg_upgrade_1() -> Result<(), anyhow:
 
     Ok(())
 }
+
+#[sim_test]
+async fn test_tree_shaking_package_system_deps() -> Result<(), anyhow::Error> {
+    let mut test = TreeShakingTest::new().await?;
+
+    // Publish package J and verify empty linkage table
+    let (package_j_id, _) = test.publish_package("J_system_deps", false).await?;
+    let move_pkg_j = fetch_move_packages(&test.client, vec![package_j_id]).await;
+    let linkage_table_j = move_pkg_j.first().unwrap().linkage_table();
+    assert!(
+        linkage_table_j.is_empty(),
+        "Package J should have no dependencies"
+    );
+
+    Ok(())
+}

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

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

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

@@ -0,0 +1,10 @@
+[package]
+name = "J_system_deps"
+edition = "2024.beta"
+
+[dependencies]
+Sui = { local = "../../sui-framework", override = true }
+
+[addresses]
+j_system_deps = "0x0"
+

crates/sui/tests/data/tree_shaking/J_system_deps/sources/j_system_deps.move

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