[FIR] Add resolution of contracts that are written using the new syntax

compiler/fir/analysis-tests/testData/resolveWithStdlib/contracts/fromSource/good/variousContracts/newSyntax/functionsWithContract.kt

@@ -0,0 +1,22 @@
+import kotlin.contracts.*
+
+fun test1(arg: Any?) contract [
+    returns() implies (arg != null)
+] {
+    require(arg != null)
+}
+
+fun test2(s: String?, block: () -> Unit) contract [
+    returns() implies (s != null),
+    callsInPlace(block, InvocationKind.EXACTLY_ONCE)
+] {
+    require(s != null)
+    block()
+}
+
+fun test3(arg: Any?): Boolean contract [
+    returns(true) implies (arg != null)
+] {
+    require(arg != null)
+    return true
+}

compiler/fir/analysis-tests/testData/resolveWithStdlib/contracts/fromSource/good/variousContracts/newSyntax/functionsWithContract.txt

@@ -0,0 +1,28 @@
+FILE: functionsWithContract.kt
+    public final fun test1(arg: R|kotlin/Any?|): R|kotlin/Unit|
+        [R|Contract description]
+         <
+            Returns(WILDCARD) -> arg != null
+        >
+     {
+        R|kotlin/require|(!=(R|<local>/arg|, Null(null)))
+    }
+    public final fun test2(s: R|kotlin/String?|, block: R|() -> kotlin/Unit|): R|kotlin/Unit|
+        [R|Contract description]
+         <
+            Returns(WILDCARD) -> s != null
+            CallsInPlace(block, EXACTLY_ONCE)
+        >
+     {
+        R|kotlin/require|(!=(R|<local>/s|, Null(null)))
+        R|<local>/block|.R|FakeOverride<kotlin/Function0.invoke: R|kotlin/Unit|>|()
+    }
+    public final fun test3(arg: R|kotlin/Any?|): R|kotlin/Boolean|
+        [R|Contract description]
+         <
+            Returns(TRUE) -> arg != null
+        >
+     {
+        R|kotlin/require|(!=(R|<local>/arg|, Null(null)))
+        ^test3 Boolean(true)
+    }

compiler/fir/analysis-tests/tests/org/jetbrains/kotlin/fir/FirDiagnosticsWithStdlibTestGenerated.java

@@ -747,6 +747,37 @@ public class FirDiagnosticsWithStdlibTestGenerated extends AbstractFirDiagnostic
                         runTest("compiler/fir/analysis-tests/testData/resolveWithStdlib/contracts/fromSource/good/returnsImplies/typePredicate.kt");
                     }
                 }
+
+                @TestMetadata("compiler/fir/analysis-tests/testData/resolveWithStdlib/contracts/fromSource/good/variousContracts")
+                @TestDataPath("$PROJECT_ROOT")
+                @RunWith(JUnit3RunnerWithInners.class)
+                public static class VariousContracts extends AbstractFirDiagnosticsWithStdlibTest {
+                    private void runTest(String testDataFilePath) throws Exception {
+                        KotlinTestUtils.runTest(this::doTest, this, testDataFilePath);
+                    }
+
+                    public void testAllFilesPresentInVariousContracts() throws Exception {
+                        KotlinTestUtils.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/fir/analysis-tests/testData/resolveWithStdlib/contracts/fromSource/good/variousContracts"), Pattern.compile("^([^.]+)\\.kt$"), null, true);
+                    }
+
+                    @TestMetadata("compiler/fir/analysis-tests/testData/resolveWithStdlib/contracts/fromSource/good/variousContracts/newSyntax")
+                    @TestDataPath("$PROJECT_ROOT")
+                    @RunWith(JUnit3RunnerWithInners.class)
+                    public static class NewSyntax extends AbstractFirDiagnosticsWithStdlibTest {
+                        private void runTest(String testDataFilePath) throws Exception {
+                            KotlinTestUtils.runTest(this::doTest, this, testDataFilePath);
+                        }
+
+                        public void testAllFilesPresentInNewSyntax() throws Exception {
+                            KotlinTestUtils.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/fir/analysis-tests/testData/resolveWithStdlib/contracts/fromSource/good/variousContracts/newSyntax"), Pattern.compile("^([^.]+)\\.kt$"), null, true);
+                        }
+
+                        @TestMetadata("functionsWithContract.kt")
+                        public void testFunctionsWithContract() throws Exception {
+                            runTest("compiler/fir/analysis-tests/testData/resolveWithStdlib/contracts/fromSource/good/variousContracts/newSyntax/functionsWithContract.kt");
+                        }
+                    }
+                }
             }
         }
     }

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

@@ -5,27 +5,44 @@
 
 package org.jetbrains.kotlin.fir.resolve.transformers.contracts
 
-import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.contracts.FirLegacyRawContractDescription
 import org.jetbrains.kotlin.fir.contracts.FirRawContractDescription
+import org.jetbrains.kotlin.fir.contracts.FirResolvedContractDescription
 import org.jetbrains.kotlin.fir.contracts.builder.buildEffectDeclaration
+import org.jetbrains.kotlin.fir.contracts.builder.buildLegacyRawContractDescription
 import org.jetbrains.kotlin.fir.contracts.builder.buildResolvedContractDescription
 import org.jetbrains.kotlin.fir.contracts.description.ConeEffectDeclaration
 import org.jetbrains.kotlin.fir.contracts.impl.FirEmptyContractDescription
 import org.jetbrains.kotlin.fir.contracts.toFirEffectDeclaration
 import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.declarations.builder.buildAnonymousFunction
+import org.jetbrains.kotlin.fir.declarations.impl.FirSimpleFunctionImpl
 import org.jetbrains.kotlin.fir.declarations.synthetic.FirSyntheticProperty
-import org.jetbrains.kotlin.fir.errorTypeFromPrototype
 import org.jetbrains.kotlin.fir.expressions.*
+import org.jetbrains.kotlin.fir.expressions.builder.*
+import org.jetbrains.kotlin.fir.references.builder.buildSimpleNamedReference
 import org.jetbrains.kotlin.fir.resolve.ResolutionMode
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
+import org.jetbrains.kotlin.fir.resolve.calls.ImplicitExtensionReceiverValue
 import org.jetbrains.kotlin.fir.resolve.diagnostics.ConeContractDescriptionError
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirBodyResolveTransformer
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirDeclarationsResolveTransformer
+import org.jetbrains.kotlin.fir.symbols.StandardClassIds
+import org.jetbrains.kotlin.fir.symbols.impl.ConeClassLikeLookupTagImpl
+import org.jetbrains.kotlin.fir.symbols.impl.FirAnonymousFunctionSymbol
+import org.jetbrains.kotlin.fir.types.ConeStarProjection.kind
+import org.jetbrains.kotlin.fir.types.builder.buildImplicitTypeRef
+import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
 import org.jetbrains.kotlin.fir.types.coneType
-import org.jetbrains.kotlin.fir.visitors.CompositeTransformResult
+import org.jetbrains.kotlin.fir.types.impl.ConeClassLikeTypeImpl
-import org.jetbrains.kotlin.fir.visitors.compose
+import org.jetbrains.kotlin.fir.types.impl.FirImplicitBuiltinTypeRef
-import org.jetbrains.kotlin.fir.visitors.transformSingle
+import org.jetbrains.kotlin.fir.visitors.*
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.name.FqNameUnsafe
+import org.jetbrains.kotlin.name.Name
+import kotlin.contracts.contract
 
 class FirContractResolveTransformer(
     session: FirSession,
@@ -116,8 +133,17 @@ class FirContractResolveTransformer(
             owner: T
         ): CompositeTransformResult<T> {
             dataFlowAnalyzer.enterContractDescription()
-            val contractDescription = owner.contractDescription
+            return when (val contractDescription = owner.contractDescription) {
-            require(contractDescription is FirLegacyRawContractDescription)
+                is FirLegacyRawContractDescription -> transformLegacyRawContractDescriptionOwner(owner, contractDescription)
+                is FirRawContractDescription -> transformRawContractDescriptionOwner(owner, contractDescription)
+                else -> throw IllegalArgumentException("$owner has a contract description of an unknown type")
+            }
+        }
+
+        private fun <T : FirContractDescriptionOwner> transformLegacyRawContractDescriptionOwner(
+            owner: T,
+            contractDescription: FirLegacyRawContractDescription
+        ): CompositeTransformResult<T> {
             val valueParameters = owner.valueParameters
             val contractCall = withNewLocalScope {
                 for (valueParameter in valueParameters) {
@@ -150,13 +176,63 @@ class FirContractResolveTransformer(
             return owner.compose()
         }
 
+        private fun <T : FirContractDescriptionOwner> transformRawContractDescriptionOwner(
+            owner: T,
+            contractDescription: FirRawContractDescription
+        ): CompositeTransformResult<T> {
+            val effectsBlock = buildAnonymousFunction {
+                session = this@FirDeclarationsContractResolveTransformer.session
+                origin = FirDeclarationOrigin.Source
+                returnTypeRef = buildImplicitTypeRef()
+                receiverTypeRef = buildImplicitTypeRef()
+                symbol = FirAnonymousFunctionSymbol()
+                isLambda = true
+
+                body = buildBlock {
+                    contractDescription.rawEffects.forEach {
+                        statements += it
+                    }
+                }
+            }
+
+            val lambdaArgument = buildLambdaArgumentExpression {
+                expression = effectsBlock
+            }
+
+            val contractCall = buildFunctionCall {
+                calleeReference = buildSimpleNamedReference {
+                    name = Name.identifier("contract")
+                }
+                argumentList = buildArgumentList {
+                    arguments += lambdaArgument
+                }
+            }
+
+            val legacyRawContractDescription = buildLegacyRawContractDescription {
+                this.contractCall = contractCall
+            }
+
+            owner.replaceContractDescription(legacyRawContractDescription)
+
+            return transformLegacyRawContractDescriptionOwner(owner, legacyRawContractDescription)
+        }
+
         private fun <T : FirContractDescriptionOwner> transformOwnerWithUnresolvedContract(owner: T): CompositeTransformResult<T> {
-            val contractDescription = owner.contractDescription as FirLegacyRawContractDescription
+            return when (val contractDescription = owner.contractDescription) {
-            val functionCall = contractDescription.contractCall
+                is FirLegacyRawContractDescription -> { // old syntax contract description
-            owner.replaceContractDescription(FirEmptyContractDescription)
+                    val functionCall = contractDescription.contractCall
-            owner.body.replaceFirstStatement(functionCall)
+                    owner.replaceContractDescription(FirEmptyContractDescription)
-            dataFlowAnalyzer.exitContractDescription()
+                    owner.body.replaceFirstStatement(functionCall)
-            return owner.compose()
+                    dataFlowAnalyzer.exitContractDescription()
+                    owner.compose()
+                }
+                is FirRawContractDescription -> { // new syntax contract description
+                    owner.replaceContractDescription(FirEmptyContractDescription)
+                    dataFlowAnalyzer.exitContractDescription()
+                    owner.compose()
+                }
+                else -> owner.compose() // TODO: change
+            }
         }
 
         override fun transformRegularClass(regularClass: FirRegularClass, data: ResolutionMode): CompositeTransformResult<FirStatement> {
@@ -200,7 +276,7 @@ class FirContractResolveTransformer(
         }
 
         private val FirContractDescriptionOwner.hasContractToResolve: Boolean
-            get() = contractDescription is FirRawContractDescription
+            get() = contractDescription is FirLegacyRawContractDescription || contractDescription is FirRawContractDescription
 
         private fun FirDeclaration.updatePhase() {
             replaceResolvePhase(FirResolvePhase.CONTRACTS)