From 171157ff784f12732cf2facd5b82ea55acc6f0f8 Mon Sep 17 00:00:00 2001
From: Ilya Kirillov <ilya.kirillov@jetbrains.com>
Date: Thu, 29 Oct 2020 23:05:40 +0300
Subject: [PATCH] FIR IDE: allow reference resolving from EDT

It may be called from some other subsystems like Find usages.
Throwing PCE every time we try to resolve some reference from EDT
may result endless try to resolving
---
 .../frontend/api/fir/utils/threadUtils.kt     | 16 ++++++++++++++
 .../idea/references/KtFirReferenceResolver.kt | 22 +++++++++----------
 2 files changed, 27 insertions(+), 11 deletions(-)
 create mode 100644 idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/utils/threadUtils.kt

diff --git a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/utils/threadUtils.kt b/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/utils/threadUtils.kt
new file mode 100644
index 00000000000..0c770150482
--- /dev/null
+++ b/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/utils/threadUtils.kt
@@ -0,0 +1,16 @@
+/*
+ * 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.frontend.api.fir.utils
+
+import com.intellij.openapi.application.ApplicationManager
+import org.jetbrains.kotlin.idea.frontend.api.HackToForceAllowRunningAnalyzeOnEDT
+import org.jetbrains.kotlin.idea.frontend.api.hackyAllowRunningOnEdt
+
+@HackToForceAllowRunningAnalyzeOnEDT
+internal inline fun <R> runInPossiblyEdtThread(action: () -> R): R = when {
+    !ApplicationManager.getApplication().isDispatchThread -> action()
+    else -> hackyAllowRunningOnEdt(action)
+}
\ No newline at end of file
diff --git a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/KtFirReferenceResolver.kt b/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/KtFirReferenceResolver.kt
index 259d136adae..545c8ebbcbf 100644
--- a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/KtFirReferenceResolver.kt
+++ b/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/KtFirReferenceResolver.kt
@@ -5,32 +5,32 @@
 
 package org.jetbrains.kotlin.idea.references
 
-import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.diagnostic.ControlFlowException
-import com.intellij.openapi.progress.ProcessCanceledException
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiElementResolveResult
 import com.intellij.psi.ResolveResult
 import com.intellij.psi.impl.source.resolve.ResolveCache
+import org.jetbrains.kotlin.idea.frontend.api.HackToForceAllowRunningAnalyzeOnEDT
 import org.jetbrains.kotlin.idea.frontend.api.analyze
+import org.jetbrains.kotlin.idea.frontend.api.fir.utils.runInPossiblyEdtThread
 import org.jetbrains.kotlin.idea.util.getElementTextInContext
 
 object KtFirReferenceResolver : ResolveCache.PolyVariantResolver<KtReference> {
     class KotlinResolveResult(element: PsiElement) : PsiElementResolveResult(element)
 
+    @OptIn(HackToForceAllowRunningAnalyzeOnEDT::class)
     override fun resolve(ref: KtReference, incompleteCode: Boolean): Array<ResolveResult> {
         check(ref is KtFirReference) { "reference should be FirKtReference, but was ${ref::class}" }
         check(ref is AbstractKtReference<*>) { "reference should be AbstractKtReference, but was ${ref::class}" }
-        if (ApplicationManager.getApplication().isDispatchThread) {
-            throw ProcessCanceledException()
+        return runInPossiblyEdtThread {
+            val resolveToPsiElements = try {
+                analyze(ref.expression) { ref.getResolvedToPsi(this) }
+            } catch (e: Throwable) {
+                if (e is ControlFlowException) throw e
+                throw KtReferenceResolveException(ref, e)
+            }
+            resolveToPsiElements.map { KotlinResolveResult(it) }.toTypedArray()
         }
-        val resolveToPsiElements = try {
-            analyze(ref.expression) { ref.getResolvedToPsi(this) }
-        } catch (e: Throwable) {
-            if (e is ControlFlowException) throw e
-            throw KtReferenceResolveException(ref, e)
-        }
-        return resolveToPsiElements.map { KotlinResolveResult(it) }.toTypedArray()
     }
 }