[FIR] Extract BodyResolveContext from FirAbstractBodyResolveTransformer

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

@@ -0,0 +1,174 @@
+/*
+ * 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.fir.resolve.transformers.body.resolve
+
+import kotlinx.collections.immutable.PersistentList
+import kotlinx.collections.immutable.persistentListOf
+import org.jetbrains.kotlin.fir.PrivateForInline
+import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.resolve.FirTowerDataContext
+import org.jetbrains.kotlin.fir.resolve.FirTowerDataContextsForClassParts
+import org.jetbrains.kotlin.fir.resolve.FirTowerDataElement
+import org.jetbrains.kotlin.fir.resolve.ImplicitReceiverStack
+import org.jetbrains.kotlin.fir.resolve.calls.ImplicitReceiverValue
+import org.jetbrains.kotlin.fir.resolve.dfa.DataFlowAnalyzerContext
+import org.jetbrains.kotlin.fir.resolve.dfa.PersistentFlow
+import org.jetbrains.kotlin.fir.resolve.transformers.ReturnTypeCalculator
+import org.jetbrains.kotlin.fir.scopes.FirScope
+import org.jetbrains.kotlin.fir.scopes.impl.FirLocalScope
+import org.jetbrains.kotlin.fir.symbols.impl.FirAnonymousFunctionSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirClassLikeSymbol
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.utils.sure
+
+class BodyResolveContext(
+    val returnTypeCalculator: ReturnTypeCalculator,
+    val dataFlowAnalyzerContext: DataFlowAnalyzerContext<PersistentFlow>,
+    val targetedLocalClasses: Set<FirClass<*>> = emptySet(),
+    val outerLocalClassForNested: MutableMap<FirClassLikeSymbol<*>, FirClassLikeSymbol<*>> = mutableMapOf()
+) {
+    val fileImportsScope: MutableList<FirScope> = mutableListOf()
+
+    @set:PrivateForInline
+    lateinit var file: FirFile
+        internal set
+
+    val implicitReceiverStack: ImplicitReceiverStack get() = towerDataContext.implicitReceiverStack
+
+    @set:PrivateForInline
+    var towerDataContext: FirTowerDataContext = FirTowerDataContext()
+
+    @set:PrivateForInline
+    var towerDataContextsForClassParts: FirTowerDataContextsForClassParts? = null
+
+    val containerIfAny: FirDeclaration?
+        get() = containers.lastOrNull()
+
+    @set:PrivateForInline
+    var containers: PersistentList<FirDeclaration> = persistentListOf()
+
+    val towerDataContextForAnonymousFunctions: MutableMap<FirAnonymousFunctionSymbol, FirTowerDataContext> = mutableMapOf()
+
+    @OptIn(PrivateForInline::class)
+    inline fun <T> withNewTowerDataForClassParts(newContexts: FirTowerDataContextsForClassParts, f: () -> T): T {
+        val old = towerDataContextsForClassParts
+
+        towerDataContextsForClassParts = newContexts
+
+        return try {
+            f()
+        } finally {
+
+            towerDataContextsForClassParts = old
+        }
+    }
+
+    fun getTowerDataContextForStaticNestedClassesUnsafe(): FirTowerDataContext =
+        firTowerDataContextsForClassParts().forNestedClasses
+
+    fun getTowerDataContextForConstructorResolution(): FirTowerDataContext =
+        firTowerDataContextsForClassParts().forConstructorHeaders
+
+    fun getPrimaryConstructorParametersScope(): FirLocalScope? =
+        towerDataContextsForClassParts?.primaryConstructorParametersScope
+
+    private fun firTowerDataContextsForClassParts() =
+        towerDataContextsForClassParts.sure { "towerDataContextForStaticNestedClasses should not be null" }
+
+    @OptIn(PrivateForInline::class)
+    inline fun <T> withContainer(declaration: FirDeclaration, crossinline f: () -> T): T {
+        val oldContainers = containers
+        containers = containers.add(declaration)
+        return try {
+            f()
+        } finally {
+            containers = oldContainers
+        }
+    }
+
+    inline fun <T> withTowerDataContext(context: FirTowerDataContext, f: () -> T): T {
+        return withTowerDataCleanup {
+            replaceTowerDataContext(context)
+            f()
+        }
+    }
+
+    @OptIn(PrivateForInline::class)
+    inline fun <R> withTowerDataCleanup(l: () -> R): R {
+        val initialContext = towerDataContext
+        return try {
+            l()
+        } finally {
+            towerDataContext = initialContext
+        }
+    }
+
+    @OptIn(PrivateForInline::class)
+    fun replaceTowerDataContext(newContext: FirTowerDataContext) {
+        towerDataContext = newContext
+    }
+
+    fun addNonLocalTowerDataElement(element: FirTowerDataElement) {
+        replaceTowerDataContext(towerDataContext.addNonLocalTowerDataElements(listOf(element)))
+    }
+
+    fun addNonLocalTowerDataElements(newElements: List<FirTowerDataElement>) {
+        replaceTowerDataContext(towerDataContext.addNonLocalTowerDataElements(newElements))
+    }
+
+    fun addLocalScope(localScope: FirLocalScope) {
+        replaceTowerDataContext(towerDataContext.addLocalScope(localScope))
+    }
+
+    fun addReceiver(name: Name?, implicitReceiverValue: ImplicitReceiverValue<*>) {
+        replaceTowerDataContext(towerDataContext.addReceiver(name, implicitReceiverValue))
+    }
+
+    fun storeClassIfNotNested(klass: FirRegularClass) {
+        if (containerIfAny is FirClass<*>) return
+        updateLastScope { storeClass(klass) }
+    }
+
+    fun storeFunction(function: FirSimpleFunction) {
+        updateLastScope { storeFunction(function) }
+    }
+
+    fun storeVariable(variable: FirVariable<*>) {
+        updateLastScope { storeVariable(variable) }
+    }
+
+    fun saveContextForAnonymousFunction(anonymousFunction: FirAnonymousFunction) {
+        towerDataContextForAnonymousFunctions[anonymousFunction.symbol] = towerDataContext
+    }
+
+    fun dropContextForAnonymousFunction(anonymousFunction: FirAnonymousFunction) {
+        towerDataContextForAnonymousFunctions.remove(anonymousFunction.symbol)
+    }
+
+    fun cleanContextForAnonymousFunction() {
+        towerDataContextForAnonymousFunctions.clear()
+    }
+
+    fun cleanDataFlowContext() {
+        dataFlowAnalyzerContext.reset()
+    }
+
+    private inline fun updateLastScope(transform: FirLocalScope.() -> FirLocalScope) {
+        val lastScope = towerDataContext.localScopes.lastOrNull() ?: return
+        replaceTowerDataContext(towerDataContext.setLastLocalScope(lastScope.transform()))
+    }
+
+    @OptIn(PrivateForInline::class)
+    fun createSnapshotForLocalClasses(
+        returnTypeCalculator: ReturnTypeCalculator,
+        targetedLocalClasses: Set<FirClass<*>>
+    ) = BodyResolveContext(returnTypeCalculator, dataFlowAnalyzerContext, targetedLocalClasses, outerLocalClassForNested).apply {
+        file = this@BodyResolveContext.file
+        towerDataContextForAnonymousFunctions.putAll(this@BodyResolveContext.towerDataContextForAnonymousFunctions)
+        containers = this@BodyResolveContext.containers
+        towerDataContext = this@BodyResolveContext.towerDataContext
+    }
+}

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

@@ -5,28 +5,19 @@
 
 package org.jetbrains.kotlin.fir.resolve.transformers.body.resolve
 
-import kotlinx.collections.immutable.PersistentList
-import kotlinx.collections.immutable.persistentListOf
 import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.resolve.*
-import org.jetbrains.kotlin.fir.resolve.calls.ImplicitReceiverValue
 import org.jetbrains.kotlin.fir.resolve.calls.ResolutionStageRunner
-import org.jetbrains.kotlin.fir.resolve.dfa.DataFlowAnalyzerContext
 import org.jetbrains.kotlin.fir.resolve.dfa.FirDataFlowAnalyzer
-import org.jetbrains.kotlin.fir.resolve.dfa.PersistentFlow
 import org.jetbrains.kotlin.fir.resolve.inference.FirCallCompleter
 import org.jetbrains.kotlin.fir.resolve.inference.InferenceComponents
 import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider
 import org.jetbrains.kotlin.fir.resolve.transformers.*
 import org.jetbrains.kotlin.fir.scopes.FirScope
 import org.jetbrains.kotlin.fir.scopes.impl.FirLocalScope
-import org.jetbrains.kotlin.fir.symbols.impl.FirAnonymousFunctionSymbol
-import org.jetbrains.kotlin.fir.symbols.impl.FirClassLikeSymbol
 import org.jetbrains.kotlin.fir.types.FirTypeRef
 import org.jetbrains.kotlin.fir.types.builder.buildImplicitTypeRef
-import org.jetbrains.kotlin.name.Name
-import org.jetbrains.kotlin.utils.sure
 
 abstract class FirAbstractBodyResolveTransformer(phase: FirResolvePhase) : FirAbstractPhaseTransformer<ResolutionMode>(phase) {
     abstract val context: BodyResolveContext
@@ -95,155 +86,6 @@ abstract class FirAbstractBodyResolveTransformer(phase: FirResolvePhase) : FirAb
             else -> null
         }
 
-    class BodyResolveContext(
-        val returnTypeCalculator: ReturnTypeCalculator,
-        val dataFlowAnalyzerContext: DataFlowAnalyzerContext<PersistentFlow>,
-        val targetedLocalClasses: Set<FirClass<*>> = emptySet(),
-        val outerLocalClassForNested: MutableMap<FirClassLikeSymbol<*>, FirClassLikeSymbol<*>> = mutableMapOf()
-    ) {
-        val fileImportsScope: MutableList<FirScope> = mutableListOf()
-
-        @set:PrivateForInline
-        lateinit var file: FirFile
-            internal set
-
-        val implicitReceiverStack: ImplicitReceiverStack get() = towerDataContext.implicitReceiverStack
-
-        @set:PrivateForInline
-        var towerDataContext: FirTowerDataContext = FirTowerDataContext()
-
-        @set:PrivateForInline
-        var towerDataContextsForClassParts: FirTowerDataContextsForClassParts? = null
-
-        val containerIfAny: FirDeclaration?
-            get() = containers.lastOrNull()
-
-        @set:PrivateForInline
-        var containers: PersistentList<FirDeclaration> = persistentListOf()
-
-        val towerDataContextForAnonymousFunctions: MutableMap<FirAnonymousFunctionSymbol, FirTowerDataContext> = mutableMapOf()
-
-        @OptIn(PrivateForInline::class)
-        inline fun <T> withNewTowerDataForClassParts(newContexts: FirTowerDataContextsForClassParts, f: () -> T): T {
-            val old = towerDataContextsForClassParts
-
-            towerDataContextsForClassParts = newContexts
-
-            return try {
-                f()
-            } finally {
-
-                towerDataContextsForClassParts = old
-            }
-        }
-
-        fun getTowerDataContextForStaticNestedClassesUnsafe(): FirTowerDataContext =
-            firTowerDataContextsForClassParts().forNestedClasses
-
-        fun getTowerDataContextForConstructorResolution(): FirTowerDataContext =
-            firTowerDataContextsForClassParts().forConstructorHeaders
-
-        fun getPrimaryConstructorParametersScope(): FirLocalScope? =
-            towerDataContextsForClassParts?.primaryConstructorParametersScope
-
-        private fun firTowerDataContextsForClassParts() =
-            towerDataContextsForClassParts.sure { "towerDataContextForStaticNestedClasses should not be null" }
-
-        @OptIn(PrivateForInline::class)
-        inline fun <T> withContainer(declaration: FirDeclaration, crossinline f: () -> T): T {
-            val oldContainers = containers
-            containers = containers.add(declaration)
-            return try {
-                f()
-            } finally {
-                containers = oldContainers
-            }
-        }
-
-        inline fun <T> withTowerDataContext(context: FirTowerDataContext, f: () -> T): T {
-            return withTowerDataCleanup {
-                replaceTowerDataContext(context)
-                f()
-            }
-        }
-
-        @OptIn(PrivateForInline::class)
-        inline fun <R> withTowerDataCleanup(l: () -> R): R {
-            val initialContext = towerDataContext
-            return try {
-                l()
-            } finally {
-                towerDataContext = initialContext
-            }
-        }
-
-        @OptIn(PrivateForInline::class)
-        fun replaceTowerDataContext(newContext: FirTowerDataContext) {
-            towerDataContext = newContext
-        }
-
-        fun addNonLocalTowerDataElement(element: FirTowerDataElement) {
-            replaceTowerDataContext(towerDataContext.addNonLocalTowerDataElements(listOf(element)))
-        }
-
-        fun addNonLocalTowerDataElements(newElements: List<FirTowerDataElement>) {
-            replaceTowerDataContext(towerDataContext.addNonLocalTowerDataElements(newElements))
-        }
-
-        fun addLocalScope(localScope: FirLocalScope) {
-            replaceTowerDataContext(towerDataContext.addLocalScope(localScope))
-        }
-
-        fun addReceiver(name: Name?, implicitReceiverValue: ImplicitReceiverValue<*>) {
-            replaceTowerDataContext(towerDataContext.addReceiver(name, implicitReceiverValue))
-        }
-
-        fun storeClassIfNotNested(klass: FirRegularClass) {
-            if (containerIfAny is FirClass<*>) return
-            updateLastScope { storeClass(klass) }
-        }
-
-        fun storeFunction(function: FirSimpleFunction) {
-            updateLastScope { storeFunction(function) }
-        }
-
-        fun storeVariable(variable: FirVariable<*>) {
-            updateLastScope { storeVariable(variable) }
-        }
-
-        fun saveContextForAnonymousFunction(anonymousFunction: FirAnonymousFunction) {
-            towerDataContextForAnonymousFunctions[anonymousFunction.symbol] = towerDataContext
-        }
-
-        fun dropContextForAnonymousFunction(anonymousFunction: FirAnonymousFunction) {
-            towerDataContextForAnonymousFunctions.remove(anonymousFunction.symbol)
-        }
-
-        fun cleanContextForAnonymousFunction() {
-            towerDataContextForAnonymousFunctions.clear()
-        }
-
-        fun cleanDataFlowContext() {
-            dataFlowAnalyzerContext.reset()
-        }
-
-        private inline fun updateLastScope(transform: FirLocalScope.() -> FirLocalScope) {
-            val lastScope = towerDataContext.localScopes.lastOrNull() ?: return
-            replaceTowerDataContext(towerDataContext.setLastLocalScope(lastScope.transform()))
-        }
-
-        @OptIn(PrivateForInline::class)
-        fun createSnapshotForLocalClasses(
-            returnTypeCalculator: ReturnTypeCalculator,
-            targetedLocalClasses: Set<FirClass<*>>
-        ) = BodyResolveContext(returnTypeCalculator, dataFlowAnalyzerContext, targetedLocalClasses, outerLocalClassForNested).apply {
-            file = this@BodyResolveContext.file
-            towerDataContextForAnonymousFunctions.putAll(this@BodyResolveContext.towerDataContextForAnonymousFunctions)
-            containers = this@BodyResolveContext.containers
-            towerDataContext = this@BodyResolveContext.towerDataContext
-        }
-    }
-
     class BodyResolveTransformerComponents(
         override val session: FirSession,
         override val scopeSession: ScopeSession,

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

@@ -170,7 +170,7 @@ private class ReturnTypeCalculatorWithJump(
     val designationMapForLocalClasses: Map<FirCallableMemberDeclaration<*>, List<FirClass<*>>> = mapOf()
 ) : ReturnTypeCalculator {
 
-    var outerBodyResolveContext: FirAbstractBodyResolveTransformer.BodyResolveContext? = null
+    var outerBodyResolveContext: BodyResolveContext? = null
 
     override fun tryCalculateReturnType(declaration: FirTypedDeclaration): FirResolvedTypeRef {
         if (declaration is FirValueParameter && declaration.returnTypeRef is FirImplicitTypeRef) {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/contracts/FirContractResolveTransformer.kt

@@ -22,6 +22,7 @@ import org.jetbrains.kotlin.fir.references.builder.buildSimpleNamedReference
 import org.jetbrains.kotlin.fir.resolve.ResolutionMode
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.diagnostics.ConeContractDescriptionError
+import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.BodyResolveContext
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirBodyResolveTransformer
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirDeclarationsResolveTransformer
 import org.jetbrains.kotlin.fir.symbols.impl.FirAnonymousFunctionSymbol

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/contracts/FirContractResolveTransformerAdapter.kt

@@ -15,6 +15,7 @@ import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.transformers.AdapterForResolveProcessor
 import org.jetbrains.kotlin.fir.resolve.transformers.FirTransformerBasedResolveProcessor
 import org.jetbrains.kotlin.fir.resolve.transformers.ReturnTypeCalculatorForFullBodyResolve
+import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.BodyResolveContext
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirAbstractBodyResolveTransformer
 import org.jetbrains.kotlin.fir.visitors.CompositeTransformResult
 import org.jetbrains.kotlin.fir.visitors.FirTransformer
@@ -41,7 +42,7 @@ class FirContractResolveTransformerAdapter(session: FirSession, scopeSession: Sc
 fun <F : FirClass<F>> F.runContractResolveForLocalClass(
     session: FirSession,
     scopeSession: ScopeSession,
-    outerBodyResolveContext: FirAbstractBodyResolveTransformer.BodyResolveContext,
+    outerBodyResolveContext: BodyResolveContext,
     targetedClasses: Set<FirClass<*>>
 ): F {
     val newContext = outerBodyResolveContext.createSnapshotForLocalClasses(
@@ -51,4 +52,4 @@ fun <F : FirClass<F>> F.runContractResolveForLocalClass(
     val transformer = FirContractResolveTransformer(session, scopeSession, newContext)
 
     return this.transformSingle(transformer, ResolutionMode.ContextIndependent)
-}
+}

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/plugin/FirAnnotationArgumentsResolveTransformer.kt

@@ -13,6 +13,7 @@ import org.jetbrains.kotlin.fir.extensions.registeredPluginAnnotations
 import org.jetbrains.kotlin.fir.references.FirErrorNamedReference
 import org.jetbrains.kotlin.fir.resolve.ResolutionMode
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
+import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.BodyResolveContext
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirBodyResolveTransformer
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirDeclarationsResolveTransformer
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirExpressionsResolveTransformer