[devkit] IDEA-309523 Extend the inspection to other loop types and add a quick fix to insert a cancellation check

GitOrigin-RevId: 7548f7bce82a119670885fc2de88f326ec8d1662

plugins/devkit/devkit-core/resources/intellij.devkit.core.xml

@@ -515,6 +515,7 @@
   <extensions defaultExtensionNs="DevKit.lang">
     <goodCodeRedVisitor language="JAVA" implementationClass="org.jetbrains.idea.devkit.inspections.internal.JavaGoodCodeRedVisitor"/>
     <cancellationCheckProvider language="JAVA" implementationClass="org.jetbrains.idea.devkit.inspections.JavaCancellationCheckProvider"/>
+    <cancellationCheckInLoopsFixProvider language="JAVA" implementationClass="org.jetbrains.idea.devkit.inspections.quickfix.JavaCancellationCheckInLoopsFixProvider"/>
     <visitorProviderForRBCInspection language="JAVA"
       implementationClass="org.jetbrains.idea.devkit.inspections.CallingMethodShouldBeRequiresBlockingContextInspection$VisitorProviderForJava"/>
     <extensionClassShouldNotBePublicProvider language="JAVA"
@@ -541,6 +542,11 @@
                     dynamic="true">
     <with attribute="implementationClass" implements="org.jetbrains.idea.devkit.inspections.CancellationCheckProvider"/>
     </extensionPoint>
+    <extensionPoint qualifiedName="DevKit.lang.cancellationCheckInLoopsFixProvider"
+                    beanClass="com.intellij.lang.LanguageExtensionPoint"
+                    dynamic="true">
+      <with attribute="implementationClass" implements="org.jetbrains.idea.devkit.inspections.quickfix.CancellationCheckInLoopsFixProvider"/>
+    </extensionPoint>
     <extensionPoint qualifiedName="DevKit.lang.visitorProviderForRBCInspection"
                     beanClass="com.intellij.lang.LanguageExtensionPoint"
                     dynamic="true">

plugins/devkit/devkit-core/resources/messages/DevKitBundle.properties

@@ -670,6 +670,7 @@ inspection.extension.registered.as.component.message=A class must not be registe
 
 inspection.cancellation.check.in.loops.display.name=Cancellation check in loops
 inspection.cancellation.check.in.loops.message=Cancellation check ''{0}'' should be placed in the first line
+inspection.insert.cancellation.check.fix.message=Insert cancellation check
 
 inspections.application.service.as.static.final.field.display.name=Application service assigned to a static final field/property
 inspections.application.service.as.static.final.field.message=Application service must not be assigned to a static final field

plugins/devkit/devkit-core/src/inspections/CancellationCheckInLoopsInspection.kt

@@ -3,14 +3,13 @@ package org.jetbrains.idea.devkit.inspections
 
 import com.intellij.codeInspection.ProblemsHolder
 import com.intellij.psi.PsiElementVisitor
-import com.intellij.psi.PsiMember
-import com.intellij.psi.util.PsiUtil.getMemberQualifiedName
 import com.intellij.uast.UastHintedVisitorAdapter
+import com.intellij.util.concurrency.annotations.RequiresReadLock
 import org.jetbrains.idea.devkit.DevKitBundle
+import org.jetbrains.idea.devkit.inspections.quickfix.CancellationCheckInLoopsFixProviders
 import org.jetbrains.uast.*
 import org.jetbrains.uast.visitor.AbstractUastNonRecursiveVisitor
 
-private const val REQUIRES_READ_LOCK_FQN = "com.intellij.util.concurrency.annotations.RequiresReadLock"
 
 class CancellationCheckInLoopsInspection : DevKitUastInspectionBase() {
 
@@ -20,12 +19,27 @@ class CancellationCheckInLoopsInspection : DevKitUastInspectionBase() {
     return UastHintedVisitorAdapter.create(
       holder.file.language,
       object : AbstractUastNonRecursiveVisitor() {
+        override fun visitForExpression(node: UForExpression): Boolean {
+          inspectLoopExpression(node, checkProvider, holder)
+          return true
+        }
+
         override fun visitForEachExpression(node: UForEachExpression): Boolean {
-          inspectForEachExpression(node, checkProvider, holder)
+          inspectLoopExpression(node, checkProvider, holder)
+          return true
+        }
+
+        override fun visitWhileExpression(node: UWhileExpression): Boolean {
+          inspectLoopExpression(node, checkProvider, holder)
+          return true
+        }
+
+        override fun visitDoWhileExpression(node: UDoWhileExpression): Boolean {
+          inspectLoopExpression(node, checkProvider, holder)
           return true
         }
       },
-      arrayOf(UForEachExpression::class.java)
+      arrayOf(ULoopExpression::class.java)
     )
   }
 
@@ -33,44 +47,47 @@ class CancellationCheckInLoopsInspection : DevKitUastInspectionBase() {
    * If the first expression in a loop is not another loop, finds the right cancellation check based on the context of a loop,
    * and if it's missing, registers the problem.
    */
-  private fun inspectForEachExpression(forEachExpression: UForEachExpression,
-                                       checkProvider: CancellationCheckProvider,
-                                       holder: ProblemsHolder) {
-    val sourcePsi = forEachExpression.sourcePsi ?: return
+  private fun inspectLoopExpression(loopExpression: ULoopExpression,
+                                    checkProvider: CancellationCheckProvider,
+                                    holder: ProblemsHolder) {
+    val sourcePsi = loopExpression.sourcePsi ?: return
 
-    if (!shouldBeRunOn(forEachExpression)) return
+    if (!shouldBeRunOn(loopExpression)) return
 
-    val firstExpressionInLoop = getFirstExpressionInLoop(forEachExpression)
+    val firstExpressionInLoop = loopExpression.bodyExpressions.firstOrNull()
 
     // Don't insert a check between nested loops if there is nothing in between
-    if (firstExpressionInLoop is UForEachExpression) return
+    if (firstExpressionInLoop is ULoopExpression) return
 
-    val cancellationCheckFqn = checkProvider.findCancellationCheckFqn(sourcePsi)
-    if (firstExpressionInLoop.isCancellationCheck(cancellationCheckFqn)) return
+    val cancellationCheckFqn = checkProvider.findCancellationCheckCall(sourcePsi)
+    if (checkProvider.isCancellationCheckCall(firstExpressionInLoop?.sourcePsi, cancellationCheckFqn)) return
 
     val anchor = sourcePsi.firstChild
-    holder.registerProblem(anchor, DevKitBundle.message("inspection.cancellation.check.in.loops.message", cancellationCheckFqn))
+    val fixProvider = CancellationCheckInLoopsFixProviders.forLanguage(holder.file.language) ?: return
+    val fixes = fixProvider.getFixes(anchor, cancellationCheckFqn)
+    holder.registerProblem(
+      anchor,
+      DevKitBundle.message("inspection.cancellation.check.in.loops.message", cancellationCheckFqn),
+      *fixes.toTypedArray()
+    )
   }
 
   /**
-   * For now, insert a cancellation check in loops with [com.intellij.util.concurrency.annotations.RequiresReadLock]
+   * For now, insert a cancellation check in loops with [com.intellij.util.concurrency.annotations.RequiresReadLock].
    */
   private fun shouldBeRunOn(uElement: UElement): Boolean {
     val containingMethod = uElement.getParentOfType<UMethod>() ?: return false
     val superMethods = containingMethod.javaPsi.findSuperMethods()
-    return superMethods.plus(containingMethod).any { it.hasAnnotation(REQUIRES_READ_LOCK_FQN) }
+    return superMethods.plus(containingMethod).any { it.hasAnnotation(RequiresReadLock::class.java.canonicalName) }
   }
 
-  private fun getFirstExpressionInLoop(loop: UForEachExpression): UExpression? {
-    return when (val body = loop.body) {
-      is UBlockExpression -> body.expressions.firstOrNull()
-      else -> body
+  private val ULoopExpression.bodyExpressions: List<UExpression>
+    get() {
+      return when (val loopBody = body) {
+        is UBlockExpression -> loopBody.expressions
+        else -> listOf(body)
+      }
     }
-  }
-
-  private fun UExpression?.isCancellationCheck(cancellationCheckFqn: String): Boolean {
-    val resolved = this?.tryResolve() as? PsiMember ?: return false
-    return getMemberQualifiedName(resolved) == cancellationCheckFqn
-  }
 
 }
+

plugins/devkit/devkit-core/src/inspections/CancellationCheckProvider.kt

@@ -4,39 +4,44 @@ package org.jetbrains.idea.devkit.inspections
 import com.intellij.lang.LanguageExtension
 import com.intellij.openapi.extensions.ExtensionPointName
 import com.intellij.openapi.extensions.ExtensionPointName.Companion.create
+import com.intellij.openapi.util.IntellijInternalApi
 import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiMethodCallExpression
+import com.intellij.psi.util.PsiUtil
+import org.jetbrains.annotations.ApiStatus
+
+
+private const val PROGRESS_MANAGER_CHECKED_CANCELED = "com.intellij.openapi.progress.ProgressManager.checkCanceled"
 
 private val EP_NAME: ExtensionPointName<CancellationCheckProvider> = create("DevKit.lang.cancellationCheckProvider")
+
 internal object CancellationCheckProviders : LanguageExtension<CancellationCheckProvider>(EP_NAME.name)
 
-private const val CANCELLATION_CHECK_BLOCKING_FQN = "com.intellij.openapi.progress.ProgressManager.checkCanceled"
-private const val CANCELLATION_CHECK_SUSPENDING_FQN = "com.intellij.openapi.progress.CoroutinesKt.checkCancelled"
 
 /**
  * Provides the right cancellation check based on the context
- * (see [com.intellij.openapi.progress.ProgressManager.checkCanceled],
- * [com.intellij.openapi.progress.checkCancelled]).
+ * (see [com.intellij.openapi.progress.ProgressManager.checkCanceled], [com.intellij.openapi.progress.checkCancelled])
+ * and checks expressions for cancellation check calls.
  */
+@IntellijInternalApi
+@ApiStatus.Internal
 interface CancellationCheckProvider {
 
-  enum class Context {
-    BLOCKING, SUSPENDING
-  }
+  fun findCancellationCheckCall(element: PsiElement): String
 
-  fun findCancellationCheckFqn(element: PsiElement): String {
-    return when (findContext(element)) {
-      Context.BLOCKING -> CANCELLATION_CHECK_BLOCKING_FQN
-      Context.SUSPENDING -> CANCELLATION_CHECK_SUSPENDING_FQN
-    }
-  }
-
-  fun findContext(element: PsiElement): Context
+  fun isCancellationCheckCall(element: PsiElement?, cancellationCheckFqn: String): Boolean
 
 }
 
 
-class JavaCancellationCheckProvider : CancellationCheckProvider {
+internal class JavaCancellationCheckProvider : CancellationCheckProvider {
 
-  override fun findContext(element: PsiElement) = CancellationCheckProvider.Context.BLOCKING
+  override fun findCancellationCheckCall(element: PsiElement): String {
+    return PROGRESS_MANAGER_CHECKED_CANCELED
+  }
 
+  override fun isCancellationCheckCall(element: PsiElement?, cancellationCheckFqn: String): Boolean {
+    val resolvedMethod = (element as? PsiMethodCallExpression)?.resolveMethod() ?: return false
+    return PsiUtil.getMemberQualifiedName(resolvedMethod) == cancellationCheckFqn
+  }
 }

plugins/devkit/devkit-core/src/inspections/quickfix/CancellationCheckInLoopsFixProvider.kt

@@ -0,0 +1,73 @@
+// Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package org.jetbrains.idea.devkit.inspections.quickfix
+
+import com.intellij.codeInsight.BlockUtils.expandSingleStatementToBlockStatement
+import com.intellij.codeInspection.LocalQuickFix
+import com.intellij.codeInspection.LocalQuickFixOnPsiElement
+import com.intellij.lang.LanguageExtension
+import com.intellij.openapi.extensions.ExtensionPointName
+import com.intellij.openapi.project.Project
+import com.intellij.openapi.util.IntellijInternalApi
+import com.intellij.psi.*
+import com.intellij.psi.codeStyle.JavaCodeStyleManager
+import com.intellij.psi.util.PsiTreeUtil
+import com.intellij.psi.util.parentOfType
+import org.jetbrains.annotations.ApiStatus
+import org.jetbrains.idea.devkit.DevKitBundle
+
+
+private val EP_NAME: ExtensionPointName<CancellationCheckInLoopsFixProvider> = ExtensionPointName.create(
+  "DevKit.lang.cancellationCheckInLoopsFixProvider")
+
+internal object CancellationCheckInLoopsFixProviders : LanguageExtension<CancellationCheckInLoopsFixProvider>(EP_NAME.name)
+
+
+@IntellijInternalApi
+@ApiStatus.Internal
+interface CancellationCheckInLoopsFixProvider {
+
+  fun getFixes(loopKeyword: PsiElement, cancellationCheckFqn: String): List<LocalQuickFix>
+
+}
+
+
+internal class JavaCancellationCheckInLoopsFixProvider : CancellationCheckInLoopsFixProvider {
+
+  override fun getFixes(loopKeyword: PsiElement, cancellationCheckFqn: String): List<LocalQuickFix> {
+    return listOf(InsertCancellationCheckFix(cancellationCheckFqn, loopKeyword))
+  }
+
+}
+
+
+class InsertCancellationCheckFix(
+  private val cancellationCheckCallFqn: String,
+  loopKeyword: PsiElement,
+) : LocalQuickFixOnPsiElement(loopKeyword) {
+
+  override fun getFamilyName(): String = DevKitBundle.message("inspection.insert.cancellation.check.fix.message")
+
+  override fun getText(): String = familyName
+
+  override fun isAvailable(project: Project, file: PsiFile, startElement: PsiElement, endElement: PsiElement): Boolean {
+    return PsiTreeUtil.getParentOfType(startElement, PsiLoopStatement::class.java) != null
+  }
+
+  override fun invoke(project: Project, file: PsiFile, startElement: PsiElement, endElement: PsiElement) {
+    val loopStatement = PsiTreeUtil.getParentOfType(startElement, PsiLoopStatement::class.java) ?: return
+    val factory = PsiElementFactory.getInstance(project)
+
+    val cancellationCheckText = "${cancellationCheckCallFqn}();"
+    val cancellationCheckStatement = factory.createStatementFromText(cancellationCheckText, loopStatement)
+
+    val body = loopStatement.body ?: return
+    val bodyBlock = expandSingleStatementToBlockStatement(body).parentOfType<PsiBlockStatement>(withSelf = true)
+
+    val insertedElement = bodyBlock?.codeBlock?.addBefore(cancellationCheckStatement, bodyBlock.codeBlock.firstBodyElement)
+
+    insertedElement?.let {
+      JavaCodeStyleManager.getInstance(project).shortenClassReferences(it)
+    }
+  }
+
+}

plugins/devkit/devkit-java-tests/testData/inspections/cancellationCheckInLoops/DoWhileLoops.java

@@ -0,0 +1,44 @@
+package inspections.cancellationCheckInLoops;
+
+import com.intellij.util.concurrency.annotations.RequiresReadLock;
+
+import static inspections.cancellationCheckInLoops.Foo.doSomething;
+
+
+class Clazz {
+
+  @RequiresReadLock
+  public static void foo() {
+    int i = 0;
+    int j = 0;
+    do {
+      // nested loops with something in between
+      <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">do</warning> {
+        doSomething();
+        <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">do</warning> {
+          doSomething();
+          i++;
+        } while (i < 10);
+        i++;
+      } while (i < 5);
+
+      // nested loops with a block in between
+      <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">do</warning> {
+        if (i != 3) {
+          //empty loop
+          <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">do</warning> {
+          } while (j > 5);
+        }
+        i++;
+      } while (i < 15);
+
+      // single-line loop
+      <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">do</warning> System.out.println(i); while (i < 20);
+
+      // no body loop
+      <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">do</warning>; while(i < 0);
+
+      i++;
+    } while (i < 100);
+  }
+}

plugins/devkit/devkit-java-tests/testData/inspections/cancellationCheckInLoops/ForEachLoops.java

@@ -32,6 +32,9 @@ class Clazz {
 
       // single-line loop
       <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">for</warning> (String item6: items) System.out.println(item);
+
+      // no body loop
+      <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">for</warning>(String item5 : items);
     }
   }
 }

plugins/devkit/devkit-java-tests/testData/inspections/cancellationCheckInLoops/ForLoops.java

@@ -0,0 +1,38 @@
+package inspections.cancellationCheckInLoops;
+
+import com.intellij.util.concurrency.annotations.RequiresReadLock;
+
+import static inspections.cancellationCheckInLoops.Foo.doSomething;
+
+
+class Clazz {
+
+  @RequiresReadLock
+  public static void foo() {
+    for (int i = 0; i < 5; i++) {
+
+      // nested loops with something in between
+      <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">for</warning> (int j = 0; j < 5; j++) {
+        doSomething();
+        <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">for</warning> (int k = 0; k < 5; k++) {
+          doSomething();
+        }
+      }
+
+      // nested loops with a block in between
+      <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">for</warning> (int j = 0; j < 5; j++) {
+        if (j != 3) {
+          //empty loop
+          <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">for</warning> (int k = 0; k < 5; k++) {
+          }
+        }
+      }
+
+      // single-line loop
+      <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">for</warning> (int j = 0; j < 5; j++) System.out.println(j);
+
+      // no body loop
+      <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">for</warning> (int y = 0; y < 10; y++);
+    }
+  }
+}

plugins/devkit/devkit-java-tests/testData/inspections/cancellationCheckInLoops/NestedLoops.java

@@ -0,0 +1,32 @@
+package inspections.cancellationCheckInLoops;
+
+import com.intellij.util.concurrency.annotations.RequiresReadLock;
+
+import static inspections.cancellationCheckInLoops.Foo.doSomething;
+
+class Clazz {
+
+  @RequiresReadLock
+  public static void foo() {
+    int j = 0;
+    String[] items = {""};
+
+    // nested loops of different kinds
+    while (j < 100) {
+      for (int i = 0; i < 5; i++) {
+        for (String item : items) {
+          <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">do</warning> {
+            doSomething();
+            j++;
+          } while (j < 5);
+        }
+      }
+      j++;
+    }
+
+    // single line nested loops
+    for (int i = 0; i < 5; i++) <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">for</warning> (String item : items) {
+      doSomething();
+    }
+  }
+}

plugins/devkit/devkit-java-tests/testData/inspections/cancellationCheckInLoops/PresentCancellationCheck.java

@@ -15,5 +15,23 @@ class Clazz {
       ProgressManager.checkCanceled();
       doSomething();
     }
+
+    for (int i = 0; i < 5; i++) {
+      ProgressManager.checkCanceled();
+      doSomething();
+    }
+
+    int i = 0;
+    while (i < 5) {
+      ProgressManager.checkCanceled();
+      doSomething();
+      i++;
+    }
+
+    do {
+      ProgressManager.checkCanceled();
+      doSomething();
+      i++;
+    } while (i < 10);
   }
 }

plugins/devkit/devkit-java-tests/testData/inspections/cancellationCheckInLoops/WhileLoops.java

@@ -0,0 +1,44 @@
+package inspections.cancellationCheckInLoops;
+
+import com.intellij.util.concurrency.annotations.RequiresReadLock;
+
+import static inspections.cancellationCheckInLoops.Foo.doSomething;
+
+
+class Clazz {
+
+  @RequiresReadLock
+  public static void foo() {
+    int i = 0;
+    int j = 0;
+    while (i < 100) {
+      // nested loops with something in between
+      <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">while</warning> (i < 5) {
+        doSomething();
+        <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">while</warning> (i < 10) {
+          doSomething();
+          i++;
+        }
+        i++;
+      }
+
+      // nested loops with a block in between
+      <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">while</warning> (i < 15) {
+        if (i != 3) {
+          //empty loop
+          <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">while</warning> (j > 5) {
+          }
+        }
+        i++;
+      }
+
+      // single-line loop
+      <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">while</warning> (i < 20) System.out.println(i);
+
+      // no body loop
+      <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">while</warning>(i < 0);
+
+      i++;
+    }
+  }
+}

plugins/devkit/devkit-java-tests/testData/inspections/insertCancellationCheckFix/BlockDoWhileLoop.java

@@ -0,0 +1,19 @@
+package inspections.cancellationCheckInLoops;
+
+import com.intellij.util.concurrency.annotations.RequiresReadLock;
+
+import static inspections.cancellationCheckInLoops.Foo.doSomething;
+
+
+class Clazz {
+
+    @RequiresReadLock
+    public static void foo() {
+        int i = 0;
+        <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">do<caret></warning> {
+            doSomething();
+            i++;
+        }
+        while (i < 100);
+    }
+}

plugins/devkit/devkit-java-tests/testData/inspections/insertCancellationCheckFix/BlockDoWhileLoop_after.java

@@ -0,0 +1,21 @@
+package inspections.cancellationCheckInLoops;
+
+import com.intellij.openapi.progress.ProgressManager;
+import com.intellij.util.concurrency.annotations.RequiresReadLock;
+
+import static inspections.cancellationCheckInLoops.Foo.doSomething;
+
+
+class Clazz {
+
+    @RequiresReadLock
+    public static void foo() {
+        int i = 0;
+        do {
+            ProgressManager.checkCanceled();
+            doSomething();
+            i++;
+        }
+        while (i < 100);
+    }
+}

plugins/devkit/devkit-java-tests/testData/inspections/insertCancellationCheckFix/BlockForEachLoop.java

@@ -0,0 +1,18 @@
+package inspections.cancellationCheckInLoops;
+
+import com.intellij.util.concurrency.annotations.RequiresReadLock;
+
+import static inspections.cancellationCheckInLoops.Foo.doSomething;
+
+
+class Clazz {
+
+    @RequiresReadLock
+    public static void foo() {
+        String[] items = {""};
+
+        <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">for<caret></warning> (String item : items) {
+            doSomething();
+        }
+    }
+}

plugins/devkit/devkit-java-tests/testData/inspections/insertCancellationCheckFix/BlockForEachLoop_after.java

@@ -0,0 +1,20 @@
+package inspections.cancellationCheckInLoops;
+
+import com.intellij.openapi.progress.ProgressManager;
+import com.intellij.util.concurrency.annotations.RequiresReadLock;
+
+import static inspections.cancellationCheckInLoops.Foo.doSomething;
+
+
+class Clazz {
+
+    @RequiresReadLock
+    public static void foo() {
+        String[] items = {""};
+
+        for (String item : items) {
+            ProgressManager.checkCanceled();
+            doSomething();
+        }
+    }
+}

plugins/devkit/devkit-java-tests/testData/inspections/insertCancellationCheckFix/BlockForLoop.java

@@ -0,0 +1,16 @@
+package inspections.cancellationCheckInLoops;
+
+import com.intellij.util.concurrency.annotations.RequiresReadLock;
+
+import static inspections.cancellationCheckInLoops.Foo.doSomething;
+
+
+class Clazz {
+
+    @RequiresReadLock
+    public static void foo() {
+        <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">for<caret></warning> (int i = 0; i < 5; i++) {
+            doSomething();
+        }
+    }
+}

plugins/devkit/devkit-java-tests/testData/inspections/insertCancellationCheckFix/BlockForLoop_after.java

@@ -0,0 +1,18 @@
+package inspections.cancellationCheckInLoops;
+
+import com.intellij.openapi.progress.ProgressManager;
+import com.intellij.util.concurrency.annotations.RequiresReadLock;
+
+import static inspections.cancellationCheckInLoops.Foo.doSomething;
+
+
+class Clazz {
+
+    @RequiresReadLock
+    public static void foo() {
+        for (int i = 0; i < 5; i++) {
+            ProgressManager.checkCanceled();
+            doSomething();
+        }
+    }
+}

plugins/devkit/devkit-java-tests/testData/inspections/insertCancellationCheckFix/BlockWhileLoop.java

@@ -0,0 +1,18 @@
+package inspections.cancellationCheckInLoops;
+
+import com.intellij.util.concurrency.annotations.RequiresReadLock;
+
+import static inspections.cancellationCheckInLoops.Foo.doSomething;
+
+
+class Clazz {
+
+    @RequiresReadLock
+    public static void foo() {
+        int i = 0;
+        <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">while<caret></warning> (i < 5) {
+            doSomething();
+            i++;
+        }
+    }
+}

plugins/devkit/devkit-java-tests/testData/inspections/insertCancellationCheckFix/BlockWhileLoop_after.java

@@ -0,0 +1,20 @@
+package inspections.cancellationCheckInLoops;
+
+import com.intellij.openapi.progress.ProgressManager;
+import com.intellij.util.concurrency.annotations.RequiresReadLock;
+
+import static inspections.cancellationCheckInLoops.Foo.doSomething;
+
+
+class Clazz {
+
+    @RequiresReadLock
+    public static void foo() {
+        int i = 0;
+        while (i < 5) {
+            ProgressManager.checkCanceled();
+            doSomething();
+            i++;
+        }
+    }
+}

plugins/devkit/devkit-java-tests/testData/inspections/insertCancellationCheckFix/EmptyDoWhileLoop.java

@@ -0,0 +1,16 @@
+package inspections.cancellationCheckInLoops;
+
+import com.intellij.util.concurrency.annotations.RequiresReadLock;
+
+
+class Clazz {
+
+    @RequiresReadLock
+    public static void foo() {
+        int i = 0;
+        <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">do<caret></warning> {
+            // check comments
+        }
+        while (i < 100);
+    }
+}

plugins/devkit/devkit-java-tests/testData/inspections/insertCancellationCheckFix/EmptyDoWhileLoop_after.java

@@ -0,0 +1,18 @@
+package inspections.cancellationCheckInLoops;
+
+import com.intellij.openapi.progress.ProgressManager;
+import com.intellij.util.concurrency.annotations.RequiresReadLock;
+
+
+class Clazz {
+
+    @RequiresReadLock
+    public static void foo() {
+        int i = 0;
+        do {
+            ProgressManager.checkCanceled();
+            // check comments
+        }
+        while (i < 100);
+    }
+}

plugins/devkit/devkit-java-tests/testData/inspections/insertCancellationCheckFix/EmptyForEachLoop.java

@@ -0,0 +1,16 @@
+package inspections.cancellationCheckInLoops;
+
+import com.intellij.util.concurrency.annotations.RequiresReadLock;
+
+
+class Clazz {
+
+    @RequiresReadLock
+    public static void foo() {
+        String[] items = {""};
+
+        <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">for<caret></warning> (String item : items) {
+            // check comments
+        }
+    }
+}

plugins/devkit/devkit-java-tests/testData/inspections/insertCancellationCheckFix/EmptyForEachLoop_after.java

@@ -0,0 +1,18 @@
+package inspections.cancellationCheckInLoops;
+
+import com.intellij.openapi.progress.ProgressManager;
+import com.intellij.util.concurrency.annotations.RequiresReadLock;
+
+
+class Clazz {
+
+    @RequiresReadLock
+    public static void foo() {
+        String[] items = {""};
+
+        for (String item : items) {
+            ProgressManager.checkCanceled();
+            // check comments
+        }
+    }
+}

plugins/devkit/devkit-java-tests/testData/inspections/insertCancellationCheckFix/EmptyForLoop.java

@@ -0,0 +1,14 @@
+package inspections.cancellationCheckInLoops;
+
+import com.intellij.util.concurrency.annotations.RequiresReadLock;
+
+
+class Clazz {
+
+    @RequiresReadLock
+    public static void foo() {
+        <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">for<caret></warning> (int i = 0; i < 5; i++) {
+            // check comments
+        }
+    }
+}

plugins/devkit/devkit-java-tests/testData/inspections/insertCancellationCheckFix/EmptyForLoop_after.java

@@ -0,0 +1,16 @@
+package inspections.cancellationCheckInLoops;
+
+import com.intellij.openapi.progress.ProgressManager;
+import com.intellij.util.concurrency.annotations.RequiresReadLock;
+
+
+class Clazz {
+
+    @RequiresReadLock
+    public static void foo() {
+        for (int i = 0; i < 5; i++) {
+            ProgressManager.checkCanceled();
+            // check comments
+        }
+    }
+}

plugins/devkit/devkit-java-tests/testData/inspections/insertCancellationCheckFix/EmptyWhileLoop.java

@@ -0,0 +1,17 @@
+package inspections.cancellationCheckInLoops;
+
+import com.intellij.util.concurrency.annotations.RequiresReadLock;
+
+import static inspections.cancellationCheckInLoops.Foo.doSomething;
+
+
+class Clazz {
+
+    @RequiresReadLock
+    public static void foo() {
+        int i = 0;
+        <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">while<caret></warning> (i < 5) {
+          // check comments
+        }
+    }
+}

plugins/devkit/devkit-java-tests/testData/inspections/insertCancellationCheckFix/EmptyWhileLoop_after.java

@@ -0,0 +1,19 @@
+package inspections.cancellationCheckInLoops;
+
+import com.intellij.openapi.progress.ProgressManager;
+import com.intellij.util.concurrency.annotations.RequiresReadLock;
+
+import static inspections.cancellationCheckInLoops.Foo.doSomething;
+
+
+class Clazz {
+
+    @RequiresReadLock
+    public static void foo() {
+        int i = 0;
+        while (i < 5) {
+            ProgressManager.checkCanceled();
+            // check comments
+        }
+    }
+}

plugins/devkit/devkit-java-tests/testData/inspections/insertCancellationCheckFix/NoBodyForEachLoop.java

@@ -0,0 +1,16 @@
+package inspections.cancellationCheckInLoops;
+
+import com.intellij.util.concurrency.annotations.RequiresReadLock;
+
+import static inspections.cancellationCheckInLoops.Foo.doSomething;
+
+
+class Clazz {
+
+    @RequiresReadLock
+    public static void foo() {
+        String[] items = {""};
+
+        <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">for<caret></warning> (String item : items);
+    }
+}

plugins/devkit/devkit-java-tests/testData/inspections/insertCancellationCheckFix/NoBodyForEachLoop_after.java

@@ -0,0 +1,20 @@
+package inspections.cancellationCheckInLoops;
+
+import com.intellij.openapi.progress.ProgressManager;
+import com.intellij.util.concurrency.annotations.RequiresReadLock;
+
+import static inspections.cancellationCheckInLoops.Foo.doSomething;
+
+
+class Clazz {
+
+    @RequiresReadLock
+    public static void foo() {
+        String[] items = {""};
+
+        for<caret> (String item : items) {
+            ProgressManager.checkCanceled();
+            ;
+        }
+    }
+}

plugins/devkit/devkit-java-tests/testData/inspections/insertCancellationCheckFix/NoBodyForLoop.java

@@ -0,0 +1,14 @@
+package inspections.cancellationCheckInLoops;
+
+import com.intellij.util.concurrency.annotations.RequiresReadLock;
+
+import static inspections.cancellationCheckInLoops.Foo.doSomething;
+
+
+class Clazz {
+
+    @RequiresReadLock
+    public static void foo() {
+        <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">for<caret></warning> (int i = 0; i < 5; i++);
+    }
+}

plugins/devkit/devkit-java-tests/testData/inspections/insertCancellationCheckFix/NoBodyForLoop_after.java

@@ -0,0 +1,18 @@
+package inspections.cancellationCheckInLoops;
+
+import com.intellij.openapi.progress.ProgressManager;
+import com.intellij.util.concurrency.annotations.RequiresReadLock;
+
+import static inspections.cancellationCheckInLoops.Foo.doSomething;
+
+
+class Clazz {
+
+    @RequiresReadLock
+    public static void foo() {
+        for<caret> (int i = 0; i < 5; i++) {
+            ProgressManager.checkCanceled();
+            ;
+        }
+    }
+}

plugins/devkit/devkit-java-tests/testData/inspections/insertCancellationCheckFix/NoBodyWhileLoop.java

@@ -0,0 +1,15 @@
+package inspections.cancellationCheckInLoops;
+
+import com.intellij.util.concurrency.annotations.RequiresReadLock;
+
+import static inspections.cancellationCheckInLoops.Foo.doSomething;
+
+
+class Clazz {
+
+    @RequiresReadLock
+    public static void foo() {
+        int i = 0;
+        <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">while<caret></warning> (i < 5);
+    }
+}

plugins/devkit/devkit-java-tests/testData/inspections/insertCancellationCheckFix/NoBodyWhileLoop_after.java

@@ -0,0 +1,19 @@
+package inspections.cancellationCheckInLoops;
+
+import com.intellij.openapi.progress.ProgressManager;
+import com.intellij.util.concurrency.annotations.RequiresReadLock;
+
+import static inspections.cancellationCheckInLoops.Foo.doSomething;
+
+
+class Clazz {
+
+    @RequiresReadLock
+    public static void foo() {
+        int i = 0;
+        while<caret> (i < 5) {
+            ProgressManager.checkCanceled();
+            ;
+        }
+    }
+}

plugins/devkit/devkit-java-tests/testData/inspections/insertCancellationCheckFix/SingleLineDoWhileLoop.java

@@ -0,0 +1,16 @@
+package inspections.cancellationCheckInLoops;
+
+import com.intellij.util.concurrency.annotations.RequiresReadLock;
+
+import static inspections.cancellationCheckInLoops.Foo.doSomething;
+
+
+class Clazz {
+
+    @RequiresReadLock
+    public static void foo() {
+        int i = 0;
+        <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">do<caret></warning> i++;
+        while (i < 100);
+    }
+}

plugins/devkit/devkit-java-tests/testData/inspections/insertCancellationCheckFix/SingleLineDoWhileLoop_after.java

@@ -0,0 +1,20 @@
+package inspections.cancellationCheckInLoops;
+
+import com.intellij.openapi.progress.ProgressManager;
+import com.intellij.util.concurrency.annotations.RequiresReadLock;
+
+import static inspections.cancellationCheckInLoops.Foo.doSomething;
+
+
+class Clazz {
+
+    @RequiresReadLock
+    public static void foo() {
+        int i = 0;
+        do {
+            ProgressManager.checkCanceled();
+            i++;
+        }
+        while (i < 100);
+    }
+}

plugins/devkit/devkit-java-tests/testData/inspections/insertCancellationCheckFix/SingleLineForEachLoop.java

@@ -0,0 +1,16 @@
+package inspections.cancellationCheckInLoops;
+
+import com.intellij.util.concurrency.annotations.RequiresReadLock;
+
+import static inspections.cancellationCheckInLoops.Foo.doSomething;
+
+
+class Clazz {
+
+    @RequiresReadLock
+    public static void foo() {
+        String[] items = {""};
+
+        <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">for<caret></warning> (String item : items) doSomething();
+    }
+}

plugins/devkit/devkit-java-tests/testData/inspections/insertCancellationCheckFix/SingleLineForEachLoop_after.java

@@ -0,0 +1,20 @@
+package inspections.cancellationCheckInLoops;
+
+import com.intellij.openapi.progress.ProgressManager;
+import com.intellij.util.concurrency.annotations.RequiresReadLock;
+
+import static inspections.cancellationCheckInLoops.Foo.doSomething;
+
+
+class Clazz {
+
+    @RequiresReadLock
+    public static void foo() {
+        String[] items = {""};
+
+        for (String item : items) {
+            ProgressManager.checkCanceled();
+            doSomething();
+        }
+    }
+}

plugins/devkit/devkit-java-tests/testData/inspections/insertCancellationCheckFix/SingleLineForLoop.java

@@ -0,0 +1,14 @@
+package inspections.cancellationCheckInLoops;
+
+import com.intellij.util.concurrency.annotations.RequiresReadLock;
+
+import static inspections.cancellationCheckInLoops.Foo.doSomething;
+
+
+class Clazz {
+
+    @RequiresReadLock
+    public static void foo() {
+        <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">for<caret></warning> (int i = 0; i < 5; i++) doSomething();
+    }
+}

plugins/devkit/devkit-java-tests/testData/inspections/insertCancellationCheckFix/SingleLineForLoop_after.java

@@ -0,0 +1,18 @@
+package inspections.cancellationCheckInLoops;
+
+import com.intellij.openapi.progress.ProgressManager;
+import com.intellij.util.concurrency.annotations.RequiresReadLock;
+
+import static inspections.cancellationCheckInLoops.Foo.doSomething;
+
+
+class Clazz {
+
+    @RequiresReadLock
+    public static void foo() {
+        for (int i = 0; i < 5; i++) {
+            ProgressManager.checkCanceled();
+            doSomething();
+        }
+    }
+}

plugins/devkit/devkit-java-tests/testData/inspections/insertCancellationCheckFix/SingleLineWhileLoop.java

@@ -0,0 +1,15 @@
+package inspections.cancellationCheckInLoops;
+
+import com.intellij.util.concurrency.annotations.RequiresReadLock;
+
+import static inspections.cancellationCheckInLoops.Foo.doSomething;
+
+
+class Clazz {
+
+    @RequiresReadLock
+    public static void foo() {
+        int i = 0;
+        <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">while<caret></warning> (i < 5) i++;
+    }
+}

plugins/devkit/devkit-java-tests/testData/inspections/insertCancellationCheckFix/SingleLineWhileLoop_after.java

@@ -0,0 +1,19 @@
+package inspections.cancellationCheckInLoops;
+
+import com.intellij.openapi.progress.ProgressManager;
+import com.intellij.util.concurrency.annotations.RequiresReadLock;
+
+import static inspections.cancellationCheckInLoops.Foo.doSomething;
+
+
+class Clazz {
+
+    @RequiresReadLock
+    public static void foo() {
+        int i = 0;
+        while (i < 5) {
+            ProgressManager.checkCanceled();
+            i++;
+        }
+    }
+}

plugins/devkit/devkit-java-tests/testSrc/org/jetbrains/idea/devkit/inspections/CancellationCheckInLoopsInspectionTest.kt

@@ -7,23 +7,36 @@ import org.jetbrains.idea.devkit.DevkitJavaTestsUtil
 @TestDataPath("\$CONTENT_ROOT/testData/inspections/cancellationCheckInLoops")
 class CancellationCheckInLoopsInspectionTest : CancellationCheckInLoopsInspectionTestBase() {
 
-  override val fileType
-    get() = "java"
+  override fun getFileExtension(): String = "java"
 
-  override fun getBasePath(): String {
-    return DevkitJavaTestsUtil.TESTDATA_PATH + "inspections/cancellationCheckInLoops"
-  }
+  override fun getBasePath() = DevkitJavaTestsUtil.TESTDATA_PATH + "inspections/cancellationCheckInLoops"
 
   fun testRunCondition() {
     doTest()
   }
 
-  fun testMultipleNestedLoops() {
-    doTest()
-  }
-
   fun testPresentCancellationCheck() {
     doTest()
   }
 
+  fun testForEachLoops() {
+    doTest()
+  }
+
+  fun testForLoops() {
+    doTest()
+  }
+
+  fun testWhileLoops() {
+    doTest()
+  }
+
+  fun testDoWhileLoops() {
+    doTest()
+  }
+
+  fun testNestedLoops() {
+    doTest()
+  }
+
 }

plugins/devkit/devkit-java-tests/testSrc/org/jetbrains/idea/devkit/inspections/quickfix/InsertCancellationCheckFixTest.kt

@@ -0,0 +1,79 @@
+// Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package org.jetbrains.idea.devkit.inspections.quickfix
+
+import com.intellij.testFramework.TestDataPath
+import org.jetbrains.idea.devkit.DevKitBundle
+import org.jetbrains.idea.devkit.DevkitJavaTestsUtil
+import org.jetbrains.idea.devkit.inspections.CancellationCheckInLoopsInspectionTestBase
+
+
+@TestDataPath("\$CONTENT_ROOT/testData/inspections/insertCancellationCheckFix")
+class InsertCancellationCheckFixTest : CancellationCheckInLoopsInspectionTestBase() {
+
+  override fun getFileExtension(): String = "java"
+
+  override fun getBasePath() = DevkitJavaTestsUtil.TESTDATA_PATH + "inspections/insertCancellationCheckFix"
+
+  private val fixName = DevKitBundle.message("inspection.insert.cancellation.check.fix.message")
+
+  fun testBlockDoWhileLoop() {
+    doTest(fixName)
+  }
+
+  fun testBlockForEachLoop() {
+    doTest(fixName)
+  }
+
+  fun testBlockForLoop() {
+    doTest(fixName)
+  }
+
+  fun testBlockWhileLoop() {
+    doTest(fixName)
+  }
+  
+  fun testEmptyDoWhileLoop() {
+    doTest(fixName)
+  }
+
+  fun testEmptyForEachLoop() {
+    doTest(fixName)
+  }
+
+  fun testEmptyForLoop() {
+    doTest(fixName)
+  }
+
+  fun testEmptyWhileLoop() {
+    doTest(fixName)
+  }
+
+  fun testSingleLineDoWhileLoop() {
+    doTest(fixName)
+  }
+
+  fun testSingleLineForEachLoop() {
+    doTest(fixName)
+  }
+
+  fun testSingleLineForLoop() {
+    doTest(fixName)
+  }
+
+  fun testSingleLineWhileLoop() {
+    doTest(fixName)
+  }
+
+  fun testNoBodyForEachLoop() {
+    doTest(fixName)
+  }
+
+  fun testNoBodyForLoop() {
+    doTest(fixName)
+  }
+
+  fun testNoBodyWhileLoop() {
+    doTest(fixName)
+  }
+
+}

plugins/devkit/devkit-kotlin-tests/testData/inspections/cancellationCheckInLoops/DoWhileLoops.kt

@@ -0,0 +1,41 @@
+package inspections.cancellationCheckInLoops
+
+import com.intellij.util.concurrency.annotations.RequiresReadLock
+
+import inspections.cancellationCheckInLoops.Foo.doSomething
+
+
+@RequiresReadLock
+fun main() {
+  var i = 0
+  do {
+    // nested loops with something in between
+    <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">do</warning> {
+      doSomething()
+      <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">do</warning> {
+        doSomething()
+        i++
+      } while (i < 20)
+      i++
+    } while (i < 10)
+
+    // sibling loop
+    <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">do</warning> {
+      doSomething()
+      i++
+    } while (i < 30)
+
+    // nested loops with a block in between
+    <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">do</warning> {
+      if (i < 5) {
+        // empty loop
+        <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">do</warning> {
+        } while (i < 3)
+      }
+    } while (i < 40)
+
+    // single-line loop
+    <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">do</warning> doSomething() while (i < 50)
+    i++
+  } while (i < 100)
+}

plugins/devkit/devkit-kotlin-tests/testData/inspections/cancellationCheckInLoops/ForEachLoops.kt

@@ -25,7 +25,7 @@ fun main() {
     <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">for</warning> (j in 1..10) {
       if (i < 5) {
         // empty loop
-        <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">for</warning> (<warning descr="[NAME_SHADOWING] Name shadowed: j">j</warning> in 1..10) {
+        <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">for</warning> (k in 1..10) {
         }
       }
     }

plugins/devkit/devkit-kotlin-tests/testData/inspections/cancellationCheckInLoops/NestedLoops.kt

@@ -0,0 +1,28 @@
+package inspections.cancellationCheckInLoops
+
+import com.intellij.util.concurrency.annotations.RequiresReadLock
+
+import inspections.cancellationCheckInLoops.Foo.doSomething
+
+
+@RequiresReadLock
+fun main() {
+  var j = 0
+  // nested loops of different kinds
+  for (i in 1..10) {
+    while (j < 5) {
+      <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">do</warning> {
+        doSomething()
+        j++
+      } while (j < 3)
+      j++
+    }
+  }
+
+  // single line nested loops
+  for (i in 1..10) <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">while</warning> (j < 5) {
+    doSomething()
+    j++
+  }
+
+}

plugins/devkit/devkit-kotlin-tests/testData/inspections/cancellationCheckInLoops/PresentCancellationCheck.kt

@@ -10,6 +10,17 @@ fun main() {
   for (i in 1..10) {
     ProgressManager.checkCanceled()
   }
+
+  var i = 0;
+  while (i < 5) {
+    ProgressManager.checkCanceled()
+    i++
+  }
+
+  do {
+    ProgressManager.checkCanceled()
+    i++
+  } while (i < 10)
 }
 
 @RequiresReadLock
@@ -20,7 +31,7 @@ suspend fun foo() {
   }
 
   // wrong cancellation check
-  <warning descr="Cancellation check 'com.intellij.openapi.progress.CoroutinesKt.checkCancelled' should be placed in the first line">for</warning> (i in 1..10) {
+  <warning descr="Cancellation check 'com.intellij.openapi.progress.checkCancelled' should be placed in the first line">for</warning> (i in 1..10) {
     ProgressManager.checkCanceled()
   }
 }

plugins/devkit/devkit-kotlin-tests/testData/inspections/cancellationCheckInLoops/SuspendingContext.kt

@@ -5,11 +5,11 @@ import com.intellij.util.concurrency.annotations.RequiresReadLock
 import inspections.cancellationCheckInLoops.Foo.doSomething
 
 @Suppress("UNUSED_PARAMETER")
-fun withSuspendLambda(l: suspend ()-> Any) { }
+fun withSuspendLambda(l: suspend () -> Any) { }
 
 @RequiresReadLock
 suspend fun mySuspendFun() {
-  <warning descr="Cancellation check 'com.intellij.openapi.progress.CoroutinesKt.checkCancelled' should be placed in the first line">for</warning> (i in 1..10) {
+  <warning descr="Cancellation check 'com.intellij.openapi.progress.checkCancelled' should be placed in the first line">for</warning> (i in 1..10) {
     doSomething()
   }
 }
@@ -17,9 +17,9 @@ suspend fun mySuspendFun() {
 @RequiresReadLock
 fun main() {
   withSuspendLambda {
-    <warning descr="Cancellation check 'com.intellij.openapi.progress.CoroutinesKt.checkCancelled' should be placed in the first line">for</warning> (i in 1..10) {
-    doSomething()
-  }
+    <warning descr="Cancellation check 'com.intellij.openapi.progress.checkCancelled' should be placed in the first line">for</warning> (i in 1..10) {
+      doSomething()
+    }
   }
 }
 

plugins/devkit/devkit-kotlin-tests/testData/inspections/cancellationCheckInLoops/WhileLoops.kt

@@ -0,0 +1,41 @@
+package inspections.cancellationCheckInLoops
+
+import com.intellij.util.concurrency.annotations.RequiresReadLock
+
+import inspections.cancellationCheckInLoops.Foo.doSomething
+
+
+@RequiresReadLock
+fun main() {
+  var i = 0
+  while (i < 100) {
+    // nested loops with something in between
+    <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">while</warning> (i < 10) {
+      doSomething()
+      <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">while</warning> (i < 20) {
+        doSomething()
+        i++
+      }
+      i++
+    }
+
+    // sibling loop
+    <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">while</warning> (i < 30) {
+      doSomething()
+      i++
+    }
+
+    // nested loops with a block in between
+    <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">while</warning> (i < 40) {
+      if (i < 5) {
+        // empty loop
+        <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">while</warning> (i < 3) {
+        }
+      }
+    }
+
+    // single-line loop
+    <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">while</warning> (i < 50) doSomething()
+    i++
+  }
+}

plugins/devkit/devkit-kotlin-tests/testData/inspections/insertCancellationCheckFix/BlockDoWhileLoop.kt

@@ -0,0 +1,16 @@
+package inspections.insertCancellationCheckFix
+
+import com.intellij.util.concurrency.annotations.RequiresReadLock
+
+import inspections.cancellationCheckInLoops.Foo.doSomething
+
+
+@RequiresReadLock
+fun main() {
+  var i = 0
+  <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">do<caret></warning> {
+    // comments
+    doSomething()
+    i++
+  } while (i < 10)
+}

plugins/devkit/devkit-kotlin-tests/testData/inspections/insertCancellationCheckFix/BlockDoWhileLoop_after.kt

@@ -0,0 +1,18 @@
+package inspections.insertCancellationCheckFix
+
+import com.intellij.openapi.progress.ProgressManager
+import com.intellij.util.concurrency.annotations.RequiresReadLock
+
+import inspections.cancellationCheckInLoops.Foo.doSomething
+
+
+@RequiresReadLock
+fun main() {
+  var i = 0
+  do {
+      ProgressManager.checkCanceled()
+      // comments
+      doSomething()
+      i++
+  } while (i < 10)
+}

plugins/devkit/devkit-kotlin-tests/testData/inspections/insertCancellationCheckFix/BlockForEachLoop.kt

@@ -0,0 +1,14 @@
+package inspections.insertCancellationCheckFix
+
+import com.intellij.util.concurrency.annotations.RequiresReadLock
+
+import inspections.cancellationCheckInLoops.Foo.doSomething
+
+
+@RequiresReadLock
+fun main() {
+  <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">for<caret></warning> (i in 1..10) {
+    // comments
+    doSomething()
+  }
+}

plugins/devkit/devkit-kotlin-tests/testData/inspections/insertCancellationCheckFix/BlockForEachLoop_after.kt

@@ -0,0 +1,16 @@
+package inspections.insertCancellationCheckFix
+
+import com.intellij.openapi.progress.ProgressManager
+import com.intellij.util.concurrency.annotations.RequiresReadLock
+
+import inspections.cancellationCheckInLoops.Foo.doSomething
+
+
+@RequiresReadLock
+fun main() {
+  for (i in 1..10) {
+      ProgressManager.checkCanceled()
+      // comments
+      doSomething()
+  }
+}

plugins/devkit/devkit-kotlin-tests/testData/inspections/insertCancellationCheckFix/BlockWhileLoop.kt

@@ -0,0 +1,16 @@
+package inspections.insertCancellationCheckFix
+
+import com.intellij.util.concurrency.annotations.RequiresReadLock
+
+import inspections.cancellationCheckInLoops.Foo.doSomething
+
+
+@RequiresReadLock
+fun main() {
+  var i = 0
+  <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">while<caret></warning> (i < 10) {
+    // comments
+    doSomething()
+    i++
+  }
+}

plugins/devkit/devkit-kotlin-tests/testData/inspections/insertCancellationCheckFix/BlockWhileLoop_after.kt

@@ -0,0 +1,18 @@
+package inspections.insertCancellationCheckFix
+
+import com.intellij.openapi.progress.ProgressManager
+import com.intellij.util.concurrency.annotations.RequiresReadLock
+
+import inspections.cancellationCheckInLoops.Foo.doSomething
+
+
+@RequiresReadLock
+fun main() {
+  var i = 0
+  while (i < 10) {
+      ProgressManager.checkCanceled()
+      // comments
+      doSomething()
+      i++
+  }
+}

plugins/devkit/devkit-kotlin-tests/testData/inspections/insertCancellationCheckFix/EmptyDoWhileLoop.kt

@@ -0,0 +1,12 @@
+package inspections.insertCancellationCheckFix
+
+import com.intellij.util.concurrency.annotations.RequiresReadLock
+
+
+@RequiresReadLock
+fun main() {
+  var i = 0
+  <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">do<caret></warning> {
+    // comments
+  } while (i < 10)
+}

plugins/devkit/devkit-kotlin-tests/testData/inspections/insertCancellationCheckFix/EmptyDoWhileLoop_after.kt

@@ -0,0 +1,14 @@
+package inspections.insertCancellationCheckFix
+
+import com.intellij.openapi.progress.ProgressManager
+import com.intellij.util.concurrency.annotations.RequiresReadLock
+
+
+@RequiresReadLock
+fun main() {
+  var i = 0
+  do {
+      ProgressManager.checkCanceled()
+      // comments
+  } while (i < 10)
+}

plugins/devkit/devkit-kotlin-tests/testData/inspections/insertCancellationCheckFix/EmptyForEachLoop.kt

@@ -0,0 +1,11 @@
+package inspections.insertCancellationCheckFix
+
+import com.intellij.util.concurrency.annotations.RequiresReadLock
+
+
+@RequiresReadLock
+fun main() {
+  <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">for<caret></warning> (i in 1..10) {
+    // comments
+  }
+}

plugins/devkit/devkit-kotlin-tests/testData/inspections/insertCancellationCheckFix/EmptyForEachLoop_after.kt

@@ -0,0 +1,13 @@
+package inspections.insertCancellationCheckFix
+
+import com.intellij.openapi.progress.ProgressManager
+import com.intellij.util.concurrency.annotations.RequiresReadLock
+
+
+@RequiresReadLock
+fun main() {
+  for (i in 1..10) {
+      ProgressManager.checkCanceled()
+      // comments
+  }
+}

plugins/devkit/devkit-kotlin-tests/testData/inspections/insertCancellationCheckFix/EmptyWhileLoop.kt

@@ -0,0 +1,12 @@
+package inspections.insertCancellationCheckFix
+
+import com.intellij.util.concurrency.annotations.RequiresReadLock
+
+
+@RequiresReadLock
+fun main() {
+  var i = 0
+  <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">while<caret></warning> (i < 10) {
+    // comments
+  }
+}

plugins/devkit/devkit-kotlin-tests/testData/inspections/insertCancellationCheckFix/EmptyWhileLoop_after.kt

@@ -0,0 +1,14 @@
+package inspections.insertCancellationCheckFix
+
+import com.intellij.openapi.progress.ProgressManager
+import com.intellij.util.concurrency.annotations.RequiresReadLock
+
+
+@RequiresReadLock
+fun main() {
+  var i = 0
+  while (i < 10) {
+      ProgressManager.checkCanceled()
+      // comments
+  }
+}

plugins/devkit/devkit-kotlin-tests/testData/inspections/insertCancellationCheckFix/NoBodyForEachLoop.kt

@@ -0,0 +1,9 @@
+package inspections.insertCancellationCheckFix
+
+import com.intellij.util.concurrency.annotations.RequiresReadLock
+
+
+@RequiresReadLock
+fun main() {
+  <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">for<caret></warning> (i in 1..10); // comments
+}

plugins/devkit/devkit-kotlin-tests/testData/inspections/insertCancellationCheckFix/NoBodyForEachLoop_after.kt

@@ -0,0 +1,12 @@
+package inspections.insertCancellationCheckFix
+
+import com.intellij.openapi.progress.ProgressManager
+import com.intellij.util.concurrency.annotations.RequiresReadLock
+
+
+@RequiresReadLock
+fun main() {
+  for (i in 1..10) {
+      ProgressManager.checkCanceled()
+  } // comments
+}

plugins/devkit/devkit-kotlin-tests/testData/inspections/insertCancellationCheckFix/NoBodyWhileLoop.kt

@@ -0,0 +1,10 @@
+package inspections.insertCancellationCheckFix
+
+import com.intellij.util.concurrency.annotations.RequiresReadLock
+
+
+@RequiresReadLock
+fun main() {
+  var i = 0
+  <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">while<caret></warning> (i < 10);
+}

plugins/devkit/devkit-kotlin-tests/testData/inspections/insertCancellationCheckFix/NoBodyWhileLoop_after.kt

@@ -0,0 +1,13 @@
+package inspections.insertCancellationCheckFix
+
+import com.intellij.openapi.progress.ProgressManager
+import com.intellij.util.concurrency.annotations.RequiresReadLock
+
+
+@RequiresReadLock
+fun main() {
+  var i = 0
+  while (i < 10) {
+      ProgressManager.checkCanceled()
+  }
+}

plugins/devkit/devkit-kotlin-tests/testData/inspections/insertCancellationCheckFix/SingleLineDoWhileLoop.kt

@@ -0,0 +1,10 @@
+package inspections.insertCancellationCheckFix
+
+import com.intellij.util.concurrency.annotations.RequiresReadLock
+
+
+@RequiresReadLock
+fun main() {
+  var i = 0
+  <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">do<caret></warning> i++ while (i < 10) // comments
+}

plugins/devkit/devkit-kotlin-tests/testData/inspections/insertCancellationCheckFix/SingleLineDoWhileLoop_after.kt

@@ -0,0 +1,14 @@
+package inspections.insertCancellationCheckFix
+
+import com.intellij.openapi.progress.ProgressManager
+import com.intellij.util.concurrency.annotations.RequiresReadLock
+
+
+@RequiresReadLock
+fun main() {
+  var i = 0
+  do {
+      ProgressManager.checkCanceled()
+      i++
+  } while (i < 10) // comments
+}

plugins/devkit/devkit-kotlin-tests/testData/inspections/insertCancellationCheckFix/SingleLineForEachLoop.kt

@@ -0,0 +1,11 @@
+package inspections.insertCancellationCheckFix
+
+import com.intellij.util.concurrency.annotations.RequiresReadLock
+
+import inspections.cancellationCheckInLoops.Foo.doSomething
+
+
+@RequiresReadLock
+fun main() {
+  <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">for<caret></warning> (i in 1..10) doSomething() // comments
+}

plugins/devkit/devkit-kotlin-tests/testData/inspections/insertCancellationCheckFix/SingleLineForEachLoop_after.kt

@@ -0,0 +1,15 @@
+package inspections.insertCancellationCheckFix
+
+import com.intellij.openapi.progress.ProgressManager
+import com.intellij.util.concurrency.annotations.RequiresReadLock
+
+import inspections.cancellationCheckInLoops.Foo.doSomething
+
+
+@RequiresReadLock
+fun main() {
+  for (i in 1..10) {
+      ProgressManager.checkCanceled()
+      doSomething() // comments
+  }
+}

plugins/devkit/devkit-kotlin-tests/testData/inspections/insertCancellationCheckFix/SingleLineWhileLoop.kt

@@ -0,0 +1,9 @@
+package inspections.insertCancellationCheckFix
+
+import com.intellij.util.concurrency.annotations.RequiresReadLock
+
+@RequiresReadLock
+fun main() {
+  var i = 0
+  <warning descr="Cancellation check 'com.intellij.openapi.progress.ProgressManager.checkCanceled' should be placed in the first line">while<caret></warning> (i < 10) i++ // comments
+}

plugins/devkit/devkit-kotlin-tests/testData/inspections/insertCancellationCheckFix/SingleLineWhileLoop_after.kt

@@ -0,0 +1,13 @@
+package inspections.insertCancellationCheckFix
+
+import com.intellij.openapi.progress.ProgressManager
+import com.intellij.util.concurrency.annotations.RequiresReadLock
+
+@RequiresReadLock
+fun main() {
+  var i = 0
+  while (i < 10) {
+      ProgressManager.checkCanceled()
+      i++ // comments
+  }
+}

plugins/devkit/devkit-kotlin-tests/testData/inspections/insertCancellationCheckFix/SuspendingDoWhileLoop.kt

@@ -0,0 +1,16 @@
+package inspections.insertCancellationCheckFix
+
+import com.intellij.util.concurrency.annotations.RequiresReadLock
+
+import inspections.cancellationCheckInLoops.Foo.doSomething
+
+
+@RequiresReadLock
+suspend fun main() {
+  var i = 0
+  <warning descr="Cancellation check 'com.intellij.openapi.progress.checkCancelled' should be placed in the first line">do<caret></warning> {
+    // comments
+    doSomething()
+    i++
+  } while (i < 10)
+}

plugins/devkit/devkit-kotlin-tests/testData/inspections/insertCancellationCheckFix/SuspendingDoWhileLoop_after.kt

@@ -0,0 +1,17 @@
+package inspections.insertCancellationCheckFix
+
+import com.intellij.util.concurrency.annotations.RequiresReadLock
+
+import inspections.cancellationCheckInLoops.Foo.doSomething
+
+
+@RequiresReadLock
+suspend fun main() {
+  var i = 0
+  do {
+      com.intellij.openapi.progress.checkCancelled()
+      // comments
+      doSomething()
+      i++
+  } while (i < 10)
+}

plugins/devkit/devkit-kotlin-tests/testData/inspections/insertCancellationCheckFix/SuspendingForEachLoop.kt

@@ -0,0 +1,14 @@
+package inspections.insertCancellationCheckFix
+
+import com.intellij.util.concurrency.annotations.RequiresReadLock
+
+import inspections.cancellationCheckInLoops.Foo.doSomething
+
+
+@RequiresReadLock
+suspend fun main() {
+  <warning descr="Cancellation check 'com.intellij.openapi.progress.checkCancelled' should be placed in the first line">for<caret></warning> (i in 1..10) {
+    // comments
+    doSomething()
+  }
+}

plugins/devkit/devkit-kotlin-tests/testData/inspections/insertCancellationCheckFix/SuspendingForEachLoop_after.kt

@@ -0,0 +1,15 @@
+package inspections.insertCancellationCheckFix
+
+import com.intellij.util.concurrency.annotations.RequiresReadLock
+
+import inspections.cancellationCheckInLoops.Foo.doSomething
+
+
+@RequiresReadLock
+suspend fun main() {
+  for (i in 1..10) {
+      com.intellij.openapi.progress.checkCancelled()
+      // comments
+      doSomething()
+  }
+}

plugins/devkit/devkit-kotlin-tests/testData/inspections/insertCancellationCheckFix/SuspendingWhileLoop.kt

@@ -0,0 +1,16 @@
+package inspections.insertCancellationCheckFix
+
+import com.intellij.util.concurrency.annotations.RequiresReadLock
+
+import inspections.cancellationCheckInLoops.Foo.doSomething
+
+
+@RequiresReadLock
+suspend fun main() {
+  var i = 0
+  <warning descr="Cancellation check 'com.intellij.openapi.progress.checkCancelled' should be placed in the first line">while<caret></warning> (i < 10) {
+    // comments
+    doSomething()
+    i++
+  }
+}

plugins/devkit/devkit-kotlin-tests/testData/inspections/insertCancellationCheckFix/SuspendingWhileLoop_after.kt

@@ -0,0 +1,17 @@
+package inspections.insertCancellationCheckFix
+
+import com.intellij.util.concurrency.annotations.RequiresReadLock
+
+import inspections.cancellationCheckInLoops.Foo.doSomething
+
+
+@RequiresReadLock
+suspend fun main() {
+  var i = 0
+  while (i < 10) {
+      com.intellij.openapi.progress.checkCancelled()
+      // comments
+      doSomething()
+      i++
+  }
+}

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

@@ -8,32 +8,45 @@ import org.jetbrains.idea.devkit.kotlin.DevkitKtTestsUtil
 
 @TestDataPath("\$CONTENT_ROOT/testData/inspections/cancellationCheckInLoops")
 class KtCancellationCheckInLoopsInspectionTest : CancellationCheckInLoopsInspectionTestBase() {
-  override val fileType
-    get() = "kt"
 
-  override fun getBasePath(): String {
-    return DevkitKtTestsUtil.TESTDATA_PATH + "inspections/cancellationCheckInLoops"
-  }
+  override fun getFileExtension(): String = "kt"
+
+  override fun getBasePath() = DevkitKtTestsUtil.TESTDATA_PATH + "inspections/cancellationCheckInLoops"
 
   override fun setUp() {
     super.setUp()
-    myFixture.addFileToProject("Coroutines.kt", """
+    myFixture.addFileToProject(
+      "Coroutines.kt",
+      //language=kotlin
+      """
       package com.intellij.openapi.progress
 
       suspend fun checkCancelled() { }
-    """.trimIndent())
+      """.trimIndent())
   }
 
   fun testSuspendingContext() {
     doTest()
   }
 
-  fun testMultipleNestedLoops() {
-    doTest()
-  }
-
   fun testPresentCancellationCheck() {
     doTest()
   }
 
+  fun testForEachLoops() {
+    doTest()
+  }
+
+  fun testWhileLoops() {
+    doTest()
+  }
+
+  fun testDoWhileLoops() {
+    doTest()
+  }
+
+  fun testNestedLoops() {
+    doTest()
+  }
+
 }

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

@@ -0,0 +1,75 @@
+// Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package org.jetbrains.idea.devkit.kotlin.inspections.quickfix
+
+import com.intellij.testFramework.TestDataPath
+import org.jetbrains.idea.devkit.DevKitBundle
+import org.jetbrains.idea.devkit.inspections.CancellationCheckInLoopsInspectionTestBase
+import org.jetbrains.idea.devkit.kotlin.DevkitKtTestsUtil
+
+
+@TestDataPath("\$CONTENT_ROOT/testData/inspections/insertCancellationCheckFix")
+class KtInsertCancellationCheckFixTest : CancellationCheckInLoopsInspectionTestBase() {
+
+  override fun getFileExtension(): String = "kt"
+
+  override fun getBasePath() = DevkitKtTestsUtil.TESTDATA_PATH + "inspections/insertCancellationCheckFix"
+
+  private val fixName = DevKitBundle.message("inspection.insert.cancellation.check.fix.message")
+  
+  fun testBlockDoWhileLoop() {
+    doTest(fixName)
+  }
+
+  fun testBlockForEachLoop() {
+    doTest(fixName)
+  }
+
+  fun testBlockWhileLoop() {
+    doTest(fixName)
+  }
+
+  fun testEmptyDoWhileLoop() {
+    doTest(fixName)
+  }
+
+  fun testEmptyForEachLoop() {
+    doTest(fixName)
+  }
+
+  fun testEmptyWhileLoop() {
+    doTest(fixName)
+  }
+
+  fun testSingleLineDoWhileLoop() {
+    doTest(fixName)
+  }
+
+  fun testSingleLineForEachLoop() {
+    doTest(fixName)
+  }
+
+  fun testSingleLineWhileLoop() {
+    doTest(fixName)
+  }
+
+  fun testNoBodyForEachLoop() {
+    doTest(fixName)
+  }
+
+  fun testNoBodyWhileLoop() {
+    doTest(fixName)
+  }
+
+  fun testSuspendingDoWhileLoop() {
+    doTest(fixName)
+  }
+
+  fun testSuspendingForEachLoop() {
+    doTest(fixName)
+  }
+
+  fun testSuspendingWhileLoop() {
+    doTest(fixName)
+  }
+
+}

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

@@ -1,11 +1,9 @@
 // Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
 package org.jetbrains.idea.devkit.inspections
 
-import com.intellij.testFramework.fixtures.LightJavaCodeInsightFixtureTestCase
+import org.jetbrains.idea.devkit.inspections.quickfix.LightDevKitInspectionFixTestBase
 
-abstract class CancellationCheckInLoopsInspectionTestBase : LightJavaCodeInsightFixtureTestCase() {
-
-  protected abstract val fileType: String
+abstract class CancellationCheckInLoopsInspectionTestBase : LightDevKitInspectionFixTestBase() {
 
   override fun setUp() {
     super.setUp()
@@ -19,7 +17,7 @@ abstract class CancellationCheckInLoopsInspectionTestBase : LightJavaCodeInsight
     myFixture.addClass("""
       package inspections.cancellationCheckInLoops;
       
-      public class Foo {
+      public final class Foo {
         public static void doSomething() { }
       }
     """.trimIndent()
@@ -34,11 +32,7 @@ abstract class CancellationCheckInLoopsInspectionTestBase : LightJavaCodeInsight
     """.trimIndent()
     )
 
-    myFixture.enableInspections(CancellationCheckInLoopsInspection::class.java)
-  }
-
-  protected open fun doTest() {
-    myFixture.testHighlighting("${getTestName(false)}.$fileType")
+    myFixture.enableInspections(CancellationCheckInLoopsInspection())
   }
 
 }

plugins/devkit/intellij.kotlin.devkit/intellij.kotlin.devkit.iml

@@ -17,5 +17,7 @@
     <orderEntry type="module" module-name="intellij.java.psi" />
     <orderEntry type="module" module-name="intellij.xml.dom" />
     <orderEntry type="module" module-name="intellij.xml.psi" />
+    <orderEntry type="module" module-name="kotlin.code-insight.intentions-shared" />
+    <orderEntry type="module" module-name="kotlin.code-insight.utils" />
   </component>
 </module>

plugins/devkit/intellij.kotlin.devkit/resources/intellij.kotlin.devkit.xml

@@ -41,7 +41,10 @@
   </extensions>
 
   <extensions defaultExtensionNs="DevKit.lang">
-    <cancellationCheckProvider language="kotlin" implementationClass="org.jetbrains.idea.devkit.kotlin.inspections.KotlinCancellationCheckProvider"/>
+    <cancellationCheckProvider language="kotlin"
+      implementationClass="org.jetbrains.idea.devkit.kotlin.inspections.KtCancellationCheckProvider"/>
+    <cancellationCheckInLoopsFixProvider language="kotlin"
+      implementationClass="org.jetbrains.idea.devkit.kotlin.inspections.KtCancellationCheckInLoopsFixProvider"/>
     <visitorProviderForRBCInspection language="kotlin"
       implementationClass="org.jetbrains.idea.devkit.kotlin.inspections.KtCallingFunctionShouldBeRequiresBlockingContextVisitorProvider"/>
     <extensionClassShouldNotBePublicProvider language="kotlin"

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

@@ -1,50 +0,0 @@
-// Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
-package org.jetbrains.idea.devkit.kotlin.inspections
-
-import com.intellij.psi.PsiElement
-import org.jetbrains.idea.devkit.inspections.CancellationCheckProvider
-import org.jetbrains.kotlin.builtins.isSuspendFunctionType
-import org.jetbrains.kotlin.idea.caches.resolve.resolveToCall
-import org.jetbrains.kotlin.lexer.KtTokens
-import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.psi.psiUtil.getParentOfType
-import org.jetbrains.kotlin.psi.psiUtil.getParentOfTypes
-import org.jetbrains.kotlin.psi.psiUtil.getStrictParentOfType
-import org.jetbrains.kotlin.resolve.calls.util.getParameterForArgument
-import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
-
-
-class KotlinCancellationCheckProvider : CancellationCheckProvider {
-
-  override fun findContext(element: PsiElement): CancellationCheckProvider.Context {
-    val containingLambda = element.getParentOfType<KtLambdaExpression>(false)
-
-    // if lambda is null, check the containing method
-    if (containingLambda == null) {
-      val containingMethod = element.getParentOfType<KtNamedFunction>(false)
-      return if (containingMethod.hasSuspendModifier()) CancellationCheckProvider.Context.SUSPENDING else CancellationCheckProvider.Context.BLOCKING
-    }
-
-    // otherwise, check containing argument (whether the corresponding parameter has `suspend` modifier)
-    val containingArgument = containingLambda.getParentOfType<KtValueArgument>(true, KtCallableDeclaration::class.java)
-    if (containingArgument != null) {
-      val callExpression = containingArgument.getStrictParentOfType<KtCallExpression>() ?: return CancellationCheckProvider.Context.BLOCKING
-      val resolvedCall = callExpression.resolveToCall(BodyResolveMode.PARTIAL) ?: return CancellationCheckProvider.Context.BLOCKING
-
-      val parameterForArgument = resolvedCall.getParameterForArgument(containingArgument) ?: return CancellationCheckProvider.Context.BLOCKING
-      val type = parameterForArgument.returnType ?: return CancellationCheckProvider.Context.BLOCKING
-      return if (type.isSuspendFunctionType) CancellationCheckProvider.Context.SUSPENDING else CancellationCheckProvider.Context.BLOCKING
-    }
-
-    // otherwise, check if it's a property or a function
-    val containingPropertyOrFunction: KtCallableDeclaration? =
-      containingLambda.getParentOfTypes(true, KtProperty::class.java, KtNamedFunction::class.java)
-    if (containingPropertyOrFunction?.typeReference.hasSuspendModifier()) return CancellationCheckProvider.Context.SUSPENDING
-    return if (containingPropertyOrFunction.hasSuspendModifier()) CancellationCheckProvider.Context.SUSPENDING else CancellationCheckProvider.Context.BLOCKING
-  }
-
-  private fun KtModifierListOwner?.hasSuspendModifier(): Boolean {
-    return this?.hasModifier(KtTokens.SUSPEND_KEYWORD) == true
-  }
-
-}

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

@@ -0,0 +1,88 @@
+// Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package org.jetbrains.idea.devkit.kotlin.inspections
+
+import com.intellij.codeInspection.LocalQuickFix
+import com.intellij.codeInspection.LocalQuickFixOnPsiElement
+import com.intellij.openapi.project.Project
+import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiFile
+import com.intellij.psi.codeStyle.CodeStyleManager
+import com.intellij.psi.util.PsiTreeUtil
+import org.jetbrains.idea.devkit.DevKitBundle
+import org.jetbrains.idea.devkit.inspections.quickfix.CancellationCheckInLoopsFixProvider
+import org.jetbrains.kotlin.KtNodeTypes
+import org.jetbrains.kotlin.idea.base.codeInsight.ShortenReferencesFacility
+import org.jetbrains.kotlin.idea.codeInsight.intentions.shared.AddBracesIntention
+import org.jetbrains.kotlin.idea.codeinsight.utils.isRedundantSemicolon
+import org.jetbrains.kotlin.lexer.KtTokens
+import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.KtPsiUtil.findChildByType
+import org.jetbrains.kotlin.utils.addToStdlib.safeAs
+
+
+internal class KtCancellationCheckInLoopsFixProvider : CancellationCheckInLoopsFixProvider {
+
+  override fun getFixes(loopKeyword: PsiElement, cancellationCheckFqn: String): List<LocalQuickFix> {
+    return listOf(KtInsertCancellationCheckFix(cancellationCheckFqn, loopKeyword))
+  }
+
+}
+
+
+internal class KtInsertCancellationCheckFix(
+  private val cancellationCheckCallFqn: String,
+  loopKeyword: PsiElement,
+) : LocalQuickFixOnPsiElement(loopKeyword) {
+
+  override fun getFamilyName(): String = DevKitBundle.message("inspection.insert.cancellation.check.fix.message")
+
+  override fun getText(): String = familyName
+
+  override fun isAvailable(project: Project, file: PsiFile, startElement: PsiElement, endElement: PsiElement): Boolean {
+    return PsiTreeUtil.getParentOfType(startElement, KtLoopExpression::class.java) != null
+  }
+
+  override fun invoke(project: Project, file: PsiFile, startElement: PsiElement, endElement: PsiElement) {
+    val loopStatement = PsiTreeUtil.getParentOfType(startElement, KtLoopExpression::class.java) ?: return
+
+    val factory = KtPsiFactory(project)
+    val cancellationCheckExpression = factory.createExpression("${cancellationCheckCallFqn}()")
+
+    val bodyBlock = loopStatement.getOrCreateBodyBlock(project)
+
+    bodyBlock?.addExpressionToFirstLine(cancellationCheckExpression)?.let {
+      ShortenReferencesFacility.getInstance().shorten(it)
+    }
+  }
+
+  private fun KtLoopExpression.getOrCreateBodyBlock(project: Project): KtBlockExpression? {
+    val factory = KtPsiFactory(project)
+
+    return when (val loopBody = body) {
+      is KtBlockExpression -> loopBody
+      is KtExpression -> {
+        AddBracesIntention.addBraces(this, loopBody)
+        body as KtBlockExpression
+      }
+      else -> {
+        val containerNode = findChildByType(this, KtNodeTypes.BODY) ?: return null
+        containerNode.add(factory.createEmptyBody())
+        deleteRedundantSemicolon(this)
+        body as KtBlockExpression
+      }
+    }
+  }
+
+  private fun deleteRedundantSemicolon(loop: KtLoopExpression) {
+    val nextSibling = loop.nextSibling
+    if (nextSibling.node.elementType == KtTokens.SEMICOLON && isRedundantSemicolon(nextSibling)) {
+      nextSibling.delete()
+    }
+  }
+
+  private fun KtBlockExpression.addExpressionToFirstLine(expression: KtExpression): KtExpression? {
+    CodeStyleManager.getInstance(project).reformat(this)
+    return addAfter(expression, lBrace).safeAs<KtExpression>()
+  }
+
+}

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

@@ -0,0 +1,82 @@
+// Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package org.jetbrains.idea.devkit.kotlin.inspections
+
+import com.intellij.psi.PsiElement
+import org.jetbrains.idea.devkit.inspections.CancellationCheckProvider
+import org.jetbrains.kotlin.analysis.api.analyze
+import org.jetbrains.kotlin.analysis.api.calls.singleFunctionCallOrNull
+import org.jetbrains.kotlin.analysis.api.calls.symbol
+import org.jetbrains.kotlin.builtins.isSuspendFunctionType
+import org.jetbrains.kotlin.idea.caches.resolve.resolveToCall
+import org.jetbrains.kotlin.lexer.KtTokens
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.getChildOfType
+import org.jetbrains.kotlin.psi.psiUtil.getParentOfType
+import org.jetbrains.kotlin.psi.psiUtil.getParentOfTypes
+import org.jetbrains.kotlin.psi.psiUtil.getStrictParentOfType
+import org.jetbrains.kotlin.resolve.calls.util.getParameterForArgument
+import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
+
+
+private const val PROGRESS_MANAGER_CHECKED_CANCELED = "com.intellij.openapi.progress.ProgressManager.checkCanceled"
+private const val COROUTINE_CHECK_CANCELLED = "com.intellij.openapi.progress.checkCancelled"
+
+class KtCancellationCheckProvider : CancellationCheckProvider {
+
+  enum class Context {
+    BLOCKING, SUSPENDING
+  }
+
+  override fun findCancellationCheckCall(element: PsiElement): String {
+    return when (findContext(element)) {
+      Context.BLOCKING -> PROGRESS_MANAGER_CHECKED_CANCELED
+      Context.SUSPENDING -> COROUTINE_CHECK_CANCELLED
+    }
+  }
+
+  override fun isCancellationCheckCall(element: PsiElement?, cancellationCheckFqn: String): Boolean {
+    val callExpression = when (element) {
+      is KtCallExpression -> element
+      is KtDotQualifiedExpression -> element.getChildOfType<KtCallExpression>() ?: return false
+      else -> return false
+    }
+
+    analyze(callExpression) {
+      val functionCalledSymbol = callExpression.resolveCall().singleFunctionCallOrNull()?.symbol ?: return false
+      return functionCalledSymbol.callableIdIfNonLocal?.asSingleFqName() == FqName(cancellationCheckFqn)
+    }
+  }
+
+  private fun findContext(element: PsiElement): Context {
+    val containingLambda = element.getParentOfType<KtLambdaExpression>(false)
+
+    // if lambda is null, check the containing method
+    if (containingLambda == null) {
+      val containingMethod = element.getParentOfType<KtNamedFunction>(false)
+      return if (containingMethod.hasSuspendModifier()) Context.SUSPENDING else Context.BLOCKING
+    }
+
+    // otherwise, check containing argument (whether the corresponding parameter has `suspend` modifier)
+    val containingArgument = containingLambda.getParentOfType<KtValueArgument>(true, KtCallableDeclaration::class.java)
+    if (containingArgument != null) {
+      val callExpression = containingArgument.getStrictParentOfType<KtCallExpression>() ?: return Context.BLOCKING
+      val resolvedCall = callExpression.resolveToCall(BodyResolveMode.PARTIAL) ?: return Context.BLOCKING
+
+      val parameterForArgument = resolvedCall.getParameterForArgument(containingArgument) ?: return Context.BLOCKING
+      val type = parameterForArgument.returnType ?: return Context.BLOCKING
+      return if (type.isSuspendFunctionType) Context.SUSPENDING else Context.BLOCKING
+    }
+
+    // otherwise, check if it's a property or a function
+    val containingPropertyOrFunction: KtCallableDeclaration? = containingLambda.getParentOfTypes(true, KtProperty::class.java,
+                                                                                                 KtNamedFunction::class.java)
+    if (containingPropertyOrFunction?.typeReference.hasSuspendModifier()) return Context.SUSPENDING
+    return if (containingPropertyOrFunction.hasSuspendModifier()) Context.SUSPENDING else Context.BLOCKING
+  }
+
+  private fun KtModifierListOwner?.hasSuspendModifier(): Boolean {
+    return this?.hasModifier(KtTokens.SUSPEND_KEYWORD) == true
+  }
+
+}

plugins/kotlin/code-insight/intentions-shared/src/org/jetbrains/kotlin/idea/codeInsight/intentions/shared/AddBracesIntention.kt

@@ -3,8 +3,10 @@
 package org.jetbrains.kotlin.idea.codeInsight.intentions.shared
 
 import com.intellij.openapi.editor.Editor
+import com.intellij.openapi.util.IntellijInternalApi
 import com.intellij.psi.PsiComment
 import com.intellij.psi.PsiWhiteSpace
+import org.jetbrains.annotations.ApiStatus.Internal
 import org.jetbrains.kotlin.idea.base.psi.getLineNumber
 import org.jetbrains.kotlin.idea.base.resources.KotlinBundle
 import org.jetbrains.kotlin.idea.codeinsight.api.classic.intentions.SelfTargetingIntention
@@ -14,7 +16,9 @@ import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.*
 
-internal class AddBracesIntention : SelfTargetingIntention<KtElement>(KtElement::class.java, KotlinBundle.lazyMessage("add.braces")) {
+@Internal
+@IntellijInternalApi
+class AddBracesIntention : SelfTargetingIntention<KtElement>(KtElement::class.java, KotlinBundle.lazyMessage("add.braces")) {
     override fun isApplicableTo(element: KtElement, caretOffset: Int): Boolean {
         val expression = element.getTargetExpression(caretOffset) ?: return false
         if (expression is KtBlockExpression) return false