FIR IDE: precalculate completion context on dependent analysis session creation

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

@@ -111,7 +111,7 @@ private class KotlinAvailableScopesCompletionProvider(prefixMatcher: PrefixMatch
 
         val explicitReceiver = nameExpression.getReceiverExpression()
 
-        with(getAnalysisSessionFor(originalFile).createContextDependentCopy()) {
+        with(getAnalysisSessionFor(originalFile).createContextDependentCopy(originalFile, nameExpression)) {
             val expectedType = nameExpression.getExpectedType()
 
             val (implicitScopes, _) = originalFile.getScopeContextForPosition(nameExpression)

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

@@ -43,6 +43,7 @@ abstract class KtAnalysisSession(final override val token: ValidityToken) : Vali
     protected abstract val callResolver: KtCallResolver
     protected abstract val completionCandidateChecker: KtCompletionCandidateChecker
     protected abstract val symbolDeclarationOverridesProvider: KtSymbolDeclarationOverridesProvider
+
     @Suppress("LeakingThis")
 
     protected open val typeRenderer: KtTypeRenderer = KtDefaultTypeRenderer(this, token)
@@ -51,16 +52,15 @@ abstract class KtAnalysisSession(final override val token: ValidityToken) : Vali
     protected abstract val subtypingComponent: KtSubtypingComponent
     protected abstract val expressionHandlingComponent: KtExpressionHandlingComponent
 
-    /// TODO: get rid of
+    abstract fun createContextDependentCopy(originalKtFile: KtFile, fakeKtElement: KtElement): KtAnalysisSession
-    @Deprecated("Used only in completion now, temporary")
-    abstract fun createContextDependentCopy(): KtAnalysisSession
 
     fun KtCallableSymbol.getOverriddenSymbols(containingDeclaration: KtClassOrObjectSymbol): List<KtCallableSymbol> =
         symbolDeclarationOverridesProvider.getOverriddenSymbols(this, containingDeclaration)
 
     fun KtExpression.getSmartCasts(): Collection<KtType> = smartCastProvider.getSmartCastedToTypes(this)
 
-    fun KtExpression.getImplicitReceiverSmartCasts(): Collection<ImplicitReceiverSmartCast> = smartCastProvider.getImplicitReceiverSmartCasts(this)
+    fun KtExpression.getImplicitReceiverSmartCasts(): Collection<ImplicitReceiverSmartCast> =
+        smartCastProvider.getImplicitReceiverSmartCasts(this)
 
     fun KtExpression.getKtType(): KtType = expressionTypeProvider.getKtExpressionType(this)
 
@@ -158,6 +158,6 @@ abstract class KtAnalysisSession(final override val token: ValidityToken) : Vali
     fun KtType.render(options: KtTypeRendererOptions = KtTypeRendererOptions.DEFAULT): String =
         typeRenderer.render(this, options)
 
-    fun KtReturnExpression.getReturnTargetSymbol(): KtFunctionLikeSymbol? =
+    fun KtReturnExpression.getReturnTargetSymbol(): KtCallableSymbol? =
         expressionHandlingComponent.getReturnExpressionTargetSymbol(this)
 }

idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/components/KtExpressionHandlingComponent.kt

@@ -5,9 +5,10 @@
 
 package org.jetbrains.kotlin.idea.frontend.api.components
 
+import org.jetbrains.kotlin.idea.frontend.api.symbols.KtCallableSymbol
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtFunctionLikeSymbol
 import org.jetbrains.kotlin.psi.KtReturnExpression
 
 abstract class KtExpressionHandlingComponent : KtAnalysisSessionComponent() {
-   abstract fun getReturnExpressionTargetSymbol(returnExpression: KtReturnExpression): KtFunctionLikeSymbol?
+   abstract fun getReturnExpressionTargetSymbol(returnExpression: KtReturnExpression): KtCallableSymbol?
 }

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

@@ -6,9 +6,11 @@
 package org.jetbrains.kotlin.idea.frontend.api.fir
 
 import com.intellij.openapi.project.Project
+import org.jetbrains.kotlin.fir.declarations.FirFile
 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.LowLevelFirApiFacadeForCompletion
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.getFirFile
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.getResolveState
 import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.ReadActionConfinementValidityToken
@@ -17,9 +19,12 @@ import org.jetbrains.kotlin.idea.frontend.api.assertIsValid
 import org.jetbrains.kotlin.idea.frontend.api.components.*
 import org.jetbrains.kotlin.idea.frontend.api.fir.components.*
 import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.KtFirSymbolProvider
+import org.jetbrains.kotlin.idea.frontend.api.fir.utils.EnclosingDeclarationContext
+import org.jetbrains.kotlin.idea.frontend.api.fir.utils.buildCompletionContext
 import org.jetbrains.kotlin.idea.frontend.api.fir.utils.threadLocal
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtSymbolProvider
 import org.jetbrains.kotlin.psi.KtElement
+import org.jetbrains.kotlin.psi.KtFile
 
 internal class KtFirAnalysisSession
 private constructor(
@@ -27,12 +32,14 @@ private constructor(
     val firResolveState: FirModuleResolveState,
     internal val firSymbolBuilder: KtSymbolByFirBuilder,
     token: ValidityToken,
-    val isContextSession: Boolean,
+    val context: KtFirAnalysisSessionContext,
 ) : KtAnalysisSession(token) {
     init {
         assertIsValid()
     }
 
+    internal val towerDataContextProvider: TowerDataContextProvider = TowerDataContextProvider(this)
+
     override val smartCastProvider: KtSmartCastProvider = KtFirSmartcastProvider(this, token)
     override val expressionTypeProvider: KtExpressionTypeProvider = KtFirExpressionTypeProvider(this, token)
     override val diagnosticProvider: KtDiagnosticProvider = KtFirDiagnosticProvider(this, token)
@@ -49,15 +56,19 @@ private constructor(
     override val typeProvider: KtTypeProvider = KtFirTypeProvider(this, token)
     override val subtypingComponent: KtSubtypingComponent = KtFirSubtypingComponent(this, token)
 
-    override fun createContextDependentCopy(): KtAnalysisSession {
+    override fun createContextDependentCopy(originalKtFile: KtFile, fakeKtElement: KtElement): KtAnalysisSession {
-        check(!isContextSession) { "Cannot create context-dependent copy of KtAnalysis session from a context dependent one" }
+        check(context == KtFirAnalysisSessionContext.DefaultContext) {
+            "Cannot create context-dependent copy of KtAnalysis session from a context dependent one"
+        }
         val contextResolveState = LowLevelFirApiFacadeForCompletion.getResolveStateForCompletion(firResolveState)
+        val originalFirFile = originalKtFile.getFirFile(firResolveState)
+        val context = KtFirAnalysisSessionContext.FakeFileContext(originalKtFile, originalFirFile, fakeKtElement, contextResolveState)
         return KtFirAnalysisSession(
             project,
             contextResolveState,
             firSymbolBuilder.createReadOnlyCopy(contextResolveState),
             token,
-            isContextSession = true
+            context
         )
     }
 
@@ -82,9 +93,31 @@ private constructor(
                 firResolveState,
                 firSymbolBuilder,
                 token,
-                isContextSession = false
+                KtFirAnalysisSessionContext.DefaultContext
             )
         }
     }
 }
 
+internal sealed class KtFirAnalysisSessionContext {
+    object DefaultContext : KtFirAnalysisSessionContext()
+
+    class FakeFileContext(
+        originalFile: KtFile,
+        firFile: FirFile,
+        fakeContextElement: KtElement,
+        fakeModuleResolveState: FirModuleResolveState
+    ) : KtFirAnalysisSessionContext() {
+        init {
+            require(!fakeContextElement.isPhysical)
+        }
+
+        val completionContext = run {
+            val enclosingContext = EnclosingDeclarationContext.detect(originalFile, fakeContextElement)
+            enclosingContext.buildCompletionContext(firFile, fakeModuleResolveState)
+        }
+
+        val fakeKtFile = fakeContextElement.containingKtFile
+    }
+}
+

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

@@ -92,13 +92,7 @@ internal class KtFirCompletionCandidateChecker(
         firFile: FirFile,
         fakeNameExpression: KtSimpleNameExpression
     ): Sequence<ImplicitReceiverValue<*>?> {
-        val enclosingContext = EnclosingDeclarationContext.detect(originalFile, fakeNameExpression)
+        val towerDataContext = analysisSession.towerDataContextProvider.getTowerDataContext(fakeNameExpression)
-
-        val completionContext = completionContextCache.computeIfAbsent(firFile to enclosingContext.fakeEnclosingDeclaration) {
-            enclosingContext.buildCompletionContext(firFile, firResolveState)
-        }
-
-        val towerDataContext = completionContext.getTowerDataContext(fakeNameExpression)
 
         return sequence {
             yield(null) // otherwise explicit receiver won't be checked when there are no implicit receivers in completion position

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

@@ -18,9 +18,9 @@ internal class KtFirExpressionHandlingComponent(
     override val analysisSession: KtFirAnalysisSession,
     override val token: ValidityToken,
 ) : KtExpressionHandlingComponent(), KtFirAnalysisSessionComponent {
-    override fun getReturnExpressionTargetSymbol(returnExpression: KtReturnExpression): KtFunctionLikeSymbol? {
+    override fun getReturnExpressionTargetSymbol(returnExpression: KtReturnExpression): KtCallableSymbol? {
         val fir = returnExpression.getOrBuildFirSafe<FirReturnExpression>(firResolveState) ?: return null
         val firTargetSymbol = fir.target.labeledElement
-        return firSymbolBuilder.buildCallableSymbol(firTargetSymbol) as KtFunctionLikeSymbol
+        return firSymbolBuilder.buildCallableSymbol(firTargetSymbol)
     }
 }

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

@@ -22,6 +22,7 @@ import org.jetbrains.kotlin.idea.frontend.api.components.KtBuiltinTypes
 import org.jetbrains.kotlin.idea.frontend.api.components.KtExpressionTypeProvider
 import org.jetbrains.kotlin.idea.frontend.api.fir.KtFirAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.fir.types.KtFirType
+import org.jetbrains.kotlin.idea.frontend.api.symbols.markers.KtTypedSymbol
 import org.jetbrains.kotlin.idea.frontend.api.types.KtType
 import org.jetbrains.kotlin.idea.frontend.api.withValidityAssertion
 import org.jetbrains.kotlin.psi.*
@@ -65,7 +66,7 @@ internal class KtFirExpressionTypeProvider(
     private fun getExpectedTypeByReturnExpression(expression: PsiElement): KtType? {
         val returnParent = expression.getReturnExpressionWithThisType() ?: return null
         val targetSymbol = with(analysisSession) { returnParent.getReturnTargetSymbol() } ?: return null
-        return targetSymbol.type
+        return (targetSymbol as? KtTypedSymbol)?.type
     }
 
     private fun PsiElement.getReturnExpressionWithThisType(): KtReturnExpression? =

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

@@ -42,13 +42,13 @@ import org.jetbrains.kotlin.psi.KtFile
 import java.util.*
 
 internal class KtFirScopeProvider(
-    analysisSession: KtAnalysisSession,
+    analysisSession: KtFirAnalysisSession,
     builder: KtSymbolByFirBuilder,
     private val project: Project,
     firResolveState: FirModuleResolveState,
     override val token: ValidityToken,
 ) : KtScopeProvider(), ValidityTokenOwner {
-    override val analysisSession: KtAnalysisSession by weakRef(analysisSession)
+    override val analysisSession: KtFirAnalysisSession by weakRef(analysisSession)
     private val builder by weakRef(builder)
     private val firResolveState by weakRef(firResolveState)
     private val firScopeStorage = FirScopeRegistry()
@@ -136,9 +136,7 @@ internal class KtFirScopeProvider(
         originalFile: KtFile,
         positionInFakeFile: KtElement
     ): KtScopeContext = withValidityAssertion {
-        val completionContext = buildCompletionContextForEnclosingDeclaration(originalFile, positionInFakeFile)
+        val towerDataContext = analysisSession.towerDataContextProvider.getTowerDataContext(positionInFakeFile)
-
-        val towerDataContext = completionContext.getTowerDataContext(positionInFakeFile)
 
         val implicitReceivers = towerDataContext.nonLocalTowerDataElements.mapNotNull { it.implicitReceiver }.distinct()
         val implicitReceiversTypes = implicitReceivers.map { builder.buildKtType(it.type) }
@@ -160,16 +158,6 @@ internal class KtFirScopeProvider(
         )
     }
 
-    private fun buildCompletionContextForEnclosingDeclaration(
-        ktFile: KtFile,
-        positionInFakeFile: KtElement
-    ): LowLevelFirApiFacadeForCompletion.FirCompletionContext {
-        val firFile = ktFile.getFirFile(firResolveState)
-        val declarationContext = EnclosingDeclarationContext.detect(ktFile, positionInFakeFile)
-
-        return declarationContext.buildCompletionContext(firFile, firResolveState)
-    }
-
     private fun convertToKtScope(firScope: FirScope): KtScope {
         firScopeStorage.register(firScope)
         return when (firScope) {

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

@@ -0,0 +1,19 @@
+/*
+ * 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.components
+
+import org.jetbrains.kotlin.fir.resolve.FirTowerDataContext
+import org.jetbrains.kotlin.idea.frontend.api.fir.KtFirAnalysisSession
+import org.jetbrains.kotlin.idea.frontend.api.fir.KtFirAnalysisSessionContext
+import org.jetbrains.kotlin.psi.KtElement
+
+internal class TowerDataContextProvider(private val analysisSession: KtFirAnalysisSession) {
+    fun getTowerDataContext(statement: KtElement): FirTowerDataContext {
+        val fakeContext =  analysisSession.context as? KtFirAnalysisSessionContext.FakeFileContext
+            ?: error("Getting data context for non context-dependent session is not supported yet")
+        return fakeContext.completionContext.getTowerDataContext(statement)
+    }
+}