diff --git a/compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/RawFirBuilder.kt b/compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/RawFirBuilder.kt index 271aef7dad5..e16032654d9 100644 --- a/compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/RawFirBuilder.kt +++ b/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() - this.body = body + val outerContractDescription = this@toFirPropertyAccessor.obtainContractDescription() + 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() - this.body = body + val outerContractDescription = function.obtainContractDescription() + 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) { + 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() diff --git a/compiler/psi/src/org/jetbrains/kotlin/psi/KtDeclarationWithBody.java b/compiler/psi/src/org/jetbrains/kotlin/psi/KtDeclarationWithBody.java index 724d15e128d..abfa45dcd76 100644 --- a/compiler/psi/src/org/jetbrains/kotlin/psi/KtDeclarationWithBody.java +++ b/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();