From 496e4dd0e3d1bccddecd6addf4805a12790e7038 Mon Sep 17 00:00:00 2001
From: Roman Golyshev <roman.golyshev@jetbrains.com>
Date: Fri, 1 Sep 2023 12:33:56 +0200
Subject: [PATCH] KTIJ-26867 [AA] Allow resolve from write action in
 `KtFirReferenceResolver`

`KtReference` might be resolved from anywhere, even
from the write action - the call site might not know
anything about Kotlin, and might not be able to
somehow prevent the `inaccessible analysis session`
exception from being thrown. That's why we
have to be permissive here.

^KTIJ-26867 Fixed
---
 analysis/analysis-api-fir/build.gradle.kts    |  1 +
 .../fir/references/KtFirReferenceResolver.kt  | 20 +++++++++++--------
 2 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/analysis/analysis-api-fir/build.gradle.kts b/analysis/analysis-api-fir/build.gradle.kts
index 703250f985d..247af117bf9 100644
--- a/analysis/analysis-api-fir/build.gradle.kts
+++ b/analysis/analysis-api-fir/build.gradle.kts
@@ -72,6 +72,7 @@ allprojects {
     tasks.withType<org.jetbrains.kotlin.gradle.dsl.KotlinCompile<*>> {
         kotlinOptions {
             freeCompilerArgs += "-opt-in=org.jetbrains.kotlin.fir.symbols.SymbolInternals"
+            freeCompilerArgs += "-opt-in=org.jetbrains.kotlin.analysis.api.lifetime.KtAllowProhibitedAnalyzeFromWriteAction"
         }
     }
 }
diff --git a/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/references/KtFirReferenceResolver.kt b/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/references/KtFirReferenceResolver.kt
index a48789e1f03..51a75db4893 100644
--- a/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/references/KtFirReferenceResolver.kt
+++ b/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/references/KtFirReferenceResolver.kt
@@ -9,8 +9,10 @@ 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.analysis.api.KtAllowAnalysisFromWriteAction
 import org.jetbrains.kotlin.analysis.api.KtAllowAnalysisOnEdt
 import org.jetbrains.kotlin.analysis.api.analyze
+import org.jetbrains.kotlin.analysis.api.lifetime.allowAnalysisFromWriteAction
 import org.jetbrains.kotlin.analysis.api.lifetime.allowAnalysisOnEdt
 import org.jetbrains.kotlin.utils.exceptions.errorWithAttachment
 import org.jetbrains.kotlin.utils.exceptions.shouldIjPlatformExceptionBeRethrown
@@ -19,21 +21,23 @@ import org.jetbrains.kotlin.utils.exceptions.withPsiEntry
 object KtFirReferenceResolver : ResolveCache.PolyVariantResolver<KtReference> {
     class KotlinResolveResult(element: PsiElement) : PsiElementResolveResult(element)
 
-    @OptIn(KtAllowAnalysisOnEdt::class)
+    @OptIn(KtAllowAnalysisOnEdt::class, KtAllowAnalysisFromWriteAction::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}" }
         return allowAnalysisOnEdt {
-            val resolveToPsiElements = try {
-                analyze(ref.expression) { ref.getResolvedToPsi(this) }
-            } catch (e: Exception) {
-                if (shouldIjPlatformExceptionBeRethrown(e)) throw e
+            allowAnalysisFromWriteAction {
+                val resolveToPsiElements = try {
+                    analyze(ref.expression) { ref.getResolvedToPsi(this) }
+                } catch (e: Exception) {
+                    if (shouldIjPlatformExceptionBeRethrown(e)) throw e
 
-                errorWithAttachment("Unable to resolve reference ${ref.element::class}", cause = e) {
-                    withPsiEntry("reference", ref.element)
+                    errorWithAttachment("Unable to resolve reference ${ref.element::class}", cause = e) {
+                        withPsiEntry("reference", ref.element)
+                    }
                 }
+                resolveToPsiElements.map { KotlinResolveResult(it) }.toTypedArray()
             }
-            resolveToPsiElements.map { KotlinResolveResult(it) }.toTypedArray()
         }
     }
 }
\ No newline at end of file