Use trace references in library desugaring tests

Fixes: b/397554277
Change-Id: Ifbe628667198304ca9b6cf1d122101a79319fa7f
diff --git a/src/main/java/com/android/tools/r8/tracereferences/Formatter.java b/src/main/java/com/android/tools/r8/tracereferences/Formatter.java
index e3978f3..cff8cf9 100644
--- a/src/main/java/com/android/tools/r8/tracereferences/Formatter.java
+++ b/src/main/java/com/android/tools/r8/tracereferences/Formatter.java
@@ -50,7 +50,7 @@
     StringUtils.append(
         output,
         ListUtils.map(method.getFormalTypes(), TypeReference::getTypeName),
-        ",",
+        ", ",
         BraceType.PARENS);
   }
 
diff --git a/src/test/java/com/android/tools/r8/desugar/desugaredlibrary/DontKeepBootstrapClassesTest.java b/src/test/java/com/android/tools/r8/desugar/desugaredlibrary/DontKeepBootstrapClassesTest.java
index 058f6cf..72a20c7 100644
--- a/src/test/java/com/android/tools/r8/desugar/desugaredlibrary/DontKeepBootstrapClassesTest.java
+++ b/src/test/java/com/android/tools/r8/desugar/desugaredlibrary/DontKeepBootstrapClassesTest.java
@@ -5,13 +5,15 @@
 
 import static com.android.tools.r8.desugar.desugaredlibrary.test.CompilationSpecification.DEFAULT_SPECIFICATIONS;
 import static com.android.tools.r8.desugar.desugaredlibrary.test.LibraryDesugaringSpecification.getJdk8Jdk11;
-import static org.junit.Assert.assertTrue;
+import static org.hamcrest.CoreMatchers.containsString;
+import static org.hamcrest.MatcherAssert.assertThat;
 
 import com.android.tools.r8.TestParameters;
 import com.android.tools.r8.desugar.desugaredlibrary.test.CompilationSpecification;
 import com.android.tools.r8.desugar.desugaredlibrary.test.CustomLibrarySpecification;
 import com.android.tools.r8.desugar.desugaredlibrary.test.LibraryDesugaringSpecification;
 import com.android.tools.r8.utils.AndroidApiLevel;
+import com.android.tools.r8.utils.StringUtils;
 import java.util.Arrays;
 import java.util.List;
 import java.util.function.Consumer;
@@ -56,11 +58,10 @@
             keepRule -> {
               if (libraryDesugaringSpecification.hasEmulatedInterfaceDesugaring(parameters)) {
                 String prefix = libraryDesugaringSpecification.functionPrefix(parameters);
-                assertTrue(
-                    keepRule.stream()
-                        .anyMatch(
-                            kr ->
-                                kr.contains("-keep class " + prefix + ".util.function.Consumer")));
+                String allKeepRules = StringUtils.lines(keepRule);
+                assertThat(
+                    allKeepRules,
+                    containsString("-keep interface " + prefix + ".util.function.Consumer"));
               }
             });
   }
diff --git a/src/test/java/com/android/tools/r8/desugar/desugaredlibrary/FeatureSplitTest.java b/src/test/java/com/android/tools/r8/desugar/desugaredlibrary/FeatureSplitTest.java
index 6799d90..1b8aa3b 100644
--- a/src/test/java/com/android/tools/r8/desugar/desugaredlibrary/FeatureSplitTest.java
+++ b/src/test/java/com/android/tools/r8/desugar/desugaredlibrary/FeatureSplitTest.java
@@ -9,6 +9,8 @@
 import static com.android.tools.r8.desugar.desugaredlibrary.test.LibraryDesugaringSpecification.getJdk8Jdk11;
 import static junit.framework.TestCase.assertEquals;
 import static junit.framework.TestCase.assertTrue;
+import static org.hamcrest.CoreMatchers.containsString;
+import static org.hamcrest.MatcherAssert.assertThat;
 import static org.junit.Assume.assumeTrue;
 
 import com.android.tools.r8.TestParameters;
@@ -19,7 +21,6 @@
 import com.android.tools.r8.desugar.desugaredlibrary.test.CompilationSpecification;
 import com.android.tools.r8.desugar.desugaredlibrary.test.DesugaredLibraryTestCompileResult;
 import com.android.tools.r8.desugar.desugaredlibrary.test.LibraryDesugaringSpecification;
-import com.android.tools.r8.dexsplitter.SplitterTestBase;
 import com.android.tools.r8.utils.AndroidApiLevel;
 import com.android.tools.r8.utils.StringUtils;
 import com.android.tools.r8.utils.codeinspector.CodeInspector;
@@ -88,17 +89,19 @@
     }
     // Ensure count, toArray and forEach are kept.
     String prefix = libraryDesugaringSpecification.functionPrefix(parameters);
-    assertTrue(
-        keepRules.contains(
+    assertThat(
+        keepRules,
+        containsString(
             "-keep class j$.lang.Iterable$-EL {\n"
-                + "    void forEach(java.lang.Iterable, "
+                + "  public static void forEach(java.lang.Iterable, "
                 + prefix
                 + ".util.function.Consumer);"));
-    assertTrue(
-        keepRules.contains(
-            "-keep class j$.util.stream.Stream {\n"
-                + "    long count();\n"
-                + "    java.lang.Object[] toArray();"));
+    assertThat(
+        keepRules,
+        containsString(
+            "-keep interface j$.util.stream.Stream {\n"
+                + "  public long count();\n"
+                + "  public java.lang.Object[] toArray();"));
   }
 
   private void assertClassPresent(Path appPath, Class<?> present) throws IOException {
@@ -244,27 +247,18 @@
     }
 
     public CompiledWithFeature invoke(FeatureSplitTest tester) throws Throwable {
-
-      basePath = temp.newFile("base.zip").toPath();
-      feature1Path = temp.newFile("feature1.zip").toPath();
-      feature2Path = temp.newFile("feature2.zip").toPath();
-
       DesugaredLibraryTestCompileResult<?> compileResult =
           tester
               .testForDesugaredLibrary(
                   parameters, libraryDesugaringSpecification, compilationSpecification)
               .addProgramClasses(BaseClass.class, RunInterface.class, SplitRunner.class)
-              .addFeatureSplit(
-                  builder ->
-                      SplitterTestBase.simpleSplitProvider(
-                          builder, feature1Path, temp, FeatureClass.class))
-              .addFeatureSplit(
-                  builder ->
-                      SplitterTestBase.simpleSplitProvider(
-                          builder, feature2Path, temp, FeatureClass2.class))
+              .addFeatureSplit(FeatureClass.class)
+              .addFeatureSplit(FeatureClass2.class)
               .addKeepAllClassesRule()
-              .compile()
-              .writeToZip(basePath);
+              .compile();
+      basePath = compileResult.writeToZip();
+      feature1Path = compileResult.getFeature(0);
+      feature2Path = compileResult.getFeature(1);
 
       // Stream desugaring is not needed >= N.
       if (parameters.getApiLevel().getLevel() >= AndroidApiLevel.N.getLevel()) {
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 db9a6f1..0d661b2 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
@@ -6,9 +6,7 @@
 
 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;
@@ -44,7 +42,7 @@
     return buildParameters(
         getTestParameters().withAllRuntimes().withAllApiLevelsAlsoForCf().build(),
         getJdk8Jdk11(),
-        ImmutableSet.of(D8_L8DEBUG, D8_L8SHRINK, R8_L8SHRINK, D8_L8SHRINK_TR, R8_L8SHRINK_TR));
+        ImmutableSet.of(D8_L8DEBUG, D8_L8SHRINK, R8_L8SHRINK));
   }
 
   public JavaUtilFunctionTest(
diff --git a/src/test/java/com/android/tools/r8/desugar/desugaredlibrary/ProgramRewritingTest.java b/src/test/java/com/android/tools/r8/desugar/desugaredlibrary/ProgramRewritingTest.java
index ca43f85..dd9fe54 100644
--- a/src/test/java/com/android/tools/r8/desugar/desugaredlibrary/ProgramRewritingTest.java
+++ b/src/test/java/com/android/tools/r8/desugar/desugaredlibrary/ProgramRewritingTest.java
@@ -189,35 +189,51 @@
   }
 
   private void assertGeneratedKeepRulesAreCorrect(String keepRules) {
-
     String prefix = libraryDesugaringSpecification.functionPrefix(parameters);
     String expectedResult =
         StringUtils.lines(
             "-keep class j$.util.Collection$-EL {",
-            "    j$.util.stream.Stream stream(java.util.Collection);",
+            "  public static j$.util.stream.Stream stream(java.util.Collection);",
             "}",
             "-keep class j$.util.Comparator$-CC {",
-            "    java.util.Comparator comparingInt(" + prefix + ".util.function.ToIntFunction);",
+            "  public static java.util.Comparator comparingInt("
+                + prefix
+                + ".util.function.ToIntFunction);",
             "}",
             "-keep class j$.util.DesugarArrays {",
-            "    j$.util.Spliterator spliterator(java.lang.Object[]);",
-            "    j$.util.Spliterator spliterator(java.lang.Object[], int, int);",
-            "    j$.util.stream.Stream stream(java.lang.Object[]);",
-            "    j$.util.stream.Stream stream(java.lang.Object[], int, int);",
+            "  public static j$.util.Spliterator spliterator(java.lang.Object[]);",
+            "  public static j$.util.Spliterator spliterator(java.lang.Object[], int, int);",
+            "  public static j$.util.stream.Stream stream(java.lang.Object[]);",
+            "  public static j$.util.stream.Stream stream(java.lang.Object[], int, int);",
             "}",
             "-keep class j$.util.List$-EL {",
-            "    void sort(java.util.List, java.util.Comparator);",
+            "  public static void sort(java.util.List, java.util.Comparator);",
             "}",
             "-keep class j$.util.Set$-EL {",
-            "    j$.util.Spliterator spliterator(java.util.Set);",
+            "  public static j$.util.Spliterator spliterator(java.util.Set);",
             "}",
-            "-keep class j$.util.Spliterator",
-            "-keep class " + prefix + ".util.function.ToIntFunction { *; }",
+            "-keep interface j$.util.Spliterator {",
+            "}");
+    if (prefix.equals("j$")) {
+      expectedResult +=
+          StringUtils.lines(
+              "-keep interface j$.util.function.ToIntFunction {",
+              "  public int applyAsInt(java.lang.Object);",
+              "}");
+    }
+    expectedResult +=
+        StringUtils.lines(
             "-keep class j$.util.stream.IntStream$-CC {",
-            "    j$.util.stream.IntStream range(int, int);",
+            "  public static j$.util.stream.IntStream range(int, int);",
             "}",
-            "-keep class j$.util.stream.IntStream",
-            "-keep class j$.util.stream.Stream");
+            "-keep interface j$.util.stream.IntStream {",
+            "}",
+            "-keep interface j$.util.stream.Stream {",
+            "}");
+    if (prefix.equals("java")) {
+      expectedResult +=
+          StringUtils.lines("-keep interface java.util.function.ToIntFunction {", "}");
+    }
     assertEquals(expectedResult.trim(), keepRules.trim());
   }
 }
diff --git a/src/test/java/com/android/tools/r8/desugar/desugaredlibrary/jdk11/FileChannelSizeTest.java b/src/test/java/com/android/tools/r8/desugar/desugaredlibrary/jdk11/FileChannelSizeTest.java
index f3f8b01..5b55897 100644
--- a/src/test/java/com/android/tools/r8/desugar/desugaredlibrary/jdk11/FileChannelSizeTest.java
+++ b/src/test/java/com/android/tools/r8/desugar/desugaredlibrary/jdk11/FileChannelSizeTest.java
@@ -35,7 +35,7 @@
   private static final String EXPECTED_KEEP_RULES =
       StringUtils.lines(
           "-keep class j$.nio.channels.DesugarChannels {",
-          "    java.nio.channels.FileChannel"
+          "  public static java.nio.channels.FileChannel"
               + " convertMaybeLegacyFileChannelFromLibrary(java.nio.channels.FileChannel);",
           "}");
 
diff --git a/src/test/java/com/android/tools/r8/optimize/serviceloader/ServiceLoaderDesugaredLibraryTest.java b/src/test/java/com/android/tools/r8/optimize/serviceloader/ServiceLoaderDesugaredLibraryTest.java
index 09b18be..cc46480 100644
--- a/src/test/java/com/android/tools/r8/optimize/serviceloader/ServiceLoaderDesugaredLibraryTest.java
+++ b/src/test/java/com/android/tools/r8/optimize/serviceloader/ServiceLoaderDesugaredLibraryTest.java
@@ -4,8 +4,8 @@
 
 package com.android.tools.r8.optimize.serviceloader;
 
-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_TR;
+import static com.android.tools.r8.desugar.desugaredlibrary.test.CompilationSpecification.D8_L8SHRINK;
+import static com.android.tools.r8.desugar.desugaredlibrary.test.CompilationSpecification.R8_L8SHRINK;
 import static com.android.tools.r8.desugar.desugaredlibrary.test.LibraryDesugaringSpecification.JDK11;
 import static org.junit.Assume.assumeTrue;
 
@@ -68,7 +68,7 @@
     return buildParameters(
         getTestParameters().withAllRuntimes().withAllApiLevelsAlsoForCf().build(),
         ImmutableList.of(JDK11),
-        ImmutableList.of(D8_L8SHRINK_TR, R8_L8SHRINK_TR));
+        ImmutableList.of(D8_L8SHRINK, R8_L8SHRINK));
   }
 
   private static final String EXPECTED_OUTPUT = StringUtils.lines("true", "true", "1", "1");
@@ -114,7 +114,7 @@
   @Test
   public void testD8WithoutDesugaredLibrary() throws Throwable {
     parameters.assumeR8TestParameters();
-    assumeTrue(compilationSpecification == D8_L8SHRINK_TR);
+    assumeTrue(compilationSpecification == D8_L8SHRINK);
     testForD8(parameters.getBackend())
         .addLibraryFiles(ToolHelper.getAndroidJar(apiLevelWithJavaTime()))
         .addInnerClasses(getClass())
@@ -130,7 +130,7 @@
   @Test
   public void testR8WithoutDesugaredLibrary() throws Throwable {
     parameters.assumeR8TestParameters();
-    assumeTrue(compilationSpecification == D8_L8SHRINK_TR);
+    assumeTrue(compilationSpecification == D8_L8SHRINK);
     testForR8(parameters.getBackend())
         .addLibraryFiles(ToolHelper.getAndroidJar(apiLevelWithJavaTime()))
         .addInnerClasses(getClass())
diff --git a/src/test/testbase/java/com/android/tools/r8/R8PartialTestCompileResult.java b/src/test/testbase/java/com/android/tools/r8/R8PartialTestCompileResult.java
index 18f75fb..13e277b 100644
--- a/src/test/testbase/java/com/android/tools/r8/R8PartialTestCompileResult.java
+++ b/src/test/testbase/java/com/android/tools/r8/R8PartialTestCompileResult.java
@@ -48,6 +48,16 @@
   }
 
   @Override
+  public boolean isR8PartialCompileResult() {
+    return true;
+  }
+
+  @Override
+  public R8PartialTestCompileResult asR8PartialCompileResult() {
+    return this;
+  }
+
+  @Override
   public R8PartialTestCompileResult self() {
     return this;
   }
diff --git a/src/test/testbase/java/com/android/tools/r8/R8TestCompileResult.java b/src/test/testbase/java/com/android/tools/r8/R8TestCompileResult.java
index 9354c62..707760d 100644
--- a/src/test/testbase/java/com/android/tools/r8/R8TestCompileResult.java
+++ b/src/test/testbase/java/com/android/tools/r8/R8TestCompileResult.java
@@ -47,6 +47,16 @@
   }
 
   @Override
+  public boolean isR8CompileResult() {
+    return true;
+  }
+
+  @Override
+  public R8TestCompileResult asR8CompileResult() {
+    return this;
+  }
+
+  @Override
   public R8TestCompileResult self() {
     return this;
   }
diff --git a/src/test/testbase/java/com/android/tools/r8/TestCompileResult.java b/src/test/testbase/java/com/android/tools/r8/TestCompileResult.java
index a2c139a..8fd2315 100644
--- a/src/test/testbase/java/com/android/tools/r8/TestCompileResult.java
+++ b/src/test/testbase/java/com/android/tools/r8/TestCompileResult.java
@@ -90,6 +90,22 @@
     this.libraryDesugaringTestConfiguration = libraryDesugaringTestConfiguration;
   }
 
+  public boolean isR8CompileResult() {
+    return false;
+  }
+
+  public R8TestCompileResult asR8CompileResult() {
+    return null;
+  }
+
+  public boolean isR8PartialCompileResult() {
+    return false;
+  }
+
+  public R8PartialTestCompileResult asR8PartialCompileResult() {
+    return null;
+  }
+
   public CR addVmArguments(Collection<String> arguments) {
     vmArguments.addAll(arguments);
     return self();
diff --git a/src/test/testbase/java/com/android/tools/r8/desugar/desugaredlibrary/test/CompilationSpecification.java b/src/test/testbase/java/com/android/tools/r8/desugar/desugaredlibrary/test/CompilationSpecification.java
index 0a8d35b..7938093 100644
--- a/src/test/testbase/java/com/android/tools/r8/desugar/desugaredlibrary/test/CompilationSpecification.java
+++ b/src/test/testbase/java/com/android/tools/r8/desugar/desugaredlibrary/test/CompilationSpecification.java
@@ -12,17 +12,14 @@
 import java.util.Set;
 
 public enum CompilationSpecification {
-  D8_L8DEBUG(false, false, false, false, DEBUG),
-  D8_L8SHRINK(false, true, false, false, RELEASE),
+  D8_L8DEBUG(false, false, false, DEBUG),
+  D8_L8SHRINK(false, true, false, RELEASE),
   // In theory no build system uses R8_L8DEBUG, for local debugging only.
-  R8_L8DEBUG(true, false, false, false, RELEASE),
-  R8_L8SHRINK(true, true, false, false, RELEASE),
+  R8_L8DEBUG(true, false, false, RELEASE),
+  R8_L8SHRINK(true, true, false, RELEASE),
   // The D8CFTOCF specifications can run either in CF or be dexed afterwards.
-  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);
+  D8CF2CF_L8DEBUG(false, false, true, DEBUG),
+  D8CF2CF_L8SHRINK(false, true, true, RELEASE);
 
   public static Set<CompilationSpecification> DEFAULT_SPECIFICATIONS =
       ImmutableSet.of(D8_L8DEBUG, D8_L8SHRINK, R8_L8SHRINK);
@@ -32,19 +29,16 @@
   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,
-      boolean traceReferences,
       CompilationMode mode) {
     this.programShrink = programShrink;
     this.l8Shrink = l8Shrink;
     this.cfToCf = cfToCf;
-    this.traceReferences = traceReferences;
     this.programCompilationMode = mode;
   }
 
@@ -63,8 +57,4 @@
   public boolean isCfToCf() {
     return cfToCf;
   }
-
-  public boolean isTraceReferences() {
-    return traceReferences;
-  }
 }
diff --git a/src/test/testbase/java/com/android/tools/r8/desugar/desugaredlibrary/test/DesugaredLibraryTestBuilder.java b/src/test/testbase/java/com/android/tools/r8/desugar/desugaredlibrary/test/DesugaredLibraryTestBuilder.java
index 7470580..1c2b231 100644
--- a/src/test/testbase/java/com/android/tools/r8/desugar/desugaredlibrary/test/DesugaredLibraryTestBuilder.java
+++ b/src/test/testbase/java/com/android/tools/r8/desugar/desugaredlibrary/test/DesugaredLibraryTestBuilder.java
@@ -9,7 +9,6 @@
 import com.android.tools.r8.CompilationMode;
 import com.android.tools.r8.D8TestBuilder;
 import com.android.tools.r8.D8TestCompileResult;
-import com.android.tools.r8.FeatureSplit;
 import com.android.tools.r8.L8TestBuilder;
 import com.android.tools.r8.L8TestCompileResult;
 import com.android.tools.r8.LibraryDesugaringTestConfiguration;
@@ -24,7 +23,6 @@
 import com.android.tools.r8.TestRuntime;
 import com.android.tools.r8.TestShrinkerBuilder;
 import com.android.tools.r8.desugar.desugaredlibrary.DesugaredLibraryTestBase;
-import com.android.tools.r8.desugar.desugaredlibrary.DesugaredLibraryTestBase.KeepRuleConsumer;
 import com.android.tools.r8.ir.desugar.desugaredlibrary.DesugaredLibrarySpecificationParser;
 import com.android.tools.r8.profile.art.ArtProfileConsumer;
 import com.android.tools.r8.profile.art.ArtProfileForRewriting;
@@ -35,6 +33,7 @@
 import com.android.tools.r8.utils.ConsumerUtils;
 import com.android.tools.r8.utils.FileUtils;
 import com.android.tools.r8.utils.InternalOptions;
+import com.android.tools.r8.utils.ThrowingConsumer;
 import com.android.tools.r8.utils.codeinspector.VerticallyMergedClassesInspector;
 import com.google.common.base.Charsets;
 import java.io.IOException;
@@ -42,10 +41,10 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
+import java.util.Collections;
 import java.util.List;
 import java.util.concurrent.ExecutionException;
 import java.util.function.Consumer;
-import java.util.function.Function;
 import org.junit.Assume;
 
 public class DesugaredLibraryTestBuilder<T extends DesugaredLibraryTestBase> {
@@ -61,7 +60,6 @@
   private boolean l8FinalPrefixVerification = true;
   private boolean overrideDefaultLibrary = false;
   private CustomLibrarySpecification customLibrarySpecification = null;
-  private TestingKeepRuleConsumer keepRuleConsumer = null;
   private List<ExternalArtProfile> l8ResidualArtProfiles = new ArrayList<>();
   private boolean managedPostPrefix = false;
 
@@ -84,10 +82,6 @@
         LibraryDesugaringTestConfiguration.builder()
             .addDesugaredLibraryConfiguration(
                 StringResource.fromFile(libraryDesugaringSpecification.getSpecification()));
-    if (compilationSpecification.isL8Shrink() && !compilationSpecification.isCfToCf()) {
-      keepRuleConsumer = new TestingKeepRuleConsumer();
-      libraryConfBuilder.setKeepRuleConsumer(keepRuleConsumer);
-    }
     builder.enableCoreLibraryDesugaring(libraryConfBuilder.build());
   }
 
@@ -224,17 +218,16 @@
   }
 
   private void withD8TestBuilder(Consumer<D8TestBuilder> consumer) {
-    if (!builder.isD8TestBuilder()) {
-      return;
+    if (builder.isD8TestBuilder()) {
+      consumer.accept((D8TestBuilder) builder);
     }
-    consumer.accept((D8TestBuilder) builder);
   }
 
-  private void withR8TestBuilder(Consumer<R8TestBuilder<?, ?, ?>> consumer) {
-    if (!builder.isTestShrinkerBuilder()) {
-      return;
+  private <E extends Throwable> void withR8TestBuilder(
+      ThrowingConsumer<R8TestBuilder<?, ?, ?>, E> consumer) throws E {
+    if (builder.isTestShrinkerBuilder()) {
+      consumer.accept((R8TestBuilder<?, ?, ?>) builder);
     }
-    consumer.accept((R8TestBuilder<?, ?, ?>) builder);
   }
 
   public DesugaredLibraryTestBuilder<T> allowUnusedDontWarnPatterns() {
@@ -267,8 +260,8 @@
     return this;
   }
 
-  public DesugaredLibraryTestBuilder<T> applyIfR8TestBuilder(
-      Consumer<R8TestBuilder<?, ?, ?>> consumer) {
+  public <E extends Throwable> DesugaredLibraryTestBuilder<T> applyIfR8TestBuilder(
+      ThrowingConsumer<R8TestBuilder<?, ?, ?>, E> consumer) throws E {
     withR8TestBuilder(consumer);
     return this;
   }
@@ -330,9 +323,8 @@
     return this;
   }
 
-  public DesugaredLibraryTestBuilder<T> addFeatureSplit(
-      Function<FeatureSplit.Builder, FeatureSplit> featureSplitBuilder) {
-    withR8TestBuilder(b -> b.addFeatureSplit(featureSplitBuilder));
+  public DesugaredLibraryTestBuilder<T> addFeatureSplit(Class<?>... classes) throws IOException {
+    withR8TestBuilder(b -> b.addFeatureSplit(classes));
     return this;
   }
 
@@ -417,8 +409,8 @@
 
   private DesugaredLibraryTestCompileResult<T> internalCompile(
       TestCompileResult<?, ? extends SingleTestRunResult<?>> compile)
-      throws CompilationFailedException, IOException, ExecutionException {
-    L8TestCompileResult l8Compile = compileDesugaredLibrary(compile, keepRuleConsumer);
+      throws CompilationFailedException, IOException {
+    L8TestCompileResult l8Compile = compileDesugaredLibrary(compile);
     D8TestCompileResult customLibCompile = compileCustomLib();
     if (managedPostPrefix) {
       System.clearProperty("com.android.tools.r8.desugaredLibraryPostPrefix");
@@ -427,7 +419,6 @@
         test,
         compile,
         parameters,
-        libraryDesugaringSpecification,
         compilationSpecification,
         customLibCompile,
         l8Compile,
@@ -442,29 +433,22 @@
   }
 
   private L8TestCompileResult compileDesugaredLibrary(
-      TestCompileResult<?, ? extends SingleTestRunResult<?>> compile,
-      KeepRuleConsumer keepRuleConsumer)
-      throws CompilationFailedException, IOException, ExecutionException {
+      TestCompileResult<?, ? extends SingleTestRunResult<?>> compile)
+      throws CompilationFailedException, IOException {
     if (!compilationSpecification.isL8Shrink()) {
-      return compileDesugaredLibrary(null);
-    }
-    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());
+      return internalCompileDesugaredLibrary(null);
     }
     L8TestCompileResult nonShrunk =
         test.testForL8(parameters.getApiLevel(), Backend.CF)
             .apply(libraryDesugaringSpecification::configureL8TestBuilder)
             .apply(b -> configure(b, Backend.CF))
             .compile();
-    String keepRules =
-        collectKeepRulesWithTraceReferences(compile.writeToZip(), nonShrunk.writeToZip());
-    return compileDesugaredLibrary(keepRules);
+    String keepRules = collectKeepRulesWithTraceReferences(compile, nonShrunk);
+    return internalCompileDesugaredLibrary(keepRules);
   }
 
-  private L8TestCompileResult compileDesugaredLibrary(String keepRule)
-      throws CompilationFailedException, IOException, ExecutionException {
+  private L8TestCompileResult internalCompileDesugaredLibrary(String keepRule)
+      throws CompilationFailedException, IOException {
     assert !compilationSpecification.isL8Shrink() || keepRule != null;
     return test.testForL8(parameters.getApiLevel(), parameters.getBackend())
         .apply(
@@ -476,21 +460,24 @@
   }
 
   private void configure(L8TestBuilder l8Builder, Backend backend) {
-    l8Builder
-        .applyIf(!l8FinalPrefixVerification, L8TestBuilder::ignoreFinalPrefixVerification)
-        .applyIf(
-            compilationSpecification.isL8Shrink() && !backend.isCf() && !l8ExtraKeepRules.isEmpty(),
-            b -> b.addKeepRules(l8ExtraKeepRules))
-        .addOptionsModifier(l8OptionModifier);
-    for (ArtProfileForRewriting artProfileForRewriting : l8ArtProfilesForRewriting) {
-      l8Builder.addArtProfileForRewriting(
-          artProfileForRewriting.getArtProfileProvider(),
-          artProfileForRewriting.getResidualArtProfileConsumer());
+    l8Builder.applyIf(!l8FinalPrefixVerification, L8TestBuilder::ignoreFinalPrefixVerification);
+    if (backend.isDex()) {
+      l8Builder
+          .applyIf(
+              compilationSpecification.isL8Shrink() && !l8ExtraKeepRules.isEmpty(),
+              b -> b.addKeepRules(l8ExtraKeepRules))
+          .addOptionsModifier(l8OptionModifier);
+      for (ArtProfileForRewriting artProfileForRewriting : l8ArtProfilesForRewriting) {
+        l8Builder.addArtProfileForRewriting(
+            artProfileForRewriting.getArtProfileProvider(),
+            artProfileForRewriting.getResidualArtProfileConsumer());
+      }
     }
   }
 
   public String collectKeepRulesWithTraceReferences(
-      Path desugaredProgramClassFile, Path desugaredLibraryClassFile)
+      TestCompileResult<?, ? extends SingleTestRunResult<?>> compileResult,
+      L8TestCompileResult l8TestCompileResult)
       throws CompilationFailedException, IOException {
     Path generatedKeepRules = test.temp.newFile().toPath();
     ArrayList<String> args = new ArrayList<>();
@@ -500,9 +487,19 @@
       args.add(libraryFile.toString());
     }
     args.add("--target");
-    args.add(desugaredLibraryClassFile.toString());
+    args.add(l8TestCompileResult.writeToZip().toString());
     args.add("--source");
-    args.add(desugaredProgramClassFile.toString());
+    args.add(compileResult.writeToZip().toString());
+    List<Path> features = Collections.emptyList();
+    if (compileResult.isR8CompileResult()) {
+      features = compileResult.asR8CompileResult().getFeatures();
+    } else if (compileResult.isR8PartialCompileResult()) {
+      features = compileResult.asR8PartialCompileResult().getFeatures();
+    }
+    for (Path feature : features) {
+      args.add("--source");
+      args.add(feature.toString());
+    }
     args.add("--output");
     args.add(generatedKeepRules.toString());
     args.add("--map-diagnostics");
diff --git a/src/test/testbase/java/com/android/tools/r8/desugar/desugaredlibrary/test/DesugaredLibraryTestCompileResult.java b/src/test/testbase/java/com/android/tools/r8/desugar/desugaredlibrary/test/DesugaredLibraryTestCompileResult.java
index 0f11fc1..2fb1cc6 100644
--- a/src/test/testbase/java/com/android/tools/r8/desugar/desugaredlibrary/test/DesugaredLibraryTestCompileResult.java
+++ b/src/test/testbase/java/com/android/tools/r8/desugar/desugaredlibrary/test/DesugaredLibraryTestCompileResult.java
@@ -5,6 +5,7 @@
 package com.android.tools.r8.desugar.desugaredlibrary.test;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.fail;
 
 import com.android.tools.r8.CompilationFailedException;
 import com.android.tools.r8.D8TestCompileResult;
@@ -31,7 +32,6 @@
   private final T test;
   private final TestCompileResult<?, ? extends SingleTestRunResult<?>> compileResult;
   private final TestParameters parameters;
-  private final LibraryDesugaringSpecification libraryDesugaringSpecification;
   private final CompilationSpecification compilationSpecification;
   private final D8TestCompileResult customLibCompile;
   private final L8TestCompileResult l8Compile;
@@ -44,7 +44,6 @@
       T test,
       TestCompileResult<?, ? extends SingleTestRunResult<?>> compileResult,
       TestParameters parameters,
-      LibraryDesugaringSpecification libraryDesugaringSpecification,
       CompilationSpecification compilationSpecification,
       D8TestCompileResult customLibCompile,
       L8TestCompileResult l8Compile,
@@ -53,7 +52,6 @@
     this.test = test;
     this.compileResult = compileResult;
     this.parameters = parameters;
-    this.libraryDesugaringSpecification = libraryDesugaringSpecification;
     this.compilationSpecification = compilationSpecification;
     this.customLibCompile = customLibCompile;
     this.l8Compile = l8Compile;
@@ -66,11 +64,15 @@
     return this;
   }
 
-  public <E extends Throwable> DesugaredLibraryTestCompileResult<T> inspectCustomLib(
-      ThrowingConsumer<CodeInspector, E> consumer) throws IOException, E {
-    assert customLibCompile != null;
-    customLibCompile.inspect(consumer);
-    return this;
+  public Path getFeature(int index) {
+    if (compileResult.isR8CompileResult()) {
+      return compileResult.asR8CompileResult().getFeature(index);
+    } else if (compileResult.isR8PartialCompileResult()) {
+      return compileResult.asR8PartialCompileResult().getFeature(index);
+    } else {
+      fail();
+      return null;
+    }
   }
 
   public <E extends Throwable> DesugaredLibraryTestCompileResult<T> inspectL8(
@@ -80,12 +82,6 @@
   }
 
   public <E extends Throwable> DesugaredLibraryTestCompileResult<T> inspectL8ResidualArtProfile(
-      ThrowingConsumer<ArtProfileInspector, E> consumer) throws E, IOException {
-    return inspectL8ResidualArtProfile(
-        (rewrittenArtProfile, inspector) -> consumer.accept(rewrittenArtProfile));
-  }
-
-  public <E extends Throwable> DesugaredLibraryTestCompileResult<T> inspectL8ResidualArtProfile(
       ThrowingBiConsumer<ArtProfileInspector, CodeInspector, E> consumer) throws E, IOException {
     assertEquals(1, l8ResidualArtProfiles.size());
     consumer.accept(