Check for not extending main-dex file in second round

This also adds a reproduction for horizontal class merger adding
classes not traced as roots in the first round.

Bug: 177333791
Bug: 177847090
Change-Id: I6548d4431004008b13171880bca49d2a01217e57
diff --git a/src/main/java/com/android/tools/r8/GenerateMainDexList.java b/src/main/java/com/android/tools/r8/GenerateMainDexList.java
index d260bcb..5774e10 100644
--- a/src/main/java/com/android/tools/r8/GenerateMainDexList.java
+++ b/src/main/java/com/android/tools/r8/GenerateMainDexList.java
@@ -96,7 +96,8 @@
     }
 
     Enqueuer enqueuer =
-        EnqueuerFactory.createForMainDexTracing(appView, subtypingInfo, graphConsumer);
+        EnqueuerFactory.createForFinalMainDexTracing(
+            appView, subtypingInfo, graphConsumer, MainDexTracingResult.NONE);
     Set<DexProgramClass> liveTypes = enqueuer.traceMainDex(mainDexRootSet, executor, timing);
     // LiveTypes is the result.
     MainDexTracingResult mainDexTracingResult = new MainDexListBuilder(liveTypes, appView).run();
diff --git a/src/main/java/com/android/tools/r8/R8.java b/src/main/java/com/android/tools/r8/R8.java
index 9e7a967..7f56ab4 100644
--- a/src/main/java/com/android/tools/r8/R8.java
+++ b/src/main/java/com/android/tools/r8/R8.java
@@ -435,7 +435,7 @@
                 .run(executorService);
         // Live types is the tracing result.
         Set<DexProgramClass> mainDexBaseClasses =
-            EnqueuerFactory.createForMainDexTracing(appView, subtypingInfo)
+            EnqueuerFactory.createForInitialMainDexTracing(appView, subtypingInfo)
                 .traceMainDex(mainDexRootSet, executorService, timing);
         // Calculate the automatic main dex list according to legacy multidex constraints.
         mainDexTracingResult = new MainDexListBuilder(mainDexBaseClasses, appView).run();
@@ -644,8 +644,11 @@
         }
 
         Enqueuer enqueuer =
-            EnqueuerFactory.createForMainDexTracing(
-                appView, new SubtypingInfo(appView), mainDexKeptGraphConsumer);
+            EnqueuerFactory.createForFinalMainDexTracing(
+                appView,
+                new SubtypingInfo(appView),
+                mainDexKeptGraphConsumer,
+                mainDexTracingResult);
         // Find classes which may have code executed before secondary dex files installation.
         // Live types is the tracing result.
         Set<DexProgramClass> mainDexBaseClasses =
@@ -1071,8 +1074,8 @@
       SubtypingInfo subtypingInfo = new SubtypingInfo(appView);
       if (forMainDex) {
         enqueuer =
-            EnqueuerFactory.createForMainDexTracing(
-                appView, subtypingInfo, whyAreYouKeepingConsumer);
+            EnqueuerFactory.createForFinalMainDexTracing(
+                appView, subtypingInfo, whyAreYouKeepingConsumer, MainDexTracingResult.NONE);
         enqueuer.traceMainDex(rootSet, executorService, timing);
       } else {
         enqueuer =
diff --git a/src/main/java/com/android/tools/r8/shaking/Enqueuer.java b/src/main/java/com/android/tools/r8/shaking/Enqueuer.java
index c15fdb8..404006a 100644
--- a/src/main/java/com/android/tools/r8/shaking/Enqueuer.java
+++ b/src/main/java/com/android/tools/r8/shaking/Enqueuer.java
@@ -5,6 +5,7 @@
 
 import static com.android.tools.r8.graph.DexProgramClass.asProgramClassOrNull;
 import static com.android.tools.r8.graph.FieldAccessInfoImpl.MISSING_FIELD_ACCESS_INFO;
+import static com.android.tools.r8.ir.optimize.enums.UnboxedEnumMemberRelocator.ENUM_UNBOXING_UTILITY_CLASS_SUFFIX;
 import static com.android.tools.r8.naming.IdentifierNameStringUtils.identifyIdentifier;
 import static com.android.tools.r8.naming.IdentifierNameStringUtils.isReflectionMethod;
 import static com.android.tools.r8.shaking.AnnotationRemover.shouldKeepAnnotation;
@@ -103,6 +104,7 @@
 import com.android.tools.r8.utils.Action;
 import com.android.tools.r8.utils.InternalOptions;
 import com.android.tools.r8.utils.InternalOptions.DesugarState;
+import com.android.tools.r8.utils.InternalOptions.OutlineOptions;
 import com.android.tools.r8.utils.IteratorUtils;
 import com.android.tools.r8.utils.MethodSignatureEquivalence;
 import com.android.tools.r8.utils.OptionalBool;
@@ -162,7 +164,8 @@
   public enum Mode {
     INITIAL_TREE_SHAKING,
     FINAL_TREE_SHAKING,
-    MAIN_DEX_TRACING,
+    INITIAL_MAIN_DEX_TRACING,
+    FINAL_MAIN_DEX_TRACING,
     WHY_ARE_YOU_KEEPING;
 
     public boolean isInitialTreeShaking() {
@@ -177,8 +180,16 @@
       return isInitialTreeShaking() || isFinalTreeShaking();
     }
 
-    public boolean isTracingMainDex() {
-      return this == MAIN_DEX_TRACING;
+    public boolean isInitialMainDexTracing() {
+      return this == INITIAL_MAIN_DEX_TRACING;
+    }
+
+    public boolean isFinalMainDexTracing() {
+      return this == FINAL_MAIN_DEX_TRACING;
+    }
+
+    public boolean isMainDexTracing() {
+      return isInitialMainDexTracing() || isFinalMainDexTracing();
     }
 
     public boolean isWhyAreYouKeeping() {
@@ -366,12 +377,14 @@
       ProgramMethodMap.create();
   private final Map<DexMethod, ProgramMethod> methodsWithBackports = new IdentityHashMap<>();
   private final Set<DexProgramClass> classesWithSerializableLambdas = Sets.newIdentityHashSet();
+  private final MainDexTracingResult previousMainDexTracingResult;
 
   Enqueuer(
       AppView<? extends AppInfoWithClassHierarchy> appView,
       SubtypingInfo subtypingInfo,
       GraphConsumer keptGraphConsumer,
-      Mode mode) {
+      Mode mode,
+      MainDexTracingResult previousMainDexTracingResult) {
     assert appView.appServices() != null;
     InternalOptions options = appView.options();
     this.appInfo = appView.appInfo();
@@ -384,6 +397,7 @@
     this.options = options;
     this.useRegistryFactory = createUseRegistryFactory();
     this.workList = EnqueuerWorklist.createWorklist();
+    this.previousMainDexTracingResult = previousMainDexTracingResult;
 
     if (mode.isInitialOrFinalTreeShaking()) {
       if (options.protoShrinking().enableGeneratedMessageLiteShrinking) {
@@ -1659,6 +1673,14 @@
       return;
     }
 
+    assert !mode.isFinalMainDexTracing()
+            || !options.testing.checkForNotExpandingMainDexTracingResult
+            || previousMainDexTracingResult.isRoot(clazz)
+            || clazz.toSourceString().contains(ENUM_UNBOXING_UTILITY_CLASS_SUFFIX)
+            // TODO(b/177847090): Consider not outlining anything in main dex.
+            || clazz.toSourceString().contains(OutlineOptions.CLASS_NAME)
+        : "Class " + clazz.toSourceString() + " was not a main dex root in the first round";
+
     // Mark types in inner-class attributes referenced.
     for (InnerClassAttribute innerClassAttribute : clazz.getInnerClasses()) {
       recordTypeReference(innerClassAttribute.getInner(), clazz, this::ignoreMissingClass);
@@ -1766,7 +1788,7 @@
 
     if (!appView.options().enableUnusedInterfaceRemoval
         || rootSet.noUnusedInterfaceRemoval.contains(type)
-        || mode.isTracingMainDex()) {
+        || mode.isMainDexTracing()) {
       markTypeAsLive(clazz, implementer);
     } else {
       if (liveTypes.contains(clazz)) {
@@ -2047,7 +2069,7 @@
   }
 
   private void ensureFromLibraryOrThrow(DexType type, DexLibraryClass context) {
-    if (mode.isTracingMainDex()) {
+    if (mode.isMainDexTracing()) {
       // b/72312389: android.jar contains parts of JUnit and most developers include JUnit in
       // their programs. This leads to library classes extending program classes. When tracing
       // main dex lists we allow this.
@@ -2385,7 +2407,7 @@
   private void markLibraryAndClasspathMethodOverridesAsLive(
       InstantiatedObject instantiation, DexClass libraryClass) {
     assert libraryClass.isNotProgramClass();
-    if (mode.isTracingMainDex()) {
+    if (mode.isMainDexTracing()) {
       // Library roots must be specified for tracing of library methods. For classpath the expected
       // use case is that the classes will be classloaded, thus they should have no bearing on the
       // content of the main dex file.
@@ -2891,7 +2913,7 @@
   public Set<DexProgramClass> traceMainDex(
       RootSet rootSet, ExecutorService executorService, Timing timing) throws ExecutionException {
     assert analyses.isEmpty();
-    assert mode.isTracingMainDex();
+    assert mode.isMainDexTracing();
     this.rootSet = rootSet;
     // Translate the result of root-set computation into enqueuer actions.
     enqueueRootItems(rootSet.noShrinking);
diff --git a/src/main/java/com/android/tools/r8/shaking/EnqueuerFactory.java b/src/main/java/com/android/tools/r8/shaking/EnqueuerFactory.java
index 34a8750..3c9101a 100644
--- a/src/main/java/com/android/tools/r8/shaking/EnqueuerFactory.java
+++ b/src/main/java/com/android/tools/r8/shaking/EnqueuerFactory.java
@@ -17,7 +17,8 @@
   public static Enqueuer createForInitialTreeShaking(
       AppView<? extends AppInfoWithClassHierarchy> appView,
       SubtypingInfo subtypingInfo) {
-    return new Enqueuer(appView, subtypingInfo, null, Mode.INITIAL_TREE_SHAKING);
+    return new Enqueuer(
+        appView, subtypingInfo, null, Mode.INITIAL_TREE_SHAKING, MainDexTracingResult.NONE);
   }
 
   public static Enqueuer createForFinalTreeShaking(
@@ -26,30 +27,46 @@
       GraphConsumer keptGraphConsumer,
       Set<DexType> initialPrunedTypes) {
     Enqueuer enqueuer =
-        new Enqueuer(appView, subtypingInfo, keptGraphConsumer, Mode.FINAL_TREE_SHAKING);
+        new Enqueuer(
+            appView,
+            subtypingInfo,
+            keptGraphConsumer,
+            Mode.FINAL_TREE_SHAKING,
+            MainDexTracingResult.NONE);
     appView.withProtoShrinker(
         shrinker -> enqueuer.setInitialDeadProtoTypes(shrinker.getDeadProtoTypes()));
     enqueuer.setInitialPrunedTypes(initialPrunedTypes);
     return enqueuer;
   }
 
-  public static Enqueuer createForMainDexTracing(
-      AppView<? extends AppInfoWithClassHierarchy> appView,
-      SubtypingInfo subtypingInfo) {
-    return createForMainDexTracing(appView, subtypingInfo, null);
+  public static Enqueuer createForInitialMainDexTracing(
+      AppView<? extends AppInfoWithClassHierarchy> appView, SubtypingInfo subtypingInfo) {
+    return new Enqueuer(
+        appView, subtypingInfo, null, Mode.INITIAL_MAIN_DEX_TRACING, MainDexTracingResult.NONE);
   }
 
-  public static Enqueuer createForMainDexTracing(
+  public static Enqueuer createForFinalMainDexTracing(
       AppView<? extends AppInfoWithClassHierarchy> appView,
       SubtypingInfo subtypingInfo,
-      GraphConsumer keptGraphConsumer) {
-    return new Enqueuer(appView, subtypingInfo, keptGraphConsumer, Mode.MAIN_DEX_TRACING);
+      GraphConsumer keptGraphConsumer,
+      MainDexTracingResult previousMainDexTracingResult) {
+    return new Enqueuer(
+        appView,
+        subtypingInfo,
+        keptGraphConsumer,
+        Mode.FINAL_MAIN_DEX_TRACING,
+        previousMainDexTracingResult);
   }
 
   public static Enqueuer createForWhyAreYouKeeping(
       AppView<? extends AppInfoWithClassHierarchy> appView,
       SubtypingInfo subtypingInfo,
       GraphConsumer keptGraphConsumer) {
-    return new Enqueuer(appView, subtypingInfo, keptGraphConsumer, Mode.WHY_ARE_YOU_KEEPING);
+    return new Enqueuer(
+        appView,
+        subtypingInfo,
+        keptGraphConsumer,
+        Mode.WHY_ARE_YOU_KEEPING,
+        MainDexTracingResult.NONE);
   }
 }
diff --git a/src/main/java/com/android/tools/r8/utils/InternalOptions.java b/src/main/java/com/android/tools/r8/utils/InternalOptions.java
index 42ccd3c..84b76fd 100644
--- a/src/main/java/com/android/tools/r8/utils/InternalOptions.java
+++ b/src/main/java/com/android/tools/r8/utils/InternalOptions.java
@@ -1473,6 +1473,8 @@
     public boolean disableMappingToOriginalProgramVerification = false;
     public boolean allowInvalidCfAccessFlags =
         System.getProperty("com.android.tools.r8.allowInvalidCfAccessFlags") != null;
+    // TODO(b/177333791): Set to true
+    public boolean checkForNotExpandingMainDexTracingResult = false;
 
     // Flag to allow processing of resources in D8. A data resource consumer still needs to be
     // specified.
diff --git a/src/test/java/com/android/tools/r8/graph/MissingClassThrowingTest.java b/src/test/java/com/android/tools/r8/graph/MissingClassThrowingTest.java
index 213db97..7277139 100644
--- a/src/test/java/com/android/tools/r8/graph/MissingClassThrowingTest.java
+++ b/src/test/java/com/android/tools/r8/graph/MissingClassThrowingTest.java
@@ -65,7 +65,7 @@
   }
 
   @Test
-  public void testSuperTypeOfExceptions() throws Exception {
+  public void testSuperTypeOfExceptions() {
     AssertUtils.assertFailsCompilation(
         () ->
             testForR8(parameters.getBackend())
diff --git a/src/test/java/com/android/tools/r8/maindexlist/MainDexListMergeInRootTest.java b/src/test/java/com/android/tools/r8/maindexlist/MainDexListMergeInRootTest.java
new file mode 100644
index 0000000..7041c75
--- /dev/null
+++ b/src/test/java/com/android/tools/r8/maindexlist/MainDexListMergeInRootTest.java
@@ -0,0 +1,113 @@
+// Copyright (c) 2021, the R8 project authors. Please see the AUTHORS file
+// for details. All rights reserved. Use of this source code is governed by a
+// BSD-style license that can be found in the LICENSE file.
+
+package com.android.tools.r8.maindexlist;
+
+import static com.android.tools.r8.DiagnosticsMatcher.diagnosticMessage;
+import static com.android.tools.r8.utils.codeinspector.AssertUtils.assertFailsCompilation;
+import static org.hamcrest.CoreMatchers.containsString;
+import static org.junit.Assume.assumeTrue;
+
+import com.android.tools.r8.NeverClassInline;
+import com.android.tools.r8.NeverInline;
+import com.android.tools.r8.NoHorizontalClassMerging;
+import com.android.tools.r8.TestBase;
+import com.android.tools.r8.TestParameters;
+import com.android.tools.r8.TestParametersCollection;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+import org.junit.runners.Parameterized.Parameters;
+
+@RunWith(Parameterized.class)
+public class MainDexListMergeInRootTest extends TestBase {
+
+  private final TestParameters parameters;
+
+  @Parameters(name = "{0}")
+  public static TestParametersCollection data() {
+    return getTestParameters().withDexRuntimes().withAllApiLevels().build();
+  }
+
+  public MainDexListMergeInRootTest(TestParameters parameters) {
+    this.parameters = parameters;
+  }
+
+  @Test
+  public void testMainDexTracing() {
+    assumeTrue(parameters.getDexRuntimeVersion().isDalvik());
+    assertFailsCompilation(
+        () ->
+            testForR8(parameters.getBackend())
+                .addProgramClasses(OutsideMainDex.class, InsideA.class, InsideB.class, Main.class)
+                .addKeepClassAndMembersRules(Main.class)
+                .setMinApi(parameters.getApiLevel())
+                .enableNeverClassInliningAnnotations()
+                .enableNoHorizontalClassMergingAnnotations()
+                .enableInliningAnnotations()
+                .noMinification()
+                .addMainDexRules(
+                    "-keep class "
+                        + Main.class.getTypeName()
+                        + " { public static void main(***); }")
+                .addOptionsModification(
+                    options -> {
+                      options.testing.checkForNotExpandingMainDexTracingResult = true;
+                    })
+                .compileWithExpectedDiagnostics(
+                    diagnostics -> {
+                      diagnostics.assertErrorsMatch(
+                          diagnosticMessage(
+                              containsString(
+                                  "Class com.android.tools.r8.maindexlist"
+                                      + ".MainDexListMergeInRootTest$OutsideMainDex"
+                                      + " was not a main dex root in the first round")));
+                    }));
+  }
+
+  @NoHorizontalClassMerging
+  @NeverClassInline
+  public static class OutsideMainDex {
+
+    @NeverInline
+    public void print(int i) {
+      System.out.println("OutsideMainDex::print" + i);
+    }
+  }
+
+  @NeverClassInline
+  public static class InsideA {
+
+    public void bar() {
+      System.out.println("A::live");
+    }
+
+    /* Not a traced root */
+    @NeverInline
+    public void foo(int i) {
+      new OutsideMainDex().print(i);
+    }
+  }
+
+  @NeverClassInline
+  public static class InsideB {
+
+    @NeverInline
+    public void foo(int i) {
+      System.out.println("InsideB::live" + i);
+    }
+  }
+
+  public static class Main {
+
+    public static void main(String[] args) {
+      new InsideB().foo(args.length);
+      new InsideA().bar();
+    }
+
+    public void keptToKeepInsideANotLive() {
+      new InsideA().foo(System.currentTimeMillis() > 0 ? 0 : 1);
+    }
+  }
+}
diff --git a/src/test/java/com/android/tools/r8/utils/codeinspector/AssertUtils.java b/src/test/java/com/android/tools/r8/utils/codeinspector/AssertUtils.java
index 87eb583..e9634a5 100644
--- a/src/test/java/com/android/tools/r8/utils/codeinspector/AssertUtils.java
+++ b/src/test/java/com/android/tools/r8/utils/codeinspector/AssertUtils.java
@@ -13,9 +13,14 @@
 
 public class AssertUtils {
 
-  public static <E extends Throwable> void assertFailsCompilation(ThrowingAction<E> action)
-      throws E {
-    assertFailsCompilationIf(true, action);
+  public static void assertFailsCompilation(ThrowingAction<CompilationFailedException> action) {
+    try {
+      assertFailsCompilationIf(true, action);
+      return;
+    } catch (CompilationFailedException e) {
+      // Should have been caught
+    }
+    fail("Should have failed with a CompilationFailedException");
   }
 
   public static <E extends Throwable> void assertFailsCompilation(