Improve memory efficiency of incremental dexing and introduce flag to error out if DexArchiveAspect misses .jar files

--
PiperOrigin-RevId: 148835665
MOS_MIGRATED_REVID=148835665

3 files changed, 69 insertions, 26 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/rules/android/AndroidBinary.java b/src/main/java/com/google/devtools/build/lib/rules/android/AndroidBinary.java
index ff22844118..21722e9f9b 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/android/AndroidBinary.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/android/AndroidBinary.java
@@ -1459,10 +1459,10 @@ public abstract class AndroidBinary implements RuleConfiguredTargetFactory {
   }
 
   /**
-   * Returns a {@link DexArchiveProvider} of all transitively generated dex archives as well as dex
-   * archives for the Jars produced by the binary target itself.
+   * Returns a {@link Map} of all transitively generated dex archives as well as dex archives for
+   * the Jars produced by the binary target itself.
    */
-  private static Function<Artifact, Artifact> collectDexArchives(
+  private static Map<Artifact, Artifact> collectDexArchives(
       RuleContext ruleContext,
       AndroidCommon common,
       List<String> dexopts,
@@ -1507,7 +1507,7 @@ public abstract class AndroidBinary implements RuleConfiguredTargetFactory {
       JavaTargetAttributes attributes,
       Function<Artifact, Artifact> derivedJarFunction,
       @Nullable Artifact mainDexList)
-      throws InterruptedException {
+      throws InterruptedException, RuleErrorException {
     checkArgument(mainDexList == null || shards.size() > 1);
     checkArgument(proguardedJar == null || inclusionFilterJar == null);
     Artifact javaResourceJar =
@@ -1548,8 +1548,25 @@ public abstract class AndroidBinary implements RuleConfiguredTargetFactory {
         // Use dex archives instead of their corresponding Jars wherever we can.  At this point
         // there should be very few or no Jar files that still end up in shards.  The dexing
         // step below will have to deal with those in addition to merging .dex files together.
-        classpath = Iterables.transform(classpath,
-            collectDexArchives(ruleContext, common, dexopts, semantics, derivedJarFunction));
+        Map<Artifact, Artifact> dexArchives =
+            collectDexArchives(ruleContext, common, dexopts, semantics, derivedJarFunction);
+        ImmutableList.Builder<Artifact> dexedClasspath = ImmutableList.builder();
+        boolean reportMissing =
+            AndroidCommon.getAndroidConfig(ruleContext).incrementalDexingErrorOnMissedJars();
+        for (Artifact jar : classpath) {
+          Artifact dexArchive = dexArchives.get(jar);
+          if (reportMissing && dexArchive == null) {
+            // Users can create this situation by directly depending on a .jar artifact (checked in
+            // or coming from a genrule or similar, b/11285003).  This will also catch new  implicit
+            // dependencies that incremental dexing would need to be extended to (b/34949364).
+            ruleContext.throwWithAttributeError("deps", "Dependencies on .jar artifacts are not "
+                + "allowed in Android binaries, please use a java_import to depend on "
+                + jar.prettyPrint() + ". If this is an implicit dependency then Bazel will need "
+                + "to be fixed to account for it correctly.");
+          }
+          dexedClasspath.add(dexArchive != null ? dexArchive : jar);
+        }
+        classpath = dexedClasspath.build();
         shardCommandLine.add("--split_dexed_classes");
       } else {
         classpath = Iterables.transform(classpath, derivedJarFunction);
diff --git a/src/main/java/com/google/devtools/build/lib/rules/android/AndroidConfiguration.java b/src/main/java/com/google/devtools/build/lib/rules/android/AndroidConfiguration.java
index eaae47df85..bfdd604f41 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/android/AndroidConfiguration.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/android/AndroidConfiguration.java
@@ -331,6 +331,18 @@ public class AndroidConfiguration extends BuildConfiguration.Fragment {
       help = "Do not use.")
     public boolean incrementalDexingForLiteProtos;
 
+    /**
+     * Whether to error out when we find Jar files when building binaries that weren't converted to
+     * a dex archive. Once this option works, we'll flip the default value in a config file, then
+     * once it is proven that it works, remove it from Bazel and said config file.
+     */
+    @Option(
+      name = "experimental_incremental_dexing_error_on_missed_jars",
+      defaultValue = "false",
+      category = "experimental",
+      help = "Do not use.")
+    public boolean incrementalDexingErrorOnMissedJars;
+
     @Option(name = "non_incremental_per_target_dexopts",
         converter = Converters.CommaSeparatedOptionListConverter.class,
         defaultValue = "--no-locals",
@@ -452,6 +464,7 @@ public class AndroidConfiguration extends BuildConfiguration.Fragment {
       host.incrementalDexing = incrementalDexing;
       host.incrementalDexingBinaries = incrementalDexingBinaries;
       host.incrementalDexingForLiteProtos = incrementalDexingForLiteProtos;
+      host.incrementalDexingErrorOnMissedJars = incrementalDexingErrorOnMissedJars;
       host.nonIncrementalPerTargetDexopts = nonIncrementalPerTargetDexopts;
       host.dexoptsSupportedInIncrementalDexing = dexoptsSupportedInIncrementalDexing;
       host.manifestMerger = manifestMerger;
@@ -500,6 +513,7 @@ public class AndroidConfiguration extends BuildConfiguration.Fragment {
   private final boolean jackSanityChecks;
   private final ImmutableSet<AndroidBinaryType> incrementalDexingBinaries;
   private final boolean incrementalDexingForLiteProtos;
+  private final boolean incrementalDexingErrorOnMissedJars;
   private final ImmutableList<String> dexoptsSupportedInIncrementalDexing;
   private final ImmutableList<String> targetDexoptsThatPreventIncrementalDexing;
   private final boolean desugarJava8;
@@ -527,6 +541,7 @@ public class AndroidConfiguration extends BuildConfiguration.Fragment {
       this.incrementalDexingBinaries = ImmutableSet.of();
     }
     this.incrementalDexingForLiteProtos = options.incrementalDexingForLiteProtos;
+    this.incrementalDexingErrorOnMissedJars = options.incrementalDexingErrorOnMissedJars;
     this.dexoptsSupportedInIncrementalDexing =
         ImmutableList.copyOf(options.dexoptsSupportedInIncrementalDexing);
     this.targetDexoptsThatPreventIncrementalDexing =
@@ -588,6 +603,14 @@ public class AndroidConfiguration extends BuildConfiguration.Fragment {
   }
 
   /**
+   * Returns whether to report an error when Jars that weren't converted to dex archives are part
+   * of an android binary.
+   */
+  public boolean incrementalDexingErrorOnMissedJars() {
+    return incrementalDexingErrorOnMissedJars;
+  }
+
+  /**
    * dx flags supported in incremental dexing actions.
    */
   public ImmutableList<String> getDexoptsSupportedInIncrementalDexing() {
@@ -637,7 +660,7 @@ public class AndroidConfiguration extends BuildConfiguration.Fragment {
   public boolean useSingleJarForMultidex() {
     return useSingleJarForMultidex;
   }
-  
+
   public boolean useResourcePrefiltering() {
     return useResourcePrefiltering;
   }
diff --git a/src/main/java/com/google/devtools/build/lib/rules/android/DexArchiveProvider.java b/src/main/java/com/google/devtools/build/lib/rules/android/DexArchiveProvider.java
index 237f2a16ed..f04594193b 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/android/DexArchiveProvider.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/android/DexArchiveProvider.java
@@ -18,18 +18,18 @@ import static com.google.common.base.Preconditions.checkNotNull;
 
 import com.google.common.base.Function;
 import com.google.common.collect.HashBasedTable;
-import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableTable;
 import com.google.common.collect.Table;
 import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.analysis.TransitiveInfoProvider;
+import com.google.devtools.build.lib.collect.nestedset.NestedSet;
+import com.google.devtools.build.lib.collect.nestedset.NestedSetBuilder;
 import com.google.devtools.build.lib.concurrent.ThreadSafety.Immutable;
-
+import java.util.HashMap;
+import java.util.Map;
 import java.util.Set;
 
-import javax.annotation.Nullable;
-
 /**
  * Provider of transitively available dex archives corresponding to Jars.  A dex archive is a zip of
  * {@code .dex} files that each encode exactly one {@code .class} file in an Android-readable form.
@@ -56,6 +56,8 @@ public class DexArchiveProvider implements TransitiveInfoProvider {
 
     private final Table<ImmutableSet<String>, Artifact, Artifact> dexArchives =
         HashBasedTable.create();
+    private final NestedSetBuilder<ImmutableTable<ImmutableSet<String>, Artifact, Artifact>>
+        transitiveDexArchives = NestedSetBuilder.stableOrder();
 
     public Builder() {
     }
@@ -66,7 +68,7 @@ public class DexArchiveProvider implements TransitiveInfoProvider {
      */
     public Builder addTransitiveProviders(Iterable<DexArchiveProvider> providers) {
       for (DexArchiveProvider provider : providers) {
-        dexArchives.putAll(provider.dexArchives);
+        transitiveDexArchives.addTransitive(provider.dexArchives);
       }
       return this;
     }
@@ -95,28 +97,29 @@ public class DexArchiveProvider implements TransitiveInfoProvider {
      * Returns the finished {@link DexArchiveProvider}.
      */
     public DexArchiveProvider build() {
-      return new DexArchiveProvider(ImmutableTable.copyOf(dexArchives));
+      return new DexArchiveProvider(
+          transitiveDexArchives.add(ImmutableTable.copyOf(dexArchives)).build());
     }
   }
 
   /** Map from Jar artifacts to the corresponding dex archives. */
-  private final ImmutableTable<ImmutableSet<String>, Artifact, Artifact> dexArchives;
+  private final NestedSet<ImmutableTable<ImmutableSet<String>, Artifact, Artifact>> dexArchives;
 
-  private DexArchiveProvider(ImmutableTable<ImmutableSet<String>, Artifact, Artifact> dexArchives) {
+  private DexArchiveProvider(
+      NestedSet<ImmutableTable<ImmutableSet<String>, Artifact, Artifact>> dexArchives) {
     this.dexArchives = dexArchives;
   }
 
-  public Function<Artifact, Artifact> archivesForDexopts(ImmutableSet<String> dexopts) {
-    final ImmutableMap<Artifact, Artifact> dexArchivesForDexopts = dexArchives.row(dexopts);
-    return new Function<Artifact, Artifact>() {
-      /** Maps Jars to available dex archives and returns the given Jar otherwise. */
-      @Override
-      @Nullable
-      public Artifact apply(@Nullable Artifact jar) {
-        Artifact dexArchive = dexArchivesForDexopts.get(jar);
-        return dexArchive != null ? dexArchive : jar; // return null iff input == null
-      }
-    };
+  /**
+   * Returns a flat map from Jars to dex archives transitively produced for the given dexopts.
+   */
+  public Map<Artifact, Artifact> archivesForDexopts(ImmutableSet<String> dexopts) {
+    // Can't use ImmutableMap because we can encounter the same key-value pair multiple times
+    HashMap<Artifact, Artifact> result = new HashMap<>();
+    for (ImmutableTable<ImmutableSet<String>, Artifact, Artifact> partialMapping : dexArchives) {
+      result.putAll(partialMapping.row(dexopts));
+    }
+    return result;
   }
 
   @Override