[FIR] Add extracting contracts in raw fir builder and light tree builder

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

@@ -9,6 +9,8 @@ import com.intellij.psi.tree.IElementType
 import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.descriptors.Visibilities
 import org.jetbrains.kotlin.fir.*
+import org.jetbrains.kotlin.fir.contracts.FirContractDescription
+import org.jetbrains.kotlin.fir.contracts.builder.buildRawContractDescription
 import org.jetbrains.kotlin.fir.declarations.FirProperty
 import org.jetbrains.kotlin.fir.declarations.FirVariable
 import org.jetbrains.kotlin.fir.declarations.builder.*
@@ -17,7 +19,9 @@ import org.jetbrains.kotlin.fir.declarations.impl.FirDefaultPropertyAccessor
 import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.expressions.builder.*
 import org.jetbrains.kotlin.fir.expressions.impl.FirSingleExpressionBlock
+import org.jetbrains.kotlin.fir.expressions.impl.FirStubStatement
 import org.jetbrains.kotlin.fir.expressions.impl.buildSingleExpressionBlock
+import org.jetbrains.kotlin.fir.references.FirNamedReference
 import org.jetbrains.kotlin.fir.references.builder.*
 import org.jetbrains.kotlin.fir.symbols.StandardClassIds
 import org.jetbrains.kotlin.fir.symbols.constructStarProjectedType
@@ -35,6 +39,8 @@ import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.types.Variance
 import org.jetbrains.kotlin.types.expressions.OperatorConventions
 import org.jetbrains.kotlin.util.OperatorNameConventions
+import kotlin.contracts.ExperimentalContracts
+import kotlin.contracts.contract
 
 fun String.parseCharacter(): Char? {
     // Strip the quotes
@@ -458,3 +464,32 @@ private val GET_VALUE = Name.identifier("getValue")
 private val SET_VALUE = Name.identifier("setValue")
 private val PROVIDE_DELEGATE = Name.identifier("provideDelegate")
 private val DELEGATED_SETTER_PARAM = Name.special("<set-?>")
+
+fun FirBlock?.extractContractDescriptionIfPossible(): Pair<FirBlock?, FirContractDescription?> {
+    if (this == null) return null to null
+    if (!isContractPresentFirCheck()) return this to null
+    val contractCall = replaceFirstStatement(FirStubStatement) as FirFunctionCall
+    return this to buildRawContractDescription {
+        source = contractCall.source
+        this.contractCall = contractCall
+    }
+}
+
+fun FirBlock.isContractPresentFirCheck(): Boolean {
+    val firstStatement = statements.firstOrNull() ?: return false
+    val contractCall = firstStatement as? FirFunctionCall ?: return false
+    if (contractCall.calleeReference.name.asString() != "contract") return false
+    val receiver = contractCall.explicitReceiver as? FirQualifiedAccessExpression ?: return true
+    if (!contractCall.checkReceiver("contracts")) return false
+    if (!receiver.checkReceiver("kotlin")) return false
+    val receiverOfReceiver = receiver.explicitReceiver as? FirQualifiedAccessExpression ?: return false
+    if (receiverOfReceiver.explicitReceiver != null) return false
+    return true
+}
+
+private fun FirExpression.checkReceiver(name: String?): Boolean {
+    if (this !is FirQualifiedAccessExpression) return false
+    val receiver = explicitReceiver as? FirQualifiedAccessExpression ?: return false
+    val receiverName = (receiver.calleeReference as? FirNamedReference)?.name?.asString() ?: return false
+    return receiverName == name
+}

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

@@ -15,7 +15,9 @@ import org.jetbrains.kotlin.descriptors.Visibilities
 import org.jetbrains.kotlin.descriptors.annotations.AnnotationUseSiteTarget
 import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.builder.Context
+import org.jetbrains.kotlin.fir.builder.extractContractDescriptionIfPossible
 import org.jetbrains.kotlin.fir.builder.generateAccessorsByDelegate
+import org.jetbrains.kotlin.fir.contracts.FirContractDescription
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.builder.*
 import org.jetbrains.kotlin.fir.declarations.impl.FirDeclarationStatusImpl
@@ -757,7 +759,8 @@ class DeclarationsConverter(
             annotations += modifiers.annotations
             typeParameters += constructorTypeParametersFromConstructedClass(classWrapper.classBuilder.typeParameters)
             valueParameters += firValueParameters.map { it.firValueParameter }
-            body = convertFunctionBody(block, null)
+            val (body, _) = convertFunctionBody(block, null)
+            this.body = body
             context.firFunctionTargets.removeLast()
         }.also {
             target.bind(it)
@@ -1053,7 +1056,11 @@ class DeclarationsConverter(
                 valueParameters += firValueParameters
             }
 
-            body = convertFunctionBody(block, expression)
+            val (body, contractDescription) = convertFunctionBody(block, expression)
+            this.body = body
+            contractDescription?.let {
+                this.contractDescription = contractDescription
+            }
             context.firFunctionTargets.removeLast()
         }.also {
             target.bind(it)
@@ -1178,7 +1185,14 @@ class DeclarationsConverter(
                     addCapturedTypeParameters(typeParameters)
                 }
                 valueParametersList?.let { list -> valueParameters += convertValueParameters(list).map { it.firValueParameter } }
-                body = convertFunctionBody(block, expression)
+                val (body, contractDescription) = convertFunctionBody(block, expression)
+                this.body = body
+                contractDescription?.let {
+                    // TODO: add error reporting for contracts on lambdas
+                    if (this is FirSimpleFunctionBuilder) {
+                        this.contractDescription = it
+                    }
+                }
             }
             context.firFunctionTargets.removeLast()
         }.build().also {
@@ -1190,13 +1204,13 @@ class DeclarationsConverter(
      * @see org.jetbrains.kotlin.parsing.KotlinParsing.parseFunctionBody
      * @see org.jetbrains.kotlin.fir.builder.RawFirBuilder.Visitor.buildFirBody
      */
-    private fun convertFunctionBody(blockNode: LighterASTNode?, expression: LighterASTNode?): FirBlock? {
+    private fun convertFunctionBody(blockNode: LighterASTNode?, expression: LighterASTNode?): Pair<FirBlock?, FirContractDescription?> {
         return when {
-            blockNode != null -> return convertBlock(blockNode)
+            blockNode != null -> convertBlock(blockNode).extractContractDescriptionIfPossible()
             expression != null -> FirSingleExpressionBlock(
                 expressionConverter.getAsFirExpression<FirExpression>(expression, "Function has no body (but should)").toReturn()
-            )
+            ) to null
-            else -> null
+            else -> null to null
         }
     }
 

compiler/fir/raw-fir/light-tree2fir/tests/org/jetbrains/kotlin/fir/lightTree/LightTree2FirConverterTestCaseGenerated.java

@@ -55,6 +55,11 @@ public class LightTree2FirConverterTestCaseGenerated extends AbstractLightTree2F
             runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/constructorInObject.kt");
         }
 
+        @TestMetadata("contractDescription.kt")
+        public void testContractDescription() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/contractDescription.kt");
+        }
+
         @TestMetadata("derivedClass.kt")
         public void testDerivedClass() throws Exception {
             runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/derivedClass.kt");

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

@@ -14,6 +14,8 @@ import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.descriptors.Visibilities
 import org.jetbrains.kotlin.descriptors.Visibility
 import org.jetbrains.kotlin.fir.*
+import org.jetbrains.kotlin.fir.contracts.FirContractDescription
+import org.jetbrains.kotlin.fir.contracts.builder.buildRawContractDescription
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.builder.*
 import org.jetbrains.kotlin.fir.declarations.impl.FirDeclarationStatusImpl
@@ -26,6 +28,7 @@ import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.expressions.builder.*
 import org.jetbrains.kotlin.fir.expressions.impl.FirModifiableQualifiedAccess
 import org.jetbrains.kotlin.fir.expressions.impl.FirSingleExpressionBlock
+import org.jetbrains.kotlin.fir.expressions.impl.FirStubStatement
 import org.jetbrains.kotlin.fir.references.FirNamedReference
 import org.jetbrains.kotlin.fir.references.builder.*
 import org.jetbrains.kotlin.fir.scopes.FirScopeProvider
@@ -42,10 +45,24 @@ import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.getStrictParentOfType
 import org.jetbrains.kotlin.psi.psiUtil.hasActualModifier
 import org.jetbrains.kotlin.psi.psiUtil.hasExpectModifier
+import org.jetbrains.kotlin.psi.psiUtil.isContractDescriptionCallPsiCheck
 import org.jetbrains.kotlin.types.Variance
 import org.jetbrains.kotlin.types.expressions.OperatorConventions
 import org.jetbrains.kotlin.util.OperatorNameConventions
 import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstanceOrNull
+import kotlin.collections.MutableList
+import kotlin.collections.any
+import kotlin.collections.contains
+import kotlin.collections.filterIsInstance
+import kotlin.collections.first
+import kotlin.collections.firstOrNull
+import kotlin.collections.forEach
+import kotlin.collections.isNotEmpty
+import kotlin.collections.lastOrNull
+import kotlin.collections.plusAssign
+import kotlin.collections.reverse
+import kotlin.collections.withIndex
+import kotlin.collections.zip
 
 class RawFirBuilder(
     session: FirSession, val baseScopeProvider: FirScopeProvider, val stubMode: Boolean
@@ -199,19 +216,20 @@ class RawFirBuilder(
                     FirSingleExpressionBlock(convert())
             }
 
-        private fun KtDeclarationWithBody.buildFirBody(): FirBlock? =
+        private fun KtDeclarationWithBody.buildFirBody(): Pair<FirBlock?, FirContractDescription?> =
             when {
                 !hasBody() ->
-                    null
+                    null to null
                 hasBlockBody() -> if (!stubMode) {
-                    bodyBlockExpression?.accept(this@Visitor, Unit) as? FirBlock
+                    val block = bodyBlockExpression?.accept(this@Visitor, Unit) as? FirBlock
+                    block.extractContractDescriptionIfPossible()
                 } else {
-                    FirSingleExpressionBlock(buildExpressionStub { source = this@buildFirBody.toFirSourceElement() }.toReturn())
+                    FirSingleExpressionBlock(buildExpressionStub { source = this@buildFirBody.toFirSourceElement() }.toReturn()) to null
                 }
                 else -> {
                     val result = { bodyExpression }.toFirExpression("Function has no body (but should)")
                     // basePsi is null, because 'return' is synthetic & should not be bound to some PSI
-                    FirSingleExpressionBlock(result.toReturn(baseSource = null))
+                    FirSingleExpressionBlock(result.toReturn(baseSource = null)) to null
                 }
             }
 
@@ -289,7 +307,11 @@ class RawFirBuilder(
                     }
                 }
                 symbol = FirPropertyAccessorSymbol()
-                body = this@toFirPropertyAccessor.buildFirBody()
+                val (body, contractDescription) = this@toFirPropertyAccessor.buildFirBody()
+                this.body = body
+                contractDescription?.let {
+                    this.contractDescription = it
+                }
             }.also {
                 accessorTarget.bind(it)
                 this@RawFirBuilder.context.firFunctionTargets.removeLast()
@@ -824,7 +846,14 @@ class RawFirBuilder(
                 }
                 withCapturedTypeParameters {
                     if (this is FirSimpleFunctionBuilder) addCapturedTypeParameters(this.typeParameters)
-                    body = function.buildFirBody()
+                    val (body, contractDescription) = function.buildFirBody()
+                    this.body = body
+                    contractDescription?.let {
+                        // TODO: add error reporting for contracts on lambdas
+                        if (this is FirSimpleFunctionBuilder) {
+                            this.contractDescription = it
+                        }
+                    }
                 }
                 context.firFunctionTargets.removeLast()
             }.build().also {
@@ -945,7 +974,8 @@ class RawFirBuilder(
                 extractAnnotationsTo(this)
                 typeParameters += constructorTypeParametersFromConstructedClass(ownerTypeParameters)
                 extractValueParametersTo(this)
-                body = buildFirBody()
+                val (body, _) = buildFirBody()
+                this.body = body
                 this@RawFirBuilder.context.firFunctionTargets.removeLast()
             }.also {
                 target.bind(it)

compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/contractDescription.kt

@@ -0,0 +1,62 @@
+// Should have raw description
+fun test_1() {
+    contract {
+        callsInPlace()
+    }
+    test_1()
+}
+
+fun test_2() {
+    kotlin.contracts.contract {
+        callsInPlace()
+        callsInPlace()
+    }
+    test_2()
+}
+
+var test_3: Int = 1
+    get() {
+        contract {
+            callsInPlace()
+        }
+        return 1
+    }
+    set(value) {
+        kotlin.contracts.contract {
+            callsInPlace()
+            callsInPlace()
+        }
+    }
+
+fun test_4() {
+    contract()
+    test_4()
+}
+
+// should not have raw description
+
+fun test_5() {
+    test_5()
+    contract()
+}
+
+fun test_6() {
+    aaa.bbb.ccc.contract {
+
+    }
+    test_6()
+}
+
+fun test_7() {
+    contracts.contract {
+
+    }
+    test_7()
+}
+
+fun test_8() {
+    aaa.kotlin.contracts.contract {
+
+    }
+    test_8()
+}

compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/contractDescription.txt

@@ -0,0 +1,80 @@
+FILE: contractDescription.kt
+    public? final? fun test_1(): R|kotlin/Unit|
+        [Contract description] <
+            contract#(<L> = contract@fun <implicit>.<anonymous>(): <implicit> {
+                callsInPlace#()
+            }
+            )
+        >
+     {
+        [StubStatement]
+        test_1#()
+    }
+    public? final? fun test_2(): R|kotlin/Unit|
+        [Contract description] <
+            kotlin#.contracts#.contract#(<L> = contract@fun <implicit>.<anonymous>(): <implicit> {
+                callsInPlace#()
+                callsInPlace#()
+            }
+            )
+        >
+     {
+        [StubStatement]
+        test_2#()
+    }
+    public? final? var test_3: Int = IntegerLiteral(1)
+        public? get(): Int
+            [Contract description] <
+                contract#(<L> = contract@fun <implicit>.<anonymous>(): <implicit> {
+                    callsInPlace#()
+                }
+                )
+            >
+         {
+            [StubStatement]
+            ^ IntegerLiteral(1)
+        }
+        public? set(value: Int): R|kotlin/Unit|
+            [Contract description] <
+                kotlin#.contracts#.contract#(<L> = contract@fun <implicit>.<anonymous>(): <implicit> {
+                    callsInPlace#()
+                    callsInPlace#()
+                }
+                )
+            >
+         {
+            [StubStatement]
+        }
+    public? final? fun test_4(): R|kotlin/Unit|
+        [Contract description] <
+            contract#()
+        >
+     {
+        [StubStatement]
+        test_4#()
+    }
+    public? final? fun test_5(): R|kotlin/Unit| {
+        test_5#()
+        contract#()
+    }
+    public? final? fun test_6(): R|kotlin/Unit| {
+        aaa#.bbb#.ccc#.contract#(<L> = contract@fun <implicit>.<anonymous>(): <implicit> {
+            Unit
+        }
+        )
+        test_6#()
+    }
+    public? final? fun test_7(): R|kotlin/Unit| {
+        contracts#.contract#(<L> = contract@fun <implicit>.<anonymous>(): <implicit> {
+            Unit
+        }
+        )
+        test_7#()
+    }
+    public? final? fun test_8(): R|kotlin/Unit| {
+        aaa#.kotlin#.contracts#.contract#(<L> = contract@fun <implicit>.<anonymous>(): <implicit> {
+            Unit
+        }
+        )
+        test_8#()
+    }

compiler/fir/raw-fir/psi2fir/tests/org/jetbrains/kotlin/fir/builder/RawFirBuilderTestCaseGenerated.java

@@ -55,6 +55,11 @@ public class RawFirBuilderTestCaseGenerated extends AbstractRawFirBuilderTestCas
             runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/constructorInObject.kt");
         }
 
+        @TestMetadata("contractDescription.kt")
+        public void testContractDescription() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/contractDescription.kt");
+        }
+
         @TestMetadata("derivedClass.kt")
         public void testDerivedClass() throws Exception {
             runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/derivedClass.kt");

compiler/fir/tree/src/org/jetbrains/kotlin/fir/FirRenderer.kt

@@ -544,7 +544,11 @@ class FirRenderer(builder: StringBuilder, private val mode: RenderMode = RenderM
     }
 
     override fun visitStatement(statement: FirStatement) {
-        visitElement(statement)
+        if (statement is FirStubStatement) {
+            print("[StubStatement]")
+        } else {
+            visitElement(statement)
+        }
     }
 
     override fun visitReturnExpression(returnExpression: FirReturnExpression) {

compiler/fir/tree/src/org/jetbrains/kotlin/fir/expressions/FirExpressionUtil.kt

@@ -84,4 +84,13 @@ fun <D> FirBlock.transformAllStatementsExceptLast(transformer: FirTransformer<D>
             TransformData.Nothing
         }
     }
+}
+
+fun FirBlock.replaceFirstStatement(statement: FirStatement): FirStatement {
+    require(this is FirBlockImpl) {
+        "replaceFirstStatement should not be called for ${this::class.simpleName}"
+    }
+    val existed = statements[0]
+    statements[0] = statement
+    return existed
 }

compiler/fir/tree/src/org/jetbrains/kotlin/fir/expressions/impl/FirStubStatement.kt

@@ -0,0 +1,32 @@
+/*
+ * 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.expressions.impl
+
+import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.FirPureAbstractElement
+import org.jetbrains.kotlin.fir.FirSourceElement
+import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
+import org.jetbrains.kotlin.fir.expressions.FirStatement
+import org.jetbrains.kotlin.fir.visitors.FirTransformer
+import org.jetbrains.kotlin.fir.visitors.FirVisitor
+
+object FirStubStatement : FirPureAbstractElement(), FirStatement {
+    override val source: FirSourceElement?
+        get() = null
+
+    override val annotations: List<FirAnnotationCall>
+        get() = emptyList()
+
+    override fun <D> transformAnnotations(transformer: FirTransformer<D>, data: D): FirStatement {
+        return this
+    }
+
+    override fun <R, D> acceptChildren(visitor: FirVisitor<R, D>, data: D) {}
+
+    override fun <D> transformChildren(transformer: FirTransformer<D>, data: D): FirElement {
+        return this
+    }
+}

compiler/visualizer/tests/org/jetbrains/kotlin/visualizer/fir/FirVisualizerForRawFirDataGenerated.java

@@ -55,6 +55,11 @@ public class FirVisualizerForRawFirDataGenerated extends AbstractFirVisualizer {
             runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/constructorInObject.kt");
         }
 
+        @TestMetadata("contractDescription.kt")
+        public void testContractDescription() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/contractDescription.kt");
+        }
+
         @TestMetadata("derivedClass.kt")
         public void testDerivedClass() throws Exception {
             runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/derivedClass.kt");

compiler/visualizer/tests/org/jetbrains/kotlin/visualizer/psi/PsiVisualizerForRawFirDataGenerated.java

@@ -55,6 +55,11 @@ public class PsiVisualizerForRawFirDataGenerated extends AbstractPsiVisualizer {
             runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/constructorInObject.kt");
         }
 
+        @TestMetadata("contractDescription.kt")
+        public void testContractDescription() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/contractDescription.kt");
+        }
+
         @TestMetadata("derivedClass.kt")
         public void testDerivedClass() throws Exception {
             runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/derivedClass.kt");