Rewrite EnclosingMethod referring to renamed lambda methods

If a lambda method is renamed as part of lambda desugaring the
EnclosingMethod if any anonymous or local classes of that lambda
method is now rewritten to match the renaming.

This requires that the class with the lambda method and the classes
for the anonymous and local classes of the lambda method are in the
same compilation unit.

Updated the tests to use testForDesugaring and added a tests for R8.

Bug: 158752316
Change-Id: Id41cf812d9aee85fe632f66f3ac5fab764f4b690
diff --git a/src/main/java/com/android/tools/r8/graph/DexClass.java b/src/main/java/com/android/tools/r8/graph/DexClass.java
index 528b79b..e3bf2d8 100644
--- a/src/main/java/com/android/tools/r8/graph/DexClass.java
+++ b/src/main/java/com/android/tools/r8/graph/DexClass.java
@@ -723,6 +723,10 @@
     return enclosingMethod;
   }
 
+  public void setEnclosingMethodAttribute(EnclosingMethodAttribute enclosingMethod) {
+    this.enclosingMethod = enclosingMethod;
+  }
+
   public void clearEnclosingMethodAttribute() {
     enclosingMethod = null;
   }
diff --git a/src/main/java/com/android/tools/r8/ir/desugar/LambdaClass.java b/src/main/java/com/android/tools/r8/ir/desugar/LambdaClass.java
index 104fc97..ecad620 100644
--- a/src/main/java/com/android/tools/r8/ir/desugar/LambdaClass.java
+++ b/src/main/java/com/android/tools/r8/ir/desugar/LambdaClass.java
@@ -617,7 +617,7 @@
                             encodedMethod.getCode(),
                             true);
                     newMethod.copyMetadata(encodedMethod);
-                    rewriter.originalMethodSignatures.put(callTarget, encodedMethod.method);
+                    rewriter.lensBuilder.move(encodedMethod.method, callTarget);
 
                     DexEncodedMethod.setDebugInfoWithFakeThisParameter(
                         newMethod.getCode(), callTarget.getArity(), appView);
@@ -689,7 +689,7 @@
                             encodedMethod.getCode(),
                             true);
                     newMethod.copyMetadata(encodedMethod);
-                    rewriter.originalMethodSignatures.put(callTarget, encodedMethod.method);
+                    rewriter.lensBuilder.move(encodedMethod.method, callTarget);
                     return newMethod;
                   });
       return new ProgramMethod(implMethodHolder, replacement);
diff --git a/src/main/java/com/android/tools/r8/ir/desugar/LambdaRewriter.java b/src/main/java/com/android/tools/r8/ir/desugar/LambdaRewriter.java
index dd7889d..fad4980 100644
--- a/src/main/java/com/android/tools/r8/ir/desugar/LambdaRewriter.java
+++ b/src/main/java/com/android/tools/r8/ir/desugar/LambdaRewriter.java
@@ -7,11 +7,13 @@
 import com.android.tools.r8.graph.AppView;
 import com.android.tools.r8.graph.DexApplication.Builder;
 import com.android.tools.r8.graph.DexCallSite;
+import com.android.tools.r8.graph.DexField;
 import com.android.tools.r8.graph.DexItemFactory;
 import com.android.tools.r8.graph.DexMethod;
 import com.android.tools.r8.graph.DexProgramClass;
 import com.android.tools.r8.graph.DexString;
 import com.android.tools.r8.graph.DexType;
+import com.android.tools.r8.graph.EnclosingMethodAttribute;
 import com.android.tools.r8.graph.GraphLens;
 import com.android.tools.r8.graph.GraphLens.NestedGraphLens;
 import com.android.tools.r8.graph.ProgramMethod;
@@ -31,8 +33,6 @@
 import com.android.tools.r8.utils.DescriptorUtils;
 import com.android.tools.r8.utils.collections.SortedProgramMethodSet;
 import com.google.common.collect.BiMap;
-import com.google.common.collect.HashBiMap;
-import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Sets;
 import java.util.ArrayList;
@@ -63,7 +63,7 @@
 
   final DexString instanceFieldName;
 
-  final BiMap<DexMethod, DexMethod> originalMethodSignatures = HashBiMap.create();
+  final LambdaRewriterLens.Builder lensBuilder = LambdaRewriterLens.builder();
 
   // Maps call sites seen so far to inferred lambda descriptor. It is intended
   // to help avoid re-matching call sites we already seen. Note that same call
@@ -160,6 +160,7 @@
       appView.appInfo().addSynthesizedClass(synthesizedClass, lambdaClass.addToMainDexList.get());
       builder.addSynthesizedClass(synthesizedClass);
     }
+    fixup();
     optimizeSynthesizedClasses(converter, executorService);
   }
 
@@ -345,34 +346,78 @@
     return knownLambdaClasses;
   }
 
-  public NestedGraphLens buildMappingLens(AppView<?> appView) {
-    if (originalMethodSignatures.isEmpty()) {
+  public NestedGraphLens fixup() {
+    LambdaRewriterLens lens = lensBuilder.build(appView.graphLens(), appView.dexItemFactory());
+    if (lens == null) {
       return null;
     }
-    return new LambdaRewriterLens(
-        originalMethodSignatures, appView.graphLens(), appView.dexItemFactory());
+    for (DexProgramClass clazz : appView.appInfo().classes()) {
+      EnclosingMethodAttribute enclosingMethod = clazz.getEnclosingMethodAttribute();
+      if (enclosingMethod != null) {
+        DexMethod mappedEnclosingMethod = lens.lookupMethod(enclosingMethod.getEnclosingMethod());
+        if (mappedEnclosingMethod != null
+            && mappedEnclosingMethod != enclosingMethod.getEnclosingMethod()) {
+          clazz.setEnclosingMethodAttribute(new EnclosingMethodAttribute(mappedEnclosingMethod));
+        }
+      }
+    }
+    ;
+    // Return lens without method map (but still retaining originalMethodSignatures), as the
+    // generated lambdas classes are generated with the an invoke to the new method, so no
+    // code rewriting is required.
+    return lens.withoutMethodMap();
   }
 
   static class LambdaRewriterLens extends NestedGraphLens {
 
-    public LambdaRewriterLens(
+    LambdaRewriterLens(
+        Map<DexType, DexType> typeMap,
+        Map<DexMethod, DexMethod> methodMap,
+        Map<DexField, DexField> fieldMap,
+        BiMap<DexField, DexField> originalFieldSignatures,
         BiMap<DexMethod, DexMethod> originalMethodSignatures,
-        GraphLens graphLens,
-        DexItemFactory factory) {
+        GraphLens previousLens,
+        DexItemFactory dexItemFactory) {
       super(
-          ImmutableMap.of(),
-          ImmutableMap.of(),
-          ImmutableMap.of(),
-          null,
+          typeMap,
+          methodMap,
+          fieldMap,
+          originalFieldSignatures,
           originalMethodSignatures,
-          graphLens,
-          factory);
+          previousLens,
+          dexItemFactory);
     }
 
     @Override
     protected boolean isLegitimateToHaveEmptyMappings() {
       return true;
     }
+
+    private LambdaRewriterLens withoutMethodMap() {
+      methodMap.clear();
+      return this;
+    }
+
+    public static LambdaRewriterLens.Builder builder() {
+      return new LambdaRewriterLens.Builder();
+    }
+
+    public static class Builder extends NestedGraphLens.Builder {
+      public LambdaRewriterLens build(GraphLens previousLens, DexItemFactory dexItemFactory) {
+        if (typeMap.isEmpty() && methodMap.isEmpty() && fieldMap.isEmpty()) {
+          return null;
+        }
+        assert typeMap.isEmpty();
+        assert fieldMap.isEmpty();
+        return new LambdaRewriterLens(
+            typeMap,
+            methodMap,
+            fieldMap,
+            originalFieldSignatures,
+            originalMethodSignatures,
+            previousLens,
+            dexItemFactory);
+      }
+    }
   }
 }
-
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 97fd5a7..771f226 100644
--- a/src/main/java/com/android/tools/r8/shaking/Enqueuer.java
+++ b/src/main/java/com/android/tools/r8/shaking/Enqueuer.java
@@ -2743,7 +2743,7 @@
   }
 
   public NestedGraphLens buildGraphLens(AppView<?> appView) {
-    return lambdaRewriter != null ? lambdaRewriter.buildMappingLens(appView) : null;
+    return lambdaRewriter != null ? lambdaRewriter.fixup() : null;
   }
 
   private void keepClassWithRules(DexProgramClass clazz, Set<ProguardKeepRuleBase> rules) {
diff --git a/src/test/java/com/android/tools/r8/desugar/DesugarLambdaWithAnonymousClass.java b/src/test/java/com/android/tools/r8/desugar/DesugarLambdaWithAnonymousClass.java
index 7192449..b2d5b54 100644
--- a/src/test/java/com/android/tools/r8/desugar/DesugarLambdaWithAnonymousClass.java
+++ b/src/test/java/com/android/tools/r8/desugar/DesugarLambdaWithAnonymousClass.java
@@ -5,22 +5,22 @@
 package com.android.tools.r8.desugar;
 
 import static com.android.tools.r8.utils.codeinspector.Matchers.isPresent;
-import static org.hamcrest.CoreMatchers.not;
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 
-import com.android.tools.r8.D8TestRunResult;
+import com.android.tools.r8.DesugarTestConfiguration;
 import com.android.tools.r8.TestBase;
 import com.android.tools.r8.TestParameters;
 import com.android.tools.r8.TestParametersCollection;
 import com.android.tools.r8.ToolHelper;
-import com.android.tools.r8.ToolHelper.DexVm.Version;
 import com.android.tools.r8.graph.DexMethod;
 import com.android.tools.r8.utils.codeinspector.ClassSubject;
 import com.android.tools.r8.utils.codeinspector.CodeInspector;
-import java.nio.file.Path;
+import com.google.common.collect.ImmutableList;
 import java.util.ArrayList;
+import java.util.List;
 import org.junit.BeforeClass;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -29,6 +29,14 @@
 @RunWith(Parameterized.class)
 public class DesugarLambdaWithAnonymousClass extends TestBase {
 
+  private List<String> EXPECTED_JAVAC_RESULT =
+      ImmutableList.of("Hello from inside lambda$test$0", "Hello from inside lambda$testStatic$1");
+
+  private List<String> EXPECTED_DESUGARED_RESULT =
+      ImmutableList.of(
+          "Hello from inside lambda$test$0$DesugarLambdaWithAnonymousClass$TestClass",
+          "Hello from inside lambda$testStatic$1");
+
   @Parameterized.Parameters(name = "{0}")
   public static TestParametersCollection data() {
     return getTestParameters().withAllRuntimes().withAllApiLevelsAlsoForCf().build();
@@ -73,111 +81,55 @@
     assertEquals(2, counter.getCount());
   }
 
-  // TODO(158752316): There should be no use of this check.
-  private void checkEnclosingMethodWrong(CodeInspector inspector) {
-    Counter counter = new Counter();
-    inspector.forAllClasses(
-        clazz -> {
-          if (clazz.getFinalName().endsWith("$TestClass$1")
-              || clazz.getFinalName().endsWith("$TestClass$2")) {
-            counter.increment();
-            assertTrue(clazz.isAnonymousClass());
-            DexMethod enclosingMethod = clazz.getFinalEnclosingMethod();
-            ClassSubject testClassSubject =
-                inspector.clazz(DesugarLambdaWithAnonymousClass.TestClass.class);
-            assertEquals(
-                testClassSubject, inspector.clazz(enclosingMethod.holder.toSourceString()));
-            if (enclosingMethod.name.toString().contains("Static")) {
-              assertThat(
-                  testClassSubject.uniqueMethodWithName(enclosingMethod.name.toString()),
-                  isPresent());
-            } else {
-              assertThat(
-                  testClassSubject.uniqueMethodWithName(enclosingMethod.name.toString()),
-                  not(isPresent()));
-            }
-          }
-        });
-    assertEquals(2, counter.getCount());
-  }
-
-  private void checkArtResult(D8TestRunResult result) {
-    // TODO(158752316): This should neither return null nor fail.
-    if (parameters.getRuntime().asDex().getVm().getVersion().isOlderThanOrEqual(Version.V4_4_4)
-        || parameters.getRuntime().asDex().getVm().getVersion().isNewerThan(Version.V6_0_1)) {
-      result.assertSuccessWithOutputLines(
-          "Hello from inside <null>", "Hello from inside lambda$testStatic$1");
-    } else {
-      result.assertFailureWithErrorThatThrows(IncompatibleClassChangeError.class);
-    }
-  }
-
   @BeforeClass
   public static void checkExpectedJavacNames() throws Exception {
     CodeInspector inspector =
         new CodeInspector(
-            ToolHelper.getClassFilesForInnerClasses(DesugarLambdaWithLocalClass.class));
-    String outer = DesugarLambdaWithLocalClass.class.getTypeName();
+            ToolHelper.getClassFilesForInnerClasses(DesugarLambdaWithAnonymousClass.class));
+    String outer = DesugarLambdaWithAnonymousClass.class.getTypeName();
     ClassSubject testClass = inspector.clazz(outer + "$TestClass");
     assertThat(testClass, isPresent());
     assertThat(testClass.uniqueMethodWithName("lambda$test$0"), isPresent());
     assertThat(testClass.uniqueMethodWithName("lambda$testStatic$1"), isPresent());
-    assertThat(inspector.clazz(outer + "$TestClass$1MyConsumerImpl"), isPresent());
-    assertThat(inspector.clazz(outer + "$TestClass$2MyConsumerImpl"), isPresent());
+    assertThat(inspector.clazz(outer + "$TestClass$1"), isPresent());
+    assertThat(inspector.clazz(outer + "$TestClass$2"), isPresent());
   }
 
   @Test
-  public void testDefault() throws Exception {
-    if (parameters.getRuntime().isCf()) {
-      // Run on the JVM.
-      testForJvm()
+  public void testDesugar() throws Exception {
+    testForDesugaring(parameters)
+        .addInnerClasses(DesugarLambdaWithAnonymousClass.class)
+        .run(parameters.getRuntime(), TestClass.class)
+        .inspect(this::checkEnclosingMethod)
+        .applyIf(
+            DesugarTestConfiguration::isNotDesugared,
+            r -> r.assertSuccessWithOutputLines(EXPECTED_JAVAC_RESULT))
+        .applyIf(
+            DesugarTestConfiguration::isDesugared,
+            r -> r.assertSuccessWithOutputLines(EXPECTED_DESUGARED_RESULT));
+  }
+
+  @Test
+  public void testR8() throws Exception {
+    try {
+      testForR8(parameters.getBackend())
+          // TODO(b/158752316, b/157700141): Disable inlining to keep the synthetic lambda methods.
+          .addOptionsModification(options -> options.enableInlining = false)
           .addInnerClasses(DesugarLambdaWithAnonymousClass.class)
-          .run(parameters.getRuntime(), TestClass.class)
+          .setMinApi(parameters.getApiLevel())
+          .addKeepRules("-keep class * { *; }")
+          // Keeping the synthetic lambda methods is currently not supported - they are
+          // forcefully unpinned. The following rule has no effect. See b/b/157700141.
+          .addKeepRules("-keep class **.*$TestClass { synthetic *; }")
+          .addKeepAttributes("InnerClasses", "EnclosingMethod")
+          .compile()
           .inspect(this::checkEnclosingMethod)
-          .assertSuccessWithOutputLines(
-              "Hello from inside lambda$test$0", "Hello from inside lambda$testStatic$1");
-    } else {
-      assert parameters.getRuntime().isDex();
-      // Run on Art.
-      checkArtResult(
-          testForD8()
-              .addInnerClasses(DesugarLambdaWithAnonymousClass.class)
-              .setMinApi(parameters.getApiLevel())
-              .compile()
-              .inspect(this::checkEnclosingMethodWrong)
-              .run(parameters.getRuntime(), TestClass.class));
-    }
-  }
-
-  @Test
-  public void testCfToCf() throws Exception {
-    // Use D8 to desugar with Java classfile output.
-    Path jar =
-        testForD8(Backend.CF)
-            .addInnerClasses(DesugarLambdaWithAnonymousClass.class)
-            .setMinApi(parameters.getApiLevel())
-            .compile()
-            .inspect(this::checkEnclosingMethodWrong)
-            .writeToZip();
-
-    if (parameters.getRuntime().isCf()) {
-      // Run on the JVM.
-      testForJvm()
-          .addProgramFiles(jar)
           .run(parameters.getRuntime(), TestClass.class)
-          // TODO(158752316): This should not fail.
-          .assertFailureWithErrorThatThrows(InternalError.class);
-    } else {
-      assert parameters.getRuntime().isDex();
-      // Compile to DEX without desugaring and run on Art.
-      checkArtResult(
-          testForD8()
-              .addProgramFiles(jar)
-              .setMinApi(parameters.getApiLevel())
-              .disableDesugaring()
-              .compile()
-              .inspect(this::checkEnclosingMethodWrong)
-              .run(parameters.getRuntime(), TestClass.class));
+          .assertSuccessWithOutputLines(
+              parameters.isCfRuntime() ? EXPECTED_JAVAC_RESULT : EXPECTED_DESUGARED_RESULT);
+      assertFalse(parameters.isDexRuntime());
+    } catch (AssertionError e) {
+      assertTrue(parameters.isDexRuntime());
     }
   }
 
diff --git a/src/test/java/com/android/tools/r8/desugar/DesugarLambdaWithLocalClass.java b/src/test/java/com/android/tools/r8/desugar/DesugarLambdaWithLocalClass.java
index 68fd89e..8c539d2 100644
--- a/src/test/java/com/android/tools/r8/desugar/DesugarLambdaWithLocalClass.java
+++ b/src/test/java/com/android/tools/r8/desugar/DesugarLambdaWithLocalClass.java
@@ -5,22 +5,21 @@
 package com.android.tools.r8.desugar;
 
 import static com.android.tools.r8.utils.codeinspector.Matchers.isPresent;
-import static org.hamcrest.CoreMatchers.not;
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
 
-import com.android.tools.r8.D8TestRunResult;
+import com.android.tools.r8.DesugarTestConfiguration;
 import com.android.tools.r8.TestBase;
 import com.android.tools.r8.TestParameters;
 import com.android.tools.r8.TestParametersCollection;
 import com.android.tools.r8.ToolHelper;
-import com.android.tools.r8.ToolHelper.DexVm.Version;
 import com.android.tools.r8.graph.DexMethod;
 import com.android.tools.r8.utils.codeinspector.ClassSubject;
 import com.android.tools.r8.utils.codeinspector.CodeInspector;
-import java.nio.file.Path;
+import com.google.common.collect.ImmutableList;
 import java.util.ArrayList;
+import java.util.List;
 import org.junit.BeforeClass;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -29,6 +28,17 @@
 @RunWith(Parameterized.class)
 public class DesugarLambdaWithLocalClass extends TestBase {
 
+  private List<String> EXPECTED_JAVAC_RESULT =
+      ImmutableList.of("Hello from inside lambda$test$0", "Hello from inside lambda$testStatic$1");
+
+  private List<String> EXPECTED_DESUGARED_RESULT =
+      ImmutableList.of(
+          "Hello from inside lambda$test$0$DesugarLambdaWithLocalClass$TestClass",
+          "Hello from inside lambda$testStatic$1");
+
+  private List<String> EXPECTED_DESUGARED_RESULT_R8_WITHOUT_INLINING =
+      ImmutableList.of("Hello from inside lambda$test$0", "Hello from inside lambda$testStatic$1");
+
   @Parameterized.Parameters(name = "{0}")
   public static TestParametersCollection data() {
     return getTestParameters().withAllRuntimes().withAllApiLevelsAlsoForCf().build();
@@ -56,7 +66,8 @@
     Counter counter = new Counter();
     inspector.forAllClasses(
         clazz -> {
-          if (clazz.getFinalName().endsWith("MyConsumerImpl")) {
+          if (clazz.getFinalName().endsWith("$TestClass$1MyConsumerImpl")
+              || clazz.getFinalName().endsWith("$TestClass$2MyConsumerImpl")) {
             counter.increment();
             assertTrue(clazz.isLocalClass());
             DexMethod enclosingMethod = clazz.getFinalEnclosingMethod();
@@ -71,111 +82,53 @@
     assertEquals(2, counter.getCount());
   }
 
-  // TODO(158752316): There should be no use of this check.
-  private void checkEnclosingMethodWrong(CodeInspector inspector) {
-    Counter counter = new Counter();
-    inspector.forAllClasses(
-        clazz -> {
-          if (clazz.getFinalName().endsWith("$TestClass$1MyConsumerImpl")
-              || clazz.getFinalName().endsWith("$TestClass$2MyConsumerImpl")) {
-            counter.increment();
-            assertTrue(clazz.isLocalClass());
-            DexMethod enclosingMethod = clazz.getFinalEnclosingMethod();
-            ClassSubject testClassSubject = inspector.clazz(TestClass.class);
-            assertEquals(
-                testClassSubject, inspector.clazz(enclosingMethod.holder.toSourceString()));
-            if (enclosingMethod.name.toString().contains("Static")) {
-              assertThat(
-                  testClassSubject.uniqueMethodWithName(enclosingMethod.name.toString()),
-                  isPresent());
-            } else {
-              assertThat(
-                  testClassSubject.uniqueMethodWithName(enclosingMethod.name.toString()),
-                  not(isPresent()));
-            }
-          }
-        });
-    assertEquals(2, counter.getCount());
-  }
-
-  private void checkArtResult(D8TestRunResult result) {
-    // TODO(158752316): This should neither return null nor fail.
-    if (parameters.getRuntime().asDex().getVm().getVersion().isOlderThanOrEqual(Version.V4_4_4)
-        || parameters.getRuntime().asDex().getVm().getVersion().isNewerThan(Version.V6_0_1)) {
-      result.assertSuccessWithOutputLines(
-          "Hello from inside <null>", "Hello from inside lambda$testStatic$1");
-    } else {
-      result.assertFailureWithErrorThatThrows(IncompatibleClassChangeError.class);
-    }
-  }
-
   @BeforeClass
   public static void checkExpectedJavacNames() throws Exception {
     CodeInspector inspector =
         new CodeInspector(
-            ToolHelper.getClassFilesForInnerClasses(DesugarLambdaWithAnonymousClass.class));
-    String outer = DesugarLambdaWithAnonymousClass.class.getTypeName();
+            ToolHelper.getClassFilesForInnerClasses(DesugarLambdaWithLocalClass.class));
+    String outer = DesugarLambdaWithLocalClass.class.getTypeName();
     ClassSubject testClass = inspector.clazz(outer + "$TestClass");
     assertThat(testClass, isPresent());
     assertThat(testClass.uniqueMethodWithName("lambda$test$0"), isPresent());
     assertThat(testClass.uniqueMethodWithName("lambda$testStatic$1"), isPresent());
-    assertThat(inspector.clazz(outer + "$TestClass$1"), isPresent());
-    assertThat(inspector.clazz(outer + "$TestClass$2"), isPresent());
+    assertThat(inspector.clazz(outer + "$TestClass$1MyConsumerImpl"), isPresent());
+    assertThat(inspector.clazz(outer + "$TestClass$2MyConsumerImpl"), isPresent());
   }
 
   @Test
-  public void testDefault() throws Exception {
-    if (parameters.getRuntime().isCf()) {
-      // Run on the JVM.
-      testForJvm()
-          .addInnerClasses(DesugarLambdaWithLocalClass.class)
-          .run(parameters.getRuntime(), TestClass.class)
-          .inspect(this::checkEnclosingMethod)
-          .assertSuccessWithOutputLines(
-              "Hello from inside lambda$test$0", "Hello from inside lambda$testStatic$1");
-    } else {
-      assert parameters.getRuntime().isDex();
-      // Run on Art.
-      checkArtResult(
-          testForD8()
-              .addInnerClasses(DesugarLambdaWithLocalClass.class)
-              .setMinApi(parameters.getApiLevel())
-              .compile()
-              .inspect(this::checkEnclosingMethodWrong)
-              .run(parameters.getRuntime(), TestClass.class));
-    }
+  public void testDesugar() throws Exception {
+    testForDesugaring(parameters)
+        .addInnerClasses(DesugarLambdaWithLocalClass.class)
+        .run(parameters.getRuntime(), TestClass.class)
+        .inspect(this::checkEnclosingMethod)
+        .applyIf(
+            DesugarTestConfiguration::isNotDesugared,
+            r -> r.assertSuccessWithOutputLines(EXPECTED_JAVAC_RESULT))
+        .applyIf(
+            DesugarTestConfiguration::isDesugared,
+            r -> r.assertSuccessWithOutputLines(EXPECTED_DESUGARED_RESULT));
   }
 
   @Test
-  public void testCfToCf() throws Exception {
-    // Use D8 to desugar with Java classfile output.
-    Path jar =
-        testForD8(Backend.CF)
-            .addInnerClasses(DesugarLambdaWithLocalClass.class)
-            .setMinApi(parameters.getApiLevel())
-            .compile()
-            .inspect(this::checkEnclosingMethodWrong)
-            .writeToZip();
-
-    if (parameters.getRuntime().isCf()) {
-      // Run on the JVM.
-      testForJvm()
-          .addProgramFiles(jar)
-          .run(parameters.getRuntime(), TestClass.class)
-          // TODO(158752316): This should not fail.
-          .assertFailureWithErrorThatThrows(InternalError.class);
-    } else {
-      assert parameters.getRuntime().isDex();
-      // Compile to DEX without desugaring and run on Art.
-      checkArtResult(
-          testForD8()
-              .addProgramFiles(jar)
-              .setMinApi(parameters.getApiLevel())
-              .disableDesugaring()
-              .compile()
-              .inspect(this::checkEnclosingMethodWrong)
-              .run(parameters.getRuntime(), TestClass.class));
-    }
+  public void testR8() throws Exception {
+    testForR8(parameters.getBackend())
+        // TODO(b/158752316, b/157700141): Disable inlining to keep the synthetic lambda methods.
+        .addOptionsModification(options -> options.enableInlining = false)
+        .addInnerClasses(DesugarLambdaWithLocalClass.class)
+        .setMinApi(parameters.getApiLevel())
+        .addKeepRules("-keep class * { *; }")
+        // Keeping the synthetic lambda methods is currently not supported - they are
+        // forcefully unpinned. The following rule has no effect. See b/b/157700141.
+        .addKeepRules("-keep class **.*$TestClass { synthetic *; }")
+        .addKeepAttributes("InnerClasses", "EnclosingMethod")
+        .compile()
+        .inspect(this::checkEnclosingMethod)
+        .run(parameters.getRuntime(), TestClass.class)
+        .assertSuccessWithOutputLines(
+            parameters.isCfRuntime()
+                ? EXPECTED_JAVAC_RESULT
+                : EXPECTED_DESUGARED_RESULT_R8_WITHOUT_INLINING);
   }
 
   public interface MyConsumer<T> {