[FIR] Fix Disappeared CONTRACT_NOT_ALLOWED

https://youtrack.jetbrains.com/issue/KT-60004/K2-Disappeared-CONTRACTNOTALLOWED

Merge-request: KT-MR-12613
Merged-by: Vladimir Sukharev <Vladimir.Sukharev@jetbrains.com>

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/diagnostic/compiler/based/DiagnosticCompilerTestFE10TestdataTestGenerated.java

@@ -6645,6 +6645,54 @@ public class DiagnosticCompilerTestFE10TestdataTestGenerated extends AbstractDia
             }
         }
 
+        @Nested
+        @TestMetadata("compiler/testData/diagnostics/tests/contracts")
+        @TestDataPath("$PROJECT_ROOT")
+        public class Contracts {
+            @Test
+            public void testAllFilesPresentInContracts() throws Exception {
+                KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/diagnostics/tests/contracts"), Pattern.compile("^(.+)\\.(kt|kts)$"), Pattern.compile("^(.+)\\.(reversed|fir|ll)\\.kts?$"), true, "multiplatform");
+            }
+
+            @Nested
+            @TestMetadata("compiler/testData/diagnostics/tests/contracts/neg")
+            @TestDataPath("$PROJECT_ROOT")
+            public class Neg {
+                @Test
+                @TestMetadata("14.kt")
+                public void test14() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/contracts/neg/14.kt");
+                }
+
+                @Test
+                public void testAllFilesPresentInNeg() throws Exception {
+                    KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/diagnostics/tests/contracts/neg"), Pattern.compile("^(.+)\\.(kt|kts)$"), Pattern.compile("^(.+)\\.(reversed|fir|ll)\\.kts?$"), true, "multiplatform");
+                }
+            }
+
+            @Nested
+            @TestMetadata("compiler/testData/diagnostics/tests/contracts/pos")
+            @TestDataPath("$PROJECT_ROOT")
+            public class Pos {
+                @Test
+                @TestMetadata("1.kt")
+                public void test1() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/contracts/pos/1.kt");
+                }
+
+                @Test
+                @TestMetadata("14UserDefined.kt")
+                public void test14UserDefined() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/contracts/pos/14UserDefined.kt");
+                }
+
+                @Test
+                public void testAllFilesPresentInPos() throws Exception {
+                    KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/diagnostics/tests/contracts/pos"), Pattern.compile("^(.+)\\.(kt|kts)$"), Pattern.compile("^(.+)\\.(reversed|fir|ll)\\.kts?$"), true, "multiplatform");
+                }
+            }
+        }
+
         @Nested
         @TestMetadata("compiler/testData/diagnostics/tests/controlFlowAnalysis")
         @TestDataPath("$PROJECT_ROOT")

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/diagnostic/compiler/based/LLFirPreresolvedReversedDiagnosticCompilerFE10TestDataTestGenerated.java

@@ -6645,6 +6645,54 @@ public class LLFirPreresolvedReversedDiagnosticCompilerFE10TestDataTestGenerated
             }
         }
 
+        @Nested
+        @TestMetadata("compiler/testData/diagnostics/tests/contracts")
+        @TestDataPath("$PROJECT_ROOT")
+        public class Contracts {
+            @Test
+            public void testAllFilesPresentInContracts() throws Exception {
+                KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/diagnostics/tests/contracts"), Pattern.compile("^(.+)\\.(kt|kts)$"), Pattern.compile("^(.+)\\.(reversed|fir|ll)\\.kts?$"), true, "multiplatform");
+            }
+
+            @Nested
+            @TestMetadata("compiler/testData/diagnostics/tests/contracts/neg")
+            @TestDataPath("$PROJECT_ROOT")
+            public class Neg {
+                @Test
+                @TestMetadata("14.kt")
+                public void test14() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/contracts/neg/14.kt");
+                }
+
+                @Test
+                public void testAllFilesPresentInNeg() throws Exception {
+                    KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/diagnostics/tests/contracts/neg"), Pattern.compile("^(.+)\\.(kt|kts)$"), Pattern.compile("^(.+)\\.(reversed|fir|ll)\\.kts?$"), true, "multiplatform");
+                }
+            }
+
+            @Nested
+            @TestMetadata("compiler/testData/diagnostics/tests/contracts/pos")
+            @TestDataPath("$PROJECT_ROOT")
+            public class Pos {
+                @Test
+                @TestMetadata("1.kt")
+                public void test1() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/contracts/pos/1.kt");
+                }
+
+                @Test
+                @TestMetadata("14UserDefined.kt")
+                public void test14UserDefined() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/contracts/pos/14UserDefined.kt");
+                }
+
+                @Test
+                public void testAllFilesPresentInPos() throws Exception {
+                    KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/diagnostics/tests/contracts/pos"), Pattern.compile("^(.+)\\.(kt|kts)$"), Pattern.compile("^(.+)\\.(reversed|fir|ll)\\.kts?$"), true, "multiplatform");
+                }
+            }
+        }
+
         @Nested
         @TestMetadata("compiler/testData/diagnostics/tests/controlFlowAnalysis")
         @TestDataPath("$PROJECT_ROOT")

compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirLightTreeOldFrontendDiagnosticsTestGenerated.java

@@ -6639,6 +6639,54 @@ public class FirLightTreeOldFrontendDiagnosticsTestGenerated extends AbstractFir
             }
         }
 
+        @Nested
+        @TestMetadata("compiler/testData/diagnostics/tests/contracts")
+        @TestDataPath("$PROJECT_ROOT")
+        public class Contracts {
+            @Test
+            public void testAllFilesPresentInContracts() throws Exception {
+                KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/diagnostics/tests/contracts"), Pattern.compile("^(.+)\\.kt$"), Pattern.compile("^(.+)\\.(reversed|fir|ll)\\.kts?$"), true, "multiplatform");
+            }
+
+            @Nested
+            @TestMetadata("compiler/testData/diagnostics/tests/contracts/neg")
+            @TestDataPath("$PROJECT_ROOT")
+            public class Neg {
+                @Test
+                @TestMetadata("14.kt")
+                public void test14() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/contracts/neg/14.kt");
+                }
+
+                @Test
+                public void testAllFilesPresentInNeg() throws Exception {
+                    KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/diagnostics/tests/contracts/neg"), Pattern.compile("^(.+)\\.kt$"), Pattern.compile("^(.+)\\.(reversed|fir|ll)\\.kts?$"), true, "multiplatform");
+                }
+            }
+
+            @Nested
+            @TestMetadata("compiler/testData/diagnostics/tests/contracts/pos")
+            @TestDataPath("$PROJECT_ROOT")
+            public class Pos {
+                @Test
+                @TestMetadata("1.kt")
+                public void test1() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/contracts/pos/1.kt");
+                }
+
+                @Test
+                @TestMetadata("14UserDefined.kt")
+                public void test14UserDefined() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/contracts/pos/14UserDefined.kt");
+                }
+
+                @Test
+                public void testAllFilesPresentInPos() throws Exception {
+                    KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/diagnostics/tests/contracts/pos"), Pattern.compile("^(.+)\\.kt$"), Pattern.compile("^(.+)\\.(reversed|fir|ll)\\.kts?$"), true, "multiplatform");
+                }
+            }
+        }
+
         @Nested
         @TestMetadata("compiler/testData/diagnostics/tests/controlFlowAnalysis")
         @TestDataPath("$PROJECT_ROOT")

compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirPsiOldFrontendDiagnosticsTestGenerated.java

@@ -6645,6 +6645,54 @@ public class FirPsiOldFrontendDiagnosticsTestGenerated extends AbstractFirPsiDia
             }
         }
 
+        @Nested
+        @TestMetadata("compiler/testData/diagnostics/tests/contracts")
+        @TestDataPath("$PROJECT_ROOT")
+        public class Contracts {
+            @Test
+            public void testAllFilesPresentInContracts() throws Exception {
+                KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/diagnostics/tests/contracts"), Pattern.compile("^(.*)\\.kts?$"), Pattern.compile("^(.+)\\.(reversed|fir|ll)\\.kts?$"), true, "multiplatform");
+            }
+
+            @Nested
+            @TestMetadata("compiler/testData/diagnostics/tests/contracts/neg")
+            @TestDataPath("$PROJECT_ROOT")
+            public class Neg {
+                @Test
+                @TestMetadata("14.kt")
+                public void test14() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/contracts/neg/14.kt");
+                }
+
+                @Test
+                public void testAllFilesPresentInNeg() throws Exception {
+                    KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/diagnostics/tests/contracts/neg"), Pattern.compile("^(.*)\\.kts?$"), Pattern.compile("^(.+)\\.(reversed|fir|ll)\\.kts?$"), true, "multiplatform");
+                }
+            }
+
+            @Nested
+            @TestMetadata("compiler/testData/diagnostics/tests/contracts/pos")
+            @TestDataPath("$PROJECT_ROOT")
+            public class Pos {
+                @Test
+                @TestMetadata("1.kt")
+                public void test1() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/contracts/pos/1.kt");
+                }
+
+                @Test
+                @TestMetadata("14UserDefined.kt")
+                public void test14UserDefined() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/contracts/pos/14UserDefined.kt");
+                }
+
+                @Test
+                public void testAllFilesPresentInPos() throws Exception {
+                    KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/diagnostics/tests/contracts/pos"), Pattern.compile("^(.*)\\.kts?$"), Pattern.compile("^(.+)\\.(reversed|fir|ll)\\.kts?$"), true, "multiplatform");
+                }
+            }
+        }
+
         @Nested
         @TestMetadata("compiler/testData/diagnostics/tests/controlFlowAnalysis")
         @TestDataPath("$PROJECT_ROOT")

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/expression/FirContractNotFirstStatementChecker.kt

@@ -9,7 +9,11 @@ import org.jetbrains.kotlin.diagnostics.DiagnosticReporter
 import org.jetbrains.kotlin.diagnostics.reportOn
 import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors
+import org.jetbrains.kotlin.fir.contracts.FirContractDescription
+import org.jetbrains.kotlin.fir.contracts.FirResolvedContractDescription
+import org.jetbrains.kotlin.fir.declarations.FirContractDescriptionOwner
 import org.jetbrains.kotlin.fir.declarations.FirFunction
+import org.jetbrains.kotlin.fir.diagnostics.ConeContractShouldBeFirstStatement
 import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
 import org.jetbrains.kotlin.fir.expressions.impl.FirContractCallBlock
 import org.jetbrains.kotlin.fir.expressions.impl.FirSingleExpressionBlock
@@ -35,5 +39,9 @@ object FirContractNotFirstStatementChecker : FirFunctionCallChecker() {
     private fun FirFunctionCall.isCorrectlyPlacedIn(functionDeclaration: FirFunction): Boolean {
         val firstStatement = functionDeclaration.body?.statements?.first()
         return firstStatement is FirContractCallBlock && firstStatement.call == this
+                && !(functionDeclaration is FirContractDescriptionOwner && functionDeclaration.contractDescription.isNonFirstStatement)
     }
+
+    private val FirContractDescription.isNonFirstStatement: Boolean
+        get() = this is FirResolvedContractDescription && diagnostic == ConeContractShouldBeFirstStatement
 }

compiler/fir/raw-fir/light-tree2fir/src/org/jetbrains/kotlin/fir/lightTree/converter/LightTreeRawFirDeclarationBuilder.kt

@@ -19,6 +19,7 @@ import org.jetbrains.kotlin.descriptors.Visibility
 import org.jetbrains.kotlin.descriptors.annotations.AnnotationUseSiteTarget
 import org.jetbrains.kotlin.descriptors.annotations.AnnotationUseSiteTarget.*
 import org.jetbrains.kotlin.fir.*
+import org.jetbrains.kotlin.fir.analysis.isCallTheFirstStatement
 import org.jetbrains.kotlin.fir.builder.*
 import org.jetbrains.kotlin.fir.contracts.FirContractDescription
 import org.jetbrains.kotlin.fir.contracts.builder.buildRawContractDescription
@@ -53,6 +54,7 @@ import org.jetbrains.kotlin.lexer.KtModifierKeywordToken
 import org.jetbrains.kotlin.lexer.KtTokens.*
 import org.jetbrains.kotlin.name.*
 import org.jetbrains.kotlin.psi.stubs.elements.KtStubElementTypes
+import org.jetbrains.kotlin.util.getChildren
 import org.jetbrains.kotlin.utils.addToStdlib.runIf
 import org.jetbrains.kotlin.utils.addToStdlib.runUnless
 
@@ -1818,7 +1820,14 @@ class LightTreeRawFirDeclarationBuilder(
         return when {
             blockNode != null -> {
                 val block = convertBlock(blockNode)
-                val contractDescription = runIf(allowLegacyContractDescription) { processLegacyContractDescription(block) }
+                val contractDescription = runIf(allowLegacyContractDescription) {
+                    val blockSource = block.source
+                    val diagnostic = when {
+                        blockSource == null || !isCallTheFirstStatement(blockSource) -> ConeContractShouldBeFirstStatement
+                        else -> null
+                    }
+                    processLegacyContractDescription(block, diagnostic)
+                }
                 block to contractDescription
             }
             expression != null -> FirSingleExpressionBlock(
@@ -1828,6 +1837,9 @@ class LightTreeRawFirDeclarationBuilder(
         }
     }
 
+    private fun isCallTheFirstStatement(sourceElement: KtSourceElement): Boolean =
+        isCallTheFirstStatement(sourceElement.lighterASTNode, { it.elementType }, { it.getChildren(sourceElement.treeStructure) })
+
     /**
      * @see org.jetbrains.kotlin.parsing.KotlinParsing.parseBlock
      */

compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/PsiRawFirBuilder.kt

@@ -14,6 +14,7 @@ import org.jetbrains.kotlin.builtins.StandardNames.BACKING_FIELD
 import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.descriptors.annotations.AnnotationUseSiteTarget.*
 import org.jetbrains.kotlin.fir.*
+import org.jetbrains.kotlin.fir.analysis.isCallTheFirstStatement
 import org.jetbrains.kotlin.fir.contracts.FirContractDescription
 import org.jetbrains.kotlin.fir.contracts.builder.buildRawContractDescription
 import org.jetbrains.kotlin.fir.declarations.*
@@ -419,7 +420,14 @@ open class PsiRawFirBuilder(
                     if (hasBlockBody()) {
                         val block = bodyBlockExpression?.accept(this@Visitor, null) as? FirBlock
                         val contractDescription = when {
-                            !hasContractEffectList() -> block?.let(::processLegacyContractDescription)
+                            !hasContractEffectList() -> block?.let {
+                                val blockSourcePsi = it.source?.psi
+                                val diagnostic = when {
+                                    blockSourcePsi == null || !isCallTheFirstStatement(blockSourcePsi) -> ConeContractShouldBeFirstStatement
+                                    else -> null
+                                }
+                                processLegacyContractDescription(block, diagnostic)
+                            }
                             else -> null
                         }
                         return@buildFirBody block to contractDescription
@@ -432,6 +440,9 @@ open class PsiRawFirBuilder(
                 null to null
             }
 
+        private fun isCallTheFirstStatement(psi: PsiElement): Boolean =
+            isCallTheFirstStatement(psi, { it.elementType }, { it.allChildren.toList() })
+
         private fun ValueArgument?.toFirExpression(): FirExpression {
             if (this == null) {
                 return buildErrorExpression(

compiler/fir/raw-fir/raw-fir.common/src/org/jetbrains/kotlin/fir/analysis/Traversals.kt

@@ -0,0 +1,33 @@
+/*
+ * 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.
+ */
+
+package org.jetbrains.kotlin.fir.analysis
+
+import com.intellij.psi.tree.IElementType
+import org.jetbrains.kotlin.KtNodeTypes
+import org.jetbrains.kotlin.lexer.KtTokens
+import org.jetbrains.kotlin.utils.addToStdlib.popLast
+
+inline fun <T> isCallTheFirstStatement(
+    root: T,
+    getElementType: (T) -> IElementType,
+    getChildren: (T) -> List<T>,
+): Boolean {
+    val stack = getChildren(root).asReversed().toMutableList()
+
+    while (stack.isNotEmpty()) {
+        val child = stack.popLast()
+        when (getElementType(child)) {
+            KtTokens.LBRACE, KtTokens.WHITE_SPACE, KtTokens.DOT, KtTokens.EOL_COMMENT -> {}
+            KtNodeTypes.CALL_EXPRESSION -> return true
+            KtNodeTypes.REFERENCE_EXPRESSION -> {}
+            KtNodeTypes.DOT_QUALIFIED_EXPRESSION -> {
+                stack += getChildren(child).asReversed()
+            }
+            else -> return false
+        }
+    }
+    return false
+}

compiler/fir/raw-fir/raw-fir.common/src/org/jetbrains/kotlin/fir/builder/ConversionUtils.kt

@@ -9,11 +9,11 @@ import com.intellij.psi.tree.IElementType
 import org.jetbrains.kotlin.KtFakeSourceElementKind
 import org.jetbrains.kotlin.KtRealSourceElementKind
 import org.jetbrains.kotlin.KtSourceElement
+import org.jetbrains.kotlin.*
 import org.jetbrains.kotlin.contracts.description.LogicOperationKind
 import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.descriptors.Visibilities
 import org.jetbrains.kotlin.descriptors.annotations.AnnotationUseSiteTarget
-import org.jetbrains.kotlin.fakeElement
 import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.contracts.FirContractDescription
 import org.jetbrains.kotlin.fir.contracts.FirLegacyRawContractDescription
@@ -56,7 +56,6 @@ import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.name.SpecialNames
 import org.jetbrains.kotlin.name.StandardClassIds
-import org.jetbrains.kotlin.realElement
 import org.jetbrains.kotlin.types.ConstantValueKind
 import org.jetbrains.kotlin.types.Variance
 import org.jetbrains.kotlin.types.expressions.OperatorConventions
@@ -556,19 +555,20 @@ fun <T> FirPropertyBuilder.generateAccessorsByDelegate(
     }
 }
 
-fun processLegacyContractDescription(block: FirBlock): FirContractDescription? {
+fun processLegacyContractDescription(block: FirBlock, diagnostic: ConeDiagnostic?): FirContractDescription? {
     if (block.isContractPresentFirCheck()) {
         val contractCall = block.replaceFirstStatement<FirFunctionCall> { FirContractCallBlock(it) }
-        return contractCall.toLegacyRawContractDescription()
+        return contractCall.toLegacyRawContractDescription(diagnostic)
     }
 
     return null
 }
 
-fun FirFunctionCall.toLegacyRawContractDescription(): FirLegacyRawContractDescription {
+fun FirFunctionCall.toLegacyRawContractDescription(diagnostic: ConeDiagnostic? = null): FirLegacyRawContractDescription {
     return buildLegacyRawContractDescription {
         this.source = this@toLegacyRawContractDescription.source
         this.contractCall = this@toLegacyRawContractDescription
+        this.diagnostic = diagnostic
     }
 }
 

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

@@ -196,6 +196,7 @@ abstract class FirAbstractContractResolveTransformerDispatcher(
                     }
                 }
                 this.source = contractDescription.source
+                this.diagnostic = contractDescription.diagnostic
             }
             owner.replaceContractDescription(resolvedContractDescription)
             dataFlowAnalyzer.exitContractDescription()

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/contracts/FirLegacyRawContractDescription.kt

@@ -10,6 +10,7 @@ package org.jetbrains.kotlin.fir.contracts
 
 import org.jetbrains.kotlin.KtSourceElement
 import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.diagnostics.ConeDiagnostic
 import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
 import org.jetbrains.kotlin.fir.visitors.FirTransformer
 import org.jetbrains.kotlin.fir.visitors.FirVisitor
@@ -20,6 +21,7 @@ import org.jetbrains.kotlin.fir.visitors.FirVisitor
 abstract class FirLegacyRawContractDescription : FirContractDescription() {
     abstract override val source: KtSourceElement?
     abstract val contractCall: FirFunctionCall
+    abstract val diagnostic: ConeDiagnostic?
 
     override fun <R, D> accept(visitor: FirVisitor<R, D>, data: D): R =
         visitor.visitLegacyRawContractDescription(this, data)

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/contracts/FirResolvedContractDescription.kt

@@ -10,6 +10,7 @@ package org.jetbrains.kotlin.fir.contracts
 
 import org.jetbrains.kotlin.KtSourceElement
 import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.diagnostics.ConeDiagnostic
 import org.jetbrains.kotlin.fir.visitors.FirTransformer
 import org.jetbrains.kotlin.fir.visitors.FirVisitor
 
@@ -20,6 +21,7 @@ abstract class FirResolvedContractDescription : FirContractDescription() {
     abstract override val source: KtSourceElement?
     abstract val effects: List<FirEffectDeclaration>
     abstract val unresolvedEffects: List<FirContractElementDeclaration>
+    abstract val diagnostic: ConeDiagnostic?
 
     override fun <R, D> accept(visitor: FirVisitor<R, D>, data: D): R =
         visitor.visitResolvedContractDescription(this, data)

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/contracts/builder/FirLegacyRawContractDescriptionBuilder.kt

@@ -15,17 +15,20 @@ import org.jetbrains.kotlin.KtSourceElement
 import org.jetbrains.kotlin.fir.builder.FirBuilderDsl
 import org.jetbrains.kotlin.fir.contracts.FirLegacyRawContractDescription
 import org.jetbrains.kotlin.fir.contracts.impl.FirLegacyRawContractDescriptionImpl
+import org.jetbrains.kotlin.fir.diagnostics.ConeDiagnostic
 import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
 
 @FirBuilderDsl
 class FirLegacyRawContractDescriptionBuilder {
     var source: KtSourceElement? = null
     lateinit var contractCall: FirFunctionCall
+    var diagnostic: ConeDiagnostic? = null
 
     fun build(): FirLegacyRawContractDescription {
         return FirLegacyRawContractDescriptionImpl(
             source,
             contractCall,
+            diagnostic,
         )
     }
 

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/contracts/builder/FirResolvedContractDescriptionBuilder.kt

@@ -17,18 +17,21 @@ import org.jetbrains.kotlin.fir.contracts.FirContractElementDeclaration
 import org.jetbrains.kotlin.fir.contracts.FirEffectDeclaration
 import org.jetbrains.kotlin.fir.contracts.FirResolvedContractDescription
 import org.jetbrains.kotlin.fir.contracts.impl.FirResolvedContractDescriptionImpl
+import org.jetbrains.kotlin.fir.diagnostics.ConeDiagnostic
 
 @FirBuilderDsl
 class FirResolvedContractDescriptionBuilder {
     var source: KtSourceElement? = null
     val effects: MutableList<FirEffectDeclaration> = mutableListOf()
     val unresolvedEffects: MutableList<FirContractElementDeclaration> = mutableListOf()
+    var diagnostic: ConeDiagnostic? = null
 
     fun build(): FirResolvedContractDescription {
         return FirResolvedContractDescriptionImpl(
             source,
             effects,
             unresolvedEffects,
+            diagnostic,
         )
     }
 

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/contracts/impl/FirLegacyRawContractDescriptionImpl.kt

@@ -12,6 +12,7 @@ package org.jetbrains.kotlin.fir.contracts.impl
 
 import org.jetbrains.kotlin.KtSourceElement
 import org.jetbrains.kotlin.fir.contracts.FirLegacyRawContractDescription
+import org.jetbrains.kotlin.fir.diagnostics.ConeDiagnostic
 import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
 import org.jetbrains.kotlin.fir.visitors.FirTransformer
 import org.jetbrains.kotlin.fir.visitors.FirVisitor
@@ -19,6 +20,7 @@ import org.jetbrains.kotlin.fir.visitors.FirVisitor
 internal class FirLegacyRawContractDescriptionImpl(
     override val source: KtSourceElement?,
     override var contractCall: FirFunctionCall,
+    override val diagnostic: ConeDiagnostic?,
 ) : FirLegacyRawContractDescription() {
     override fun <R, D> acceptChildren(visitor: FirVisitor<R, D>, data: D) {
         contractCall.accept(visitor, data)

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/contracts/impl/FirResolvedContractDescriptionImpl.kt

@@ -14,6 +14,7 @@ import org.jetbrains.kotlin.KtSourceElement
 import org.jetbrains.kotlin.fir.contracts.FirContractElementDeclaration
 import org.jetbrains.kotlin.fir.contracts.FirEffectDeclaration
 import org.jetbrains.kotlin.fir.contracts.FirResolvedContractDescription
+import org.jetbrains.kotlin.fir.diagnostics.ConeDiagnostic
 import org.jetbrains.kotlin.fir.visitors.FirTransformer
 import org.jetbrains.kotlin.fir.visitors.FirVisitor
 import org.jetbrains.kotlin.fir.visitors.transformInplace
@@ -22,6 +23,7 @@ internal class FirResolvedContractDescriptionImpl(
     override val source: KtSourceElement?,
     override val effects: MutableList<FirEffectDeclaration>,
     override val unresolvedEffects: MutableList<FirContractElementDeclaration>,
+    override val diagnostic: ConeDiagnostic?,
 ) : FirResolvedContractDescription() {
     override fun <R, D> acceptChildren(visitor: FirVisitor<R, D>, data: D) {
         effects.forEach { it.accept(visitor, data) }

compiler/fir/tree/src/org/jetbrains/kotlin/fir/diagnostics/ConeSimpleDiagnostic.kt

@@ -87,6 +87,10 @@ object ConeMissingConstructorKeyword : ConeDiagnostic {
     override val reason: String get() = "Use the 'constructor' keyword after the modifiers of the primary constructor."
 }
 
+object ConeContractShouldBeFirstStatement : ConeDiagnostic {
+    override val reason: String get() = "Contract should be the first statement."
+}
+
 enum class DiagnosticKind {
     ExpressionExpected,
     NotLoopLabel,

compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/NodeConfigurator.kt

@@ -777,10 +777,12 @@ object NodeConfigurator : AbstractFieldConfigurator<FirTreeBuilder>(FirTreeBuild
         resolvedContractDescription.configure {
             +fieldList("effects", effectDeclaration)
             +fieldList("unresolvedEffects", contractElementDeclaration)
+            +field("diagnostic", coneDiagnosticType, nullable = true)
         }
 
         legacyRawContractDescription.configure {
             +field("contractCall", functionCall)
+            +field("diagnostic", coneDiagnosticType, nullable = true)
         }
     }
 }

compiler/testData/diagnostics/tests/contracts/neg/14.fir.kt

@@ -0,0 +1,23 @@
+// !DIAGNOSTICS: -UNUSED_VARIABLE
+// !OPT_IN: kotlin.contracts.ExperimentalContracts
+
+/*
+ * ADDITION TO `KOTLIN DIAGNOSTICS NOT LINKED SPEC TEST (NEGATIVE)`
+ *
+ * SECTIONS: contracts, declarations, contractBuilder, common
+ * NUMBER: 14
+ * DESCRIPTION: Contract is first statement in control flow terms, but not in tokens order terms.
+ * ISSUES: KT-26153
+ * HELPERS: functions
+ */
+
+import kotlin.contracts.*
+
+// TESTCASE NUMBER: 8
+val myProp = 8
+inline fun case_8(block: () -> Unit) {
+    myProp
+    <!CONTRACT_NOT_ALLOWED!>contract<!> {
+        callsInPlace(<!USAGE_IS_NOT_INLINABLE!>block<!>, InvocationKind.EXACTLY_ONCE)
+    }
+}

compiler/testData/diagnostics/tests/contracts/neg/14.kt

@@ -0,0 +1,23 @@
+// !DIAGNOSTICS: -UNUSED_VARIABLE
+// !OPT_IN: kotlin.contracts.ExperimentalContracts
+
+/*
+ * ADDITION TO `KOTLIN DIAGNOSTICS NOT LINKED SPEC TEST (NEGATIVE)`
+ *
+ * SECTIONS: contracts, declarations, contractBuilder, common
+ * NUMBER: 14
+ * DESCRIPTION: Contract is first statement in control flow terms, but not in tokens order terms.
+ * ISSUES: KT-26153
+ * HELPERS: functions
+ */
+
+import kotlin.contracts.*
+
+// TESTCASE NUMBER: 8
+val myProp = 8
+inline fun case_8(block: () -> Unit) {
+    myProp
+    <!CONTRACT_NOT_ALLOWED!>contract<!> {
+        callsInPlace(block, InvocationKind.EXACTLY_ONCE)
+    }
+}

compiler/testData/diagnostics/tests/contracts/pos/1.kt

@@ -0,0 +1,20 @@
+// FIR_IDENTICAL
+// !OPT_IN: kotlin.contracts.ExperimentalContracts
+
+/*
+ * ADDITION TO `KOTLIN DIAGNOSTICS NOT LINKED SPEC TEST (POSITIVE)`
+ *
+ * SECTIONS: contracts, declarations, contractBuilder, common
+ * NUMBER: 1
+ * DESCRIPTION: Functions with simple contracts.
+ */
+
+import kotlin.contracts.*
+
+// The following is an useful addition to testcases in compiler/tests-spec/testData/diagnostics/notLinked/contracts/declarations/contractBuilder/common/pos/1.kt
+
+// TESTCASE NUMBER: 4
+inline fun case_4(block: () -> Unit) {
+    kotlin.contracts.contract { callsInPlace(block, InvocationKind.EXACTLY_ONCE) }
+    return block()
+}

compiler/testData/diagnostics/tests/contracts/pos/14UserDefined.fir.kt

@@ -0,0 +1,82 @@
+// KT-60004: Ensure diagnostics for `kotlin.contracts.contract()` are not raised for user-defined fun `contract()`
+
+// !DIAGNOSTICS: -UNUSED_VARIABLE
+// !OPT_IN: kotlin.contracts.ExperimentalContracts
+
+/*
+ * POSITIVE ADDITION TO `KOTLIN DIAGNOSTICS NOT LINKED SPEC TEST (NEGATIVE)`
+ *
+ * SECTIONS: contracts, declarations, contractBuilder, common
+ * NUMBER: 14UserDefined
+ * DESCRIPTION: Invocation of user-defined fun `contract()` is first statement in control flow terms, but not in tokens order terms.
+ * ISSUES: KT-60004
+ * HELPERS: functions
+ */
+
+import kotlin.contracts.ContractBuilder
+import kotlin.contracts.InvocationKind
+
+// User-defined fun `contract()` must not be mixed up with `kotlin.contracts.contract()`
+inline fun contract(builder: ContractBuilder.() -> Unit) {}
+
+// TESTCASE NUMBER: 1
+inline fun case_1(block: () -> Unit) {
+    return contract {
+        callsInPlace(<!USAGE_IS_NOT_INLINABLE!>block<!>, InvocationKind.EXACTLY_ONCE)
+    }
+}
+
+// TESTCASE NUMBER: 2
+fun case_2() = contract { }
+
+// TESTCASE NUMBER: 3
+inline fun case_3(block: () -> Unit) {
+    val value_1 = contract {
+        callsInPlace(<!USAGE_IS_NOT_INLINABLE!>block<!>, InvocationKind.EXACTLY_ONCE)
+    }
+    block()
+}
+
+// TESTCASE NUMBER: 4
+inline fun case_4(block: () -> Unit) {
+    (contract {
+        callsInPlace(<!USAGE_IS_NOT_INLINABLE!>block<!>, InvocationKind.EXACTLY_ONCE)
+    })
+    return block()
+}
+
+// TESTCASE NUMBER: 5
+inline fun case_5(block: () -> Unit) {
+    test@ contract {
+        callsInPlace(<!USAGE_IS_NOT_INLINABLE!>block<!>, InvocationKind.EXACTLY_ONCE)
+    }
+    return block()
+}
+
+// TESTCASE NUMBER: 6
+inline fun case_6(block: () -> Unit) {
+    throw Exception(contract {
+        callsInPlace(<!USAGE_IS_NOT_INLINABLE!>block<!>, InvocationKind.EXACTLY_ONCE)
+    }.toString())
+}
+
+// TESTCASE NUMBER: 7
+inline fun case_7(block: () -> Unit) {
+    funWithAnyArg(contract {
+        callsInPlace(<!USAGE_IS_NOT_INLINABLE!>block<!>, InvocationKind.EXACTLY_ONCE)
+    })
+}
+
+// TESTCASE NUMBER: 8
+val myProp = 8
+inline fun case_8(block: () -> Unit) {
+    myProp
+    contract {
+        callsInPlace(<!USAGE_IS_NOT_INLINABLE!>block<!>, InvocationKind.EXACTLY_ONCE)
+    }
+}
+
+// Copy-pasted from compiler/tests-spec/testData/diagnostics/helpers/functions.kt
+fun funWithAnyArg(value_1: Any): Int {
+    return value_1.hashCode()
+}

compiler/testData/diagnostics/tests/contracts/pos/14UserDefined.kt

@@ -0,0 +1,82 @@
+// KT-60004: Ensure diagnostics for `kotlin.contracts.contract()` are not raised for user-defined fun `contract()`
+
+// !DIAGNOSTICS: -UNUSED_VARIABLE
+// !OPT_IN: kotlin.contracts.ExperimentalContracts
+
+/*
+ * POSITIVE ADDITION TO `KOTLIN DIAGNOSTICS NOT LINKED SPEC TEST (NEGATIVE)`
+ *
+ * SECTIONS: contracts, declarations, contractBuilder, common
+ * NUMBER: 14UserDefined
+ * DESCRIPTION: Invocation of user-defined fun `contract()` is first statement in control flow terms, but not in tokens order terms.
+ * ISSUES: KT-60004
+ * HELPERS: functions
+ */
+
+import kotlin.contracts.ContractBuilder
+import kotlin.contracts.InvocationKind
+
+// User-defined fun `contract()` must not be mixed up with `kotlin.contracts.contract()`
+inline fun contract(builder: ContractBuilder.() -> Unit) {}
+
+// TESTCASE NUMBER: 1
+inline fun case_1(block: () -> Unit) {
+    return contract {
+        callsInPlace(block, InvocationKind.EXACTLY_ONCE)
+    }
+}
+
+// TESTCASE NUMBER: 2
+fun case_2() = contract { }
+
+// TESTCASE NUMBER: 3
+inline fun case_3(block: () -> Unit) {
+    val value_1 = contract {
+        callsInPlace(block, InvocationKind.EXACTLY_ONCE)
+    }
+    block()
+}
+
+// TESTCASE NUMBER: 4
+inline fun case_4(block: () -> Unit) {
+    (contract {
+        callsInPlace(block, InvocationKind.EXACTLY_ONCE)
+    })
+    return block()
+}
+
+// TESTCASE NUMBER: 5
+inline fun case_5(block: () -> Unit) {
+    <!REDUNDANT_LABEL_WARNING!>test@<!> contract {
+        callsInPlace(block, InvocationKind.EXACTLY_ONCE)
+    }
+    return block()
+}
+
+// TESTCASE NUMBER: 6
+inline fun case_6(block: () -> Unit) {
+    throw Exception(contract {
+        callsInPlace(block, InvocationKind.EXACTLY_ONCE)
+    }.toString())
+}
+
+// TESTCASE NUMBER: 7
+inline fun case_7(block: () -> Unit) {
+    funWithAnyArg(contract {
+        callsInPlace(block, InvocationKind.EXACTLY_ONCE)
+    })
+}
+
+// TESTCASE NUMBER: 8
+val myProp = 8
+inline fun case_8(block: () -> Unit) {
+    myProp
+    contract {
+        callsInPlace(block, InvocationKind.EXACTLY_ONCE)
+    }
+}
+
+// Copy-pasted from compiler/tests-spec/testData/diagnostics/helpers/functions.kt
+fun funWithAnyArg(value_1: Any): Int {
+    return value_1.hashCode()
+}

compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/DiagnosticTestGenerated.java

@@ -6645,6 +6645,54 @@ public class DiagnosticTestGenerated extends AbstractDiagnosticTest {
             }
         }
 
+        @Nested
+        @TestMetadata("compiler/testData/diagnostics/tests/contracts")
+        @TestDataPath("$PROJECT_ROOT")
+        public class Contracts {
+            @Test
+            public void testAllFilesPresentInContracts() throws Exception {
+                KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/diagnostics/tests/contracts"), Pattern.compile("^(.*)\\.kts?$"), Pattern.compile("^(.+)\\.(reversed|fir|ll)\\.kts?$"), true);
+            }
+
+            @Nested
+            @TestMetadata("compiler/testData/diagnostics/tests/contracts/neg")
+            @TestDataPath("$PROJECT_ROOT")
+            public class Neg {
+                @Test
+                @TestMetadata("14.kt")
+                public void test14() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/contracts/neg/14.kt");
+                }
+
+                @Test
+                public void testAllFilesPresentInNeg() throws Exception {
+                    KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/diagnostics/tests/contracts/neg"), Pattern.compile("^(.*)\\.kts?$"), Pattern.compile("^(.+)\\.(reversed|fir|ll)\\.kts?$"), true);
+                }
+            }
+
+            @Nested
+            @TestMetadata("compiler/testData/diagnostics/tests/contracts/pos")
+            @TestDataPath("$PROJECT_ROOT")
+            public class Pos {
+                @Test
+                @TestMetadata("1.kt")
+                public void test1() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/contracts/pos/1.kt");
+                }
+
+                @Test
+                @TestMetadata("14UserDefined.kt")
+                public void test14UserDefined() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/contracts/pos/14UserDefined.kt");
+                }
+
+                @Test
+                public void testAllFilesPresentInPos() throws Exception {
+                    KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/diagnostics/tests/contracts/pos"), Pattern.compile("^(.*)\\.kts?$"), Pattern.compile("^(.+)\\.(reversed|fir|ll)\\.kts?$"), true);
+                }
+            }
+        }
+
         @Nested
         @TestMetadata("compiler/testData/diagnostics/tests/controlFlowAnalysis")
         @TestDataPath("$PROJECT_ROOT")

compiler/tests-spec/testData/diagnostics/notLinked/contracts/declarations/contractBuilder/common/neg/14.fir.kt

@@ -33,7 +33,7 @@ inline fun case_3(block: () -> Unit) {
 
 // TESTCASE NUMBER: 4
 inline fun case_4(block: () -> Unit) {
-    (contract {
+    (<!CONTRACT_NOT_ALLOWED!>contract<!> {
         callsInPlace(block, InvocationKind.EXACTLY_ONCE)
     })
     return block()
@@ -41,7 +41,7 @@ inline fun case_4(block: () -> Unit) {
 
 // TESTCASE NUMBER: 5
 inline fun case_5(block: () -> Unit) {
-    test@ contract {
+    test@ <!CONTRACT_NOT_ALLOWED!>contract<!> {
         callsInPlace(block, InvocationKind.EXACTLY_ONCE)
     }
     return block()