From 99c7d4e8ad91a4c64775c39db90ea230956fe090 Mon Sep 17 00:00:00 2001
From: Vladimir Kryachko <vkryachko@google.com>
Date: Thu, 24 Nov 2022 13:56:00 -0500
Subject: [PATCH] Enable dead code elimination for vendored deps.

---
 .../firebase/gradle/plugins/VendorPlugin.kt   | 44 ++++++++++++++++---
 .../firebase/gradle/plugins/VendorTests.kt    | 11 ++---
 .../transport-runtime.gradle                  |  6 +++
 3 files changed, 46 insertions(+), 15 deletions(-)

diff --git a/buildSrc/src/main/java/com/google/firebase/gradle/plugins/VendorPlugin.kt b/buildSrc/src/main/java/com/google/firebase/gradle/plugins/VendorPlugin.kt
index d7a9caf0872..da1a0659fbd 100644
--- a/buildSrc/src/main/java/com/google/firebase/gradle/plugins/VendorPlugin.kt
+++ b/buildSrc/src/main/java/com/google/firebase/gradle/plugins/VendorPlugin.kt
@@ -33,18 +33,31 @@ import org.gradle.api.Plugin
 import org.gradle.api.Project
 import org.gradle.api.artifacts.Configuration
 import org.gradle.api.logging.Logger
+import org.gradle.api.provider.Property
+import org.gradle.api.provider.Provider
 import org.gradle.kotlin.dsl.apply
+import org.gradle.kotlin.dsl.create
+
+abstract class VendorExtension {
+  /** Controls dead code elimination, enabled if true. */
+  abstract val optimize: Property<Boolean>
+
+  init {
+    optimize.convention(true)
+  }
+}
 
 class VendorPlugin : Plugin<Project> {
   override fun apply(project: Project) {
+    val vendorConfig = project.extensions.create<VendorExtension>("vendor")
     project.plugins.all {
       when (this) {
-        is LibraryPlugin -> configureAndroid(project)
+        is LibraryPlugin -> configureAndroid(project, vendorConfig)
       }
     }
   }
 
-  fun configureAndroid(project: Project) {
+  fun configureAndroid(project: Project, vendorConfig: VendorExtension) {
     project.apply(plugin = "LicenseResolverPlugin")
 
     val vendor = project.configurations.create("vendor")
@@ -71,7 +84,8 @@ class VendorPlugin : Plugin<Project> {
           },
           jarJarProvider = { jarJar.resolve() },
           project = project,
-          logger = project.logger
+          logger = project.logger,
+          optimize = vendorConfig.optimize
         ),
         logger = project.logger
       )
@@ -80,19 +94,35 @@ class VendorPlugin : Plugin<Project> {
 }
 
 interface JarTransformer {
-  fun transform(inputJar: File, outputJar: File, packagesToVendor: Set<String>)
+  fun transform(
+    inputJar: File,
+    outputJar: File,
+    ownPackages: Set<String>,
+    packagesToVendor: Set<String>
+  )
 }
 
 class JarJarTransformer(
   private val parentPackageProvider: () -> String,
   private val jarJarProvider: () -> Collection<File>,
   private val project: Project,
-  private val logger: Logger
+  private val logger: Logger,
+  private val optimize: Provider<Boolean>
 ) : JarTransformer {
-  override fun transform(inputJar: File, outputJar: File, packagesToVendor: Set<String>) {
+  override fun transform(
+    inputJar: File,
+    outputJar: File,
+    ownPackages: Set<String>,
+    packagesToVendor: Set<String>
+  ) {
     val parentPackage = parentPackageProvider()
     val rulesFile = File.createTempFile(parentPackage, ".jarjar")
     rulesFile.printWriter().use {
+      if (optimize.get()) {
+        for (packageName in ownPackages) {
+          it.println("keep $packageName.**")
+        }
+      }
       for (externalPackageName in packagesToVendor) {
         it.println("rule $externalPackageName.** $parentPackage.@0")
       }
@@ -227,7 +257,7 @@ class VendorTransform(
     zipAll(unzippedDir, jar)
     val outputJar = File(workDir, "output.jar")
 
-    jarTransformer.transform(jar, outputJar, externalPackageNames)
+    jarTransformer.transform(jar, outputJar, ownPackageNames, externalPackageNames)
     return outputJar
   }
 
diff --git a/buildSrc/src/test/kotlin/com/google/firebase/gradle/plugins/VendorTests.kt b/buildSrc/src/test/kotlin/com/google/firebase/gradle/plugins/VendorTests.kt
index ad556bf11f0..de722c12f75 100644
--- a/buildSrc/src/test/kotlin/com/google/firebase/gradle/plugins/VendorTests.kt
+++ b/buildSrc/src/test/kotlin/com/google/firebase/gradle/plugins/VendorTests.kt
@@ -86,11 +86,11 @@ class VendorPluginTests {
       )
 
     // ImmutableList is not used, so it should be stripped out.
-    assertThat(classes).doesNotContain("com/google/common/collect/ImmutableList.class")
+    assertThat(classes).doesNotContain("com/example/com/google/common/collect/ImmutableList.class")
   }
 
   @Test
-  fun `vendor dagger excluding javax transitive deps and not using it should include dagger`() {
+  fun `vendor dagger excluding javax transitive deps and not using it should not include dagger`() {
     val classes =
       buildWith(
         """
@@ -113,12 +113,7 @@ class VendorPluginTests {
         )
       )
     // expected classes
-    assertThat(classes)
-      .containsAtLeast(
-        "com/example/Hello.class",
-        "com/example/BuildConfig.class",
-        "com/example/dagger/Lazy.class"
-      )
+    assertThat(classes).containsExactly("com/example/Hello.class", "com/example/BuildConfig.class")
   }
 
   @Test
diff --git a/transport/transport-runtime/transport-runtime.gradle b/transport/transport-runtime/transport-runtime.gradle
index 2272e16ba95..334959a4f1a 100644
--- a/transport/transport-runtime/transport-runtime.gradle
+++ b/transport/transport-runtime/transport-runtime.gradle
@@ -67,6 +67,12 @@ firebaseLibrary {
     }
 }
 
+vendor {
+    // Integration tests use dagger classes that are not used in the main SDK,
+    // so we disable dead code elimination to ensure those classes are preserved.
+    optimize = false
+}
+
 android {
     compileSdkVersion project.targetSdkVersion
     defaultConfig {