Add main dex list consumer api to D8.

When merging classes for a multidex build, D8 honors the input
main dex list and also puts anything that is synthesized because
of main dex classes in the main dex file.

When building app bundles, the main dex list is put in the bundle
so that feature split code can be merged with base split code
for shipping to pre-L devices. However, since the main dex list
contains no synthesized classes, bundletool can generate an apk
with some needed classes missing from the main dex file.

In order to fix this, D8 should produce an updated main dex list
containing all the classes that it synthesized and determined
were needed in the main dex file.

R=gavra@google.com, zerny@google.com

Bug: 120039166
Change-Id: I6dd94f8c3ed8617b89fbd1c18ec69bca27172437
diff --git a/src/main/java/com/android/tools/r8/BaseCompilerCommand.java b/src/main/java/com/android/tools/r8/BaseCompilerCommand.java
index 34a3a93..584b4de 100644
--- a/src/main/java/com/android/tools/r8/BaseCompilerCommand.java
+++ b/src/main/java/com/android/tools/r8/BaseCompilerCommand.java
@@ -25,6 +25,7 @@
 
   private final CompilationMode mode;
   private final ProgramConsumer programConsumer;
+  private final StringConsumer mainDexListConsumer;
   private final int minApiLevel;
   private final Reporter reporter;
   private final boolean enableDesugaring;
@@ -33,6 +34,7 @@
   BaseCompilerCommand(boolean printHelp, boolean printVersion) {
     super(printHelp, printVersion);
     programConsumer = null;
+    mainDexListConsumer = null;
     mode = null;
     minApiLevel = 0;
     reporter = new Reporter();
@@ -44,6 +46,7 @@
       AndroidApp app,
       CompilationMode mode,
       ProgramConsumer programConsumer,
+      StringConsumer mainDexListConsumer,
       int minApiLevel,
       Reporter reporter,
       boolean enableDesugaring,
@@ -53,6 +56,7 @@
     assert mode != null;
     this.mode = mode;
     this.programConsumer = programConsumer;
+    this.mainDexListConsumer = mainDexListConsumer;
     this.minApiLevel = minApiLevel;
     this.reporter = reporter;
     this.enableDesugaring = enableDesugaring;
@@ -81,6 +85,13 @@
     return programConsumer;
   }
 
+  /**
+   * Get the main dex list consumer that will receive the final complete main dex list.
+   */
+  public StringConsumer getMainDexListConsumer() {
+    return mainDexListConsumer;
+  }
+
   /** Get the use-desugaring state. True if enabled, false otherwise. */
   public boolean getEnableDesugaring() {
     return enableDesugaring;
@@ -109,6 +120,7 @@
       extends BaseCommand.Builder<C, B> {
 
     private ProgramConsumer programConsumer = null;
+    private StringConsumer mainDexListConsumer = null;
     private Path outputPath = null;
     // TODO(b/70656566): Remove default output mode when deprecated API is removed.
     private OutputMode outputMode = OutputMode.DexIndexed;
@@ -189,6 +201,13 @@
     }
 
     /**
+     * Get the main dex list consumer that will receive the final complete main dex list.
+     */
+    public StringConsumer getMainDexListConsumer() {
+      return mainDexListConsumer;
+    }
+
+    /**
      * If set to true, legacy multidex partitioning will be optimized to reduce LinearAlloc usage
      * during Dalvik DexOpt. Has no effect when compiling for a target with native multidex support
      * or without main dex list specification.
@@ -224,6 +243,33 @@
     }
 
     /**
+     * Set an output destination to which main-dex-list content should be written.
+     *
+     * <p>This is a short-hand for setting a {@link StringConsumer.FileConsumer} using {@link
+     * #setMainDexListConsumer}. Note that any subsequent call to this method or {@link
+     * #setMainDexListConsumer} will override the previous setting.
+     *
+     * @param mainDexListOutputPath File-system path to write output at.
+     */
+    public B setMainDexListOutputPath(Path mainDexListOutputPath) {
+      mainDexListConsumer = new StringConsumer.FileConsumer(mainDexListOutputPath);
+      return self();
+    }
+
+    /**
+     * Set a consumer for receiving the main-dex-list content.
+     *
+     * <p>Note that any subsequent call to this method or {@link #setMainDexListOutputPath} will
+     * override the previous setting.
+     *
+     * @param mainDexListConsumer Consumer to receive the content once produced.
+     */
+    public B setMainDexListConsumer(StringConsumer mainDexListConsumer) {
+      this.mainDexListConsumer = mainDexListConsumer;
+      return self();
+    }
+
+    /**
      * Set the output path-and-mode.
      *
      * <p>Setting the output path-and-mode will override any previous set consumer or any previous
diff --git a/src/main/java/com/android/tools/r8/D8Command.java b/src/main/java/com/android/tools/r8/D8Command.java
index e0f330d..f0217b9 100644
--- a/src/main/java/com/android/tools/r8/D8Command.java
+++ b/src/main/java/com/android/tools/r8/D8Command.java
@@ -137,6 +137,8 @@
         if (getProgramConsumer() instanceof DexFilePerClassFileConsumer) {
           reporter.error("Option --main-dex-list cannot be used with --file-per-class");
         }
+      } else if (getMainDexListConsumer() != null) {
+        reporter.error("Option --main-dex-list-output require --main-dex-list");
       }
       super.validate();
     }
@@ -153,6 +155,7 @@
           getAppBuilder().build(),
           getMode(),
           getProgramConsumer(),
+          getMainDexListConsumer(),
           getMinApiLevel(),
           getReporter(),
           !getDisableDesugaring(),
@@ -209,6 +212,7 @@
       AndroidApp inputApp,
       CompilationMode mode,
       ProgramConsumer programConsumer,
+      StringConsumer mainDexListConsumer,
       int minApiLevel,
       Reporter diagnosticsHandler,
       boolean enableDesugaring,
@@ -218,6 +222,7 @@
         inputApp,
         mode,
         programConsumer,
+        mainDexListConsumer,
         minApiLevel,
         diagnosticsHandler,
         enableDesugaring,
@@ -235,6 +240,7 @@
     assert !internal.debug;
     internal.debug = getMode() == CompilationMode.DEBUG;
     internal.programConsumer = getProgramConsumer();
+    internal.mainDexListConsumer = getMainDexListConsumer();
     internal.minimalMainDex = internal.debug;
     internal.minApiLevel = getMinApiLevel();
     internal.intermediate = intermediate;
diff --git a/src/main/java/com/android/tools/r8/D8CommandParser.java b/src/main/java/com/android/tools/r8/D8CommandParser.java
index a933b8f..b84b398 100644
--- a/src/main/java/com/android/tools/r8/D8CommandParser.java
+++ b/src/main/java/com/android/tools/r8/D8CommandParser.java
@@ -104,6 +104,7 @@
               "  --file-per-class        # Produce a separate dex file per input class",
               "  --no-desugaring         # Force disable desugaring.",
               "  --main-dex-list <file>  # List of classes to place in the primary dex file.",
+              "  --main-dex-list-output <file> # Output resulting main dex list in <file>.",
               "  --version               # Print the version of d8.",
               "  --help                  # Print this message."));
 
@@ -198,6 +199,8 @@
           }
         } else if (arg.equals("--main-dex-list")) {
           builder.addMainDexListFiles(Paths.get(expandedArgs[++i]));
+        } else if (arg.equals("--main-dex-list-output")) {
+          builder.setMainDexListOutputPath(Paths.get(expandedArgs[++i]));
         } else if (arg.equals("--optimize-multidex-for-linearalloc")) {
           builder.setOptimizeMultidexForLinearAlloc(true);
         } else if (arg.equals("--min-api")) {
diff --git a/src/main/java/com/android/tools/r8/R8Command.java b/src/main/java/com/android/tools/r8/R8Command.java
index a1858a9..338fb4f 100644
--- a/src/main/java/com/android/tools/r8/R8Command.java
+++ b/src/main/java/com/android/tools/r8/R8Command.java
@@ -90,8 +90,6 @@
     private boolean allowPartiallyImplementedProguardOptions = false;
     private boolean allowTestProguardOptions = false;
 
-    private StringConsumer mainDexListConsumer = null;
-
     // TODO(zerny): Consider refactoring CompatProguardCommandBuilder to avoid subclassing.
     Builder() {
       this(new DefaultR8DiagnosticsHandler());
@@ -178,33 +176,6 @@
       return self();
     }
 
-    /**
-     * Set an output destination to which main-dex-list content should be written.
-     *
-     * <p>This is a short-hand for setting a {@link StringConsumer.FileConsumer} using {@link
-     * #setMainDexListConsumer}. Note that any subsequent call to this method or {@link
-     * #setMainDexListConsumer} will override the previous setting.
-     *
-     * @param mainDexListOutputPath File-system path to write output at.
-     */
-    public Builder setMainDexListOutputPath(Path mainDexListOutputPath) {
-      mainDexListConsumer = new StringConsumer.FileConsumer(mainDexListOutputPath);
-      return self();
-    }
-
-    /**
-     * Set a consumer for receiving the main-dex-list content.
-     *
-     * <p>Note that any subsequent call to this method or {@link #setMainDexListOutputPath} will
-     * override the previous setting.
-     *
-     * @param mainDexListConsumer Consumer to receive the content once produced.
-     */
-    public Builder setMainDexListConsumer(StringConsumer mainDexListConsumer) {
-      this.mainDexListConsumer = mainDexListConsumer;
-      return self();
-    }
-
     /** Add proguard configuration-file resources. */
     public Builder addProguardConfigurationFiles(Path... paths) {
       guard(() -> {
@@ -321,7 +292,7 @@
       if (getProgramConsumer() instanceof DexFilePerClassFileConsumer) {
         reporter.error("R8 does not support compiling to a single DEX file per Java class file");
       }
-      if (mainDexListConsumer != null
+      if (getMainDexListConsumer() != null
           && mainDexRules.isEmpty()
           && !getAppBuilder().hasMainDexList()) {
         reporter.error(
@@ -433,7 +404,7 @@
               getAppBuilder().build(),
               getProgramConsumer(),
               mainDexKeepRules,
-              mainDexListConsumer,
+              getMainDexListConsumer(),
               configuration,
               getMode(),
               getMinApiLevel(),
@@ -503,7 +474,6 @@
   static final String USAGE_MESSAGE = R8CommandParser.USAGE_MESSAGE;
 
   private final ImmutableList<ProguardConfigurationRule> mainDexKeepRules;
-  private final StringConsumer mainDexListConsumer;
   private final ProguardConfiguration proguardConfiguration;
   private final boolean enableTreeShaking;
   private final boolean enableMinification;
@@ -576,12 +546,11 @@
       StringConsumer proguardMapConsumer,
       Path proguardCompatibilityRulesOutput,
       boolean optimizeMultidexForLinearAlloc) {
-    super(inputApp, mode, programConsumer, minApiLevel, reporter, enableDesugaring,
-        optimizeMultidexForLinearAlloc);
+    super(inputApp, mode, programConsumer, mainDexListConsumer, minApiLevel, reporter,
+        enableDesugaring, optimizeMultidexForLinearAlloc);
     assert proguardConfiguration != null;
     assert mainDexKeepRules != null;
     this.mainDexKeepRules = mainDexKeepRules;
-    this.mainDexListConsumer = mainDexListConsumer;
     this.proguardConfiguration = proguardConfiguration;
     this.enableTreeShaking = enableTreeShaking;
     this.enableMinification = enableMinification;
@@ -594,7 +563,6 @@
   private R8Command(boolean printHelp, boolean printVersion) {
     super(printHelp, printVersion);
     mainDexKeepRules = ImmutableList.of();
-    mainDexListConsumer = null;
     proguardConfiguration = null;
     enableTreeShaking = false;
     enableMinification = false;
@@ -641,7 +609,7 @@
     assert !internal.verbose;
     internal.mainDexKeepRules = mainDexKeepRules;
     internal.minimalMainDex = getMode() == CompilationMode.DEBUG;
-    internal.mainDexListConsumer = mainDexListConsumer;
+    internal.mainDexListConsumer = getMainDexListConsumer();
     internal.lineNumberOptimization =
         !internal.debug && (proguardConfiguration.isOptimizing() || internal.enableMinification)
             ? LineNumberOptimization.ON
diff --git a/src/test/java/com/android/tools/r8/TestCompilerBuilder.java b/src/test/java/com/android/tools/r8/TestCompilerBuilder.java
index afa8e7b..9425df9 100644
--- a/src/test/java/com/android/tools/r8/TestCompilerBuilder.java
+++ b/src/test/java/com/android/tools/r8/TestCompilerBuilder.java
@@ -37,6 +37,7 @@
   // Default initialized setup. Can be overwritten if needed.
   private Path defaultLibrary;
   private ProgramConsumer programConsumer;
+  private StringConsumer mainDexListConsumer;
   private AndroidApiLevel defaultMinApiLevel = ToolHelper.getMinApiLevelForDexVm();
   private Consumer<InternalOptions> optionsConsumer = DEFAULT_OPTIONS;
 
@@ -62,6 +63,7 @@
   public CR compile() throws CompilationFailedException {
     AndroidAppConsumers sink = new AndroidAppConsumers();
     builder.setProgramConsumer(sink.wrapProgramConsumer(programConsumer));
+    builder.setMainDexListConsumer(mainDexListConsumer);
     if (defaultLibrary != null) {
       builder.addLibraryFiles(defaultLibrary);
     }
@@ -113,6 +115,17 @@
     return self();
   }
 
+  public T setMainDexListConsumer(StringConsumer consumer) {
+    assert consumer != null;
+    this.mainDexListConsumer = consumer;
+    return self();
+  }
+
+  public T addMainDexListFiles(Collection<Path> files) {
+    builder.addMainDexListFiles(files);
+    return self();
+  }
+
   @Override
   public T addProgramClassFileData(Collection<byte[]> classes) {
     for (byte[] clazz : classes) {
diff --git a/src/test/java/com/android/tools/r8/maindexlist/MainDexListOutputTest.java b/src/test/java/com/android/tools/r8/maindexlist/MainDexListOutputTest.java
index c053e62..9d9f21d 100644
--- a/src/test/java/com/android/tools/r8/maindexlist/MainDexListOutputTest.java
+++ b/src/test/java/com/android/tools/r8/maindexlist/MainDexListOutputTest.java
@@ -9,16 +9,19 @@
 
 import com.android.tools.r8.CompilationFailedException;
 import com.android.tools.r8.DexIndexedConsumer;
+import com.android.tools.r8.DexIndexedConsumer.ArchiveConsumer;
 import com.android.tools.r8.Diagnostic;
 import com.android.tools.r8.DiagnosticsHandler;
 import com.android.tools.r8.OutputMode;
 import com.android.tools.r8.R8Command;
+import com.android.tools.r8.StringConsumer;
 import com.android.tools.r8.TestBase;
 import com.android.tools.r8.ToolHelper;
 import com.android.tools.r8.origin.Origin;
 import com.android.tools.r8.utils.FileUtils;
 import com.android.tools.r8.utils.StringDiagnostic;
 import com.google.common.collect.ImmutableList;
+import java.io.IOException;
 import java.nio.file.Path;
 import java.util.stream.Collectors;
 import org.junit.Rule;
@@ -27,6 +30,34 @@
 
 public class MainDexListOutputTest extends TestBase {
 
+  interface MyConsumer<T> {
+    void accept(T element);
+  }
+
+  class TestClass {
+    public void f(MyConsumer<String> s) {
+      s.accept("asdf");
+    }
+
+    public void g() {
+      f(System.out::println);
+    }
+  }
+
+  private static String testClassMainDexName =
+      "com/android/tools/r8/maindexlist/MainDexListOutputTest$TestClass.class";
+
+  private static class TestMainDexListConsumer implements StringConsumer {
+    public boolean called = false;
+
+    @Override
+    public void accept(String string, DiagnosticsHandler handler) {
+      called = true;
+      assertTrue(string.contains(testClassMainDexName));
+      assertTrue(string.contains("Lambda"));
+    }
+  }
+
   class Reporter implements DiagnosticsHandler {
     int errorCount = 0;
 
@@ -77,4 +108,37 @@
             .filter(s -> !s.isEmpty())
             .collect(Collectors.toList()));
   }
+
+  @Test
+  public void testD8DesugaredLambdasInMainDexList() throws IOException, CompilationFailedException {
+    Path mainDexList = writeTextToTempFile(testClassMainDexName);
+    TestMainDexListConsumer consumer = new TestMainDexListConsumer();
+    testForD8()
+        .addProgramClasses(ImmutableList.of(TestClass.class, MyConsumer.class))
+        .addMainDexListFiles(ImmutableList.of(mainDexList))
+        .setMainDexListConsumer(consumer)
+        .compile();
+    assertTrue(consumer.called);
+  }
+
+  @Test
+  public void testD8DesugaredLambdasInMainDexListMerging()
+      throws IOException, CompilationFailedException {
+    Path mainDexList = writeTextToTempFile(testClassMainDexName);
+    Path dexOutput = temp.getRoot().toPath().resolve("classes.zip");
+    // Build intermediate dex code first.
+    testForD8()
+        .addProgramClasses(ImmutableList.of(TestClass.class, MyConsumer.class))
+        .setIntermediate(true)
+        .setProgramConsumer(new ArchiveConsumer(dexOutput))
+        .compile();
+    // Now test that when merging with a main dex list it is correctly updated.
+    TestMainDexListConsumer consumer = new TestMainDexListConsumer();
+    testForD8()
+        .addProgramFiles(dexOutput)
+        .addMainDexListFiles(ImmutableList.of(mainDexList))
+        .setMainDexListConsumer(consumer)
+        .compile();
+    assertTrue(consumer.called);
+  }
 }