Automated rollback of commit 8ec17fddab0bf586545570f471d4e6b61b5ee807.

PiperOrigin-RevId: 929333708
Change-Id: I12d1331b84e989e866ee52af1ee094de1268eb81
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 556abf4..aa0a257 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,9 +27,8 @@
 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.Multimap;
+import com.google.common.collect.TreeMultimap;
 import com.google.common.io.ByteStreams;
 import com.google.common.io.Closer;
 import com.google.common.util.concurrent.ListenableFuture;
@@ -46,13 +45,12 @@
 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;
@@ -83,7 +81,19 @@
         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(
@@ -111,17 +121,29 @@
   @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 (preserving input order),
-      // keeping all zips open.
+      // 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).
 
       Predicate<ZipEntry> inclusionFilter = ZipEntryPredicates.suffixes(".dex", ".class");
       if (expected != null) {
@@ -129,10 +151,10 @@
       }
 
       // Maps a dex file name to the zip file containing that dex file.
-      LinkedHashMap<String, ZipFile> dexFilesAndContainingZip = new LinkedHashMap<>();
+      TreeMap<String, ZipFile> dexFilesAndContainingZip =
+          new TreeMap<>(ZipEntryComparator::compareClassNames);
       // Maps a class to its synthetic classes, if any.
-      LinkedHashMultimap<String, String> contextClassesToSyntheticClasses =
-          LinkedHashMultimap.create();
+      TreeMultimap<String, String> contextClassesToSyntheticClasses = TreeMultimap.create();
 
       for (Path inputArchive : options.inputArchives) {
         ZipFile zip = closer.register(new ZipFile(inputArchive.toFile()));
@@ -151,8 +173,28 @@
       }
 
       // 2. Process each class in desired order, rolling from shard to shard as needed.
-      out.processDexes(
-          dexFilesAndContainingZip, contextClassesToSyntheticClasses, Predicates.alwaysTrue());
+      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());
+      }
     }
   }
 
@@ -166,7 +208,7 @@
   }
 
   private static void parseSyntheticContextsMap(
-      InputStream inputStream, Multimap<String, String> syntheticClassContexts) {
+      InputStream inputStream, TreeMultimap<String, String> syntheticClassContexts) {
     Scanner scanner = new Scanner(inputStream, UTF_8);
     scanner.useDelimiter("[;\n]");
     while (scanner.hasNext()) {
@@ -244,7 +286,7 @@
 
   private void processDexes(
       Map<String, ZipFile> dexFilesAndContainingZip,
-      Multimap<String, String> contextClassesToSyntheticClasses,
+      TreeMultimap<String, String> contextClassesToSyntheticClasses,
       Predicate<String> filter)
       throws ExecutionException, InterruptedException, IOException {
 
@@ -266,7 +308,7 @@
             // file all together as a unit, so skip them here.
             if (!syntheticClasses.contains(filename)) {
               ZipFile zipFile = entry.getValue();
-            Collection<String> synths = contextClassesToSyntheticClasses.get(filename);
+              Set<String> synths = contextClassesToSyntheticClasses.get(filename);
 
             ListenableFuture<ZipEntryAndContent> contextFuture =
                 listeningExecutor.submit(() -> readAndParseDex(zipFile, filename));
@@ -360,10 +402,11 @@
         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 8e8a42f..6bbe5db 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,7 +180,66 @@
     }
   }
 
+  @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 {
@@ -189,6 +248,8 @@
             REAL_WORLD_IDX_PER_DEX,
             SIMPLE_JAR,
             "filtered",
+            /* mainDexList= */ null,
+            /* minimalMainDex= */ false,
             multidexArchive,
             simpleDexArchive);
 
@@ -211,11 +272,7 @@
   }
 
   @Test
-  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));
-
+  public void testShuffledInputsDeterminism() throws Exception {
     // Run 1: Order A, B, C
     ImmutableList<Path> outputArchives1 =
         runDexSplitter(
@@ -225,18 +282,7 @@
             jsimpleDexArchive,
             multidexArchive);
 
-    // 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
+    // Run 2: Order C, A, B (inverted context)
     ImmutableList<Path> outputArchives2 =
         runDexSplitter(
             SMALL_IDX_PER_DEX,
@@ -245,21 +291,15 @@
             simpleDexArchive,
             jsimpleDexArchive);
 
-    // 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();
+    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);
+    }
   }
 
-
-
-
   @Test
   public void testErrorPropagation() throws Exception {
     Path corruptZip = tmp.newFile("corrupt.zip").toPath();
@@ -311,6 +351,26 @@
     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();
@@ -334,19 +394,6 @@
     }
   }
 
-  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 {
@@ -354,6 +401,8 @@
         maxNumberOfIdxPerDex,
         /*inclusionFilterJar=*/ null,
         outputRoot,
+        /*mainDexList=*/ null,
+        /*minimalMainDex=*/ false,
         dexArchives);
   }
 
@@ -361,12 +410,16 @@
       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();