Fix bug in kotlin lambda merger

Change-Id: I8c4a204ab9ae1f60df52095d0ad26c52c94b7ef6
diff --git a/src/main/java/com/android/tools/r8/ir/conversion/IRConverter.java b/src/main/java/com/android/tools/r8/ir/conversion/IRConverter.java
index ba89d0c..c271af1 100644
--- a/src/main/java/com/android/tools/r8/ir/conversion/IRConverter.java
+++ b/src/main/java/com/android/tools/r8/ir/conversion/IRConverter.java
@@ -658,7 +658,8 @@
     }
 
     // Process the application identifying outlining candidates.
-    GraphLens graphLensForIR = appView.graphLens();
+    GraphLens initialGraphLensForIR = appView.graphLens();
+    GraphLens graphLensForIR = initialGraphLensForIR;
     OptimizationFeedbackDelayed feedback = delayedOptimizationFeedback;
     PostMethodProcessor.Builder postMethodProcessorBuilder =
         new PostMethodProcessor.Builder(getOptimizationsForPostIRProcessing());
@@ -762,7 +763,7 @@
 
     printPhase("Lambda merging finalization");
     // TODO(b/127694949): Adapt to PostOptimization.
-    finalizeLambdaMerging(application, feedback, builder, executorService);
+    finalizeLambdaMerging(application, feedback, builder, executorService, initialGraphLensForIR);
 
     printPhase("Desugared library API Conversion finalization");
     generateDesugaredLibraryAPIWrappers(builder, executorService);
@@ -937,11 +938,12 @@
       DexApplication application,
       OptimizationFeedback feedback,
       Builder<?> builder,
-      ExecutorService executorService)
+      ExecutorService executorService,
+      GraphLens appliedGraphLens)
       throws ExecutionException {
     if (lambdaMerger != null) {
       lambdaMerger.applyLambdaClassMapping(
-          application, this, feedback, builder, executorService);
+          application, this, feedback, builder, executorService, appliedGraphLens);
     } else {
       appView.setHorizontallyMergedLambdaClasses(HorizontallyMergedLambdaClasses.empty());
     }
diff --git a/src/main/java/com/android/tools/r8/ir/optimize/lambda/LambdaMerger.java b/src/main/java/com/android/tools/r8/ir/optimize/lambda/LambdaMerger.java
index ee35087..9ccbbbc 100644
--- a/src/main/java/com/android/tools/r8/ir/optimize/lambda/LambdaMerger.java
+++ b/src/main/java/com/android/tools/r8/ir/optimize/lambda/LambdaMerger.java
@@ -17,6 +17,7 @@
 import com.android.tools.r8.graph.DexMethod;
 import com.android.tools.r8.graph.DexProgramClass;
 import com.android.tools.r8.graph.DexType;
+import com.android.tools.r8.graph.GraphLens;
 import com.android.tools.r8.graph.ProgramMethod;
 import com.android.tools.r8.graph.ResolutionResult;
 import com.android.tools.r8.graph.classmerging.HorizontallyMergedLambdaClasses;
@@ -339,7 +340,8 @@
       IRConverter converter,
       OptimizationFeedback feedback,
       Builder<?> builder,
-      ExecutorService executorService)
+      ExecutorService executorService,
+      GraphLens appliedGraphLens)
       throws ExecutionException {
     if (lambdas.isEmpty()) {
       appView.setHorizontallyMergedLambdaClasses(HorizontallyMergedLambdaClasses.empty());
@@ -389,7 +391,7 @@
 
     // Rewrite lambda class references into lambda group class
     // references inside methods from the processing queue.
-    rewriteLambdaReferences(converter, executorService);
+    rewriteLambdaReferences(converter, executorService, appliedGraphLens);
     this.mode = null;
 
     appView.setHorizontallyMergedLambdaClasses(
@@ -454,12 +456,13 @@
     }
   }
 
-  private void rewriteLambdaReferences(IRConverter converter, ExecutorService executorService)
+  private void rewriteLambdaReferences(
+      IRConverter converter, ExecutorService executorService, GraphLens appliedGraphLens)
       throws ExecutionException {
     if (methodsToReprocess.isEmpty()) {
       return;
     }
-    SortedProgramMethodSet methods = methodsToReprocess.build(appView);
+    SortedProgramMethodSet methods = methodsToReprocess.build(appView, appliedGraphLens);
     converter.processMethodsConcurrently(methods, executorService);
     assert methods.stream()
         .map(DexClassAndMethod::getDefinition)
diff --git a/src/main/java/com/android/tools/r8/utils/collections/LongLivedProgramMethodSetBuilder.java b/src/main/java/com/android/tools/r8/utils/collections/LongLivedProgramMethodSetBuilder.java
index 8297a9f..f4115a4 100644
--- a/src/main/java/com/android/tools/r8/utils/collections/LongLivedProgramMethodSetBuilder.java
+++ b/src/main/java/com/android/tools/r8/utils/collections/LongLivedProgramMethodSetBuilder.java
@@ -60,10 +60,10 @@
     return build(appView, null);
   }
 
-  public T build(AppView<AppInfoWithLiveness> appView, GraphLens applied) {
+  public T build(AppView<AppInfoWithLiveness> appView, GraphLens appliedGraphLens) {
     T result = factory.apply(methods.size());
     for (DexMethod oldMethod : methods) {
-      DexMethod method = appView.graphLens().getRenamedMethodSignature(oldMethod, applied);
+      DexMethod method = appView.graphLens().getRenamedMethodSignature(oldMethod, appliedGraphLens);
       DexProgramClass holder = appView.definitionForHolder(method).asProgramClass();
       result.createAndAdd(holder, holder.lookupMethod(method));
     }
diff --git a/src/test/java/com/android/tools/r8/kotlin/lambda/KotlinLambdaMergingWithSmallInliningBudgetTest.java b/src/test/java/com/android/tools/r8/kotlin/lambda/KotlinLambdaMergingWithSmallInliningBudgetTest.java
index 84ee93d..446d982 100644
--- a/src/test/java/com/android/tools/r8/kotlin/lambda/KotlinLambdaMergingWithSmallInliningBudgetTest.java
+++ b/src/test/java/com/android/tools/r8/kotlin/lambda/KotlinLambdaMergingWithSmallInliningBudgetTest.java
@@ -35,7 +35,6 @@
 
   @Test
   public void testJStyleRunnable() throws Exception {
-    expectThrowsWithHorizontalClassMergingIf(allowAccessModification);
     final String mainClassName = "lambdas_jstyle_runnable.MainKt";
     runTest("lambdas_jstyle_runnable", mainClassName, optionsModifier, null);
   }