FIR IDE: Add FirTowerDataContextCollector class

- This class just encapsulates the map that is used to collect scopes

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

@@ -7,17 +7,16 @@ package org.jetbrains.kotlin.idea.fir.low.level.api
 
 import org.jetbrains.kotlin.diagnostics.Diagnostic
 import org.jetbrains.kotlin.fir.FirElement
-import org.jetbrains.kotlin.fir.declarations.FirDeclaration
-import org.jetbrains.kotlin.fir.declarations.FirFunction
-import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.FirDeclaration
+import org.jetbrains.kotlin.fir.declarations.FirFile
+import org.jetbrains.kotlin.fir.declarations.FirFunction
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
-import org.jetbrains.kotlin.fir.expressions.FirStatement
-import org.jetbrains.kotlin.fir.resolve.FirTowerDataContext
 import org.jetbrains.kotlin.idea.caches.project.IdeaModuleInfo
 import org.jetbrains.kotlin.idea.caches.project.getModuleInfo
 import org.jetbrains.kotlin.idea.fir.low.level.api.diagnostics.DiagnosticsCollector
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirElementBuilder
+import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirTowerDataContextCollector
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.PsiToFirCache
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.FirFileBuilder
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache
@@ -52,7 +51,7 @@ abstract class FirModuleResolveState {
         containerFirFile: FirFile,
         firIdeProvider: FirIdeProvider,
         toPhase: FirResolvePhase,
-        towerDataContextForStatement: MutableMap<FirStatement, FirTowerDataContext>
+        towerDataContextCollector: FirTowerDataContextCollector
     )
 
 }
@@ -97,14 +96,14 @@ internal open class FirModuleResolveStateImpl(
         containerFirFile: FirFile,
         firIdeProvider: FirIdeProvider,
         toPhase: FirResolvePhase,
-        towerDataContextForStatement: MutableMap<FirStatement, FirTowerDataContext>
+        towerDataContextCollector: FirTowerDataContextCollector
     ) {
         firLazyDeclarationResolver.runLazyResolveWithoutLock(
             firFunction,
             containerFirFile,
             firIdeProvider,
             toPhase,
-            towerDataContextForStatement,
+            towerDataContextCollector,
             checkPCE = false
         )
     }

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

@@ -12,9 +12,8 @@ import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.declarations.FirFunction
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
-import org.jetbrains.kotlin.fir.expressions.FirStatement
-import org.jetbrains.kotlin.fir.resolve.FirTowerDataContext
 import org.jetbrains.kotlin.idea.caches.project.IdeaModuleInfo
+import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirTowerDataContextCollector
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.PsiToFirCache
 import org.jetbrains.kotlin.idea.fir.low.level.api.providers.FirIdeProvider
 import org.jetbrains.kotlin.psi.KtElement
@@ -61,9 +60,9 @@ internal class FirModuleResolveStateForCompletion(
         containerFirFile: FirFile,
         firIdeProvider: FirIdeProvider,
         toPhase: FirResolvePhase,
-        towerDataContextForStatement: MutableMap<FirStatement, FirTowerDataContext>
+        towerDataContextCollector: FirTowerDataContextCollector
     ) {
-        originalState.lazyResolveFunctionForCompletion(firFunction, containerFirFile, firIdeProvider, toPhase, towerDataContextForStatement)
+        originalState.lazyResolveFunctionForCompletion(firFunction, containerFirFile, firIdeProvider, toPhase, towerDataContextCollector)
     }
 
     override fun getDiagnostics(element: KtElement): List<Diagnostic> {

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

@@ -19,8 +19,8 @@ import org.jetbrains.kotlin.idea.caches.project.getModuleInfo
 import org.jetbrains.kotlin.idea.fir.low.level.api.providers.firIdeProvider
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtNamedFunction
-import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
+import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirTowerDataContextCollector
 import org.jetbrains.kotlin.idea.util.getElementTextInContext
 
 object LowLevelFirApiFacade {
@@ -40,7 +40,7 @@ object LowLevelFirApiFacade {
 
     class FirCompletionContext internal constructor(
         val session: FirSession,
-        private val towerDataContextForStatement: Map<FirStatement, FirTowerDataContext>,
+        private val towerDataContextCollector: FirTowerDataContextCollector,
         private val state: FirModuleResolveState,
     ) {
         fun getTowerDataContext(element: KtElement): FirTowerDataContext {
@@ -49,13 +49,7 @@ object LowLevelFirApiFacade {
                 val mappedFir = state.getCachedMappingForCompletion(current)
 
                 if (mappedFir is FirStatement) {
-                    towerDataContextForStatement[mappedFir]?.let { return it }
+                    towerDataContextCollector.getContext(mappedFir)?.let { return it }
-                    for ((key, value) in towerDataContextForStatement) {
-                        // TODO Rework that
-                        if (key.psi == mappedFir.psi) {
-                            return value
-                        }
-                    }
                 }
                 current = current.parent
             }
@@ -74,7 +68,7 @@ object LowLevelFirApiFacade {
         val firIdeProvider = firFile.session.firIdeProvider
         val originalFunction = state.getOrBuildFirFor(originalElement, phase) as FirSimpleFunction
         val builtFunction = firIdeProvider.buildFunctionWithBody(element)
-        val towerDataContextForStatement = mutableMapOf<FirStatement, FirTowerDataContext>()
+        val contextCollector = FirTowerDataContextCollector()
 
         // right now we can't resolve builtFunction header properly, as it built right in air,
         // without file, which is now required for running stages other then body resolve, so we
@@ -87,13 +81,13 @@ object LowLevelFirApiFacade {
         }
 
         val function = frankensteinFunction.apply {
-            state.lazyResolveFunctionForCompletion(this, firFile, firIdeProvider, phase, towerDataContextForStatement)
+            state.lazyResolveFunctionForCompletion(this, firFile, firIdeProvider, phase, contextCollector)
             state.recordPsiToFirMappingsForCompletionFrom(this, firFile, element.containingKtFile)
         }
 
         return FirCompletionContext(
             function.session,
-            towerDataContextForStatement,
+            contextCollector,
             state
         )
     }

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/element/builder/FirDesignatedBodyResolveTransformerForIDE.kt

@@ -10,7 +10,6 @@ import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.expressions.FirStatement
-import org.jetbrains.kotlin.fir.resolve.FirTowerDataContext
 import org.jetbrains.kotlin.fir.resolve.ResolutionMode
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirBodyResolveTransformer
@@ -23,7 +22,7 @@ internal class FirDesignatedBodyResolveTransformerForIDE(
     session: FirSession,
     scopeSession: ScopeSession,
     implicitTypeOnly: Boolean,
-    private val towerDataContextForStatement: MutableMap<FirStatement, FirTowerDataContext>? = null
+    private val towerDataContextCollector: FirTowerDataContextCollector? = null
 ) : FirBodyResolveTransformer(
     session,
     phase = FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE,
@@ -42,7 +41,6 @@ internal class FirDesignatedBodyResolveTransformerForIDE(
     }
 
     override fun onBeforeStatementResolution(statement: FirStatement) {
-        if (towerDataContextForStatement == null) return
+        towerDataContextCollector?.addStatementContext(statement, context.towerDataContext)
-        towerDataContextForStatement[statement] = context.towerDataContext
     }
 }

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/element/builder/FirTowerDataContextCollector.kt

@@ -0,0 +1,30 @@
+/*
+ * 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.element.builder
+
+import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.expressions.FirStatement
+import org.jetbrains.kotlin.fir.psi
+import org.jetbrains.kotlin.fir.resolve.FirTowerDataContext
+
+internal class FirTowerDataContextCollector {
+    private val state: MutableMap<FirElement, FirTowerDataContext> = mutableMapOf()
+
+    fun addStatementContext(statement: FirStatement, context: FirTowerDataContext) {
+        state[statement] = context
+    }
+
+    fun getContext(statement: FirElement): FirTowerDataContext? {
+        this.state[statement]?.let { return it }
+        for ((key, value) in this.state) {
+            // TODO Rework that
+            if (key.psi == statement.psi) {
+                return value
+            }
+        }
+        return null
+    }
+}

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/lazy/resolve/FirLazyDeclarationResolver.kt

@@ -6,13 +6,12 @@
 package org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve
 
 import org.jetbrains.kotlin.fir.declarations.*
-import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.render
-import org.jetbrains.kotlin.fir.resolve.FirTowerDataContext
 import org.jetbrains.kotlin.fir.resolve.ResolutionMode
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.providers.FirProvider
+import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirTowerDataContextCollector
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirDesignatedBodyResolveTransformerForIDE
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.getNonLocalContainingDeclarationWithFqName
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.FirFileBuilder
@@ -32,7 +31,7 @@ internal class FirLazyDeclarationResolver(
         declaration: FirDeclaration,
         moduleFileCache: ModuleFileCache,
         toPhase: FirResolvePhase,
-        towerDataContextForStatement: MutableMap<FirStatement, FirTowerDataContext>? = null,
+        towerDataContextCollector: FirTowerDataContextCollector? = null,
         checkPCE: Boolean = false
     ) {
         if (declaration.resolvePhase >= toPhase) return
@@ -41,12 +40,12 @@ internal class FirLazyDeclarationResolver(
         val provider = firFile.session.firIdeProvider
         if (checkPCE) {
             firFileBuilder.runCustomResolveWithPCECheck(firFile, moduleFileCache) {
-                runLazyResolveWithoutLock(declaration, firFile, provider, toPhase, towerDataContextForStatement, checkPCE = true)
+                runLazyResolveWithoutLock(declaration, firFile, provider, toPhase, towerDataContextCollector, checkPCE = true)
             }
         } else {
             firFileBuilder.runCustomResolveUnderLock(firFile, moduleFileCache) {
                 executeWithoutPCE {
-                    runLazyResolveWithoutLock(declaration, firFile, provider, toPhase, towerDataContextForStatement, checkPCE = false)
+                    runLazyResolveWithoutLock(declaration, firFile, provider, toPhase, towerDataContextCollector, checkPCE = false)
                 }
             }
         }
@@ -57,7 +56,7 @@ internal class FirLazyDeclarationResolver(
         containerFirFile: FirFile,
         provider: FirProvider,
         toPhase: FirResolvePhase,
-        towerDataContextForStatement: MutableMap<FirStatement, FirTowerDataContext>? = null,
+        towerDataContextCollector: FirTowerDataContextCollector? = null,
         checkPCE: Boolean
     ) {
         val nonLazyPhase = minOf(toPhase, FirResolvePhase.DECLARATIONS)
@@ -71,7 +70,7 @@ internal class FirLazyDeclarationResolver(
         }
         if (toPhase <= nonLazyPhase) return
         if (checkPCE) checkCanceled()
-        runLazyResolvePhase(firDeclarationToResolve, containerFirFile, provider, toPhase, towerDataContextForStatement)
+        runLazyResolvePhase(firDeclarationToResolve, containerFirFile, provider, toPhase, towerDataContextCollector)
     }
 
     private fun runLazyResolvePhase(
@@ -79,7 +78,7 @@ internal class FirLazyDeclarationResolver(
         containerFirFile: FirFile,
         provider: FirProvider,
         toPhase: FirResolvePhase,
-        towerDataContextForStatement: MutableMap<FirStatement, FirTowerDataContext>?,
+        towerDataContextCollector: FirTowerDataContextCollector?,
     ) {
         val nonLocalDeclarationToResolve = firDeclarationToResolve.getNonLocalDeclarationToResolve(provider)
 
@@ -110,7 +109,7 @@ internal class FirLazyDeclarationResolver(
             designation.iterator(), containerFirFile.session,
             scopeSession,
             implicitTypeOnly = toPhase == FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE,
-            towerDataContextForStatement
+            towerDataContextCollector
         )
         containerFirFile.transform<FirFile, ResolutionMode>(transformer, ResolutionMode.ContextDependent)
     }