Support multiple (disjoint) inputs and add --multidex=given_shard flag to DexFileMerger tool

RELNOTES: None.

PiperOrigin-RevId: 175082253

2 files changed, 64 insertions, 25 deletions

diff --git a/src/tools/android/java/com/google/devtools/build/android/dexer/DexFileMerger.java b/src/tools/android/java/com/google/devtools/build/android/dexer/DexFileMerger.java
index e892f7a94a..fded7b9562 100644
--- a/src/tools/android/java/com/google/devtools/build/android/dexer/DexFileMerger.java
+++ b/src/tools/android/java/com/google/devtools/build/android/dexer/DexFileMerger.java
@@ -47,7 +47,10 @@ import java.nio.file.Files;
 import java.nio.file.Path;
 import java.util.ArrayList;
 import java.util.Collections;
+import java.util.List;
 import java.util.concurrent.Executors;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
 import java.util.zip.ZipEntry;
 import java.util.zip.ZipFile;
 import java.util.zip.ZipOutputStream;
@@ -71,15 +74,18 @@ class DexFileMerger {
   public static class Options extends OptionsBase {
     @Option(
       name = "input",
-      defaultValue = "null",
+      allowMultiple = true,
+      defaultValue = "",
       category = "input",
       documentationCategory = OptionDocumentationCategory.UNCATEGORIZED,
       effectTags = {OptionEffectTag.UNKNOWN},
       converter = ExistingPathConverter.class,
       abbrev = 'i',
-      help = "Input file to read to aggregate."
+      help = "Input archives with .dex files to merge.  Inputs are processed in given order, so "
+          + "classes from later inputs will be added after earlier inputs.  Classes mustn't appear "
+          + "more than once."
     )
-    public Path inputArchive;
+    public List<Path> inputArchives;
 
     @Option(
       name = "output",
@@ -164,7 +170,6 @@ class DexFileMerger {
       defaultValue = "false", // dx's default
       documentationCategory = OptionDocumentationCategory.UNDOCUMENTED,
       effectTags = {OptionEffectTag.UNKNOWN},
-      allowMultiple = false,
       help = "Typically not needed flag intended to imitate dx's --forceJumbo."
     )
     public boolean forceJumbo;
@@ -175,7 +180,6 @@ class DexFileMerger {
       category = "misc",
       documentationCategory = OptionDocumentationCategory.UNCATEGORIZED,
       effectTags = {OptionEffectTag.UNKNOWN},
-      allowMultiple = false,
       help = "Dex file output prefix."
     )
     public String dexPrefix;
@@ -197,6 +201,10 @@ class DexFileMerger {
   @VisibleForTesting
   static void buildMergedDexFiles(Options options) throws IOException {
     ListeningExecutorService executor;
+    checkArgument(!options.inputArchives.isEmpty(), "Need at least one --input");
+    checkArgument(
+        options.mainDexListFile == null || options.inputArchives.size() == 1,
+        "--main-dex-list only supported with exactly one --input, use DexFileSplitter for more");
     if (options.multidexMode.isMultidexAllowed()) {
       executor = createThreadPool();
     } else {
@@ -216,31 +224,37 @@ class DexFileMerger {
         ? ImmutableSet.copyOf(Files.readAllLines(options.mainDexListFile, UTF_8))
         : null;
     PrintStream originalStdOut = System.out;
-    try (ZipFile zip = new ZipFile(options.inputArchive.toFile());
-        DexFileAggregator out = createDexFileAggregator(options, executor)) {
-      ArrayList<ZipEntry> dexFiles = filesToProcess(zip);
-
+    try (DexFileAggregator out = createDexFileAggregator(options, executor)) {
       if (!options.verbose) {
         // com.android.dx.merge.DexMerger prints status information to System.out that we silence
         // here unless it was explicitly requested.  (It also prints debug info to DxContext.out,
         // which we populate accordingly below.)
         System.setOut(Dexing.nullout);
       }
-      if (classesInMainDex == null) {
-        processDexFiles(zip, dexFiles, out);
-      } else {
-        // To honor --main_dex_list make two passes:
-        // 1. process only the classes listed in the given file
-        // 2. process the remaining files
-        Predicate<ZipEntry> mainDexFilter = ZipEntryPredicates.classFileFilter(classesInMainDex);
-        processDexFiles(zip, Iterables.filter(dexFiles, mainDexFilter), out);
-        // Fail if main_dex_list is too big, following dx's example
-        checkState(out.getDexFilesWritten() == 0, "Too many classes listed in main dex list file "
-            + "%s, main dex capacity exceeded", options.mainDexListFile);
-        if (options.minimalMainDex) {
-          out.flush(); // Start new .dex file if requested
+
+      for (Path inputArchive : options.inputArchives) {
+        // Simply merge files from inputs in order.  Doing that with a main dex list doesn't work,
+        // but we rule out more than one input with a main dex list above.
+        try (ZipFile zip = new ZipFile(inputArchive.toFile())) {
+          ArrayList<ZipEntry> dexFiles = filesToProcess(zip);
+          if (classesInMainDex == null) {
+            processDexFiles(zip, dexFiles, out);
+          } else {
+            // To honor --main_dex_list make two passes:
+            // 1. process only the classes listed in the given file
+            // 2. process the remaining files
+            Predicate<ZipEntry> mainDexFilter =
+                ZipEntryPredicates.classFileFilter(classesInMainDex);
+            processDexFiles(zip, Iterables.filter(dexFiles, mainDexFilter), out);
+            // Fail if main_dex_list is too big, following dx's example
+            checkState(out.getDexFilesWritten() == 0, "Too many classes listed in main dex list "
+                + "file %s, main dex capacity exceeded", options.mainDexListFile);
+            if (options.minimalMainDex) {
+              out.flush(); // Start new .dex file if requested
+            }
+            processDexFiles(zip, Iterables.filter(dexFiles, Predicates.not(mainDexFilter)), out);
+          }
         }
-        processDexFiles(zip, Iterables.filter(dexFiles, Predicates.not(mainDexFilter)), out);
       }
     } finally {
       // Kill threads in the pool so we don't hang
@@ -280,6 +294,21 @@ class DexFileMerger {
 
   private static DexFileAggregator createDexFileAggregator(
       Options options, ListeningExecutorService executor) throws IOException {
+    String filePrefix = options.dexPrefix;
+    if (options.multidexMode == MultidexStrategy.GIVEN_SHARD) {
+      checkArgument(options.inputArchives.size() == 1,
+          "--multidex=given_shard requires exactly one --input");
+      Pattern namingPattern = Pattern.compile("([0-9]+)\\..*");
+      Matcher matcher = namingPattern.matcher(options.inputArchives.get(0).toFile().getName());
+      checkArgument(matcher.matches(),
+          "expect input named <N>.xxx.zip for --multidex=given_shard but got %s",
+          options.inputArchives.get(0).toFile().getName());
+      int shard = Integer.parseInt(matcher.group(1));
+      checkArgument(shard > 0, "expect positive N in input named <N>.xxx.zip but got %s", shard);
+      if (shard > 1) {  // first shard conventionally isn't numbered
+        filePrefix += shard;
+      }
+    }
     return new DexFileAggregator(
         new DxContext(options.verbose ? System.out : ByteStreams.nullOutputStream(), System.err),
         new DexFileArchive(
@@ -290,7 +319,7 @@ class DexFileMerger {
         options.forceJumbo,
         options.maxNumberOfIdxPerDex,
         options.wasteThresholdPerDex,
-        options.dexPrefix);
+        filePrefix);
   }
 
   /**
diff --git a/src/tools/android/java/com/google/devtools/build/android/dexer/MultidexStrategy.java b/src/tools/android/java/com/google/devtools/build/android/dexer/MultidexStrategy.java
index f1903c7e3b..70a49e4dcd 100644
--- a/src/tools/android/java/com/google/devtools/build/android/dexer/MultidexStrategy.java
+++ b/src/tools/android/java/com/google/devtools/build/android/dexer/MultidexStrategy.java
@@ -19,6 +19,8 @@ package com.google.devtools.build.android.dexer;
 public enum MultidexStrategy {
   /** Create exactly one .dex file.  The operation will fail if .dex limits are exceeded. */
   OFF,
+  /** Create exactly one &lt;prefixN&gt;.dex file with N taken from the (single) input archive. */
+  GIVEN_SHARD,
   /**
    * Assemble .dex files similar to {@link com.android.dx.command.dexer.Main dx}, with all but one
    * file as large as possible.
@@ -31,6 +33,14 @@ public enum MultidexStrategy {
   BEST_EFFORT;
 
   public boolean isMultidexAllowed() {
-    return this != OFF;
+    switch (this) {
+      case OFF:
+      case GIVEN_SHARD:
+        return false;
+      case MINIMAL:
+      case BEST_EFFORT:
+        return true;
+    }
+    throw new AssertionError("Unknown: " + this);
   }
 }