[FIR] Properly resolve declarations inside contract calls on CONTRACT phase

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

@@ -29,14 +29,15 @@ internal class LLFirDesignatedContractsResolveTransformer(
 
     private val ideDeclarationTransformer = LLFirDeclarationTransformer(designation)
 
-    override val declarationsTransformer: FirDeclarationsResolveTransformer = object : FirDeclarationsContractResolveTransformer(this) {
-        override fun transformDeclarationContent(firClass: FirClass, data: ResolutionMode) {
-            ideDeclarationTransformer.transformDeclarationContent(this, firClass, data) {
-                super.transformDeclarationContent(firClass, data)
-                firClass
+    override val contractDeclarationsTransformer: FirDeclarationsContractResolveTransformer
+        get() = object : FirDeclarationsContractResolveTransformer() {
+            override fun transformDeclarationContent(firClass: FirClass, data: ResolutionMode) {
+                ideDeclarationTransformer.transformDeclarationContent(this, firClass, data) {
+                    super.transformDeclarationContent(firClass, data)
+                    firClass
+                }
             }
         }
-    }
 
     override fun transformDeclarationContent(declaration: FirDeclaration, data: ResolutionMode): FirDeclaration =
         ideDeclarationTransformer.transformDeclarationContent(this, declaration, data) {

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

@@ -48,6 +48,14 @@ abstract class FirAbstractContractResolveTransformerDispatcher(
     final override val expressionsTransformer: FirExpressionsResolveTransformer =
         FirExpressionsResolveTransformer(this)
 
+    final override val declarationsTransformer: FirDeclarationsResolveTransformer
+        get() = if (contractMode) contractDeclarationsTransformer else regularDeclarationsTransformer
+
+    protected abstract val contractDeclarationsTransformer: FirDeclarationsContractResolveTransformer
+    private val regularDeclarationsTransformer = FirDeclarationsResolveTransformer(this)
+
+    private var contractMode = true
+
     override fun transformAnnotation(annotation: FirAnnotation, data: ResolutionMode): FirStatement {
         return annotation
     }
@@ -56,9 +64,8 @@ abstract class FirAbstractContractResolveTransformerDispatcher(
         return annotationCall
     }
 
-    protected open class FirDeclarationsContractResolveTransformer(
-        transformer: FirAbstractBodyResolveTransformerDispatcher
-    ) : FirDeclarationsResolveTransformer(transformer) {
+    protected open inner class FirDeclarationsContractResolveTransformer :
+        FirDeclarationsResolveTransformer(this@FirAbstractContractResolveTransformerDispatcher) {
         override fun transformSimpleFunction(
             simpleFunction: FirSimpleFunction,
             data: ResolutionMode
@@ -136,7 +143,12 @@ abstract class FirAbstractContractResolveTransformerDispatcher(
             for (valueParameter in valueParameters) {
                 context.storeVariable(valueParameter, session)
             }
-            val contractCall = contractDescription.contractCall.transformSingle(transformer, ResolutionMode.ContextIndependent)
+            val contractCall = try {
+                contractMode = false
+                contractDescription.contractCall.transformSingle(transformer, ResolutionMode.ContextIndependent)
+            } finally {
+                contractMode = true
+            }
             val resolvedId = contractCall.toResolvedCallableSymbol()?.callableId ?: return transformOwnerWithUnresolvedContract(owner)
             if (resolvedId != FirContractsDslNames.CONTRACT) return transformOwnerWithUnresolvedContract(owner)
             if (contractCall.arguments.size != 1) return transformOwnerOfErrorContract(owner)

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

@@ -9,7 +9,7 @@ import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.*
 
-open class FirContractResolveTransformer(
+class FirContractResolveTransformer(
     session: FirSession,
     scopeSession: ScopeSession,
     outerBodyResolveContext: BodyResolveContext? = null
@@ -18,6 +18,6 @@ open class FirContractResolveTransformer(
     scopeSession,
     outerBodyResolveContext,
 ) {
-    final override val declarationsTransformer: FirDeclarationsResolveTransformer =
-        FirDeclarationsContractResolveTransformer(this)
+    override val contractDeclarationsTransformer: FirDeclarationsContractResolveTransformer
+        get() = FirDeclarationsContractResolveTransformer()
 }

compiler/testData/diagnostics/testsWithStdLib/contracts/dsl/errors/illegalConstructionInContractBlock.fir.kt

@@ -25,7 +25,7 @@ fun whenInContract(x: Any?, boolean: Boolean) {
 
 fun forInContract(x: Any?) {
     contract {
-        for (i in <!HAS_NEXT_MISSING!>0..1<!>) {
+        for (i in 0..1) {
             returns() implies (x is String)
         }
     }
@@ -52,4 +52,4 @@ fun localValInContract(x: Any?) {
         <!ERROR_IN_CONTRACT_DESCRIPTION!>val y: Int = 42<!>
         returns() implies (x is String)
     }<!>
-}
+}

compiler/testData/diagnostics/testsWithStdLib/contracts/dsl/errors/recursiveContractCustomContractFunction.fir.kt

@@ -1,39 +0,0 @@
-// !OPT_IN: kotlin.contracts.ExperimentalContracts
-// !LANGUAGE: +UseReturnsEffect
-// Issue: KT-26386
-
-fun myRun(block: () -> Unit) {
-    block()
-}
-
-fun contract(block: () -> Unit) {
-    block()
-}
-
-fun case_1(): Boolean? {
-    contract { case_1() }
-    return null
-}
-
-fun case_2(): Boolean? {
-    contract { case_3() }
-    return null
-}
-
-fun case_3(): Boolean? {
-    contract { case_2() }
-    return null
-}
-
-fun case4() {
-    contract {
-        myRun {
-            val s: String
-            run {
-                <!UNRESOLVED_REFERENCE!>s<!> = "hello"
-            }
-            <!UNRESOLVED_REFERENCE!>s<!>.length
-        }
-    }
-}
-

compiler/testData/diagnostics/testsWithStdLib/contracts/dsl/errors/recursiveContractCustomContractFunction.kt

@@ -1,3 +1,4 @@
+// FIR_IDENTICAL
 // !OPT_IN: kotlin.contracts.ExperimentalContracts
 // !LANGUAGE: +UseReturnsEffect
 // Issue: KT-26386