[FIR] Substitute types from arguments in evaluating of contract

#KT-41078

compiler/fir/analysis-tests/legacy-fir-tests/tests-gen/org/jetbrains/kotlin/fir/LazyBodyIsNotTouchedTilContractsPhaseTestGenerated.java

@@ -980,6 +980,24 @@ public class LazyBodyIsNotTouchedTilContractsPhaseTestGenerated extends Abstract
         }
     }
 
+    @TestMetadata("compiler/fir/analysis-tests/testData/resolve/contracts")
+    @TestDataPath("$PROJECT_ROOT")
+    @RunWith(JUnit3RunnerWithInners.class)
+    public static class Contracts extends AbstractLazyBodyIsNotTouchedTilContractsPhaseTest {
+        private void runTest(String testDataFilePath) throws Exception {
+            KotlinTestUtils.runTest(this::doTest, this, testDataFilePath);
+        }
+
+        public void testAllFilesPresentInContracts() throws Exception {
+            KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/fir/analysis-tests/testData/resolve/contracts"), Pattern.compile("^([^.]+)\\.kt$"), null, true);
+        }
+
+        @TestMetadata("genericContract.kt")
+        public void testGenericContract() throws Exception {
+            runTest("compiler/fir/analysis-tests/testData/resolve/contracts/genericContract.kt");
+        }
+    }
+
     @TestMetadata("compiler/fir/analysis-tests/testData/resolve/delegates")
     @TestDataPath("$PROJECT_ROOT")
     @RunWith(JUnit3RunnerWithInners.class)

compiler/fir/analysis-tests/testData/resolve/contracts/genericContract.fir.txt

@@ -0,0 +1,35 @@
+FILE: genericContract.kt
+    public final inline fun <reified T> requreIsInstance(value: R|kotlin/Any|): R|kotlin/Unit|
+        [R|Contract description]
+         <
+            Returns(WILDCARD) -> value is T
+        >
+     {
+        when () {
+            (R|<local>/value| !is R|T|) ->  {
+                throw R|java/lang/IllegalArgumentException.IllegalArgumentException|()
+            }
+        }
+
+    }
+    public final fun test_1(s: R|kotlin/Any|): R|kotlin/Unit| {
+        R|/requreIsInstance|<R|kotlin/String|>(R|<local>/s|)
+        R|<local>/s|.R|kotlin/String.length|
+    }
+    public final inline fun <reified T> requreIsInstanceOf(value: R|kotlin/Any|, requiredValue: R|T|): R|kotlin/Unit|
+        [R|Contract description]
+         <
+            Returns(WILDCARD) -> value is T
+        >
+     {
+        when () {
+            (R|<local>/value| !is R|T|) ->  {
+                throw R|java/lang/IllegalArgumentException.IllegalArgumentException|()
+            }
+        }
+
+    }
+    public final fun test_2(x: R|kotlin/Any|, s: R|kotlin/String|): R|kotlin/Unit| {
+        R|/requreIsInstanceOf|<R|kotlin/String|>(R|<local>/x|, R|<local>/s|)
+        R|<local>/x|.R|kotlin/String.length|
+    }

compiler/fir/analysis-tests/testData/resolve/contracts/genericContract.kt

@@ -0,0 +1,25 @@
+// WITH_STDLIB
+
+import kotlin.contracts.*
+
+inline fun <reified T> requreIsInstance(value: Any) contract [
+    returns() implies (value is T)
+] {
+    if (value !is T) throw IllegalArgumentException()
+}
+
+fun test_1(s: Any) {
+    requreIsInstance<String>(s)
+    s.length
+}
+
+inline fun <reified T> requreIsInstanceOf(value: Any, requiredValue: T) contract [
+    returns() implies (value is T)
+] {
+    if (value !is T) throw IllegalArgumentException()
+}
+
+fun test_2(x: Any, s: String) {
+    requreIsInstanceOf(x, s)
+    x.length
+}

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

@@ -1128,6 +1128,22 @@ public class FirDiagnosticTestGenerated extends AbstractFirDiagnosticTest {
             }
         }
 
+        @Nested
+        @TestMetadata("compiler/fir/analysis-tests/testData/resolve/contracts")
+        @TestDataPath("$PROJECT_ROOT")
+        public class Contracts {
+            @Test
+            public void testAllFilesPresentInContracts() throws Exception {
+                KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/fir/analysis-tests/testData/resolve/contracts"), Pattern.compile("^([^.]+)\\.kt$"), null, true);
+            }
+
+            @Test
+            @TestMetadata("genericContract.kt")
+            public void testGenericContract() throws Exception {
+                runTest("compiler/fir/analysis-tests/testData/resolve/contracts/genericContract.kt");
+            }
+        }
+
         @Nested
         @TestMetadata("compiler/fir/analysis-tests/testData/resolve/delegates")
         @TestDataPath("$PROJECT_ROOT")

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

@@ -1138,6 +1138,23 @@ public class FirDiagnosticsWithLightTreeTestGenerated extends AbstractFirDiagnos
             }
         }
 
+        @Nested
+        @TestMetadata("compiler/fir/analysis-tests/testData/resolve/contracts")
+        @TestDataPath("$PROJECT_ROOT")
+        @Execution(ExecutionMode.SAME_THREAD)
+        public class Contracts {
+            @Test
+            public void testAllFilesPresentInContracts() throws Exception {
+                KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/fir/analysis-tests/testData/resolve/contracts"), Pattern.compile("^([^.]+)\\.kt$"), null, true);
+            }
+
+            @Test
+            @TestMetadata("genericContract.kt")
+            public void testGenericContract() throws Exception {
+                runTest("compiler/fir/analysis-tests/testData/resolve/contracts/genericContract.kt");
+            }
+        }
+
         @Nested
         @TestMetadata("compiler/fir/analysis-tests/testData/resolve/delegates")
         @TestDataPath("$PROJECT_ROOT")

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/FirDataFlowAnalyzer.kt

@@ -23,9 +23,12 @@ import org.jetbrains.kotlin.fir.resolve.dfa.cfg.*
 import org.jetbrains.kotlin.fir.resolve.dfa.contracts.buildContractFir
 import org.jetbrains.kotlin.fir.resolve.dfa.contracts.createArgumentsMapping
 import org.jetbrains.kotlin.fir.resolve.inference.inferenceComponents
+import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
+import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutorByMap
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirAbstractBodyResolveTransformer
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.resultType
 import org.jetbrains.kotlin.fir.symbols.AbstractFirBasedSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirTypeParameterSymbol
 import org.jetbrains.kotlin.name.CallableId
 import org.jetbrains.kotlin.name.StandardClassIds
 import org.jetbrains.kotlin.fir.symbols.impl.FirVariableSymbol
@@ -865,12 +868,24 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
         val conditionalEffects = contractDescription.effects.map { it.effect }.filterIsInstance<ConeConditionalEffectDeclaration>()
         if (conditionalEffects.isEmpty()) return
         val argumentsMapping = createArgumentsMapping(qualifiedAccess) ?: return
+
+        val typeParameters = (owner as? FirTypeParameterRefsOwner)?.typeParameters
+        val substitutor = if (!typeParameters.isNullOrEmpty()) {
+            @Suppress("UNCHECKED_CAST")
+            val substitutionFromArguments = typeParameters.zip(qualifiedAccess.typeArguments).map { (typeParameterRef, typeArgument) ->
+                typeParameterRef.symbol to typeArgument.toConeTypeProjection().type
+            }.filter { it.second != null }.toMap() as Map<FirTypeParameterSymbol, ConeKotlinType>
+            ConeSubstitutorByMap(substitutionFromArguments)
+        } else {
+            ConeSubstitutor.Empty
+        }
+
         contractDescriptionVisitingMode = true
         graphBuilder.enterContract(qualifiedAccess).mergeIncomingFlow()
         val lastFlow = graphBuilder.lastNode.flow
         val functionCallVariable = variableStorage.getOrCreateVariable(lastFlow, qualifiedAccess)
         for (conditionalEffect in conditionalEffects) {
-            val fir = conditionalEffect.buildContractFir(argumentsMapping) ?: continue
+            val fir = conditionalEffect.buildContractFir(argumentsMapping, substitutor) ?: continue
             val effect = conditionalEffect.effect as? ConeReturnsEffectDeclaration ?: continue
             fir.transformSingle(components.transformer, ResolutionMode.ContextDependent)
             val argumentVariable = variableStorage.getOrCreateVariable(lastFlow, fir)

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/contracts/ConeConditionalEffectToFirVisitor.kt

@@ -11,16 +11,20 @@ import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.expressions.builder.*
 import org.jetbrains.kotlin.fir.expressions.impl.FirNoReceiverExpression
 import org.jetbrains.kotlin.fir.references.builder.buildSimpleNamedReference
+import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
 import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
 import org.jetbrains.kotlin.types.ConstantValueKind
 import org.jetbrains.kotlin.util.OperatorNameConventions
 
-private object ConeConditionalEffectToFirVisitor : ConeContractDescriptionVisitor<FirExpression?, Map<Int, FirExpression>>() {
+private class ConeConditionalEffectToFirVisitor(
-    override fun visitConditionalEffectDeclaration(conditionalEffect: ConeConditionalEffectDeclaration, data: Map<Int, FirExpression>): FirExpression? {
+    val valueParametersMapping: Map<Int, FirExpression>,
+    val substitutor: ConeSubstitutor
+) : ConeContractDescriptionVisitor<FirExpression?, Nothing?>() {
+    override fun visitConditionalEffectDeclaration(conditionalEffect: ConeConditionalEffectDeclaration, data: Nothing?): FirExpression? {
         return conditionalEffect.condition.accept(this, data)
     }
 
-    override fun visitConstantDescriptor(constantReference: ConeConstantReference, data: Map<Int, FirExpression>): FirExpression? {
+    override fun visitConstantDescriptor(constantReference: ConeConstantReference, data: Nothing?): FirExpression? {
         return when (constantReference) {
             ConeBooleanConstantReference.TRUE -> buildConstExpression(null, ConstantValueKind.Boolean, true)
             ConeBooleanConstantReference.FALSE -> buildConstExpression(null, ConstantValueKind.Boolean, false)
@@ -31,7 +35,7 @@ private object ConeConditionalEffectToFirVisitor : ConeContractDescriptionVisito
 
     override fun visitLogicalBinaryOperationContractExpression(
         binaryLogicExpression: ConeBinaryLogicExpression,
-        data: Map<Int, FirExpression>
+        data: Nothing?
     ): FirExpression? {
         val leftExpression = binaryLogicExpression.left.accept(this, data) ?: return null
         val rightExpression = binaryLogicExpression.right.accept(this, data) ?: return null
@@ -42,7 +46,7 @@ private object ConeConditionalEffectToFirVisitor : ConeContractDescriptionVisito
         }
     }
 
-    override fun visitLogicalNot(logicalNot: ConeLogicalNot, data: Map<Int, FirExpression>): FirExpression? {
+    override fun visitLogicalNot(logicalNot: ConeLogicalNot, data: Nothing?): FirExpression? {
         val explicitReceiver = logicalNot.arg.accept(this, data) ?: return null
         return buildFunctionCall {
             calleeReference = buildSimpleNamedReference { name = OperatorNameConventions.NOT }
@@ -50,7 +54,7 @@ private object ConeConditionalEffectToFirVisitor : ConeContractDescriptionVisito
         }
     }
 
-    override fun visitIsInstancePredicate(isInstancePredicate: ConeIsInstancePredicate, data: Map<Int, FirExpression>): FirExpression? {
+    override fun visitIsInstancePredicate(isInstancePredicate: ConeIsInstancePredicate, data: Nothing?): FirExpression? {
         val argument = isInstancePredicate.arg.accept(this@ConeConditionalEffectToFirVisitor, data) ?: return null
         return buildTypeOperatorCall {
             argumentList = buildUnaryArgumentList(argument)
@@ -59,11 +63,11 @@ private object ConeConditionalEffectToFirVisitor : ConeContractDescriptionVisito
             } else {
                 FirOperation.IS
             }
-            conversionTypeRef = buildResolvedTypeRef { type = isInstancePredicate.type }
+            conversionTypeRef = buildResolvedTypeRef { type = substitutor.substituteOrSelf(isInstancePredicate.type) }
         }
     }
 
-    override fun visitIsNullPredicate(isNullPredicate: ConeIsNullPredicate, data: Map<Int, FirExpression>): FirExpression? {
+    override fun visitIsNullPredicate(isNullPredicate: ConeIsNullPredicate, data: Nothing?): FirExpression? {
         val argument = isNullPredicate.arg.accept(this, data) ?: return null
         return buildEqualityOperatorCall {
             operation = if (isNullPredicate.isNegated) {
@@ -75,15 +79,18 @@ private object ConeConditionalEffectToFirVisitor : ConeContractDescriptionVisito
         }
     }
 
-    override fun visitValueParameterReference(valueParameterReference: ConeValueParameterReference, data: Map<Int, FirExpression>): FirExpression? {
+    override fun visitValueParameterReference(valueParameterReference: ConeValueParameterReference, data: Nothing?): FirExpression? {
-        return data[valueParameterReference.parameterIndex]
+        return valueParametersMapping[valueParameterReference.parameterIndex]
     }
 
     private fun createConstNull(): FirConstExpression<*> = buildConstExpression(null, ConstantValueKind.Null, null)
 }
 
-fun ConeConditionalEffectDeclaration.buildContractFir(argumentMapping: Map<Int, FirExpression>): FirExpression? {
+fun ConeConditionalEffectDeclaration.buildContractFir(
-    return condition.accept(ConeConditionalEffectToFirVisitor, argumentMapping)
+    argumentMapping: Map<Int, FirExpression>,
+    substitutor: ConeSubstitutor
+): FirExpression? {
+    return condition.accept(ConeConditionalEffectToFirVisitor(argumentMapping, substitutor), null)
 }
 
 fun createArgumentsMapping(qualifiedAccess: FirQualifiedAccess): Map<Int, FirExpression>? {
@@ -106,4 +113,4 @@ fun createArgumentsMapping(qualifiedAccess: FirQualifiedAccess): Map<Int, FirExp
     return argumentsMapping
 }
 
-private fun FirExpression.unwrap(): FirExpression = if (this is FirWrappedArgumentExpression) expression else this
+private fun FirExpression.unwrap(): FirExpression = if (this is FirWrappedArgumentExpression) expression else this

compiler/testData/diagnostics/testsWithStdLib/contracts/smartcasts/reifiedGeneric.fir.kt

@@ -22,10 +22,10 @@ inline fun <reified T> cast(value: Any?): T {
 
 fun test_1(x: Any) {
     requireIsInstance<String>(x)
-    x.<!UNRESOLVED_REFERENCE!>length<!>
+    x.length
 }
 
 fun test_2(x: Any) {
     val s: String = cast(x)
-    x.<!UNRESOLVED_REFERENCE!>length<!>
+    x.length
 }