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

@@ -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"
         }
     }
 }

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()
         }
     }
 }