Introduce trace references on dex for testing desugared lib

Bug: b:231287675
Change-Id: I48c3c20c05dd002bcbb7a4e6971982479ac5fe6d
diff --git a/src/test/java/com/android/tools/r8/desugar/desugaredlibrary/JavaUtilFunctionTest.java b/src/test/java/com/android/tools/r8/desugar/desugaredlibrary/JavaUtilFunctionTest.java
index d4891da..d14b845 100644
--- a/src/test/java/com/android/tools/r8/desugar/desugaredlibrary/JavaUtilFunctionTest.java
+++ b/src/test/java/com/android/tools/r8/desugar/desugaredlibrary/JavaUtilFunctionTest.java
@@ -4,7 +4,11 @@
 
 package com.android.tools.r8.desugar.desugaredlibrary;
 
-import static com.android.tools.r8.desugar.desugaredlibrary.test.CompilationSpecification.DEFAULT_SPECIFICATIONS;
+import static com.android.tools.r8.desugar.desugaredlibrary.test.CompilationSpecification.D8_L8DEBUG;
+import static com.android.tools.r8.desugar.desugaredlibrary.test.CompilationSpecification.D8_L8SHRINK;
+import static com.android.tools.r8.desugar.desugaredlibrary.test.CompilationSpecification.D8_L8SHRINK_TR;
+import static com.android.tools.r8.desugar.desugaredlibrary.test.CompilationSpecification.R8_L8SHRINK;
+import static com.android.tools.r8.desugar.desugaredlibrary.test.CompilationSpecification.R8_L8SHRINK_TR;
 import static com.android.tools.r8.desugar.desugaredlibrary.test.LibraryDesugaringSpecification.getJdk8Jdk11;
 import static com.android.tools.r8.utils.codeinspector.Matchers.isPresent;
 import static org.hamcrest.MatcherAssert.assertThat;
@@ -18,6 +22,7 @@
 import com.android.tools.r8.utils.StringUtils;
 import com.android.tools.r8.utils.codeinspector.ClassSubject;
 import com.android.tools.r8.utils.codeinspector.CodeInspector;
+import com.google.common.collect.ImmutableSet;
 import java.util.List;
 import java.util.function.Function;
 import org.junit.Test;
@@ -39,7 +44,7 @@
     return buildParameters(
         getTestParameters().withAllRuntimes().withAllApiLevelsAlsoForCf().build(),
         getJdk8Jdk11(),
-        DEFAULT_SPECIFICATIONS);
+        ImmutableSet.of(D8_L8DEBUG, D8_L8SHRINK, R8_L8SHRINK, D8_L8SHRINK_TR, R8_L8SHRINK_TR));
   }
 
   public JavaUtilFunctionTest(
diff --git a/src/test/java/com/android/tools/r8/desugar/desugaredlibrary/test/CompilationSpecification.java b/src/test/java/com/android/tools/r8/desugar/desugaredlibrary/test/CompilationSpecification.java
index 3aecfba..0a8d35b 100644
--- a/src/test/java/com/android/tools/r8/desugar/desugaredlibrary/test/CompilationSpecification.java
+++ b/src/test/java/com/android/tools/r8/desugar/desugaredlibrary/test/CompilationSpecification.java
@@ -12,14 +12,17 @@
 import java.util.Set;
 
 public enum CompilationSpecification {
-  D8_L8DEBUG(false, false, false, DEBUG),
-  D8_L8SHRINK(false, true, false, RELEASE),
+  D8_L8DEBUG(false, false, false, false, DEBUG),
+  D8_L8SHRINK(false, true, false, false, RELEASE),
   // In theory no build system uses R8_L8DEBUG, for local debugging only.
-  R8_L8DEBUG(true, false, false, RELEASE),
-  R8_L8SHRINK(true, true, false, RELEASE),
+  R8_L8DEBUG(true, false, false, false, RELEASE),
+  R8_L8SHRINK(true, true, false, false, RELEASE),
   // The D8CFTOCF specifications can run either in CF or be dexed afterwards.
-  D8CF2CF_L8DEBUG(false, false, true, DEBUG),
-  D8CF2CF_L8SHRINK(false, true, true, RELEASE);
+  D8CF2CF_L8DEBUG(false, false, true, false, DEBUG),
+  D8CF2CF_L8SHRINK(false, true, true, true, RELEASE),
+  // Variants with trace reference in dex.
+  D8_L8SHRINK_TR(false, true, false, true, RELEASE),
+  R8_L8SHRINK_TR(true, true, false, true, RELEASE);
 
   public static Set<CompilationSpecification> DEFAULT_SPECIFICATIONS =
       ImmutableSet.of(D8_L8DEBUG, D8_L8SHRINK, R8_L8SHRINK);
@@ -29,13 +32,19 @@
   private final boolean programShrink;
   private final boolean l8Shrink;
   private final boolean cfToCf;
+  private final boolean traceReferences;
   private final CompilationMode programCompilationMode;
 
   CompilationSpecification(
-      boolean programShrink, boolean l8Shrink, boolean cfToCf, CompilationMode mode) {
+      boolean programShrink,
+      boolean l8Shrink,
+      boolean cfToCf,
+      boolean traceReferences,
+      CompilationMode mode) {
     this.programShrink = programShrink;
     this.l8Shrink = l8Shrink;
     this.cfToCf = cfToCf;
+    this.traceReferences = traceReferences;
     this.programCompilationMode = mode;
   }
 
@@ -54,4 +63,8 @@
   public boolean isCfToCf() {
     return cfToCf;
   }
+
+  public boolean isTraceReferences() {
+    return traceReferences;
+  }
 }
diff --git a/src/test/java/com/android/tools/r8/desugar/desugaredlibrary/test/DesugaredLibraryTestBuilder.java b/src/test/java/com/android/tools/r8/desugar/desugaredlibrary/test/DesugaredLibraryTestBuilder.java
index 33ff933..18bf530 100644
--- a/src/test/java/com/android/tools/r8/desugar/desugaredlibrary/test/DesugaredLibraryTestBuilder.java
+++ b/src/test/java/com/android/tools/r8/desugar/desugaredlibrary/test/DesugaredLibraryTestBuilder.java
@@ -4,8 +4,6 @@
 
 package com.android.tools.r8.desugar.desugaredlibrary.test;
 
-import static com.android.tools.r8.desugar.desugaredlibrary.test.CompilationSpecification.D8CF2CF_L8SHRINK;
-
 import com.android.tools.r8.CompilationFailedException;
 import com.android.tools.r8.D8TestCompileResult;
 import com.android.tools.r8.FeatureSplit;
@@ -321,13 +319,11 @@
     if (!compilationSpecification.isL8Shrink()) {
       return compileDesugaredLibrary(null);
     }
-    if (!compilationSpecification.isCfToCf()) {
+    if (!compilationSpecification.isTraceReferences()) {
       // When going to dex we can get the generated keep rule through the keep rule consumer.
       assert keepRuleConsumer != null;
       return compileDesugaredLibrary(keepRuleConsumer.get());
     }
-    // In D8CF2CF_L8SHRINK, we use trace reference to extract the keep rules.
-    assert compilationSpecification == D8CF2CF_L8SHRINK;
     L8TestCompileResult nonShrunk =
         test.testForL8(parameters.getApiLevel(), Backend.CF)
             .apply(libraryDesugaringSpecification::configureL8TestBuilder)