Guard resolve calls with runWithCancellationCheck

Relates to #KT-35135

idea/idea-analysis/src/org/jetbrains/kotlin/idea/caches/resolve/ModuleResolutionFacadeImpl.kt

@@ -27,6 +27,7 @@ import org.jetbrains.kotlin.descriptors.ModuleDescriptor
 import org.jetbrains.kotlin.idea.caches.project.IdeaModuleInfo
 import org.jetbrains.kotlin.idea.project.ResolveElementCache
 import org.jetbrains.kotlin.idea.resolve.ResolutionFacade
+import org.jetbrains.kotlin.idea.util.application.runWithCancellationCheck
 import org.jetbrains.kotlin.psi.KtDeclaration
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtPsiUtil
@@ -58,28 +59,33 @@ internal class ModuleResolutionFacadeImpl(
 
         if (elements.isEmpty()) return BindingContext.EMPTY
         val resolveElementCache = getFrontendService(elements.first(), ResolveElementCache::class.java)
-        return resolveElementCache.resolveToElements(elements, bodyResolveMode)
+        return runWithCancellationCheck {
+            resolveElementCache.resolveToElements(elements, bodyResolveMode)
+        }
     }
 
     override fun analyzeWithAllCompilerChecks(elements: Collection<KtElement>): AnalysisResult {
         ResolveInDispatchThreadManager.assertNoResolveInDispatchThread()
 
-        return projectFacade.getAnalysisResultsForElements(elements)
-    }
-
-    override fun resolveToDescriptor(declaration: KtDeclaration, bodyResolveMode: BodyResolveMode): DeclarationDescriptor {
-        return if (KtPsiUtil.isLocal(declaration)) {
-            val bindingContext = analyze(declaration, bodyResolveMode)
-            bindingContext[BindingContext.DECLARATION_TO_DESCRIPTOR, declaration]
-                ?: getFrontendService(moduleInfo, AbsentDescriptorHandler::class.java).diagnoseDescriptorNotFound(declaration)
-        } else {
-            ResolveInDispatchThreadManager.assertNoResolveInDispatchThread()
-
-            val resolveSession = projectFacade.resolverForElement(declaration).componentProvider.get<ResolveSession>()
-            resolveSession.resolveToDescriptor(declaration)
+        return runWithCancellationCheck {
+            projectFacade.getAnalysisResultsForElements(elements)
         }
     }
 
+    override fun resolveToDescriptor(declaration: KtDeclaration, bodyResolveMode: BodyResolveMode): DeclarationDescriptor =
+        runWithCancellationCheck {
+            if (KtPsiUtil.isLocal(declaration)) {
+                val bindingContext = analyze(declaration, bodyResolveMode)
+                bindingContext[BindingContext.DECLARATION_TO_DESCRIPTOR, declaration]
+                    ?: getFrontendService(moduleInfo, AbsentDescriptorHandler::class.java).diagnoseDescriptorNotFound(declaration)
+            } else {
+                ResolveInDispatchThreadManager.assertNoResolveInDispatchThread()
+
+                val resolveSession = projectFacade.resolverForElement(declaration).componentProvider.get<ResolveSession>()
+                resolveSession.resolveToDescriptor(declaration)
+            }
+        }
+
     override fun <T : Any> getFrontendService(serviceClass: Class<T>): T = getFrontendService(moduleInfo, serviceClass)
 
     override fun <T : Any> getIdeService(serviceClass: Class<T>): T {

idea/idea-analysis/src/org/jetbrains/kotlin/idea/references/KtReference.kt

@@ -38,6 +38,7 @@ import org.jetbrains.kotlin.idea.codeInsight.DescriptorToSourceUtilsIde
 import org.jetbrains.kotlin.idea.fir.FirResolution
 import org.jetbrains.kotlin.idea.fir.firResolveState
 import org.jetbrains.kotlin.idea.fir.getOrBuildFir
+import org.jetbrains.kotlin.idea.util.application.runWithCancellationCheck
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.KtDotQualifiedExpression
@@ -282,8 +283,10 @@ abstract class AbstractKtReference<T : KtElement>(element: T) : PsiPolyVariantRe
             }
 
             override fun resolve(ref: AbstractKtReference<KtElement>, incompleteCode: Boolean): Array<ResolveResult> {
-                val resolveToPsiElements = resolveToPsiElements(ref)
-                return resolveToPsiElements.map { KotlinResolveResult(it) }.toTypedArray()
+                return runWithCancellationCheck {
+                    val resolveToPsiElements = resolveToPsiElements(ref)
+                    resolveToPsiElements.map { KotlinResolveResult(it) }.toTypedArray()
+                }
             }
         }
 

idea/idea-analysis/src/org/jetbrains/kotlin/idea/util/ApplicationUtils.kt

@@ -21,6 +21,7 @@ import com.intellij.openapi.command.CommandProcessor
 import com.intellij.openapi.progress.ProgressIndicator
 import com.intellij.openapi.progress.ProgressIndicatorProvider
 import com.intellij.openapi.progress.ProgressManager
+import com.intellij.openapi.progress.impl.CancellationCheck
 import com.intellij.openapi.project.Project
 
 fun <T> runReadAction(action: () -> T): T {
@@ -45,3 +46,5 @@ fun <T> Project.executeCommand(name: String, groupId: Any? = null, command: () -
     @Suppress("USELESS_CAST")
     return result as T
 }
+
+fun <T> runWithCancellationCheck(block: () -> T): T = CancellationCheck.runWithCancellationCheck(block)

idea/idea-analysis/src/org/jetbrains/kotlin/idea/util/ApplicationUtils.kt.192

@@ -0,0 +1,49 @@
+/*
+ * Copyright 2010-2015 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.jetbrains.kotlin.idea.util.application
+
+import com.intellij.openapi.application.ApplicationManager
+import com.intellij.openapi.command.CommandProcessor
+import com.intellij.openapi.progress.ProgressIndicator
+import com.intellij.openapi.progress.ProgressIndicatorProvider
+import com.intellij.openapi.progress.ProgressManager
+import com.intellij.openapi.project.Project
+
+fun <T> runReadAction(action: () -> T): T {
+    return ApplicationManager.getApplication().runReadAction<T>(action)
+}
+
+fun <T> runWriteAction(action: () -> T): T {
+    return ApplicationManager.getApplication().runWriteAction<T>(action)
+}
+
+fun Project.executeWriteCommand(name: String, command: () -> Unit) {
+    CommandProcessor.getInstance().executeCommand(this, { runWriteAction(command) }, name, null)
+}
+
+fun <T> Project.executeWriteCommand(name: String, groupId: Any? = null, command: () -> T): T {
+    return executeCommand<T>(name, groupId) { runWriteAction(command) }
+}
+
+fun <T> Project.executeCommand(name: String, groupId: Any? = null, command: () -> T): T {
+    @Suppress("UNCHECKED_CAST") var result: T = null as T
+    CommandProcessor.getInstance().executeCommand(this, { result = command() }, name, groupId)
+    @Suppress("USELESS_CAST")
+    return result as T
+}
+
+fun <T> runWithCancellationCheck(block: () -> T): T = block()