Fix call site optimization for targeted lambda implementation methods

Bug: 190664352
Change-Id: Id330ea5e3f0ab4f02c994d34edbe35e17ccd05b6
diff --git a/src/main/java/com/android/tools/r8/graph/AppView.java b/src/main/java/com/android/tools/r8/graph/AppView.java
index 6ef4c6a..b509c1f 100644
--- a/src/main/java/com/android/tools/r8/graph/AppView.java
+++ b/src/main/java/com/android/tools/r8/graph/AppView.java
@@ -323,6 +323,13 @@
     return callSiteOptimizationInfoPropagator;
   }
 
+  public <E extends Throwable> void withCallSiteOptimizationInfoPropagator(
+      ThrowingConsumer<CallSiteOptimizationInfoPropagator, E> consumer) throws E {
+    if (callSiteOptimizationInfoPropagator != null) {
+      consumer.accept(callSiteOptimizationInfoPropagator);
+    }
+  }
+
   public LibraryMemberOptimizer libraryMethodOptimizer() {
     return libraryMemberOptimizer;
   }
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 5799d59..38ece59 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
@@ -672,6 +672,12 @@
     // 2) Revisit DexEncodedMethods for the collected candidates.
 
     printPhase("Primary optimization pass");
+
+    appView.withCallSiteOptimizationInfoPropagator(
+        optimization ->
+            optimization.abandonCallSitePropagationForLambdaImplementationMethods(
+                executorService, timing));
+
     if (fieldAccessAnalysis != null) {
       fieldAccessAnalysis.fieldAssignmentTracker().initialize();
     }
diff --git a/src/main/java/com/android/tools/r8/ir/optimize/CallSiteOptimizationInfoPropagator.java b/src/main/java/com/android/tools/r8/ir/optimize/CallSiteOptimizationInfoPropagator.java
index 2f04f11..2387a14 100644
--- a/src/main/java/com/android/tools/r8/ir/optimize/CallSiteOptimizationInfoPropagator.java
+++ b/src/main/java/com/android/tools/r8/ir/optimize/CallSiteOptimizationInfoPropagator.java
@@ -97,10 +97,6 @@
       return;
     }
 
-    if (appView.appInfo().isMethodTargetedByInvokeDynamic(code.context().getReference())) {
-      abandonCallSitePropagationForMethodAndOverrides(code.context());
-    }
-
     ProgramMethod context = code.context();
     for (Instruction instruction : code.instructions()) {
       if (instruction.isInvokeMethod()) {
@@ -293,6 +289,30 @@
     }
   }
 
+  public void abandonCallSitePropagationForLambdaImplementationMethods(
+      ExecutorService executorService, Timing timing) throws ExecutionException {
+    if (appView.options().isGeneratingClassFiles()) {
+      timing.begin("Call site optimization: abandon lambda implementation methods");
+      ForEachable<ProgramMethod> lambdaImplementationMethods =
+          consumer ->
+              appView
+                  .appInfo()
+                  .forEachMethod(
+                      method -> {
+                        if (appView
+                            .appInfo()
+                            .isMethodTargetedByInvokeDynamic(method.getReference())) {
+                          consumer.accept(method);
+                        }
+                      });
+      ThreadUtils.processItems(
+          lambdaImplementationMethods,
+          this::abandonCallSitePropagationForMethodAndOverrides,
+          executorService);
+      timing.end();
+    }
+  }
+
   public void abandonCallSitePropagationForPinnedMethodsAndOverrides(
       ExecutorService executorService) throws ExecutionException {
     ThreadUtils.processItems(