[LL FIR] move context calculation to CONTRACTS phase

ARGUMENTS_OF_ANNOTATIONS will be dropped, so CONTRACTS is the next one

^KT-62679

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/api/LowLevelFirApiFacadeForResolveOnAir.kt

@@ -436,7 +436,7 @@ object LowLevelFirApiFacadeForResolveOnAir {
             annotationMappingRequired(container, elementToReplace) -> FirResolvePhase.ANNOTATIONS_ARGUMENTS_MAPPING
             else -> {
                 /** Currently it is a minimal phase there we can collect [FirResolveContextCollector] */
-                FirResolvePhase.ARGUMENTS_OF_ANNOTATIONS
+                FirResolvePhase.CONTRACTS
             }
         }
     }

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirAnnotationArgumentsLazyResolver.kt

@@ -15,11 +15,9 @@ import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.*
 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.FirResolveContextCollector
 import org.jetbrains.kotlin.fir.resolve.transformers.plugin.FirAnnotationArgumentsResolveTransformer
 import org.jetbrains.kotlin.fir.visitors.transformSingle
-import org.jetbrains.kotlin.util.PrivateForInline
 
 internal object LLFirAnnotationArgumentsLazyResolver : LLFirLazyResolver(FirResolvePhase.ARGUMENTS_OF_ANNOTATIONS) {
     override fun resolve(
@@ -54,101 +52,18 @@ private class LLFirAnnotationArgumentsTargetResolver(
         firResolveContextCollector = firResolveContextCollector,
     )
 
-    private inline fun actionWithContextCollector(
-        noinline action: () -> Unit,
-        crossinline collect: (FirResolveContextCollector, BodyResolveContext) -> Unit,
-    ): () -> Unit {
-        val collector = transformer.firResolveContextCollector ?: return action
-        return {
-            collect(collector, transformer.context)
-            action()
-        }
-    }
-
-    override fun withScript(firScript: FirScript, action: () -> Unit) {
-        val actionWithCollector = actionWithContextCollector(action) { collector, context ->
-            collector.addDeclarationContext(firScript, context)
-        }
-
-        super.withScript(firScript, actionWithCollector)
-    }
-
-    override fun withFile(firFile: FirFile, action: () -> Unit) {
-        val actionWithCollector = actionWithContextCollector(action) { collector, context ->
-            collector.addFileContext(firFile, context.towerDataContext)
-        }
-
-        super.withFile(firFile, actionWithCollector)
-    }
-
     @Deprecated("Should never be called directly, only for override purposes, please use withRegularClass", level = DeprecationLevel.ERROR)
     override fun withRegularClassImpl(firClass: FirRegularClass, action: () -> Unit) {
-        val actionWithCollector = actionWithContextCollector(action) { collector, context ->
-            collector.addDeclarationContext(firClass, context)
-        }
-
         transformer.declarationsTransformer.withRegularClass(firClass) {
-            actionWithCollector()
+            action()
             firClass
         }
     }
 
     override fun doLazyResolveUnderLock(target: FirElementWithResolveState) {
-        collectTowerDataContext(target)
-
         FirLazyBodiesCalculator.calculateAnnotations(target)
         transformAnnotations(target)
     }
-
-    private fun collectTowerDataContext(target: FirElementWithResolveState) {
-        val contextCollector = transformer.firResolveContextCollector
-        if (contextCollector == null || target !is FirDeclaration) return
-
-        val bodyResolveContext = transformer.context
-        withTypeParametersIfMemberDeclaration(bodyResolveContext, target) {
-            when (target) {
-                is FirRegularClass -> {
-                    contextCollector.addClassHeaderContext(target, bodyResolveContext.towerDataContext)
-                }
-
-                is FirFunction -> bodyResolveContext.forFunctionBody(target, transformer.components) {
-                    contextCollector.addDeclarationContext(target, bodyResolveContext)
-                    for (valueParameter in target.valueParameters) {
-                        bodyResolveContext.withValueParameter(valueParameter, transformer.session) {
-                            contextCollector.addDeclarationContext(valueParameter, bodyResolveContext)
-                        }
-                    }
-                }
-
-                is FirScript -> {}
-
-                else -> contextCollector.addDeclarationContext(target, bodyResolveContext)
-            }
-        }
-
-        /**
-         * [withRegularClass] and [withScript] already have [FirResolveContextCollector.addDeclarationContext] call,
-         * so we shouldn't do anything inside
-         */
-        when (target) {
-            is FirRegularClass -> withRegularClass(target) { }
-            is FirScript -> withScript(target) { }
-            else -> {}
-        }
-    }
-
-    private inline fun withTypeParametersIfMemberDeclaration(
-        context: BodyResolveContext,
-        target: FirElementWithResolveState,
-        action: () -> Unit,
-    ) {
-        if (target is FirMemberDeclaration) {
-            @OptIn(PrivateForInline::class)
-            context.withTypeParametersOf(target, action)
-        } else {
-            action()
-        }
-    }
 }
 
 internal fun LLFirAbstractBodyTargetResolver.transformAnnotations(target: FirElementWithResolveState) {

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirContractsLazyResolver.kt

@@ -19,9 +19,11 @@ import org.jetbrains.kotlin.fir.contracts.FirRawContractDescription
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.impl.FirPrimaryConstructor
 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.FirResolveContextCollector
 import org.jetbrains.kotlin.fir.resolve.transformers.contracts.FirContractResolveTransformer
 import org.jetbrains.kotlin.fir.types.FirImplicitTypeRef
+import org.jetbrains.kotlin.util.PrivateForInline
 
 internal object LLFirContractsLazyResolver : LLFirLazyResolver(FirResolvePhase.CONTRACTS) {
     override fun resolve(
@@ -31,7 +33,7 @@ internal object LLFirContractsLazyResolver : LLFirLazyResolver(FirResolvePhase.C
         scopeSession: ScopeSession,
         towerDataContextCollector: FirResolveContextCollector?,
     ) {
-        val resolver = LLFirContractsTargetResolver(target, lockProvider, session, scopeSession)
+        val resolver = LLFirContractsTargetResolver(target, lockProvider, session, scopeSession, towerDataContextCollector)
         resolver.resolveDesignation()
     }
 
@@ -46,15 +48,22 @@ private class LLFirContractsTargetResolver(
     lockProvider: LLFirLockProvider,
     session: FirSession,
     scopeSession: ScopeSession,
+    firResolveContextCollector: FirResolveContextCollector?,
 ) : LLFirAbstractBodyTargetResolver(
     target,
     lockProvider,
     scopeSession,
-    FirResolvePhase.CONTRACTS
+    FirResolvePhase.CONTRACTS,
 ) {
-    override val transformer = FirContractResolveTransformer(session, scopeSession)
+    override val transformer = FirContractResolveTransformer(
+        session,
+        scopeSession,
+        firResolveContextCollector = firResolveContextCollector,
+    )
 
     override fun doLazyResolveUnderLock(target: FirElementWithResolveState) {
+        collectTowerDataContext(target)
+
         when (target) {
             is FirPrimaryConstructor, is FirErrorPrimaryConstructor -> {
                 // No contracts here
@@ -93,6 +102,94 @@ private class LLFirContractsTargetResolver(
             else -> throwUnexpectedFirElementError(target)
         }
     }
+
+
+    private inline fun actionWithContextCollector(
+        noinline action: () -> Unit,
+        crossinline collect: (FirResolveContextCollector, BodyResolveContext) -> Unit,
+    ): () -> Unit {
+        val collector = transformer.firResolveContextCollector ?: return action
+        return {
+            collect(collector, transformer.context)
+            action()
+        }
+    }
+
+    override fun withScript(firScript: FirScript, action: () -> Unit) {
+        val actionWithCollector = actionWithContextCollector(action) { collector, context ->
+            collector.addDeclarationContext(firScript, context)
+        }
+
+        super.withScript(firScript, actionWithCollector)
+    }
+
+    override fun withFile(firFile: FirFile, action: () -> Unit) {
+        val actionWithCollector = actionWithContextCollector(action) { collector, context ->
+            collector.addFileContext(firFile, context.towerDataContext)
+        }
+
+        super.withFile(firFile, actionWithCollector)
+    }
+
+    @Deprecated("Should never be called directly, only for override purposes, please use withRegularClass", level = DeprecationLevel.ERROR)
+    override fun withRegularClassImpl(firClass: FirRegularClass, action: () -> Unit) {
+        val actionWithCollector = actionWithContextCollector(action) { collector, context ->
+            collector.addDeclarationContext(firClass, context)
+        }
+
+        @Suppress("DEPRECATION_ERROR")
+        super.withRegularClassImpl(firClass, actionWithCollector)
+    }
+
+    private fun collectTowerDataContext(target: FirElementWithResolveState) {
+        val contextCollector = transformer.firResolveContextCollector
+        if (contextCollector == null || target !is FirDeclaration) return
+
+        val bodyResolveContext = transformer.context
+        withTypeParametersIfMemberDeclaration(bodyResolveContext, target) {
+            when (target) {
+                is FirRegularClass -> {
+                    contextCollector.addClassHeaderContext(target, bodyResolveContext.towerDataContext)
+                }
+
+                is FirFunction -> bodyResolveContext.forFunctionBody(target, transformer.components) {
+                    contextCollector.addDeclarationContext(target, bodyResolveContext)
+                    for (valueParameter in target.valueParameters) {
+                        bodyResolveContext.withValueParameter(valueParameter, transformer.session) {
+                            contextCollector.addDeclarationContext(valueParameter, bodyResolveContext)
+                        }
+                    }
+                }
+
+                is FirScript -> {}
+
+                else -> contextCollector.addDeclarationContext(target, bodyResolveContext)
+            }
+        }
+
+        /**
+         * [withRegularClass] and [withScript] already have [FirResolveContextCollector.addDeclarationContext] call,
+         * so we shouldn't do anything inside
+         */
+        when (target) {
+            is FirRegularClass -> withRegularClass(target) { }
+            is FirScript -> withScript(target) { }
+            else -> {}
+        }
+    }
+
+    private inline fun withTypeParametersIfMemberDeclaration(
+        context: BodyResolveContext,
+        target: FirElementWithResolveState,
+        action: () -> Unit,
+    ) {
+        if (target is FirMemberDeclaration) {
+            @OptIn(PrivateForInline::class)
+            context.withTypeParametersOf(target, action)
+        } else {
+            action()
+        }
+    }
 }
 
 private object ContractStateKeepers {

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

@@ -30,6 +30,7 @@ import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.BodyResolveCon
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirAbstractBodyResolveTransformerDispatcher
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirDeclarationsResolveTransformer
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirExpressionsResolveTransformer
+import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirResolveContextCollector
 import org.jetbrains.kotlin.fir.symbols.impl.FirAnonymousFunctionSymbol
 import org.jetbrains.kotlin.fir.types.impl.FirImplicitTypeRefImplWithoutSource
 import org.jetbrains.kotlin.fir.utils.exceptions.withFirEntry
@@ -40,14 +41,16 @@ import org.jetbrains.kotlin.utils.exceptions.errorWithAttachment
 abstract class FirAbstractContractResolveTransformerDispatcher(
     session: FirSession,
     scopeSession: ScopeSession,
-    outerBodyResolveContext: BodyResolveContext? = null
+    outerBodyResolveContext: BodyResolveContext? = null,
+    firResolveContextCollector: FirResolveContextCollector? = null,
 ) : FirAbstractBodyResolveTransformerDispatcher(
     session,
     FirResolvePhase.CONTRACTS,
     implicitTypeOnly = false,
     scopeSession,
     returnTypeCalculator = ReturnTypeCalculatorForFullBodyResolve.Contract,
-    outerBodyResolveContext = outerBodyResolveContext
+    outerBodyResolveContext = outerBodyResolveContext,
+    firResolveContextCollector = firResolveContextCollector,
 ) {
     final override val expressionsTransformer: FirExpressionsResolveTransformer =
         FirExpressionsResolveTransformer(this)

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Copyright 2010-2023 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.
  */
 
@@ -12,11 +12,13 @@ import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.*
 class FirContractResolveTransformer(
     session: FirSession,
     scopeSession: ScopeSession,
-    outerBodyResolveContext: BodyResolveContext? = null
+    outerBodyResolveContext: BodyResolveContext? = null,
+    firResolveContextCollector: FirResolveContextCollector? = null,
 ) : FirAbstractContractResolveTransformerDispatcher(
     session,
     scopeSession,
     outerBodyResolveContext,
+    firResolveContextCollector,
 ) {
     override val contractDeclarationsTransformer: FirDeclarationsContractResolveTransformer
         get() = FirDeclarationsContractResolveTransformer()