FIR IDE: throw PCE in highlighter & reference resolver if on EDT thread

As resolve is not allowed in EDT thread

idea/idea-fir/src/org/jetbrains/kotlin/idea/fir/highlighter/KotlinFirPsiChecker.kt

@@ -6,6 +6,8 @@
 package org.jetbrains.kotlin.idea.fir.highlighter
 
 import com.intellij.lang.annotation.AnnotationHolder
+import com.intellij.openapi.application.ApplicationManager
+import com.intellij.openapi.progress.ProcessCanceledException
 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.diagnostics.Diagnostic
 import org.jetbrains.kotlin.idea.fir.highlighter.visitors.FirAfterResolveHighlightingVisitor
@@ -24,7 +26,10 @@ class KotlinFirPsiChecker : AbstractKotlinPsiChecker() {
 
     override fun annotateElement(element: PsiElement, containingFile: KtFile, holder: AnnotationHolder) {
         if (element !is KtElement) return
-        val analysisSession = FirAnalysisSession()
+        if (ApplicationManager.getApplication().isDispatchThread) {
+            throw ProcessCanceledException()
+        }
+        val analysisSession = FirAnalysisSession(element)
 
         highlightDiagnostics(element, analysisSession, holder)
 

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/KtFirReferenceResolver.kt

@@ -5,6 +5,8 @@
 
 package org.jetbrains.kotlin.idea.references
 
+import com.intellij.openapi.application.ApplicationManager
+import com.intellij.openapi.progress.ProcessCanceledException
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiElementResolveResult
 import com.intellij.psi.ResolveResult
@@ -17,7 +19,10 @@ object KtFirReferenceResolver : ResolveCache.PolyVariantResolver<KtReference> {
     override fun resolve(ref: KtReference, incompleteCode: Boolean): Array<ResolveResult> {
         check(ref is FirKtReference) { "reference should be FirKtReference, but was ${ref::class}" }
         check(ref is AbstractKtReference<*>) { "reference should be AbstractKtReference, but was ${ref::class}" }
-        val analysisSession = FirAnalysisSession()
+        if (ApplicationManager.getApplication().isDispatchThread) {
+            throw ProcessCanceledException()
+        }
+        val analysisSession = FirAnalysisSession(ref.expression)
         val resolveToPsiElements = ref.getResolvedToPsi(analysisSession)
         return resolveToPsiElements.map { KotlinResolveResult(it) }.toTypedArray()
     }

idea/tests/org/jetbrains/kotlin/idea/resolve/AbstractReferenceResolveTest.kt

 /*
- * Copyright 2010-2019 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
  * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
  */
 
 package org.jetbrains.kotlin.idea.resolve
 
 import com.google.common.collect.Lists
+import com.intellij.openapi.application.ApplicationManager
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiPolyVariantReference
 import com.intellij.psi.PsiReference
@@ -16,9 +17,11 @@ import org.jetbrains.kotlin.idea.completion.test.configureWithExtraFile
 import org.jetbrains.kotlin.idea.test.KotlinLightPlatformCodeInsightFixtureTestCase
 import org.jetbrains.kotlin.idea.test.KotlinWithJdkAndRuntimeLightProjectDescriptor
 import org.jetbrains.kotlin.idea.test.PluginTestCaseBase
+import org.jetbrains.kotlin.idea.util.application.runReadAction
 import org.jetbrains.kotlin.test.InTextDirectivesUtils
 import org.jetbrains.kotlin.test.util.renderAsGotoImplementation
 import org.junit.Assert
+import java.util.concurrent.Callable
 import kotlin.test.assertTrue
 
 abstract class AbstractReferenceResolveTest : KotlinLightPlatformCodeInsightFixtureTestCase() {
@@ -66,7 +69,9 @@ abstract class AbstractReferenceResolveTest : KotlinLightPlatformCodeInsightFixt
             assertTrue(psiReference is PsiPolyVariantReference)
             psiReference as PsiPolyVariantReference
 
-            val results = wrapReference(psiReference).multiResolve(true)
+            val results = executeOnPooledThreadInReadAction {
+                wrapReference(psiReference).multiResolve(true)
+            }
 
             val actualResolvedTo = Lists.newArrayList<String>()
             for (result in results) {
@@ -128,7 +133,7 @@ abstract class AbstractReferenceResolveTest : KotlinLightPlatformCodeInsightFixt
         ) {
             val expectedString = expectedResolveData.referenceString
             if (psiReference != null) {
-                val resolvedTo = psiReference.resolve()
+                val resolvedTo = executeOnPooledThreadInReadAction { psiReference.resolve() }
                 if (resolvedTo != null) {
                     checkResolvedTo(resolvedTo)
                     val resolvedToElementStr = replacePlaceholders(resolvedTo.renderAsGotoImplementation())
@@ -161,3 +166,6 @@ abstract class AbstractReferenceResolveTest : KotlinLightPlatformCodeInsightFixt
         }
     }
 }
+
+private fun <R> executeOnPooledThreadInReadAction(action: () -> R): R =
+    ApplicationManager.getApplication().executeOnPooledThread<R> { runReadAction(action) }.get()