[devkit] IDEA-309527: handle 'invokeLater'

Provides two fixes:
- withContext(EDT)
- if there a CoroutineScope, launch(EDT)

GitOrigin-RevId: a9f1840b08b222217d948b4876d2001d4e10c83a

platform/core-api/src/com/intellij/openapi/application/Application.java

@@ -289,6 +289,7 @@ public interface Application extends ComponentManager {
    *
    * @param runnable the runnable to execute.
    */
+  @RequiresBlockingContext
   void invokeLater(@NotNull Runnable runnable);
 
   /**
@@ -316,6 +317,7 @@ public interface Application extends ComponentManager {
    * @param runnable the runnable to execute.
    * @param state    the state in which the runnable will be executed.
    */
+  @RequiresBlockingContext
   void invokeLater(@NotNull Runnable runnable, @NotNull ModalityState state);
 
   /**

platform/util-ex/src/com/intellij/openapi/application/actions.kt

@@ -56,6 +56,7 @@ fun runInEdt(modalityState: ModalityState? = null, runnable: () -> Unit) {
   }
 }
 
+@RequiresBlockingContext
 fun invokeLater(modalityState: ModalityState? = null, runnable: () -> Unit) {
   ApplicationManager.getApplication().invokeLater({ runnable() }, modalityState ?: ModalityState.defaultModalityState())
 }

plugins/devkit/devkit-kotlin-tests/testSrc/org/jetbrains/idea/devkit/kotlin/inspections/ForbiddenInSuspectContextMethodInspectionTest.kt

@@ -48,6 +48,10 @@ class ForbiddenInSuspectContextMethodInspectionTest : LightJavaCodeInsightFixtur
 
   private val navigateToSuspendContextFix = "Navigate to suspend context"
 
+  private val invokeLaterDescr = "'invokeLater' is not idiomatic in suspend context. Use 'Dispatchers.EDT' instead"
+  private val invokeLaterFixWithContext = "(May change semantics) Replace 'invokeLater' call with 'withContext(Dispatchers.EDT) {}'"
+  private val invokeLaterFixLaunch = "Replace 'invokeLater' with 'launch(Dispatchers.EDT)'"
+
   @Test
   fun `progress manager checkCanceled in suspend function`() {
     addCheckCanceledFunctions()
@@ -703,6 +707,70 @@ class ForbiddenInSuspectContextMethodInspectionTest : LightJavaCodeInsightFixtur
     """.trimIndent())
   }
 
+  @Test
+  fun `invokeLater replace withContext`() {
+    addApplicationAndEtc()
+
+    myFixture.configureByText("file.kt", """
+      import com.intellij.openapi.application.*
+      
+      suspend fun a() {
+        ApplicationManager.getApplication().<warning descr="$invokeLaterDescr">inv<caret>okeLater</warning> {
+          println()
+        }
+      }
+    """.trimIndent())
+    myFixture.testHighlighting()
+
+    val intention = myFixture.getAvailableIntention(invokeLaterFixWithContext)
+    assertNotNullK(intention)
+    myFixture.checkPreviewAndLaunchAction(intention)
+
+    myFixture.checkResult("""
+      import com.intellij.openapi.application.*
+      import kotlinx.coroutines.Dispatchers
+      import kotlinx.coroutines.withContext
+      
+      suspend fun a() {
+        withContext(Dispatchers.EDT) {
+          println()
+        }
+      }
+    """.trimIndent())
+  }
+
+  @Test
+  fun `invokeLater with launch`() {
+    addApplicationAndEtc()
+
+    myFixture.configureByText("file.kt", """
+      import com.intellij.openapi.application.*
+      import kotlinx.coroutines.*
+      
+      suspend fun CoroutineScope.a() {
+        ApplicationManager.getApplication().<warning descr="$invokeLaterDescr">inv<caret>okeLater</warning> {
+          println()
+        }
+      }
+    """.trimIndent())
+    myFixture.testHighlighting()
+
+    val intention = myFixture.getAvailableIntention(invokeLaterFixLaunch)
+    assertNotNullK(intention)
+    myFixture.checkPreviewAndLaunchAction(intention)
+
+    myFixture.checkResult("""
+      import com.intellij.openapi.application.*
+      import kotlinx.coroutines.*
+      
+      suspend fun CoroutineScope.a() {
+          launch(Dispatchers.EDT) {
+            println()
+          }
+      }
+    """.trimIndent())
+  }
+
   private fun addApplicationAndEtc() {
     myFixture.addClass("""
         package com.intellij.openapi.progress;
@@ -727,6 +795,9 @@ class ForbiddenInSuspectContextMethodInspectionTest : LightJavaCodeInsightFixtur
             
             @RequiresBlockingContext
             ModalityState getDefaultModalityState();
+            
+            @RequiresBlockingContext
+            void invokeLater(Runnable runnable);
         }
       """.trimIndent())
 

plugins/devkit/intellij.kotlin.devkit/resources/messages/DevKitKotlinBundle.properties

@@ -7,6 +7,9 @@ inspections.forbidden.method.in.suspend.context.check.canceled.text=Do not call
 inspections.forbidden.method.in.suspend.context.check.canceled.fix.text=Replace 'ProgressManager.checkCanceled' with coroutine-friendly 'checkCancelled'
 inspections.forbidden.method.in.suspend.context.invoke.and.wait.text='invokeAndWait' can block current coroutine. Use 'Dispatchers.EDT' instead
 inspections.forbidden.method.in.suspend.context.invoke.and.wait.fix.text=Replace 'invokeAndWait' call with 'withContext(Dispatchers.EDT) {}'
+inspections.forbidden.method.in.suspend.context.invoke.later.text='invokeLater' is not idiomatic in suspend context. Use 'Dispatchers.EDT' instead
+inspections.forbidden.method.in.suspend.context.invoke.later.fix.with.context.text=(May change semantics) Replace 'invokeLater' call with 'withContext(Dispatchers.EDT) {}'
+inspections.forbidden.method.in.suspend.context.invoke.later.fix.launch.text=Replace 'invokeLater' with 'launch(Dispatchers.EDT)'
 inspections.forbidden.method.in.suspend.context.default.modality.state.text='defaultModalityState()' does not work in suspend context. If it is really necessary, use 'contextModality()'
 inspections.forbidden.method.in.suspend.context.default.modality.state.fix.text=Replace with 'contextModality()'
 inspections.forbidden.method.in.suspend.context.navigate.to.suspend.context=Navigate to suspend context

plugins/devkit/intellij.kotlin.devkit/src/inspections/ForbiddenInSuspectContextMethodInspection.kt

@@ -7,6 +7,7 @@ import com.intellij.openapi.project.Project
 import com.intellij.psi.*
 import com.intellij.psi.search.GlobalSearchScope
 import com.intellij.util.concurrency.annotations.RequiresBlockingContext
+import com.intellij.util.containers.toArray
 import org.jetbrains.idea.devkit.inspections.DevKitInspectionUtil
 import org.jetbrains.idea.devkit.kotlin.DevKitKotlinBundle
 import org.jetbrains.kotlin.analysis.api.KtAnalysisSession
@@ -21,6 +22,7 @@ import org.jetbrains.kotlin.analysis.api.symbols.KtSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.markers.KtNamedSymbol
 import org.jetbrains.kotlin.analysis.api.types.KtFunctionalType
 import org.jetbrains.kotlin.analysis.api.types.KtType
+import org.jetbrains.kotlin.analysis.api.types.KtUsualClassType
 import org.jetbrains.kotlin.idea.base.codeInsight.ShortenReferencesFacility
 import org.jetbrains.kotlin.idea.base.utils.fqname.isImported
 import org.jetbrains.kotlin.idea.util.ImportInsertHelperImpl
@@ -34,27 +36,34 @@ private val REQUIRES_SUSPEND_CONTEXT_ANNOTATION = RequiresBlockingContext::class
 private const val PROGRESS_MANAGER_CHECKED_CANCELED = "com.intellij.openapi.progress.ProgressManager.checkCanceled"
 private const val APPLICATION_INVOKE_AND_WAIT = "com.intellij.openapi.application.Application.invokeAndWait"
 private const val INVOKE_AND_WAIT_IF_NEEDED = "com.intellij.openapi.application.invokeAndWaitIfNeeded"
+private const val APPLICATION_INVOKE_LATER = "com.intellij.openapi.application.Application.invokeLater"
+private const val INVOKE_LATER_KT = "com.intellij.openapi.application.invokeLater"
 private const val MODALITY_STATE_DEFAULT_MODALITY_STATE = "com.intellij.openapi.application.ModalityState.defaultModalityState"
 private const val APPLICATION_GET_DEFAULT_MODALITY_STATE = "com.intellij.openapi.application.Application.getDefaultModalityState"
 private const val RESTRICTS_SUSPENSION = "kotlin.coroutines.RestrictsSuspension"
 private const val INTELLIJ_EDT_DISPATCHER = "com.intellij.openapi.application.EDT"
 private const val CONTEXT_MODALITY_EXT = "com.intellij.openapi.application.contextModality"
+private const val LAUNCH = "kotlinx.coroutines.launch"
 
 private val requiresSuspendContextAnnotation = FqName(REQUIRES_SUSPEND_CONTEXT_ANNOTATION)
 private val progressManagerCheckedCanceledName = FqName(PROGRESS_MANAGER_CHECKED_CANCELED)
 private val applicationInvokeAndWaitName = FqName(APPLICATION_INVOKE_AND_WAIT)
 private val invokeAndWaitIfNeeded = FqName(INVOKE_AND_WAIT_IF_NEEDED)
+private val applicationInvokeLater = FqName(APPLICATION_INVOKE_LATER)
+private val invokeLaterKt = FqName(INVOKE_LATER_KT)
 private val modalityStateDefaultModalityState = FqName(MODALITY_STATE_DEFAULT_MODALITY_STATE)
 private val applicationGetDefaultModalityState = FqName(APPLICATION_GET_DEFAULT_MODALITY_STATE)
 private val restrictsSuspensionName = FqName(RESTRICTS_SUSPENSION)
 private val intelliJEdtDispatcher = FqName(INTELLIJ_EDT_DISPATCHER)
 private val contextModalityExt = FqName(CONTEXT_MODALITY_EXT)
+private val coroutinesLaunch = FqName(LAUNCH)
 
 private const val COROUTINE_CHECK_CANCELED_FIX = "com.intellij.openapi.progress.checkCancelled"
 private const val WITH_CONTEXT = "kotlinx.coroutines.withContext"
 private const val DISPATCHERS = "kotlinx.coroutines.Dispatchers"
 private const val CURRENT_COROUTINE_CONTEXT = "kotlinx.coroutines.currentCoroutineContext"
 private const val KOTLIN_TODO = "kotlin.TODO"
+private const val COROUTINE_SCOPE = "kotlinx.coroutines.CoroutineScope"
 
 internal class ForbiddenInSuspectContextMethodInspection : LocalInspectionTool() {
   override fun buildVisitor(holder: ProblemsHolder, isOnTheFly: Boolean): PsiElementVisitor {
@@ -146,6 +155,14 @@ internal class ForbiddenInSuspectContextMethodInspection : LocalInspectionTool()
               ifInSuspend { ReplaceDefaultModalityStateWithCurrentModalityQuickFix(expression) }
             )
           }
+          applicationInvokeLater, invokeLaterKt -> {
+            holder.registerProblem(
+              extractElementToHighlight(expression),
+              DevKitKotlinBundle.message("inspections.forbidden.method.in.suspend.context.invoke.later.text"),
+              ProblemHighlightType.GENERIC_ERROR_OR_WARNING,
+              *provideFixesForInvokeLater(expression)
+            )
+          }
           else -> {
             holder.registerProblem(
               extractElementToHighlight(expression),
@@ -198,6 +215,27 @@ internal class ForbiddenInSuspectContextMethodInspection : LocalInspectionTool()
       return callingElement?.let { arrayOf(NavigateToCallInSuspendFunction(it)) } ?: LocalQuickFix.EMPTY_ARRAY
     }
 
+    private fun KtAnalysisSession.provideFixesForInvokeLater(callExpression: KtCallExpression): Array<LocalQuickFix> {
+      callingElement?.let { return arrayOf(NavigateToCallInSuspendFunction(it)) }
+
+      return buildList<LocalQuickFix> {
+        add(ReplaceInvokeLaterWithWithContextQuickFix(callExpression))
+
+        val implicitReceiverTypesAtPosition = getImplicitReceiverTypesAtPosition(callExpression)
+        val coroutineScopeClass = ClassId.topLevel(FqName(COROUTINE_SCOPE))
+        val hasCoroutineScope = implicitReceiverTypesAtPosition.any { type ->
+          type is KtUsualClassType &&
+          (
+            type.classId == coroutineScopeClass ||
+            type.getAllSuperTypes().any { superType -> superType is KtUsualClassType && superType.classId == coroutineScopeClass }
+          )
+        }
+        if (hasCoroutineScope) {
+          add(ReplaceInvokeLaterWithLaunchQuickFix(callExpression))
+        }
+      }.toArray(LocalQuickFix.EMPTY_ARRAY)
+    }
+
     override fun visitLambdaExpression(lambdaExpression: KtLambdaExpression): Unit = Unit
 
     override fun visitDeclaration(dcl: KtDeclaration) {
@@ -205,126 +243,163 @@ internal class ForbiddenInSuspectContextMethodInspection : LocalInspectionTool()
         dcl.initializer?.accept(this)
       }
     }
-  }
 
-  private class ReplaceProgressManagerCheckCanceledQuickFix(element: PsiElement) : LocalQuickFixAndIntentionActionOnPsiElement(element) {
-    override fun getFamilyName(): String = DevKitKotlinBundle.message(
-      "inspections.forbidden.method.in.suspend.context.check.canceled.fix.text")
+    private class ReplaceProgressManagerCheckCanceledQuickFix(element: PsiElement) : LocalQuickFixAndIntentionActionOnPsiElement(element) {
+      override fun getFamilyName(): String = DevKitKotlinBundle.message(
+        "inspections.forbidden.method.in.suspend.context.check.canceled.fix.text")
 
-    override fun getText(): String = familyName
+      override fun getText(): String = familyName
 
-    override fun isAvailable(project: Project, file: PsiFile, startElement: PsiElement, endElement: PsiElement): Boolean =
-      getCallExpression(startElement) != null
+      override fun isAvailable(project: Project, file: PsiFile, startElement: PsiElement, endElement: PsiElement): Boolean =
+        getCallExpression(startElement) != null
 
-    override fun invoke(project: Project, file: PsiFile, editor: Editor?, startElement: PsiElement, endElement: PsiElement) {
-      val callExpression = getCallExpression(startElement)!!
-      val factory = KtPsiFactory(project)
-      val suspendAwareCheckCanceled = factory.createExpression("$COROUTINE_CHECK_CANCELED_FIX()")
-      val qualifiedExpression = callExpression.getQualifiedExpressionForSelector()
-      val expressionToReplace = qualifiedExpression ?: callExpression
-      val resultExpression = expressionToReplace.replace(suspendAwareCheckCanceled)
-      ShortenReferencesFacility.getInstance().shorten(resultExpression as KtElement)
-    }
-  }
-
-  private class ReplaceInvokeAndWaitWithWithContextQuickFix(element: PsiElement) : LocalQuickFixAndIntentionActionOnPsiElement(element) {
-    override fun getFamilyName(): String = DevKitKotlinBundle.message(
-      "inspections.forbidden.method.in.suspend.context.invoke.and.wait.fix.text")
-
-    override fun getText(): String = familyName
-
-    override fun isAvailable(project: Project, file: PsiFile, startElement: PsiElement, endElement: PsiElement): Boolean {
-      val callExpression = getCallExpression(startElement) ?: return false
-      return getLambdaArgumentExpression(callExpression) != null
+      override fun invoke(project: Project, file: PsiFile, editor: Editor?, startElement: PsiElement, endElement: PsiElement) {
+        val callExpression = getCallExpression(startElement)!!
+        val factory = KtPsiFactory(project)
+        val suspendAwareCheckCanceled = factory.createExpression("$COROUTINE_CHECK_CANCELED_FIX()")
+        val qualifiedExpression = callExpression.getQualifiedExpressionForSelector()
+        val expressionToReplace = qualifiedExpression ?: callExpression
+        val resultExpression = expressionToReplace.replace(suspendAwareCheckCanceled)
+        ShortenReferencesFacility.getInstance().shorten(resultExpression as KtElement)
+      }
     }
 
-    override fun invoke(project: Project, file: PsiFile, editor: Editor?, startElement: PsiElement, endElement: PsiElement) {
-      val callExpression = getCallExpression(startElement)!!
+    private open class ReplaceInvokeAndWaitWithWithContextQuickFix(element: PsiElement) : LocalQuickFixAndIntentionActionOnPsiElement(
+      element) {
+      override fun getFamilyName(): String = DevKitKotlinBundle.message(
+        "inspections.forbidden.method.in.suspend.context.invoke.and.wait.fix.text")
 
-      if (!isImported(FqName("com.intellij.openapi.application.EDT"), callExpression.containingKtFile)) {
-        ImportInsertHelperImpl.addImport(project, callExpression.containingKtFile, intelliJEdtDispatcher)
+      override fun getText(): String = familyName
+
+      override fun isAvailable(project: Project, file: PsiFile, startElement: PsiElement, endElement: PsiElement): Boolean {
+        val callExpression = getCallExpression(startElement) ?: return false
+        return getLambdaArgumentExpression(callExpression) != null
       }
 
-      val factory = KtPsiFactory(project)
-      val expression = factory.createExpression("$WITH_CONTEXT($DISPATCHERS.${intelliJEdtDispatcher.shortName().asString()}) {}")
-      val argument = getLambdaArgumentExpression(callExpression)
-      (expression as KtQualifiedExpression).selectorExpression
-        .let { it as KtCallExpression }
-        .lambdaArguments
-        .first()
-        .getLambdaExpression()!!
-        .replace(argument!!)
+      override fun invoke(project: Project, file: PsiFile, editor: Editor?, startElement: PsiElement, endElement: PsiElement) {
+        val callExpression = getCallExpression(startElement)!!
 
-      val qualifiedExpression = callExpression.getQualifiedExpressionForSelector()
-      val resultExpression = if (qualifiedExpression != null) {
-        qualifiedExpression.replace(expression)
+        if (!isImported(intelliJEdtDispatcher, callExpression.containingKtFile)) {
+          ImportInsertHelperImpl.addImport(project, callExpression.containingKtFile, intelliJEdtDispatcher)
+        }
+
+        replaceMethodInCallWithLambda(callExpression, "$WITH_CONTEXT($DISPATCHERS.${intelliJEdtDispatcher.shortName().asString()}) {}")
       }
-      else {
-        callExpression.replace(expression)
+    }
+
+    private class ReplaceInvokeLaterWithWithContextQuickFix(element: PsiElement) : ReplaceInvokeAndWaitWithWithContextQuickFix(element) {
+      override fun getFamilyName(): String = DevKitKotlinBundle.message(
+        "inspections.forbidden.method.in.suspend.context.invoke.later.fix.with.context.text")
+    }
+
+    private class ReplaceInvokeLaterWithLaunchQuickFix(element: PsiElement) : LocalQuickFixAndIntentionActionOnPsiElement(element) {
+      override fun getFamilyName(): String = DevKitKotlinBundle.message(
+        "inspections.forbidden.method.in.suspend.context.invoke.later.fix.launch.text")
+
+      override fun getText(): String = familyName
+
+      override fun isAvailable(project: Project, file: PsiFile, startElement: PsiElement, endElement: PsiElement): Boolean {
+        val callExpression = getCallExpression(startElement) ?: return false
+        return getLambdaArgumentExpression(callExpression) != null
       }
 
-      ShortenReferencesFacility.getInstance().shorten(resultExpression as KtElement)
+      override fun invoke(project: Project, file: PsiFile, editor: Editor?, startElement: PsiElement, endElement: PsiElement) {
+        val callExpression = getCallExpression(startElement)!!
+
+        if (!isImported(intelliJEdtDispatcher, callExpression.containingKtFile)) {
+          ImportInsertHelperImpl.addImport(project, callExpression.containingKtFile, intelliJEdtDispatcher)
+        }
+        if (!isImported(coroutinesLaunch, callExpression.containingKtFile)) {
+          ImportInsertHelperImpl.addImport(project, callExpression.containingKtFile, coroutinesLaunch)
+        }
+
+        replaceMethodInCallWithLambda(callExpression,
+                                      "${coroutinesLaunch.shortName()}($DISPATCHERS.${intelliJEdtDispatcher.shortName()}) {}")
+      }
     }
 
-    private fun getLambdaArgumentExpression(callExpression: KtCallExpression): KtExpression? {
-      val lambdaExpression = callExpression.lambdaArguments.firstOrNull()?.getArgumentExpression() as? KtLambdaExpression
-      if (lambdaExpression != null) return lambdaExpression
-      return callExpression.valueArguments.firstOrNull()?.getArgumentExpression() as? KtLambdaExpression
+    private class ReplaceDefaultModalityStateWithCurrentModalityQuickFix(element: PsiElement) : LocalQuickFixAndIntentionActionOnPsiElement(
+      element
+    ) {
+      override fun getFamilyName(): String = DevKitKotlinBundle.message(
+        "inspections.forbidden.method.in.suspend.context.default.modality.state.fix.text")
+
+      override fun getText(): String = familyName
+
+      override fun isAvailable(project: Project, file: PsiFile, startElement: PsiElement, endElement: PsiElement): Boolean =
+        getCallExpression(startElement) != null
+
+      override fun invoke(project: Project, file: PsiFile, editor: Editor?, startElement: PsiElement, endElement: PsiElement) {
+        val callExpression = getCallExpression(startElement)!!
+        val factory = KtPsiFactory(project)
+
+        if (!isImported(contextModalityExt, callExpression.containingKtFile)) {
+          ImportInsertHelperImpl.addImport(project, callExpression.containingKtFile, contextModalityExt)
+        }
+
+        // ?: because defaultModalityState is @NotNull, so it can be unexpected to replace with something nullable
+        val contextModalityCall = factory.createExpression(
+          "$CURRENT_COROUTINE_CONTEXT().${contextModalityExt.shortName()}() ?: $KOTLIN_TODO(\"Handle absence of ModalityState\")"
+        )
+        val qualifiedExpression = callExpression.getQualifiedExpressionForSelector()
+        val expressionToReplace = qualifiedExpression ?: callExpression
+        val resultExpression = expressionToReplace.replace(contextModalityCall)
+        ShortenReferencesFacility.getInstance().shorten(resultExpression as KtElement)
+      }
     }
-  }
 
-  private class ReplaceDefaultModalityStateWithCurrentModalityQuickFix(element: PsiElement) : LocalQuickFixAndIntentionActionOnPsiElement(
-    element
-  ) {
-    override fun getFamilyName(): String = DevKitKotlinBundle.message(
-      "inspections.forbidden.method.in.suspend.context.default.modality.state.fix.text")
+    private class NavigateToCallInSuspendFunction(callingElement: KtCallExpression) : IntentionAndQuickFixAction() {
+      private val pointer: SmartPsiElementPointer<KtCallExpression> = SmartPointerManager.createPointer(callingElement)
 
-    override fun getText(): String = familyName
+      override fun getFamilyName(): String = DevKitKotlinBundle.message(
+        "inspections.forbidden.method.in.suspend.context.navigate.to.suspend.context")
 
-    override fun isAvailable(project: Project, file: PsiFile, startElement: PsiElement, endElement: PsiElement): Boolean =
-      getCallExpression(startElement) != null
-
-    override fun invoke(project: Project, file: PsiFile, editor: Editor?, startElement: PsiElement, endElement: PsiElement) {
-      val callExpression = getCallExpression(startElement)!!
-      val factory = KtPsiFactory(project)
-
-      if (!isImported(contextModalityExt, callExpression.containingKtFile)) {
-        ImportInsertHelperImpl.addImport(project, callExpression.containingKtFile, contextModalityExt)
+      override fun isAvailable(project: Project, editor: Editor?, file: PsiFile?): Boolean {
+        return editor != null && pointer.element != null
       }
 
-      // ?: because defaultModalityState is @NotNull, so it can be unexpected to replace with something nullable
-      val contextModalityCall = factory.createExpression(
-        "$CURRENT_COROUTINE_CONTEXT().${contextModalityExt.shortName()}() ?: $KOTLIN_TODO(\"Handle absence of ModalityState\")"
-      )
-      val qualifiedExpression = callExpression.getQualifiedExpressionForSelector()
-      val expressionToReplace = qualifiedExpression ?: callExpression
-      val resultExpression = expressionToReplace.replace(contextModalityCall)
-      ShortenReferencesFacility.getInstance().shorten(resultExpression as KtElement)
+      override fun applyFix(project: Project, file: PsiFile?, editor: Editor?) {
+        extractElementToHighlight(pointer.element!!).navigate(true)
+      }
+
+      override fun startInWriteAction(): Boolean = false
+
+      override fun getName(): String = familyName
     }
   }
-
-  private class NavigateToCallInSuspendFunction(callingElement: KtCallExpression) : IntentionAndQuickFixAction() {
-    private val pointer: SmartPsiElementPointer<KtCallExpression> = SmartPointerManager.createPointer(callingElement)
-
-    override fun getFamilyName(): String = DevKitKotlinBundle.message("inspections.forbidden.method.in.suspend.context.navigate.to.suspend.context")
-
-    override fun isAvailable(project: Project, editor: Editor?, file: PsiFile?): Boolean {
-      return editor != null && pointer.element != null
-    }
-
-    override fun applyFix(project: Project, file: PsiFile?, editor: Editor?) {
-      extractElementToHighlight(pointer.element!!).navigate (true)
-    }
-
-    override fun startInWriteAction(): Boolean = false
-
-    override fun getName(): String = familyName
-  }
 }
 
 private fun getCallExpression(startElement: PsiElement): KtCallExpression? =
   startElement.getParentOfType<KtCallExpression>(false)
 
+private fun getLambdaArgumentExpression(callExpression: KtCallExpression): KtLambdaExpression? {
+  val lambdaExpression = callExpression.lambdaArguments.firstOrNull()?.getArgumentExpression() as? KtLambdaExpression
+  if (lambdaExpression != null) return lambdaExpression
+  return callExpression.valueArguments.firstOrNull()?.getArgumentExpression() as? KtLambdaExpression
+}
+
+private fun replaceMethodInCallWithLambda(callExpression: KtCallExpression, newCallPattern: String) {
+  val factory = KtPsiFactory(callExpression.project)
+  val expression = factory.createExpression(newCallPattern)
+  val argument = getLambdaArgumentExpression(callExpression)
+  (expression as? KtQualifiedExpression)?.selectorExpression?.let { it as KtCallExpression }
+    .let { it ?: expression as KtCallExpression }
+    .lambdaArguments
+    .first()
+    .getLambdaExpression()!!
+    .replace(argument!!)
+
+  val qualifiedExpression = callExpression.getQualifiedExpressionForSelector()
+  val resultExpression = if (qualifiedExpression != null) {
+    qualifiedExpression.replace(expression)
+  }
+  else {
+    callExpression.replace(expression)
+  }
+
+  ShortenReferencesFacility.getInstance().shorten(resultExpression as KtElement)
+}
+
 private fun extractElementToHighlight(expression: KtCallExpression): KtElement = expression.getCallNameExpression() ?: expression
 
 private fun isImported(name: FqName, file: KtFile): Boolean {