FIR IDE: Modify KtFirAnalysisSession for completion

- This is WIP solution for having a specialized `KtFirAnalysisSession`
for the completion

Co-authored-by: Ilya Kirillov <ilya.kirillov@jetbrains.com>

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

@@ -35,4 +35,6 @@ abstract class KtAnalysisSession(project: Project) : ValidityOwner {
 
     abstract fun resolveCall(call: KtCallExpression): CallInfo?
     abstract fun resolveCall(call: KtBinaryExpression): CallInfo?
+
+    abstract fun analyzeInContext(): KtAnalysisSession
 }

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

@@ -8,10 +8,19 @@ package org.jetbrains.kotlin.idea.frontend.api.scopes
 import org.jetbrains.kotlin.idea.frontend.api.ValidityOwner
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtClassOrObjectSymbol
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtPackageSymbol
+import org.jetbrains.kotlin.idea.frontend.api.types.KtType
+import org.jetbrains.kotlin.psi.KtElement
+import org.jetbrains.kotlin.psi.KtFile
 
 abstract class KtScopeProvider : ValidityOwner {
     abstract fun getMemberScope(classSymbol: KtClassOrObjectSymbol): KtMemberScope
     abstract fun getDeclaredMemberScope(classSymbol: KtClassOrObjectSymbol): KtDeclaredMemberScope
     abstract fun getPackageScope(packageSymbol: KtPackageSymbol): KtPackageScope
     abstract fun getCompositeScope(subScopes: List<KtScope>): KtCompositeScope
-}
+
+    abstract fun getScopeForType(type: KtType): KtScope?
+
+    abstract fun getScopeContextForPosition(originalFile: KtFile, positionInFakeFile: KtElement): KtScopeContext
+}
+
+data class KtScopeContext(val scopes: KtCompositeScope, val implicitReceiversTypes: List<KtType>)

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

@@ -63,9 +63,9 @@ object LowLevelFirApiFacade {
     fun buildCompletionContextForFunction(
         firFile: FirFile,
         element: KtNamedFunction,
+        state: FirModuleResolveState,
         phase: FirResolvePhase = FirResolvePhase.BODY_RESOLVE
     ): FirCompletionContext {
-        val state = element.firResolveState()
         val firIdeProvider = firFile.session.firIdeProvider
         val builtFunction = firIdeProvider.buildFunctionWithBody(element)
         val towerDataContextForStatement = mutableMapOf<FirStatement, FirTowerDataContext>()

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

@@ -55,7 +55,11 @@ private object KotlinFirCompletionProvider : CompletionProvider<CompletionParame
 
         val parentFunction = nameExpression.getNonStrictParentOfType<KtNamedFunction>() ?: return
 
-        val completionContext = LowLevelFirApiFacade.buildCompletionContextForFunction(originalFileFir, parentFunction)
+        val completionContext = LowLevelFirApiFacade.buildCompletionContextForFunction(
+            originalFileFir,
+            parentFunction,
+            LowLevelFirApiFacade.getResolveStateFor(originalFile)
+        )
 
         val element = nameExpression.getFirOfClosestParent() as? FirQualifiedAccessExpression ?: return
         val towerDataContext = completionContext.getTowerDataContext(nameExpression)

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

@@ -19,6 +19,8 @@ import org.jetbrains.kotlin.fir.symbols.impl.FirConstructorSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.idea.fir.*
+import org.jetbrains.kotlin.idea.fir.low.level.api.FirModuleResolveState
+import org.jetbrains.kotlin.idea.fir.low.level.api.FirModuleResolveStateImpl
 import org.jetbrains.kotlin.idea.fir.low.level.api.LowLevelFirApiFacade
 import org.jetbrains.kotlin.idea.frontend.api.*
 import org.jetbrains.kotlin.idea.frontend.api.fir.scopes.KtFirScopeProvider
@@ -38,11 +40,12 @@ import org.jetbrains.kotlin.psi.*
 class KtFirAnalysisSession
 @Deprecated("Please use org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSessionProviderKt.analyze")
 constructor(
-    private val element: KtElement
+    private val element: KtElement,
+    resolveState: FirModuleResolveState
 ) : KtAnalysisSession(element.project) {
     internal val token get() = validityToken
 
-    internal val firResolveState = LowLevelFirApiFacade.getResolveStateFor(element)
+    internal val firResolveState = resolveState
 
     internal val firSession get() = LowLevelFirApiFacade.getSessionFor(element, firResolveState)
 
@@ -67,7 +70,7 @@ constructor(
             firSymbolBuilder
         )
 
-    override val scopeProvider: KtScopeProvider = KtFirScopeProvider(token, firSymbolBuilder, firSession)
+    override val scopeProvider: KtScopeProvider = KtFirScopeProvider(token, firSymbolBuilder, firSession, firResolveState)
 
     init {
         assertIsValid()
@@ -129,6 +132,13 @@ constructor(
         resolveCall(firCall, call)
     }
 
+    override fun analyzeInContext(): KtAnalysisSession {
+        return KtFirAnalysisSession(
+            element,
+            LowLevelFirApiFacade.getResolveStateForCompletion(element, firResolveState as FirModuleResolveStateImpl)
+        )
+    }
+
     override fun resolveCall(call: KtCallExpression): CallInfo? = withValidityAssertion {
         val firCall = when (val fir = call.getOrBuildFir(firResolveState)) {
             is FirFunctionCall -> fir

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.LowLevelFirApiFacade
 import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSessionProvider
 import org.jetbrains.kotlin.psi.KtElement
@@ -28,6 +29,6 @@ class KtFirAnalysisSessionProvider(project: Project) : KtAnalysisSessionProvider
     @Suppress("DEPRECATION")
     override fun getAnalysisSessionFor(contextElement: KtElement): KtAnalysisSession =
         analysisSessionByModuleInfoCache.value.getOrPut(contextElement.getModuleInfo()) {
-            KtFirAnalysisSession(contextElement)
+            KtFirAnalysisSession(contextElement, LowLevelFirApiFacade.getResolveStateFor(contextElement))
         }
 }

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

@@ -6,22 +6,37 @@
 package org.jetbrains.kotlin.idea.frontend.api.fir.scopes
 
 import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.FirFile
+import org.jetbrains.kotlin.fir.resolve.ScopeSession
+import org.jetbrains.kotlin.fir.resolve.scope
+import org.jetbrains.kotlin.fir.scopes.FirScope
+import org.jetbrains.kotlin.idea.fir.getOrBuildFirOfType
+import org.jetbrains.kotlin.idea.fir.low.level.api.FirModuleResolveState
+import org.jetbrains.kotlin.idea.fir.low.level.api.LowLevelFirApiFacade
 import org.jetbrains.kotlin.idea.frontend.api.ValidityOwner
 import org.jetbrains.kotlin.idea.frontend.api.ValidityOwnerByValidityToken
 import org.jetbrains.kotlin.idea.frontend.api.fir.KtSymbolByFirBuilder
 import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.KtFirClassOrObjectSymbol
 import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.KtFirPackageSymbol
+import org.jetbrains.kotlin.idea.frontend.api.fir.types.KtFirType
 import org.jetbrains.kotlin.idea.frontend.api.fir.utils.weakRef
 import org.jetbrains.kotlin.idea.frontend.api.scopes.*
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtClassOrObjectSymbol
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtPackageSymbol
+import org.jetbrains.kotlin.idea.frontend.api.types.KtType
 import org.jetbrains.kotlin.idea.frontend.api.withValidityAssertion
+import org.jetbrains.kotlin.idea.util.getElementTextInContext
+import org.jetbrains.kotlin.psi.KtElement
+import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.psi.KtNamedFunction
+import org.jetbrains.kotlin.psi.psiUtil.getNonStrictParentOfType
 import java.util.*
 
 internal class KtFirScopeProvider(
     override val token: ValidityOwner,
     private val builder: KtSymbolByFirBuilder,
-    session: FirSession
+    session: FirSession,
+    val firResolveState: FirModuleResolveState
 ) : KtScopeProvider(), ValidityOwnerByValidityToken {
     private val session by weakRef(session)
 
@@ -53,4 +68,43 @@ internal class KtFirScopeProvider(
     override fun getCompositeScope(subScopes: List<KtScope>): KtCompositeScope = withValidityAssertion {
         KtFirCompositeScope(subScopes, token)
     }
+
+    override fun getScopeForType(type: KtType): KtScope? {
+        check(type is KtFirType) { "KtFirScopePriovider can only work with KtFirType, but ${type::class} was provided" }
+
+        val firTypeScope = type.coneType.scope(session, ScopeSession()) ?: return null
+        return convertToKtScope(firTypeScope)
+    }
+
+    override fun getScopeContextForPosition(originalFile: KtFile, positionInFakeFile: KtElement): KtScopeContext = withValidityAssertion {
+        val originalFirFile = originalFile.getOrBuildFirOfType<FirFile>(firResolveState)
+        val fakeEnclosingFunction = positionInFakeFile.getNonStrictParentOfType<KtNamedFunction>()
+            ?: error("Cannot find enclosing function for ${positionInFakeFile.getElementTextInContext()}")
+
+        val completionContext = LowLevelFirApiFacade.buildCompletionContextForFunction(
+            originalFirFile,
+            fakeEnclosingFunction,
+            state = firResolveState
+        )
+
+        val towerDataContext = completionContext.getTowerDataContext(positionInFakeFile)
+
+        val implicitReceivers = towerDataContext.nonLocalTowerDataElements.mapNotNull { it.implicitReceiver }
+        val implicitReceiversTypes = implicitReceivers.map { builder.buildKtType(it.type) }
+
+        val implicitReceiverScopes = implicitReceivers.mapNotNull { it.implicitScope }
+        val nonLocalScopes = towerDataContext.nonLocalTowerDataElements.mapNotNull { it.scope }
+        val firLocalScopes = towerDataContext.localScopes
+
+        val allKtScopes = (implicitReceiverScopes + nonLocalScopes + firLocalScopes).map { convertToKtScope(it) }
+
+        KtScopeContext(getCompositeScope(allKtScopes), implicitReceiversTypes)
+    }
+
+    private fun convertToKtScope(firScope: FirScope): KtScope = KtFirDelegatingScopeImpl(firScope, builder, token)
+}
+
+private class KtFirDelegatingScopeImpl(firScope: FirScope, builder: KtSymbolByFirBuilder, override val token: ValidityOwner) :
+    KtFirDelegatingScope(builder), ValidityOwnerByValidityToken {
+    override val firScope: FirScope by weakRef(firScope)
 }

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

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.idea.frontend.api
 
 import com.intellij.openapi.util.io.FileUtil
 import org.jetbrains.kotlin.idea.executeOnPooledThreadInReadAction
+import org.jetbrains.kotlin.idea.fir.low.level.api.LowLevelFirApiFacade
 import org.jetbrains.kotlin.idea.frontend.api.fir.KtFirAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.symbols.DebugSymbolRenderer
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtSymbol
@@ -36,7 +37,7 @@ abstract class AbstractSymbolByFqNameTest : KotlinLightCodeInsightFixtureTestCas
                 val symbols = createSymbols(symbolData, this)
                 symbols.map { DebugSymbolRenderer.render(it) }
             }
-            val analysisSession = KtFirAnalysisSession(fakeKtFile)
+            val analysisSession = KtFirAnalysisSession(fakeKtFile, LowLevelFirApiFacade.getResolveStateFor(fakeKtFile))
             val symbols = createSymbols(symbolData, analysisSession)
             symbols.map { DebugSymbolRenderer.render(it) }
         }