Add support in Blaze for running a monolithic "optimizing" dexer action.

PiperOrigin-RevId: 538401036
Change-Id: I244bfb99bd21c21367ac88ca197a916c1d034088

src/main/java/com/google/devtools/build/lib/rules/android/AndroidBinary.java

@@ -507,13 +507,16 @@ public static RuleConfiguredTargetBuilder createAndroidBinary(
     boolean postprocessingRewritesMap = androidSemantics.postprocessClassesRewritesMap(ruleContext);
     boolean desugarJava8LibsGeneratesMap =
         AndroidCommon.getAndroidConfig(ruleContext).desugarJava8Libs();
+    boolean optimizingDexing = ruleContext.getExecutablePrerequisite(":optimizing_dexer") != null;
     if (generateProguardMap) {
       // Determine the output of the Proguard map from shrinking the app. This depends on the
       // additional steps which can process the map before the final Proguard map artifact is
       // generated.
       if (!hasProguardSpecs && !postprocessingRewritesMap) {
         // When no shrinking happens a generating rule for the output map artifact is still needed.
         proguardOutputMap = androidSemantics.getProguardOutputMap(ruleContext);
+      } else if (optimizingDexing) {
+        proguardOutputMap = ProguardHelper.getProguardTempArtifact(ruleContext, "pre_dexing.map");
       } else if (postprocessingRewritesMap) {
         // Proguard map from shrinking goes to postprocessing.
         proguardOutputMap =
@@ -545,13 +548,13 @@ public static RuleConfiguredTargetBuilder createAndroidBinary(
     if (proguardOutput.hasMapping()) {
       // Determine the Proguard map artifacts for the additional steps (if any) if shrinking of
       // the app is enabled.
-      if (postprocessingRewritesMap && desugarJava8LibsGeneratesMap) {
+      if ((optimizingDexing || postprocessingRewritesMap) && desugarJava8LibsGeneratesMap) {
         // Proguard map from preprocessing will be merged with Proguard map for desugared
         // library.
         postProcessingOutputMap =
             getDxArtifact(ruleContext, "_proguard_output_for_desugared_library.map");
         finalProguardOutputMap = androidSemantics.getProguardOutputMap(ruleContext);
-      } else if (postprocessingRewritesMap) {
+      } else if (optimizingDexing || postprocessingRewritesMap) {
         // No desugared library, Proguard map from preprocessing is the final Proguard map.
         postProcessingOutputMap = androidSemantics.getProguardOutputMap(ruleContext);
         finalProguardOutputMap = postProcessingOutputMap;
@@ -592,7 +595,10 @@ public static RuleConfiguredTargetBuilder createAndroidBinary(
             resourceApk.getMainDexProguardConfig(),
             resourceClasses,
             derivedJarFunction,
-            proguardOutputMap);
+            proguardOutputMap,
+            postProcessingOutputMap,
+            proguardOutput.getLibraryJar(),
+            !proguardSpecs.isEmpty());
 
     DexPostprocessingOutput dexPostprocessingOutput =
         androidSemantics.postprocessClassesDexZip(
@@ -1120,11 +1126,12 @@ private static ProguardOutput createEmptyProguardAction(
     ProguardOutput outputs =
         ProguardHelper.getProguardOutputs(
             proguardOutputJar,
-            /*proguardSeeds=*/ null,
-            /*proguardUsage=*/ null,
+            /* proguardSeeds= */ null,
+            /* proguardUsage= */ null,
             ruleContext,
             semantics,
-            proguardOutputMap);
+            proguardOutputMap,
+            null);
     outputs.addAllToSet(failures);
     ruleContext.registerAction(
         new FailAction(
@@ -1318,8 +1325,12 @@ private static DexingOutput dex(
       @Nullable Artifact mainDexProguardSpec,
       JavaTargetAttributes attributes,
       Function<Artifact, Artifact> derivedJarFunction,
-      @Nullable Artifact proguardOutputMap)
+      @Nullable Artifact proguardOutputMap,
+      @Nullable Artifact postProcessingOutputMap,
+      @Nullable Artifact libraryJar,
+      boolean isOptimizedBuild)
       throws InterruptedException, RuleErrorException {
+    FilesToRunProvider optimizingDexer = ruleContext.getExecutablePrerequisite(":optimizing_dexer");
     List<String> dexopts = ruleContext.getExpander().withDataLocations().tokenized("dexopts");
     MultidexMode multidexMode = getMultidexMode(ruleContext);
     if (!supportsMultidexMode(ruleContext, multidexMode)) {
@@ -1370,7 +1381,55 @@ private static DexingOutput dex(
     }
 
     Artifact classesDex = getDxArtifact(ruleContext, "classes.dex.zip");
-    if (dexShards > 1) {
+    if (optimizingDexer != null && isOptimizedBuild) {
+      SpawnAction.Builder dexAction =
+          new SpawnAction.Builder()
+              .useDefaultShellEnvironment()
+              .setExecutable(optimizingDexer)
+              .setProgressMessage("Optimized dexing for %{label}")
+              .setMnemonic("OptimizingDex")
+              .addInput(proguardedJar)
+              .addOutput(classesDex);
+
+      boolean nativeMultidex = multidexMode == MultidexMode.NATIVE;
+      CustomCommandLine.Builder dexCommand =
+          CustomCommandLine.builder()
+              .addExecPath(proguardedJar)
+              .add("--release")
+              .add("--no-desugaring")
+              .addExecPath("--output", classesDex)
+              .addAll(dexopts);
+
+      if (proguardOutputMap != null) {
+        dexAction.addInput(proguardOutputMap).addOutput(postProcessingOutputMap);
+        dexCommand
+            .addExecPath("--pg-map", proguardOutputMap)
+            .addExecPath("--pg-map-output", postProcessingOutputMap);
+      }
+
+      // TODO(b/261110876): Pass min SDK through here based on the value in the merged manifest. The
+      // current value is statically defined for the entire depot.
+      // We currently set the minimum SDK version to 21 if you are doing native multidex as that is
+      // required for native multidex to work in the first place and as a result is required for
+      // correct behavior from the dexer.
+      int sdk = nativeMultidex ? Math.max(21, minSdkVersion) : minSdkVersion;
+      if (sdk != 0) {
+        dexCommand.add("--min-api", Integer.toString(sdk));
+      }
+      if (mainDexList != null) {
+        dexCommand.addExecPath("--main-dex-list", mainDexList);
+        dexAction.addInput(mainDexList);
+      }
+      if (libraryJar != null) {
+        dexCommand.addExecPath("--lib", libraryJar);
+        dexAction.addInput(libraryJar);
+      }
+
+      dexAction.addCommandLine(dexCommand.build());
+      ruleContext.registerAction(dexAction.build(ruleContext));
+      return new DexingOutput(
+          classesDex, javaResourceSourceJar, ImmutableList.of(classesDex), mainDexList);
+    } else if (dexShards > 1) {
       ImmutableList<Artifact> shards =
           makeShardArtifacts(ruleContext, dexShards, usesDexArchives ? ".jar.dex.zip" : ".jar");
 

src/main/java/com/google/devtools/build/lib/rules/android/AndroidConfiguration.java

@@ -1060,6 +1060,15 @@ public static class Options extends FragmentOptions {
                 + " dex when compiling legacy multidex.")
     public Label legacyMainDexListGenerator;
 
+    @Option(
+        name = "optimizing_dexer",
+        defaultValue = "null",
+        converter = LabelConverter.class,
+        documentationCategory = OptionDocumentationCategory.UNCATEGORIZED,
+        effectTags = {OptionEffectTag.UNKNOWN},
+        help = "Specifies a binary to use to do dexing without sharding.")
+    public Label optimizingDexer;
+
     @Option(
         name = "experimental_disable_instrumentation_manifest_merge",
         defaultValue = "false",
@@ -1185,6 +1194,7 @@ public FragmentOptions getExec() {
   private final boolean useRTxtFromMergedResources;
   private final boolean outputLibraryMergedAssets;
   private final Label legacyMainDexListGenerator;
+  private final Label optimizingDexer;
   private final boolean disableInstrumentationManifestMerging;
   private final boolean incompatibleUseToolchainResolution;
   private final boolean hwasan;
@@ -1249,6 +1259,7 @@ public AndroidConfiguration(BuildOptions buildOptions) throws InvalidConfigurati
     this.useRTxtFromMergedResources = options.useRTxtFromMergedResources;
     this.outputLibraryMergedAssets = options.outputLibraryMergedAssets;
     this.legacyMainDexListGenerator = options.legacyMainDexListGenerator;
+    this.optimizingDexer = options.optimizingDexer;
     this.disableInstrumentationManifestMerging = options.disableInstrumentationManifestMerging;
     this.incompatibleUseToolchainResolution = options.incompatibleUseToolchainResolution;
     this.hwasan = options.hwasan;
@@ -1562,4 +1573,13 @@ boolean outputLibraryMergedAssets() {
   public Label getLegacyMainDexListGenerator() {
     return legacyMainDexListGenerator;
   }
+
+  /** Returns the label provided with --optimizing_dexer, if any. */
+  @StarlarkConfigurationField(
+      name = "optimizing_dexer",
+      doc = "Returns the label provided with --optimizing_dexer, if any.")
+  @Nullable
+  public Label getOptimizingDexer() {
+    return optimizingDexer;
+  }
 }

src/main/java/com/google/devtools/build/lib/rules/android/AndroidRuleClasses.java

@@ -223,6 +223,13 @@ public static LabelLateBoundDefault<?> getAndroidSdkLabel(Label androidSdk) {
           null,
           (rule, attributes, androidConfig) -> androidConfig.getLegacyMainDexListGenerator());
 
+  @SerializationConstant
+  static final LabelLateBoundDefault<AndroidConfiguration> OPTIMIZING_DEXER =
+      LabelLateBoundDefault.fromTargetConfiguration(
+          AndroidConfiguration.class,
+          null,
+          (rule, attributes, androidConfig) -> androidConfig.getOptimizingDexer());
+
   public static final FileType ANDROID_IDL = FileType.of(".aidl");
 
   public static final String[] ALLOWED_DEPENDENCIES = {
@@ -902,6 +909,12 @@ public Object getDefault(AttributeMap rule) {
                   .cfg(ExecutionTransitionFactory.createFactory())
                   .value(LEGACY_MAIN_DEX_LIST_GENERATOR)
                   .exec())
+          // This comes from the --optimizing_dexer flag.
+          .add(
+              attr(":optimizing_dexer", LABEL)
+                  .cfg(ExecutionTransitionFactory.createFactory())
+                  .value(OPTIMIZING_DEXER)
+                  .exec())
           .removeAttribute("data")
           .advertiseStarlarkProvider(StarlarkProviderIdentifier.forKey(ApkInfo.PROVIDER.getKey()))
           .advertiseStarlarkProvider(StarlarkProviderIdentifier.forKey(JavaInfo.PROVIDER.getKey()))

src/main/java/com/google/devtools/build/lib/rules/android/ProguardHelper.java

@@ -61,6 +61,7 @@ public static final class ProguardOutput {
     @Nullable private final Artifact seeds;
     @Nullable private final Artifact usage;
     @Nullable private final Artifact constantStringObfuscatedMapping;
+    @Nullable private final Artifact libraryJar;
     private final Artifact config;
 
     public ProguardOutput(
@@ -70,18 +71,20 @@ public ProguardOutput(
         @Nullable Artifact seeds,
         @Nullable Artifact usage,
         @Nullable Artifact constantStringObfuscatedMapping,
+        @Nullable Artifact libraryJar,
         Artifact config) {
       this.outputJar = checkNotNull(outputJar);
       this.mapping = mapping;
       this.protoMapping = protoMapping;
       this.seeds = seeds;
       this.usage = usage;
       this.constantStringObfuscatedMapping = constantStringObfuscatedMapping;
+      this.libraryJar = libraryJar;
       this.config = config;
     }
 
     public static ProguardOutput createEmpty(Artifact outputJar) {
-      return new ProguardOutput(outputJar, null, null, null, null, null, null);
+      return new ProguardOutput(outputJar, null, null, null, null, null, null, null);
     }
 
     public Artifact getOutputJar() {
@@ -117,6 +120,11 @@ public Artifact getUsage() {
       return usage;
     }
 
+    @Nullable
+    public Artifact getLibraryJar() {
+      return libraryJar;
+    }
+
     public Artifact getConfig() {
       return config;
     }
@@ -247,7 +255,8 @@ public static ProguardOutput getProguardOutputs(
       @Nullable Artifact proguardUsage,
       RuleContext ruleContext,
       JavaSemantics semantics,
-      @Nullable Artifact proguardOutputMap)
+      @Nullable Artifact proguardOutputMap,
+      @Nullable Artifact libraryJar)
       throws InterruptedException {
     boolean mappingRequested = genProguardMapping(ruleContext.attributes());
 
@@ -274,6 +283,7 @@ public static ProguardOutput getProguardOutputs(
         proguardSeeds,
         proguardUsage,
         proguardConstantStringMap,
+        libraryJar,
         proguardConfigOutput);
   }
 
@@ -312,14 +322,7 @@ public static ProguardOutput createOptimizationActions(
       throws InterruptedException {
     Preconditions.checkArgument(!proguardSpecs.isEmpty());
 
-    ProguardOutput output =
-        getProguardOutputs(
-            proguardOutputJar,
-            proguardSeeds,
-            proguardUsage,
-            ruleContext,
-            semantics,
-            proguardOutputMap);
+    Artifact libraryJar = null;
 
     if (!libraryJars.isEmpty() && !libraryJars.isSingleton()) {
       JavaTargetAttributes attributes = new JavaTargetAttributes.Builder(semantics).build();
@@ -331,28 +334,42 @@ public static ProguardOutput createOptimizationActions(
           .addRuntimeJars(libraryJars)
           .build();
       libraryJars = NestedSetBuilder.create(Order.STABLE_ORDER, combinedLibraryJar);
+      libraryJar = combinedLibraryJar;
+    } else if (libraryJars.isSingleton()) {
+      libraryJar = libraryJars.getSingleton();
     }
 
     boolean filterLibraryJarWithProgramJar =
         ruleContext.getFragment(AndroidConfiguration.class).filterLibraryJarWithProgramJar();
 
     if (filterLibraryJarWithProgramJar) {
       Preconditions.checkState(libraryJars.isSingleton());
-      Artifact libraryJar = libraryJars.getSingleton();
+      Artifact singletonLibraryJar = libraryJars.getSingleton();
 
       Artifact filteredLibraryJar =
           getProguardTempArtifact(ruleContext, "combined_library_jars_filtered.jar");
 
       new ZipFilterBuilder(ruleContext)
-          .setInputZip(libraryJar)
+          .setInputZip(singletonLibraryJar)
           .setOutputZip(filteredLibraryJar)
           .addFilterZips(ImmutableList.of(programJar))
           .setCheckHashMismatchMode(ZipFilterBuilder.CheckHashMismatchMode.NONE)
           .build();
 
       libraryJars = NestedSetBuilder.create(Order.STABLE_ORDER, filteredLibraryJar);
+      libraryJar = filteredLibraryJar;
     }
 
+    ProguardOutput output =
+        getProguardOutputs(
+            proguardOutputJar,
+            proguardSeeds,
+            proguardUsage,
+            ruleContext,
+            semantics,
+            proguardOutputMap,
+            libraryJar);
+
     JavaConfiguration javaConfiguration =
         ruleContext.getConfiguration().getFragment(JavaConfiguration.class);
     JavaConfiguration.NamedLabel optimizer = javaConfiguration.getBytecodeOptimizer();