FIR IDE: Introduce module resolve state for completion

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

@@ -19,10 +19,10 @@ internal interface FirIdeResolveStateService {
 
     val fallbackModificationTracker: ModificationTracker?
 
-    fun getResolveState(moduleInfo: IdeaModuleInfo): FirModuleResolveState
+    fun getResolveState(moduleInfo: IdeaModuleInfo): FirModuleResolveStateImpl
 }
 
-private class FirModuleData(val state: FirModuleResolveState, val modificationTracker: ModificationTracker?) {
+private class FirModuleData(val state: FirModuleResolveStateImpl, val modificationTracker: ModificationTracker?) {
     val modificationCount: Long = modificationTracker?.modificationCount ?: Long.MIN_VALUE
 
     fun isOutOfDate(): Boolean {
@@ -34,7 +34,7 @@ private class FirModuleData(val state: FirModuleResolveState, val modificationTr
 internal class FirIdeResolveStateServiceImpl(val project: Project) : FirIdeResolveStateService {
     private val stateCache = mutableMapOf<IdeaModuleInfo, FirModuleData>()
 
-    private fun createResolveState(): FirModuleResolveState {
+    private fun createResolveState(): FirModuleResolveStateImpl {
         val provider = FirProjectSessionProvider(project)
         return FirModuleResolveStateImpl(provider)
     }
@@ -47,7 +47,7 @@ internal class FirIdeResolveStateServiceImpl(val project: Project) : FirIdeResol
     }
 
     // TODO: multi thread protection
-    override fun getResolveState(moduleInfo: IdeaModuleInfo): FirModuleResolveState {
+    override fun getResolveState(moduleInfo: IdeaModuleInfo): FirModuleResolveStateImpl {
         var moduleData = stateCache.getOrPut(moduleInfo) {
             createModuleData()
         }

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

@@ -60,7 +60,7 @@ abstract class FirModuleResolveState {
         }
     }
 
-    internal abstract operator fun get(psi: KtElement): FirElement?
+    internal abstract fun getCachedMapping(psi: KtElement): FirElement?
 
     internal abstract fun getDiagnostics(psi: KtElement): List<Diagnostic>
 
@@ -73,14 +73,34 @@ abstract class FirModuleResolveState {
     internal abstract fun setDiagnosticsForFile(file: KtFile, fir: FirFile, diagnostics: Iterable<FirDiagnostic<*>> = emptyList())
 }
 
-internal class FirModuleResolveStateImpl(override val sessionProvider: FirProjectSessionProvider) : FirModuleResolveState() {
+internal abstract class PsiToFirCache {
+    abstract operator fun get(psi: KtElement): FirElement?
+    abstract operator fun set(psi: KtElement, fir: FirElement)
+    abstract fun remove(psi: PsiElement)
+}
+
+internal class PsiToFirCacheImpl : PsiToFirCache() {
     private val cache = mutableMapOf<KtElement, FirElement>()
 
+    override fun get(psi: KtElement): FirElement? = cache[psi]
+
+    override fun set(psi: KtElement, fir: FirElement) {
+        cache[psi] = fir
+    }
+
+    override fun remove(psi: PsiElement) {
+        cache.remove(psi)
+    }
+}
+
+open class FirModuleResolveStateImpl(override val sessionProvider: FirProjectSessionProvider) : FirModuleResolveState() {
+    internal open val cache: PsiToFirCache = PsiToFirCacheImpl()
+
     private val diagnosticCache = mutableMapOf<KtElement, MutableList<Diagnostic>>()
 
     private val diagnosedFiles = mutableMapOf<KtFile, Long>()
 
-    override fun get(psi: KtElement): FirElement? = cache[psi]
+    override fun getCachedMapping(psi: KtElement): FirElement? = cache[psi]
 
     override fun getDiagnostics(psi: KtElement): List<Diagnostic> {
         return diagnosticCache[psi] ?: emptyList()

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

@@ -0,0 +1,31 @@
+/*
+ * 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.fir.low.level.api
+
+import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.psi.KtElement
+
+class FirModuleResolveStateForCompletion(
+    mainState: FirModuleResolveStateImpl
+) : FirModuleResolveStateImpl(mainState.sessionProvider) {
+    override val cache: PsiToFirCache = PsiToFirCacheForCompletion(mainState.cache)
+}
+
+private class PsiToFirCacheForCompletion(private val delegate: PsiToFirCache) : PsiToFirCache() {
+    private val cache = mutableMapOf<KtElement, FirElement>()
+
+    override fun get(psi: KtElement): FirElement? =
+        cache[psi] ?: delegate[psi]
+
+    override fun set(psi: KtElement, fir: FirElement) {
+        cache[psi] = fir
+    }
+
+    override fun remove(psi: PsiElement) {
+        cache.remove(psi)
+    }
+}

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

@@ -20,7 +20,6 @@ import org.jetbrains.kotlin.fir.scopes.impl.FirPackageMemberScope
 import org.jetbrains.kotlin.fir.symbols.CallableId
 import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
 import org.jetbrains.kotlin.fir.types.FirErrorTypeRef
-import org.jetbrains.kotlin.fir.types.FirTypeRef
 import org.jetbrains.kotlin.fir.types.FirUserTypeRef
 import org.jetbrains.kotlin.fir.visitors.CompositeTransformResult
 import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
@@ -255,7 +254,7 @@ internal fun KtElement.getOrBuildFir(
         is KtPropertyDelegate -> this.expression ?: this
         else -> this
     }
-    return state[this] ?: run {
+    return state.getCachedMapping(this) ?: run {
         containerFir.accept(object : FirVisitorVoid() {
             override fun visitElement(element: FirElement) {
                 (element.realPsi as? KtElement)?.let {
@@ -288,7 +287,7 @@ internal fun KtElement.getOrBuildFir(
         })
         var current: PsiElement? = psi
         while (current is KtElement) {
-            val mappedFir = state[current]
+            val mappedFir = state.getCachedMapping(current)
             if (mappedFir != null) {
                 if (current != this) {
                     state.record(current, mappedFir)

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

@@ -15,6 +15,10 @@ object LowLevelFirApiFacade {
     fun getResolveStateFor(element: KtElement): FirModuleResolveState =
         element.firResolveState()
 
+    fun getResolveStateForCompletion(element: KtElement, mainState: FirModuleResolveStateImpl): FirModuleResolveStateForCompletion {
+        return FirModuleResolveStateForCompletion(mainState)
+    }
+
     fun getSessionFor(element: KtElement, resolveState: FirModuleResolveState): FirSession =
         resolveState.getSession(element)