[FIR IDE] Collect tower contexts for local declarations

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirBodyResolveTransformer.kt

@@ -30,7 +30,8 @@ open class FirBodyResolveTransformer(
     override var implicitTypeOnly: Boolean,
     scopeSession: ScopeSession,
     val returnTypeCalculator: ReturnTypeCalculator = ReturnTypeCalculatorForFullBodyResolve(),
-    outerBodyResolveContext: BodyResolveContext? = null
+    outerBodyResolveContext: BodyResolveContext? = null,
+    val firTowerDataContextCollector: FirTowerDataContextCollector? = null
 ) : FirAbstractBodyResolveTransformer(phase) {
 
     final override val context: BodyResolveContext =
@@ -47,7 +48,7 @@ open class FirBodyResolveTransformer(
     override fun transformFile(file: FirFile, data: ResolutionMode): FirFile {
         checkSessionConsistency(file)
         return context.withFile(file, components) {
-            onBeforeFileContentResolution(file)
+            firTowerDataContextCollector?.addFileContext(file, context.towerDataContext)
 
             file.replaceResolvePhase(transformerPhase)
             @Suppress("UNCHECKED_CAST")
@@ -77,12 +78,6 @@ open class FirBodyResolveTransformer(
         return data.expectedTypeRef
     }
 
-    open fun onBeforeFileContentResolution(file: FirFile) {}
-
-    open fun onBeforeStatementResolution(statement: FirStatement) {}
-
-    open fun onBeforeDeclarationContentResolve(declaration: FirDeclaration) {}
-
     // ------------------------------------- Expressions -------------------------------------
 
     override fun transformExpression(expression: FirExpression, data: ResolutionMode): FirStatement {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirDeclarationsResolveTransformer.kt

@@ -64,7 +64,7 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
     private fun transformDeclarationContent(
         declaration: FirDeclaration, data: ResolutionMode
     ): FirDeclaration {
-        transformer.onBeforeDeclarationContentResolve(declaration)
+        transformer.firTowerDataContextCollector?.addDeclarationContext(declaration, context.towerDataContext)
         transformer.replaceDeclarationResolvePhaseIfNeeded(declaration, transformerPhase)
         return transformer.transformDeclarationContent(declaration, data)
     }
@@ -351,7 +351,7 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
 
     override fun transformRegularClass(regularClass: FirRegularClass, data: ResolutionMode): FirStatement {
         if (regularClass.isLocal && regularClass !in context.targetedLocalClasses) {
-            return regularClass.runAllPhasesForLocalClass(transformer, components, data)
+            return regularClass.runAllPhasesForLocalClass(transformer, components, data, transformer.firTowerDataContextCollector)
         }
 
         doTransformTypeParameters(regularClass)
@@ -360,11 +360,11 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
 
     override fun transformTypeAlias(typeAlias: FirTypeAlias, data: ResolutionMode): FirDeclaration {
         if (typeAlias.isLocal && typeAlias !in context.targetedLocalClasses) {
-            return typeAlias.runAllPhasesForLocalClass(transformer, components, data)
+            return typeAlias.runAllPhasesForLocalClass(transformer, components, data, transformer.firTowerDataContextCollector)
         }
         doTransformTypeParameters(typeAlias)
         typeAlias.transformAnnotations(transformer, data)
-        transformer.onBeforeDeclarationContentResolve(typeAlias)
+        transformer.firTowerDataContextCollector?.addDeclarationContext(typeAlias, context.towerDataContext)
         typeAlias.transformExpandedTypeRef(transformer, data)
         return typeAlias
     }
@@ -400,7 +400,7 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
         data: ResolutionMode
     ): FirStatement {
         if (anonymousObject !in context.targetedLocalClasses) {
-            return anonymousObject.runAllPhasesForLocalClass(transformer, components, data)
+            return anonymousObject.runAllPhasesForLocalClass(transformer, components, data, transformer.firTowerDataContextCollector)
         }
         dataFlowAnalyzer.enterClass()
         if (anonymousObject.typeRef !is FirResolvedTypeRef) {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirExpressionsResolveTransformer.kt

@@ -319,7 +319,7 @@ open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransform
 
         block.transformStatementsIndexed(transformer) { index ->
             val value = if (index == numberOfStatements - 1) data else ResolutionMode.ContextIndependent
-            transformer.onBeforeStatementResolution(block.statements[index])
+            transformer.firTowerDataContextCollector?.addStatementContext(block.statements[index], context.towerDataContext)
             TransformData.Data(value)
         }
         block.transformOtherChildren(transformer, data)

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirImplicitBodyResolve.kt

@@ -58,6 +58,7 @@ fun <F : FirClassLikeDeclaration<F>> F.runContractAndBodiesResolutionForLocalCla
     components: FirAbstractBodyResolveTransformer.BodyResolveTransformerComponents,
     resolutionMode: ResolutionMode,
     localClassesNavigationInfo: LocalClassesNavigationInfo,
+    firTowerDataContextCollector: FirTowerDataContextCollector? = null
 ): F {
     val (designationMap, targetedClasses) = localClassesNavigationInfo.run {
         designationMap to parentForClass.keys + this@runContractAndBodiesResolutionForLocalClass
@@ -85,7 +86,8 @@ fun <F : FirClassLikeDeclaration<F>> F.runContractAndBodiesResolutionForLocalCla
         FirResolvePhase.BODY_RESOLVE,
         implicitTypeOnly = false,
         returnTypeCalculator,
-        outerBodyResolveContext = newContext
+        outerBodyResolveContext = newContext,
+        firTowerDataContextCollector = firTowerDataContextCollector
     )
 
     val graphBuilder = components.context.dataFlowAnalyzerContext.graphBuilder
@@ -124,14 +126,16 @@ open class FirImplicitAwareBodyResolveTransformer(
     phase: FirResolvePhase,
     implicitTypeOnly: Boolean,
     returnTypeCalculator: ReturnTypeCalculator,
-    outerBodyResolveContext: BodyResolveContext? = null
+    outerBodyResolveContext: BodyResolveContext? = null,
+    firTowerDataContextCollector: FirTowerDataContextCollector? = null,
 ) : FirBodyResolveTransformer(
     session,
     phase,
     implicitTypeOnly,
     scopeSession,
     returnTypeCalculator,
-    outerBodyResolveContext
+    outerBodyResolveContext,
+    firTowerDataContextCollector
 ) {
     override fun transformSimpleFunction(
         simpleFunction: FirSimpleFunction,

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirTowerDataContextCollector.kt

@@ -0,0 +1,17 @@
+/*
+ * Copyright 2010-2021 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.fir.resolve.transformers.body.resolve
+
+import org.jetbrains.kotlin.fir.declarations.FirDeclaration
+import org.jetbrains.kotlin.fir.declarations.FirFile
+import org.jetbrains.kotlin.fir.expressions.FirStatement
+import org.jetbrains.kotlin.fir.resolve.FirTowerDataContext
+
+interface FirTowerDataContextCollector {
+    fun addFileContext(file: FirFile, context: FirTowerDataContext)
+    fun addStatementContext(statement: FirStatement, context: FirTowerDataContext)
+    fun addDeclarationContext(declaration: FirDeclaration, context: FirTowerDataContext)
+}

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/LocalClassesResolution.kt

@@ -17,7 +17,8 @@ import org.jetbrains.kotlin.fir.scopes.impl.createCurrentScopeList
 fun <F : FirClassLikeDeclaration<F>> F.runAllPhasesForLocalClass(
     transformer: FirAbstractBodyResolveTransformer,
     components: FirAbstractBodyResolveTransformer.BodyResolveTransformerComponents,
-    resolutionMode: ResolutionMode
+    resolutionMode: ResolutionMode,
+    firTowerDataContextCollector: FirTowerDataContextCollector?
 ): F {
     if (this.resolvePhase > FirResolvePhase.RAW_FIR) return this
     if (this is FirRegularClass) {
@@ -48,6 +49,11 @@ fun <F : FirClassLikeDeclaration<F>> F.runAllPhasesForLocalClass(
         components.createCurrentScopeList(),
         localClassesNavigationInfo
     )
-    runContractAndBodiesResolutionForLocalClass(components, resolutionMode, localClassesNavigationInfo)
+    runContractAndBodiesResolutionForLocalClass(
+        components,
+        resolutionMode,
+        localClassesNavigationInfo,
+        firTowerDataContextCollector
+    )
     return this
 }

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

@@ -15,11 +15,11 @@ 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
+import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirTowerDataContextCollector
 import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
 import org.jetbrains.kotlin.idea.fir.low.level.api.FirModuleResolveStateDepended
 import org.jetbrains.kotlin.idea.fir.low.level.api.FirModuleResolveStateImpl
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirTowerContextProvider
-import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirTowerDataContextCollector
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirTowerDataContextAllElementsCollector
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.SingleElementTowerProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.structure.FirElementsRecorder

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

@@ -11,6 +11,7 @@ import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.resolve.FirTowerDataContext
+import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirTowerDataContextCollector
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.originalDeclaration
 import org.jetbrains.kotlin.psi.*
 
@@ -18,12 +19,6 @@ interface FirTowerContextProvider {
     fun getClosestAvailableParentContext(ktElement: KtElement): FirTowerDataContext?
 }
 
-interface FirTowerDataContextCollector {
-    fun addFileContext(file: FirFile, context: FirTowerDataContext)
-    fun addStatementContext(statement: FirStatement, context: FirTowerDataContext)
-    fun addDeclarationContext(declaration: FirDeclaration, context: FirTowerDataContext)
-}
-
 internal class SingleElementTowerProvider(
     private val element: KtElement,
     private val context: FirTowerDataContext

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/file/structure/FileStructureCache.kt

@@ -5,7 +5,6 @@
 
 package org.jetbrains.kotlin.idea.fir.low.level.api.file.structure
 
-import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirTowerDataContextCollector
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.FirFileBuilder
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache
 import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.FirLazyDeclarationResolver

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

@@ -13,9 +13,9 @@ 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.fir.resolve.symbolProvider
+import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirTowerDataContextCollector
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationDesignation
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.collectDesignation
-import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirTowerDataContextCollector
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.getNonLocalContainingOrThisDeclaration
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.FirFileBuilder
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache

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

@@ -9,22 +9,21 @@ 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.FirResolvePhase
-import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.resolve.ResolutionMode
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirBodyResolveTransformer
+import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirTowerDataContextCollector
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.ImplicitBodyResolveComputationSession
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.createReturnTypeCalculatorForIDE
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationDesignation
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirIdeDesignatedBodyResolveTransformerForReturnTypeCalculator
-import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirTowerDataContextCollector
 
 internal class FirDesignatedBodyResolveTransformerForIDE(
     private val firFile: FirFile,
     designation: FirDeclarationDesignation,
     session: FirSession,
     scopeSession: ScopeSession,
-    private val towerDataContextCollector: FirTowerDataContextCollector? = null
+    towerDataContextCollector: FirTowerDataContextCollector? = null
 ) : FirLazyTransformerForIDE, FirBodyResolveTransformer(
     session,
     phase = FirResolvePhase.BODY_RESOLVE,
@@ -35,7 +34,8 @@ internal class FirDesignatedBodyResolveTransformerForIDE(
         scopeSession,
         ImplicitBodyResolveComputationSession(),
         ::FirIdeDesignatedBodyResolveTransformerForReturnTypeCalculator
-    )
+    ),
+    firTowerDataContextCollector = towerDataContextCollector
 ) {
     private val ideDeclarationTransformer = IDEDeclarationTransformer(designation.toDesignationIterator())
 
@@ -45,18 +45,6 @@ internal class FirDesignatedBodyResolveTransformerForIDE(
             super.transformDeclarationContent(declaration, data)
         }
 
-    override fun onBeforeFileContentResolution(file: FirFile) {
-        towerDataContextCollector?.addFileContext(file, context.towerDataContext)
-    }
-
-    override fun onBeforeDeclarationContentResolve(declaration: FirDeclaration) {
-        towerDataContextCollector?.addDeclarationContext(declaration, context.towerDataContext)
-    }
-
-    override fun onBeforeStatementResolution(statement: FirStatement) {
-        towerDataContextCollector?.addStatementContext(statement, context.towerDataContext)
-    }
-
     override fun needReplacePhase(firDeclaration: FirDeclaration): Boolean =
         ideDeclarationTransformer.needReplacePhase(firDeclaration)