[FIR IDE] Collect tower contexts for local declarations

2021-04-30 19:05:28 +02:00
parent ae63ad0b6d
commit 959ff50fa7
11 changed files with 48 additions and 44 deletions
@@ -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 {
@@ -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) {
@@ -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)
@@ -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,
@@ -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)
+}
@@ -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
 }
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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)