FIR IDE: introduce analyze function for with by existing module resolve state

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/FirIdeResolveStateService.kt

@@ -33,7 +33,6 @@ internal class FirIdeResolveStateService(project: Project) {
         ConcurrentHashMap<IdeaModuleInfo, FirModuleResolveStateImpl>()
     }
 
-
     fun getResolveState(moduleInfo: IdeaModuleInfo): FirModuleResolveStateImpl =
         stateCache.computeIfAbsent(moduleInfo) { createResolveStateFor(moduleInfo, sessionProviderStorage) }
 
@@ -48,6 +47,7 @@ internal class FirIdeResolveStateService(project: Project) {
             val sessionProvider = sessionProviderStorage.getSessionProvider(moduleInfo)
             val firFileBuilder = sessionProvider.rootModuleSession.firFileBuilder
             return FirModuleResolveStateImpl(
+                moduleInfo.project,
                 moduleInfo,
                 sessionProvider,
                 firFileBuilder,

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/FirModuleResolveStateForCompletion.kt

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.idea.fir.low.level.api
 
+import com.intellij.openapi.project.Project
 import org.jetbrains.kotlin.diagnostics.Diagnostic
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirSession
@@ -21,6 +22,7 @@ import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtFile
 
 internal class FirModuleResolveStateForCompletion(
+    override val project: Project,
     private val originalState: FirModuleResolveStateImpl
 ) : FirModuleResolveState() {
     override val moduleInfo: IdeaModuleInfo get() = originalState.moduleInfo

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/FirModuleResolveStateImpl.kt

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.idea.fir.low.level.api
 
+import com.intellij.openapi.project.Project
 import org.jetbrains.kotlin.diagnostics.Diagnostic
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirSession
@@ -31,6 +32,7 @@ import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtFile
 
 internal class FirModuleResolveStateImpl(
+    override val project: Project,
     override val moduleInfo: IdeaModuleInfo,
     private val sessionProvider: FirIdeSessionProvider,
     val firFileBuilder: FirFileBuilder,

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/api/FirModuleResolveState.kt

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.idea.fir.low.level.api.api
 
 import org.jetbrains.annotations.TestOnly
+import com.intellij.openapi.project.Project
 import org.jetbrains.kotlin.diagnostics.Diagnostic
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirSession
@@ -26,9 +27,11 @@ import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtFile
 
 abstract class FirModuleResolveState {
+    abstract val project: Project
+
     abstract val rootModuleSession: FirSession
 
-    internal abstract val moduleInfo: IdeaModuleInfo
+    abstract val moduleInfo: IdeaModuleInfo
 
     abstract val firTransformerProvider: FirTransformerProvider
 

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/api/LowLevelFirApiFacadeForCompletion.kt

@@ -27,9 +27,9 @@ import org.jetbrains.kotlin.psi.KtNamedFunction
 import org.jetbrains.kotlin.psi.KtProperty
 
 object LowLevelFirApiFacadeForCompletion {
-    fun getResolveStateForCompletion(element: KtElement, originalState: FirModuleResolveState): FirModuleResolveState {
+    fun getResolveStateForCompletion(originalState: FirModuleResolveState): FirModuleResolveState {
         check(originalState is FirModuleResolveStateImpl)
-        return FirModuleResolveStateForCompletion(originalState)
+        return FirModuleResolveStateForCompletion(originalState.project, originalState)
     }
 
     class FirCompletionContext internal constructor(

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

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.idea.frontend.api.fir
 
+import com.intellij.openapi.project.Project
 import org.jetbrains.kotlin.fir.resolve.firSymbolProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirModuleResolveState
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.LowLevelFirApiFacade
@@ -22,7 +23,7 @@ import org.jetbrains.kotlin.psi.KtElement
 
 internal class KtFirAnalysisSession
 private constructor(
-    private val element: KtElement,
+    private val project: Project,
     val firResolveState: FirModuleResolveState,
     internal val firSymbolBuilder: KtSymbolByFirBuilder,
     token: ValidityToken,
@@ -32,8 +33,6 @@ private constructor(
         assertIsValid()
     }
 
-    private val project = element.project
-
     override val smartCastProvider: KtSmartCastProvider = KtFirSmartcastProvider(this, token)
     override val typeProvider: KtTypeProvider = KtFirTypeProvider(this, token)
     override val diagnosticProvider: KtDiagnosticProvider = KtFirDiagnosticProvider(this, token)
@@ -48,9 +47,9 @@ private constructor(
 
     override fun createContextDependentCopy(): KtAnalysisSession {
         check(!isContextSession) { "Cannot create context-dependent copy of KtAnalysis session from a context dependent one" }
-        val contextResolveState = LowLevelFirApiFacadeForCompletion.getResolveStateForCompletion(element, firResolveState)
+        val contextResolveState = LowLevelFirApiFacadeForCompletion.getResolveStateForCompletion(firResolveState)
         return KtFirAnalysisSession(
-            element,
+            project,
             contextResolveState,
             firSymbolBuilder.createReadOnlyCopy(contextResolveState),
             token,
@@ -62,7 +61,12 @@ private constructor(
         @Deprecated("Please use org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSessionProviderKt.analyze")
         internal fun createForElement(element: KtElement): KtFirAnalysisSession {
             val firResolveState = LowLevelFirApiFacade.getResolveStateFor(element)
-            val project = element.project
+            return createAnalysisSessionByResolveState(firResolveState)
+        }
+
+        @Deprecated("Please use org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSessionProviderKt.analyze")
+        internal fun createAnalysisSessionByResolveState(firResolveState: FirModuleResolveState): KtFirAnalysisSession {
+            val project = firResolveState.project
             val token = ReadActionConfinementValidityToken(project)
             val firSymbolBuilder = KtSymbolByFirBuilder(
                 firResolveState,
@@ -70,7 +74,7 @@ private constructor(
                 token
             )
             return KtFirAnalysisSession(
-                element,
+                project,
                 firResolveState,
                 firSymbolBuilder,
                 token,

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

@@ -11,6 +11,7 @@ import com.intellij.psi.util.CachedValuesManager
 import com.intellij.psi.util.PsiModificationTracker
 import org.jetbrains.kotlin.analyzer.ModuleInfo
 import org.jetbrains.kotlin.idea.caches.project.getModuleInfo
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirModuleResolveState
 import org.jetbrains.kotlin.idea.frontend.api.InvalidWayOfUsingAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSessionProvider
@@ -19,7 +20,7 @@ import org.jetbrains.kotlin.psi.KtElement
 import java.util.concurrent.ConcurrentHashMap
 
 @OptIn(InvalidWayOfUsingAnalysisSession::class)
-class KtFirAnalysisSessionProvider(project: Project) : KtAnalysisSessionProvider() {
+internal class KtFirAnalysisSessionProvider(project: Project) : KtAnalysisSessionProvider() {
     private val analysisSessionByModuleInfoCache =
         CachedValuesManager.getManager(project).createCachedValue {
             CachedValueProvider.Result(
@@ -36,4 +37,9 @@ class KtFirAnalysisSessionProvider(project: Project) : KtAnalysisSessionProvider
             assertIsValid()
         }
     }
+
+    internal fun getAnalysisSessionByResolveState(firModuleResolveState: FirModuleResolveState): KtAnalysisSession =
+        analysisSessionByModuleInfoCache.value.getOrPut(firModuleResolveState.moduleInfo) {
+            KtFirAnalysisSession.createAnalysisSessionByResolveState(firModuleResolveState)
+        }
 }

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

@@ -0,0 +1,23 @@
+/*
+ * 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
+
+import com.intellij.openapi.components.service
+import org.jetbrains.kotlin.idea.frontend.api.InvalidWayOfUsingAnalysisSession
+import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSession
+import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSessionProvider
+import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.KtFirSymbol
+import org.jetbrains.kotlin.idea.frontend.api.symbols.KtSymbol
+
+@OptIn(InvalidWayOfUsingAnalysisSession::class)
+internal inline fun <R> analyzeWithSymbolAsContext(contextSymbol: KtSymbol, action: KtAnalysisSession.() -> R): R {
+    require(contextSymbol is KtFirSymbol<*>)
+    val resolveState = contextSymbol.firRef.resolveState
+    val analysisSessionProvider = resolveState.project.service<KtAnalysisSessionProvider>()
+    check(analysisSessionProvider is KtFirAnalysisSessionProvider)
+    val analysisSession = analysisSessionProvider.getAnalysisSessionByResolveState(resolveState)
+    return action(analysisSession)
+}

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

@@ -21,11 +21,12 @@ internal class FirRefWithValidityCheck<D : FirDeclaration>(fir: D, resolveState:
     inline fun <R> withFir(phase: FirResolvePhase = FirResolvePhase.RAW_FIR, action: (fir: D) -> R): R {
         token.assertIsValid()
         val fir = firWeakRef.get() ?: error("FirElement was garbage collected while analysis session is still valid")
-        val resolveState =
-            resolveStateWeakRef.get() ?: error("FirModuleResolveState was garbage collected while analysis session is still valid")
         return action(LowLevelFirApiFacade.resolvedFirToPhase(fir, phase, resolveState))
     }
 
+    val resolveState
+        get() = resolveStateWeakRef.get() ?: error("FirModuleResolveState was garbage collected while analysis session is still valid")
+
     inline fun <R> withFirAndCache(phase: FirResolvePhase = FirResolvePhase.RAW_FIR, crossinline createValue: (fir: D) -> R) =
         ValidityAwareCachedValue(token) {
             withFir(phase) { fir -> createValue(fir) }