FIR IDE: add validity check of analysis component that it is valid only in analysis context

idea/idea-fir/src/org/jetbrains/kotlin/idea/completion/KotlinFirCompletionContributor.kt

@@ -17,7 +17,7 @@ import org.jetbrains.kotlin.idea.fir.low.level.api.util.originalKtFile
 import org.jetbrains.kotlin.idea.frontend.api.InvalidWayOfUsingAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.components.KtScopeContext
-import org.jetbrains.kotlin.idea.frontend.api.getAnalysisSessionFor
+import org.jetbrains.kotlin.idea.frontend.api.analyseInFakeAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.scopes.KtCompositeScope
 import org.jetbrains.kotlin.idea.frontend.api.scopes.KtScope
 import org.jetbrains.kotlin.idea.frontend.api.scopes.KtScopeNameFilter
@@ -133,7 +133,7 @@ private class KotlinCommonCompletionProvider(
 
         val explicitReceiver = nameExpression.getReceiverExpression()
 
-        with(getAnalysisSessionFor(originalFile).createContextDependentCopy(originalFile, nameExpression)) {
+        analyseInFakeAnalysisSession(originalFile, nameExpression) {
             val expectedType = nameExpression.getExpectedType()
 
             val scopesContext = originalFile.getScopeContextForPosition(nameExpression)

idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/KtAnalysisSessionProvider.kt

@@ -11,10 +11,14 @@ import com.intellij.openapi.progress.ProgressIndicator
 import com.intellij.openapi.progress.Task
 import org.jetbrains.kotlin.idea.util.application.runReadAction
 import org.jetbrains.kotlin.psi.KtElement
+import org.jetbrains.kotlin.psi.KtFile
 
 @RequiresOptIn("To use analysis session, consider using analyze/analyzeWithReadAction/analyseInModalWindow methods")
 annotation class InvalidWayOfUsingAnalysisSession
 
+@RequiresOptIn
+annotation class KtAnalysisSessionProviderInternals
+
 /**
  * Provides [KtAnalysisSession] by [contextElement]
  * Should not be used directly, consider using [analyse]/[analyseWithReadAction]/[analyseInModalWindow] instead
@@ -23,11 +27,38 @@ annotation class InvalidWayOfUsingAnalysisSession
 abstract class KtAnalysisSessionProvider {
     @InvalidWayOfUsingAnalysisSession
     abstract fun getAnalysisSessionFor(contextElement: KtElement): KtAnalysisSession
-}
 
-@InvalidWayOfUsingAnalysisSession
-fun getAnalysisSessionFor(contextElement: KtElement): KtAnalysisSession =
-    contextElement.project.service<KtAnalysisSessionProvider>().getAnalysisSessionFor(contextElement)
+    @InvalidWayOfUsingAnalysisSession
+    inline fun <R> analyseInFakeAnalysisSession(originalFile: KtFile, fakeExpresion: KtElement, action: KtAnalysisSession.() -> R): R {
+        val fakeAnalysisSession = getAnalysisSessionFor(originalFile).createContextDependentCopy(originalFile, fakeExpresion)
+        return analyse(fakeAnalysisSession, action)
+    }
+
+    @InvalidWayOfUsingAnalysisSession
+    inline fun <R> analyse(contextElement: KtElement, action: KtAnalysisSession.() -> R): R =
+        analyse(getAnalysisSessionFor(contextElement), action)
+
+    @OptIn(KtAnalysisSessionProviderInternals::class)
+    @InvalidWayOfUsingAnalysisSession
+    inline fun <R> analyse(analysisSession: KtAnalysisSession, action: KtAnalysisSession.() -> R): R {
+        currentAnalysisContextEnteringCount.set(currentAnalysisContextEnteringCount.get() + 1)
+        return try {
+            analysisSession.action()
+        } finally {
+            currentAnalysisContextEnteringCount.set(currentAnalysisContextEnteringCount.get() - 1)
+        }
+    }
+
+    companion object {
+        @KtAnalysisSessionProviderInternals
+        val currentAnalysisContextEnteringCount = object : ThreadLocal<Int>() {
+            override fun initialValue() = 0
+        }
+
+        @OptIn(KtAnalysisSessionProviderInternals::class)
+        fun isInsideAnalysisContext() = currentAnalysisContextEnteringCount.get() > 0
+    }
+}
 
 /**
  * Execute given [action] in [KtAnalysisSession] context
@@ -43,8 +74,11 @@ fun getAnalysisSessionFor(contextElement: KtElement): KtAnalysisSession =
  */
 @OptIn(InvalidWayOfUsingAnalysisSession::class)
 inline fun <R> analyse(contextElement: KtElement, action: KtAnalysisSession.() -> R): R =
-    getAnalysisSessionFor(contextElement).action()
+    contextElement.project.service<KtAnalysisSessionProvider>().analyse(contextElement, action)
 
+@OptIn(InvalidWayOfUsingAnalysisSession::class)
+inline fun <R> analyseInFakeAnalysisSession(originalFile: KtFile, fakeExpresion: KtElement, action: KtAnalysisSession.() -> R): R =
+    originalFile.project.service<KtAnalysisSessionProvider>().analyseInFakeAnalysisSession(originalFile, fakeExpresion, action)
 
 /**
  * Execute given [action] in [KtAnalysisSession] context like [analyse] does but execute it in read action

idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/ReadActionConfinementValidityToken.kt

@@ -23,16 +23,17 @@ class ReadActionConfinementValidityToken(project: Project) : ValidityToken() {
         error("Getting invalidation reason for valid validity token")
     }
 
-    @OptIn(HackToForceAllowRunningAnalyzeOnEDT::class, ForbidKtResolveInternals::class)
+    @OptIn(HackToForceAllowRunningAnalyzeOnEDT::class, ForbidKtResolveInternals::class, InvalidWayOfUsingAnalysisSession::class)
     override fun isAccessible(): Boolean {
         val application = ApplicationManager.getApplication()
         if (application.isDispatchThread && !allowOnEdt.get()) return false
         if (ForbidKtResolve.resovleIsForbidenInActionWithName.get() != null) return false
         if (!application.isReadAccessAllowed) return false
+        if (!KtAnalysisSessionProvider.isInsideAnalysisContext()) return false
         return true
     }
 
-    @OptIn(HackToForceAllowRunningAnalyzeOnEDT::class, ForbidKtResolveInternals::class)
+    @OptIn(HackToForceAllowRunningAnalyzeOnEDT::class, ForbidKtResolveInternals::class, InvalidWayOfUsingAnalysisSession::class)
     override fun getInaccessibilityReason(): String {
         val application = ApplicationManager.getApplication()
         if (application.isDispatchThread && !allowOnEdt.get()) return "Called in EDT thread"
@@ -40,6 +41,7 @@ class ReadActionConfinementValidityToken(project: Project) : ValidityToken() {
         ForbidKtResolve.resovleIsForbidenInActionWithName.get()?.let { actionName ->
             return "Resolve is forbidden in $actionName"
         }
+        if (!KtAnalysisSessionProvider.isInsideAnalysisContext()) return "Called outside analyse method"
         error("Getting inaccessibility reason for validity token when it is accessible")
     }
 

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/KtFirAnalysisSession.kt

@@ -34,9 +34,6 @@ private constructor(
     token: ValidityToken,
     val context: KtFirAnalysisSessionContext,
 ) : KtAnalysisSession(token) {
-    init {
-        assertIsValidAndAccessible()
-    }
 
     override val smartCastProviderImpl = KtFirSmartcastProvider(this, token)
 

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/KtFirAnalysisSessionProvider.kt

@@ -38,8 +38,6 @@ class KtFirAnalysisSessionProvider(project: Project) : KtAnalysisSessionProvider
         return analysisSessionByModuleInfoCache.value.getOrPut(contextElement.getModuleInfo()) {
             @Suppress("DEPRECATION")
             KtFirAnalysisSession.createForElement(contextElement)
-        }.apply {
-            assertIsValidAndAccessible()
         }
     }
 

idea/idea-frontend-fir/tests/org/jetbrains/kotlin/idea/frontend/api/fir/AbstractResolveCallTest.kt

@@ -60,8 +60,10 @@ abstract class AbstractResolveCallTest : @Suppress("DEPRECATION") LightCodeInsig
             buildString {
                 append(textWithoutLatestComments)
                 append("\n\n")
-                callInfos.joinTo(this, separator = "\n") { info ->
-                    "// CALL: ${info?.stringRepresentation()}"
+                analyse(file as KtFile) {
+                    callInfos.joinTo(this@buildString, separator = "\n") { info ->
+                        "// CALL: ${info?.stringRepresentation()}"
+                    }
                 }
             }
         }