Optimize DexFileSplitter by switching to O(N) collections and removing legacy options.

PiperOrigin-RevId: 928644246
Change-Id: Ic1883828a0c6ec48d09e903d43e7b4ba7d6781f3
diff --git a/src/tools/java/com/google/devtools/build/android/dexer/DexFileSplitter.java b/src/tools/java/com/google/devtools/build/android/dexer/DexFileSplitter.java
index aa0a257..556abf4 100644
--- a/src/tools/java/com/google/devtools/build/android/dexer/DexFileSplitter.java
+++ b/src/tools/java/com/google/devtools/build/android/dexer/DexFileSplitter.java
@@ -27,8 +27,9 @@
 import com.google.common.base.Predicates;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.LinkedHashMultimap;
 import com.google.common.collect.Lists;
-import com.google.common.collect.TreeMultimap;
+import com.google.common.collect.Multimap;
 import com.google.common.io.ByteStreams;
 import com.google.common.io.Closer;
 import com.google.common.util.concurrent.ListenableFuture;
@@ -45,12 +46,13 @@
 import java.nio.file.Path;
 import java.nio.file.StandardOpenOption;
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.HashSet;
+import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Scanner;
 import java.util.Set;
-import java.util.TreeMap;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
@@ -81,19 +83,7 @@
         description = "Directory to write dex archives to merge.")
     public Path outputDirectory = new CompatPathConverter().convert(".");
 
-    @Parameter(
-        names = "--main-dex-list",
-        converter = CompatExistingPathConverter.class,
-        description = "List of classes to be placed into \"main\" classes.dex file.")
-    public Path mainDexListFile;
 
-    @Parameter(
-        names = "--minimal-main-dex",
-        arity = 1,
-        description =
-            "If true, *only* classes listed in --main_dex_list file are placed into \"main\" "
-                + "classes.dex file.")
-    public boolean minimalMainDex;
 
     // Undocumented dx option for testing multidex logic
     @Parameter(
@@ -121,29 +111,17 @@
   @VisibleForTesting
   static void splitIntoShards(Options options)
       throws ExecutionException, InterruptedException, IOException {
-    checkArgument(
-        !options.minimalMainDex || options.mainDexListFile != null,
-        "--minimal-main-dex not allowed without --main-dex-list");
-
     if (!Files.exists(options.outputDirectory)) {
       Files.createDirectories(options.outputDirectory);
     }
-
-    ImmutableSet<String> classesInMainDex =
-        options.mainDexListFile != null
-            ? ImmutableSet.copyOf(Files.readAllLines(options.mainDexListFile, UTF_8))
-            : null;
     ImmutableSet<String> expected =
         options.inclusionFilterJar != null ? expectedEntries(options.inclusionFilterJar) : null;
     try (Closer closer = Closer.create();
         DexFileSplitter out =
             new DexFileSplitter(options.outputDirectory, options.maxNumberOfIdxPerDex)) {
 
-      // 1. Scan inputs in order and keep first occurrence of each class, keeping all zips open.
-      // We don't process anything yet so we can shard in sorted order, which is what dx would do
-      // if presented with a single jar containing all the given inputs.
-      // TODO(kmb): Abandon alphabetic sorting to process each input fully before moving on (still
-      // requires scanning inputs twice for main dex list).
+      // 1. Scan inputs in order and keep first occurrence of each class (preserving input order),
+      // keeping all zips open.
 
       Predicate<ZipEntry> inclusionFilter = ZipEntryPredicates.suffixes(".dex", ".class");
       if (expected != null) {
@@ -151,10 +129,10 @@
       }
 
       // Maps a dex file name to the zip file containing that dex file.
-      TreeMap<String, ZipFile> dexFilesAndContainingZip =
-          new TreeMap<>(ZipEntryComparator::compareClassNames);
+      LinkedHashMap<String, ZipFile> dexFilesAndContainingZip = new LinkedHashMap<>();
       // Maps a class to its synthetic classes, if any.
-      TreeMultimap<String, String> contextClassesToSyntheticClasses = TreeMultimap.create();
+      LinkedHashMultimap<String, String> contextClassesToSyntheticClasses =
+          LinkedHashMultimap.create();
 
       for (Path inputArchive : options.inputArchives) {
         ZipFile zip = closer.register(new ZipFile(inputArchive.toFile()));
@@ -173,28 +151,8 @@
       }
 
       // 2. Process each class in desired order, rolling from shard to shard as needed.
-      if (classesInMainDex == null || classesInMainDex.isEmpty()) {
-        out.processDexes(
-            dexFilesAndContainingZip, contextClassesToSyntheticClasses, Predicates.alwaysTrue());
-      } else {
-        checkArgument(classesInMainDex.stream().noneMatch(s -> s.startsWith("j$/")),
-            "%s lists classes in package 'j$', which can't be included in classes.dex and can "
-                + "cause runtime errors. Please avoid needing these classes in the main dex file.",
-            options.mainDexListFile);
-        // To honor --main_dex_list make two passes:
-        // 1. process only the classes listed in the given file
-        // 2. process the remaining files
-        Predicate<String> mainDexFilter = ZipEntryPredicates.classFileNameFilter(classesInMainDex);
-        out.processDexes(dexFilesAndContainingZip, contextClassesToSyntheticClasses, mainDexFilter);
-        // Fail if main_dex_list is too big, following dx's example
-        checkState(out.shardsWritten() == 0, "Too many classes listed in main dex list file "
-            + "%s, main dex capacity exceeded", options.mainDexListFile);
-        if (options.minimalMainDex) {
-          out.nextShard(); // Start new .dex file if requested
-        }
-        out.processDexes(
-            dexFilesAndContainingZip, contextClassesToSyntheticClasses, mainDexFilter.negate());
-      }
+      out.processDexes(
+          dexFilesAndContainingZip, contextClassesToSyntheticClasses, Predicates.alwaysTrue());
     }
   }
 
@@ -208,7 +166,7 @@
   }
 
   private static void parseSyntheticContextsMap(
-      InputStream inputStream, TreeMultimap<String, String> syntheticClassContexts) {
+      InputStream inputStream, Multimap<String, String> syntheticClassContexts) {
     Scanner scanner = new Scanner(inputStream, UTF_8);
     scanner.useDelimiter("[;\n]");
     while (scanner.hasNext()) {
@@ -286,7 +244,7 @@
 
   private void processDexes(
       Map<String, ZipFile> dexFilesAndContainingZip,
-      TreeMultimap<String, String> contextClassesToSyntheticClasses,
+      Multimap<String, String> contextClassesToSyntheticClasses,
       Predicate<String> filter)
       throws ExecutionException, InterruptedException, IOException {
 
@@ -308,7 +266,7 @@
             // file all together as a unit, so skip them here.
             if (!syntheticClasses.contains(filename)) {
               ZipFile zipFile = entry.getValue();
-              Set<String> synths = contextClassesToSyntheticClasses.get(filename);
+            Collection<String> synths = contextClassesToSyntheticClasses.get(filename);
 
             ListenableFuture<ZipEntryAndContent> contextFuture =
                 listeningExecutor.submit(() -> readAndParseDex(zipFile, filename));
@@ -402,11 +360,10 @@
         filename);
     checkState(entry.getMethod() == ZipEntry.STORED, "Expect to process STORED: %s", filename);
 
+    // We don't want to use the Dex(InputStream) constructor because it closes the stream,
+    // which will break the for loop, and it has its own bespoke way of reading the file into
+    // a byte buffer before effectively calling Dex(byte[]) anyway.
     try (InputStream entryStream = zip.getInputStream(entry)) {
-      // We don't want to use the Dex(InputStream) constructor because it closes the stream,
-      // which will break the for loop, and it has its own bespoke way of reading the file into
-      // a byte buffer before effectively calling Dex(byte[]) anyway.
-      // TODO(kmb) since entry is stored, mmap content and give to Dex(ByteBuffer) and output zip
       byte[] content = new byte[(int) entry.getSize()];
       ByteStreams.readFully(entryStream, content); // throws if file is smaller than expected
       checkState(
diff --git a/src/tools/javatests/com/google/devtools/build/android/dexer/DexFileSplitterTest.java b/src/tools/javatests/com/google/devtools/build/android/dexer/DexFileSplitterTest.java
index 6bbe5db..8e8a42f 100644
--- a/src/tools/javatests/com/google/devtools/build/android/dexer/DexFileSplitterTest.java
+++ b/src/tools/javatests/com/google/devtools/build/android/dexer/DexFileSplitterTest.java
@@ -180,66 +180,7 @@
     }
   }
 
-  @Test
-  public void testMainDexList() throws Exception {
-    Path mainDexFile = tmp.newFile("main_dex_list.txt").toPath();
-    Files.write(mainDexFile, ImmutableList.of("multidex/Class2.class"), UTF_8);
 
-    ImmutableList<Path> outputArchives =
-        runDexSplitter(
-            SMALL_IDX_PER_DEX,
-            /* inclusionFilterJar= */ null,
-            "main_dex_list",
-            mainDexFile,
-            /* minimalMainDex= */ false,
-            simpleDexArchive,
-            multidexArchive);
-
-    HashSet<String> expectedEntries = new HashSet<>();
-    expectedEntries.addAll(dexEntries(simpleDexArchive));
-    expectedEntries.addAll(dexEntries(multidexArchive));
-    assertThat(outputArchives.size()).isGreaterThan(1);
-    assertThat(dexEntries(outputArchives.get(0))).contains("multidex/Class2.class.dex");
-    assertExpectedEntries(outputArchives, expectedEntries);
-  }
-
-  @Test
-  public void testMainDexList_containsForbidden() throws Exception {
-    Path mainDexFile = tmp.newFile("main_dex_list.txt").toPath();
-    Files.write(mainDexFile, ImmutableList.of("com/google/Ok.class", "j$/my/Bad.class"), UTF_8);
-    IllegalArgumentException e =
-        assertThrows(
-            IllegalArgumentException.class,
-            () ->
-                runDexSplitter(
-                    REAL_WORLD_IDX_PER_DEX,
-                    /* inclusionFilterJar= */ null,
-                    "invalid_main_dex_list",
-                    mainDexFile,
-                    /* minimalMainDex= */ false,
-                    simpleDexArchive));
-    assertThat(e).hasMessageThat().contains("j$");
-  }
-
-  @Test
-  public void testMinimalMainDex() throws Exception {
-    Path mainDexFile = tmp.newFile("minimal_main_dex_list.txt").toPath();
-    Files.write(mainDexFile, ImmutableList.of("multidex/Class1.class"), UTF_8);
-
-    ImmutableList<Path> outputArchives =
-        runDexSplitter(
-            REAL_WORLD_IDX_PER_DEX,
-            /* inclusionFilterJar= */ null,
-            "minimal_main_dex",
-            mainDexFile,
-            /* minimalMainDex= */ true,
-            multidexArchive);
-
-    ImmutableSet<String> expectedEntries = dexEntries(multidexArchive);
-    assertThat(outputArchives.size()).isGreaterThan(1);
-    assertThat(dexEntries(outputArchives.get(0))).containsExactly("multidex/Class1.class.dex");
-    assertExpectedEntries(outputArchives, expectedEntries);
-  }
 
   @Test
   public void testInclusionFilterJar() throws Exception {
@@ -248,8 +189,6 @@
             REAL_WORLD_IDX_PER_DEX,
             SIMPLE_JAR,
             "filtered",
-            /* mainDexList= */ null,
-            /* minimalMainDex= */ false,
             multidexArchive,
             simpleDexArchive);
 
@@ -272,7 +211,11 @@
   }
 
   @Test
-  public void testShuffledInputsDeterminism() throws Exception {
+  public void testShuffledInputsOrder() throws Exception {
+    ImmutableList<String> entriesSimple = ImmutableList.copyOf(dexEntries(simpleDexArchive));
+    ImmutableList<String> entriesJSimple = ImmutableList.copyOf(dexEntries(jsimpleDexArchive));
+    ImmutableList<String> entriesMultidex = ImmutableList.copyOf(dexEntries(multidexArchive));
+
     // Run 1: Order A, B, C
     ImmutableList<Path> outputArchives1 =
         runDexSplitter(
@@ -282,7 +225,18 @@
             jsimpleDexArchive,
             multidexArchive);
 
-    // Run 2: Order C, A, B (inverted context)
+    // Expected order: Simple then JSimple then Multidex
+    ImmutableList<String> expectedOrder1 =
+        ImmutableList.<String>builder()
+            .addAll(entriesSimple)
+            .addAll(entriesJSimple)
+            .addAll(entriesMultidex)
+            .build();
+    assertThat(allDexEntries(outputArchives1))
+        .containsExactlyElementsIn(expectedOrder1)
+        .inOrder();
+
+    // Run 2: Order C, A, B
     ImmutableList<Path> outputArchives2 =
         runDexSplitter(
             SMALL_IDX_PER_DEX,
@@ -291,15 +245,21 @@
             simpleDexArchive,
             jsimpleDexArchive);
 
-    assertThat(outputArchives1).hasSize(outputArchives2.size());
-
-    for (int i = 0; i < outputArchives1.size(); i++) {
-      ImmutableSet<String> entries1 = dexEntries(outputArchives1.get(i));
-      ImmutableSet<String> entries2 = dexEntries(outputArchives2.get(i));
-      assertThat(entries1).containsExactlyElementsIn(entries2);
-    }
+    // Expected order: Multidex then Simple then JSimple
+    ImmutableList<String> expectedOrder2 =
+        ImmutableList.<String>builder()
+            .addAll(entriesMultidex)
+            .addAll(entriesSimple)
+            .addAll(entriesJSimple)
+            .build();
+    assertThat(allDexEntries(outputArchives2))
+        .containsExactlyElementsIn(expectedOrder2)
+        .inOrder();
   }
 
+
+
+
   @Test
   public void testErrorPropagation() throws Exception {
     Path corruptZip = tmp.newFile("corrupt.zip").toPath();
@@ -351,26 +311,6 @@
     assertThat(outputArchives).hasSize(4);
   }
 
-
-
-  @Test
-  public void testMultidexOffWithMultidexFlags() throws Exception {
-    IllegalArgumentException e =
-        assertThrows(
-            IllegalArgumentException.class,
-            () ->
-                runDexSplitter(
-                    SMALL_IDX_PER_DEX,
-                    /* inclusionFilterJar= */ null,
-                    "should_fail",
-                    /* mainDexList= */ null,
-                    /* minimalMainDex= */ true,
-                    simpleDexArchive));
-    assertThat(e)
-        .hasMessageThat()
-        .isEqualTo("--minimal-main-dex not allowed without --main-dex-list");
-  }
-
   private void assertExpectedEntries(
       ImmutableList<Path> outputArchives, Set<String> expectedEntries) throws IOException {
     ImmutableSet.Builder<String> actualFiles = ImmutableSet.builder();
@@ -394,6 +334,19 @@
     }
   }
 
+  private ImmutableList<String> allDexEntries(ImmutableList<Path> outputArchives) throws IOException {
+    ImmutableList.Builder<String> result = ImmutableList.builder();
+    for (Path outputArchive : outputArchives) {
+      try (ZipFile input = new ZipFile(outputArchive.toFile())) {
+        input.stream()
+            .map(ZipEntry::getName)
+            .filter(Predicates.containsPattern(".*\\.class.dex$"))
+            .forEach(result::add);
+      }
+    }
+    return result.build();
+  }
+
   private ImmutableList<Path> runDexSplitter(
       int maxNumberOfIdxPerDex, String outputRoot, Path... dexArchives)
       throws ExecutionException, InterruptedException, IOException {
@@ -401,8 +354,6 @@
         maxNumberOfIdxPerDex,
         /*inclusionFilterJar=*/ null,
         outputRoot,
-        /*mainDexList=*/ null,
-        /*minimalMainDex=*/ false,
         dexArchives);
   }
 
@@ -410,16 +361,12 @@
       int maxNumberOfIdxPerDex,
       @Nullable Path inclusionFilterJar,
       String outputRoot,
-      @Nullable Path mainDexList,
-      boolean minimalMainDex,
       Path... dexArchives)
       throws ExecutionException, InterruptedException, IOException {
     DexFileSplitter.Options options = new DexFileSplitter.Options();
     options.inputArchives = ImmutableList.copyOf(dexArchives);
     options.outputDirectory = tmp.newFolder(outputRoot).toPath();
     options.maxNumberOfIdxPerDex = maxNumberOfIdxPerDex;
-    options.mainDexListFile = mainDexList;
-    options.minimalMainDex = minimalMainDex;
     options.inclusionFilterJar = inclusionFilterJar;
     DexFileSplitter.splitIntoShards(options);
     assertThat(options.outputDirectory.toFile().exists()).isTrue();