FIR IDE: make context-dependent analysis session to use original cache on read-only mode

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

@@ -31,5 +31,5 @@ abstract class KtAnalysisSession(override val token: ValidityToken) : ValidityTo
     abstract fun resolveCall(call: KtCallExpression): CallInfo?
     abstract fun resolveCall(call: KtBinaryExpression): CallInfo?
 
-    abstract fun analyzeInContext(): KtAnalysisSession
+    abstract fun createContextDependentCopy(): KtAnalysisSession
 }

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

@@ -34,7 +34,7 @@ private object KotlinHighLevelApiContributor : CompletionProvider<CompletionPara
         val possibleReceiver = nameExpression.getQualifiedExpressionForSelector()?.receiverExpression
 
         val originalSession = getAnalysisSessionFor(originalFile)
-        val sessionForCompletion = originalSession.analyzeInContext()
+        val sessionForCompletion = originalSession.createContextDependentCopy()
         val scopeProvider = sessionForCompletion.scopeProvider
 
         val (implicitScopes, implicitReceivers) = scopeProvider.getScopeContextForPosition(originalFile, nameExpression)

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

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.idea.frontend.api.fir
 
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
 import org.jetbrains.kotlin.diagnostics.Diagnostic
+import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirCallableDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirClass
 import org.jetbrains.kotlin.fir.declarations.isSuspend
@@ -17,6 +18,7 @@ import org.jetbrains.kotlin.fir.resolve.transformers.firClassLike
 import org.jetbrains.kotlin.fir.symbols.CallableId
 import org.jetbrains.kotlin.fir.symbols.impl.FirConstructorSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
+import org.jetbrains.kotlin.fir.typeContext
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.idea.fir.*
 import org.jetbrains.kotlin.idea.fir.low.level.api.FirModuleResolveState
@@ -38,34 +40,20 @@ import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.*
 
 class KtFirAnalysisSession
-@Deprecated("Please use org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSessionProviderKt.analyze")
+private constructor(
-constructor(
     private val element: KtElement,
-    resolveState: FirModuleResolveState
+    val firSession: FirSession,
-) : KtAnalysisSession(element.project) {
+    val firResolveState: FirModuleResolveState,
-    internal val token get() = validityToken
+    internal val firSymbolBuilder: KtSymbolByFirBuilder,
+    private val typeContext: ConeTypeCheckerContext,
+    token: ValidityToken,
+    val isContextSession: Boolean,
+) : KtAnalysisSession(token) {
     private val project = element.project
 
-    internal val firResolveState = resolveState
-
-    internal val firSession get() = LowLevelFirApiFacade.getSessionFor(element, firResolveState)
-
-    private val typeCheckerContext = ConeTypeCheckerContext(
-        isErrorTypeEqualsToAnything = true,
-        isStubTypeEqualsToAnything = true,
-        firSession
-    )
-
-    internal val firSymbolBuilder = KtSymbolByFirBuilder(
-        firSession.firSymbolProvider,
-        typeCheckerContext,
-        project,
-        validityToken
-    )
-
     override val symbolProvider: KtSymbolProvider =
         KtFirSymbolProvider(
-            this,
+            token,
             firSession.firSymbolProvider,
             firResolveState,
             firSymbolBuilder
@@ -133,10 +121,17 @@ constructor(
         resolveCall(firCall, call)
     }
 
-    override fun analyzeInContext(): KtAnalysisSession {
+    override fun createContextDependentCopy(): KtAnalysisSession {
+        check(!isContextSession) { "Cannot create context-dependent copy of KtAnalysis session from a context dependent one" }
+        val contextResolveState = LowLevelFirApiFacade.getResolveStateForCompletion(element, firResolveState as FirModuleResolveStateImpl)
         return KtFirAnalysisSession(
             element,
-            LowLevelFirApiFacade.getResolveStateForCompletion(element, firResolveState as FirModuleResolveStateImpl)
+            firSession,
+            contextResolveState,
+            firSymbolBuilder,
+            typeContext,
+            token,
+            isContextSession = true
         )
     }
 
@@ -201,5 +196,29 @@ constructor(
                 CallableId(KotlinBuiltIns.getSuspendFunctionClassId(arity), Name.identifier("invoke"))
             )
         }
+
+        @Deprecated("Please use org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSessionProviderKt.analyze")
+        fun createForElement(element: KtElement): KtFirAnalysisSession {
+            val firResolveState = LowLevelFirApiFacade.getResolveStateFor(element)
+            val firSession = LowLevelFirApiFacade.getSessionFor(element, firResolveState)
+            val project = element.project
+            val typeContext = ConeTypeCheckerContext(isErrorTypeEqualsToAnything = true, isStubTypeEqualsToAnything = true, firSession)
+            val token = ReadActionConfinementValidityToken(project)
+            val firSymbolBuilder = KtSymbolByFirBuilder(
+                firSession.firSymbolProvider,
+                typeContext,
+                project,
+                token
+            )
+            return KtFirAnalysisSession(
+                element,
+                firSession,
+                firResolveState,
+                firSymbolBuilder.createReadOnlyCopy(),
+                typeContext,
+                token,
+                isContextSession = false
+            )
+        }
     }
 }

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

@@ -26,9 +26,9 @@ class KtFirAnalysisSessionProvider(project: Project) : KtAnalysisSessionProvider
             )
         }
 
-    @Suppress("DEPRECATION")
     override fun getAnalysisSessionFor(contextElement: KtElement): KtAnalysisSession =
         analysisSessionByModuleInfoCache.value.getOrPut(contextElement.getModuleInfo()) {
-            KtFirAnalysisSession(contextElement, LowLevelFirApiFacade.getResolveStateFor(contextElement))
+            @Suppress("DEPRECATION")
+            KtFirAnalysisSession.createForElement(contextElement)
         }
 }

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

@@ -18,7 +18,7 @@ import org.jetbrains.kotlin.fir.symbols.impl.FirTypeParameterSymbol
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.fir.types.impl.ConeClassLikeTypeImpl
 import org.jetbrains.kotlin.idea.frontend.api.*
-import org.jetbrains.kotlin.idea.frontend.api.ValidityOwner
+import org.jetbrains.kotlin.idea.frontend.api.ValidityTokenOwner
 import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.*
 import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.KtFirClassOrObjectSymbol
 import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.KtFirLocalVariableSymbol
@@ -36,21 +36,52 @@ import org.jetbrains.kotlin.idea.frontend.api.symbols.*
 import org.jetbrains.kotlin.idea.frontend.api.types.KtType
 import org.jetbrains.kotlin.idea.stubindex.PackageIndexUtil
 import org.jetbrains.kotlin.name.FqName
+import java.util.concurrent.ConcurrentMap
 
 /**
  * Maps FirElement to KtSymbol & ConeType to KtType, thread safe
  */
-internal class KtSymbolByFirBuilder(
+internal class KtSymbolByFirBuilder private constructor(
     firProvider: FirSymbolProvider,
     typeCheckerContext: ConeTypeCheckerContext,
     private val project: Project,
-    override val token: ValidityOwner
+    override val token: ValidityToken,
-) : ValidityOwnerByValidityToken {
+    val withReadOnlyCaching: Boolean,
+    private val symbolsCache: BuilderCache<FirDeclaration, KtSymbol>,
+    private val typesCache: BuilderCache<ConeKotlinType, KtType>
+) : ValidityTokenOwner {
+
+    constructor(
+        firProvider: FirSymbolProvider,
+        typeCheckerContext: ConeTypeCheckerContext,
+        project: Project,
+        token: ValidityToken
+    ) : this(
+        firProvider = firProvider,
+        typeCheckerContext = typeCheckerContext,
+        project = project,
+        token = token,
+        withReadOnlyCaching = false,
+        symbolsCache = BuilderCache(),
+        typesCache = BuilderCache()
+    )
+
     private val firProvider by weakRef(firProvider)
     private val typeCheckerContext by weakRef(typeCheckerContext)
 
-    private val symbolsCache = BuilderCache<FirDeclaration, KtSymbol>()
+    fun createReadOnlyCopy(): KtSymbolByFirBuilder {
-    private val typesCache = BuilderCache<ConeKotlinType, KtType>()
+        check(!withReadOnlyCaching) { "Cannot create readOnly KtSymbolByFirBuilder from a readonly one" }
+        return KtSymbolByFirBuilder(
+            firProvider,
+            typeCheckerContext,
+            project,
+            token,
+            withReadOnlyCaching = true,
+            symbolsCache.createReadOnlyCopy(),
+            typesCache.createReadOnlyCopy()
+        )
+    }
+
 
     fun buildSymbol(fir: FirDeclaration): KtSymbol = symbolsCache.cache(fir) {
         when (fir) {
@@ -153,11 +184,23 @@ internal class KtSymbolByFirBuilder(
     }
 }
 
-private class BuilderCache<From, To> {
+private class BuilderCache<From, To> private constructor(
-    private val cache = MapMaker().weakKeys().makeMap<From, To>()
+    private val cache: ConcurrentMap<From, To>,
+    private val isReadOnly: Boolean
+) {
+    constructor() : this(cache = MapMaker().weakKeys().makeMap(), isReadOnly = false)
 
-    inline fun <reified S : To> cache(key: From, calculation: () -> S): S =
+    fun createReadOnlyCopy(): BuilderCache<From, To> {
-        cache.getOrPut(key, calculation) as S
+        check(!isReadOnly) { "Cannot create readOnly BuilderCache from a readonly one" }
+        return BuilderCache(cache, isReadOnly = true)
+    }
+
+    inline fun <reified S : To> cache(key: From, calculation: () -> S): S {
+        if (isReadOnly) {
+            return (cache[key] ?: calculation()) as S
+        }
+        return cache.getOrPut(key, calculation) as S
+    }
 }
 
 internal fun FirElement.buildSymbol(builder: KtSymbolByFirBuilder) =

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

@@ -27,12 +27,14 @@ import org.jetbrains.kotlin.name.Name
 
 internal abstract class KtFirDelegatingScope(
     private val builder: KtSymbolByFirBuilder,
-    override val token: ValidityToken
+    final override val token: ValidityToken
-) : KtScope, ValidityTokenOwner {
+) : KtScope {
 
     abstract val firScope: FirScope
 
-    private val allNamesCached by cached { getCallableNames() + getClassLikeSymbolNames() }
+    private val allNamesCached by cached(token) {
+        getCallableNames() + getClassLikeSymbolNames()
+    }
 
     override fun getAllNames(): Set<Name> = allNamesCached
 

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

@@ -19,7 +19,7 @@ import org.jetbrains.kotlin.name.Name
 internal class KtFirNonStarImportingScope(
     firScope: FirAbstractSimpleImportingScope,
     builder: KtSymbolByFirBuilder,
-    override val token: ValidityToken
+    token: ValidityToken
 ) : KtFirDelegatingScope(builder, token), KtNonStarImportingScope, ValidityTokenOwner {
     override val firScope: FirAbstractSimpleImportingScope = firScope