[FIR] Add contract description blocks support to RawFirBuilder

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

@@ -11,6 +11,7 @@ import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.descriptors.Modality
 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
@@ -225,7 +226,11 @@ class RawFirBuilder(
                     null to null
                 hasBlockBody() -> if (!stubMode) {
                     val block = bodyBlockExpression?.accept(this@Visitor, Unit) as? FirBlock
-                    block.extractContractDescriptionIfPossible()
+                    if (hasContractEffectList()) {
+                        block to null
+                    } else {
+                        block.extractContractDescriptionIfPossible()
+                    }
                 } else {
                     FirSingleExpressionBlock(buildExpressionStub { source = this@buildFirBody.toFirSourceElement() }.toReturn()) to null
                 }
@@ -331,8 +336,10 @@ class RawFirBuilder(
                     }
                 }
                 symbol = FirPropertyAccessorSymbol()
-                val (body, contractDescription) = this@toFirPropertyAccessor.buildFirBody()
+                val outerContractDescription = this@toFirPropertyAccessor.obtainContractDescription()
-                this.body = body
+                val bodyWithContractDescription = this@toFirPropertyAccessor.buildFirBody()
+                this.body = bodyWithContractDescription.first
+                val contractDescription = outerContractDescription ?: bodyWithContractDescription.second
                 contractDescription?.let {
                     this.contractDescription = it
                 }
@@ -954,8 +961,10 @@ class RawFirBuilder(
                 }
                 withCapturedTypeParameters {
                     if (this is FirSimpleFunctionBuilder) addCapturedTypeParameters(this.typeParameters)
-                    val (body, contractDescription) = function.buildFirBody()
+                    val outerContractDescription = function.obtainContractDescription()
-                    this.body = body
+                    val bodyWithContractDescription = function.buildFirBody()
+                    this.body = bodyWithContractDescription.first
+                    val contractDescription = outerContractDescription ?: bodyWithContractDescription.second
                     contractDescription?.let {
                         // TODO: add error reporting for contracts on lambdas
                         if (this is FirSimpleFunctionBuilder) {
@@ -969,6 +978,21 @@ class RawFirBuilder(
             }
         }
 
+        private fun KtDeclarationWithBody.obtainContractDescription(): FirContractDescription? {
+            return when(val description = contractDescription) {
+                null -> null
+                else -> buildRawContractDescription {
+                    source = description.toFirSourceElement()
+                    description.extractRawEffects(rawEffects)
+                }
+            }
+        }
+
+        private fun KtContractEffectList.extractRawEffects(destination: MutableList<FirExpression>) {
+            getExpressions()
+                .mapTo(destination) { it.accept(this@Visitor, Unit) as FirExpression }
+        }
+
         override fun visitLambdaExpression(expression: KtLambdaExpression, data: Unit): FirElement {
             val literal = expression.functionLiteral
             val literalSource = literal.toFirSourceElement()

compiler/psi/src/org/jetbrains/kotlin/psi/KtDeclarationWithBody.java

@@ -38,6 +38,10 @@ public interface KtDeclarationWithBody extends KtDeclaration {
         return null;
     }
 
+    default boolean hasContractEffectList() {
+        return getContractDescription() != null;
+    }
+
     boolean hasBlockBody();
 
     boolean hasBody();