[FIR] Extract contracts from original functions of fake overrides

^KT-57911 Fixed

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

@@ -8,7 +8,6 @@ package org.jetbrains.kotlin.fir.resolve.dfa
 import org.jetbrains.kotlin.contracts.description.canBeRevisited
 import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.fir.*
-import org.jetbrains.kotlin.fir.contracts.FirResolvedContractDescription
 import org.jetbrains.kotlin.fir.contracts.description.ConeConditionalEffectDeclaration
 import org.jetbrains.kotlin.fir.contracts.description.ConeReturnsEffectDeclaration
 import org.jetbrains.kotlin.fir.declarations.*
@@ -20,12 +19,15 @@ import org.jetbrains.kotlin.fir.references.toResolvedPropertySymbol
 import org.jetbrains.kotlin.fir.resolve.*
 import org.jetbrains.kotlin.fir.resolve.calls.ImplicitReceiverValue
 import org.jetbrains.kotlin.fir.resolve.dfa.cfg.*
+import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutorByMap
+import org.jetbrains.kotlin.fir.resolve.substitution.chain
 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.resolve.transformers.unwrapAnonymousFunctionExpression
 import org.jetbrains.kotlin.fir.scopes.getFunctions
 import org.jetbrains.kotlin.fir.scopes.impl.declaredMemberScope
+import org.jetbrains.kotlin.fir.scopes.impl.toConeType
 import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirClassSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
@@ -834,7 +836,8 @@ abstract class FirDataFlowAnalyzer(
             return exitBooleanNot(flow, qualifiedAccess as FirFunctionCall)
         }
 
-        val contractDescription = callee.contractDescription as? FirResolvedContractDescription ?: return
+        val originalFunction = callee.originalIfFakeOverride()
+        val contractDescription = (originalFunction?.symbol ?: callee.symbol).resolvedContractDescription ?: return
         val conditionalEffects = contractDescription.effects.mapNotNull { it.effect as? ConeConditionalEffectDeclaration }
         if (conditionalEffects.isEmpty()) return
 
@@ -845,14 +848,22 @@ abstract class FirDataFlowAnalyzer(
         if (argumentVariables.all { it == null }) return
 
         val typeParameters = callee.typeParameters
-        val substitutor = if (typeParameters.isNotEmpty() && qualifiedAccess is FirQualifiedAccessExpression) {
+        val typeArgumentsSubstitutor = if (typeParameters.isNotEmpty() && qualifiedAccess is FirQualifiedAccessExpression) {
             @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, components.session)
         } else {
-            null
+            ConeSubstitutor.Empty
+        }
+
+
+        val substitutor = if (originalFunction == null) {
+            typeArgumentsSubstitutor
+        } else {
+            val map = originalFunction.symbol.typeParameterSymbols.zip(typeParameters.map { it.symbol.toConeType() }).toMap()
+            ConeSubstitutorByMap(map, components.session).chain(typeArgumentsSubstitutor)
         }
 
         for (conditionalEffect in conditionalEffects) {

compiler/testData/diagnostics/testsWithStdLib/contracts/fromStdlib/contractWithSubstitution.fir.kt

@@ -37,7 +37,7 @@ abstract class Base<T> {
 class Derived: Base<String>() {
     override fun foo(s: String?) {
         checkNotNull(s)
-        s<!UNSAFE_CALL!>.<!>length
+        s.length
     }
 
     fun test_1(s: Any) {
@@ -54,14 +54,14 @@ class Derived: Base<String>() {
 
     fun test_3(s: Any) {
         if (checkIsReifiedR<String>(s)) {
-            s.<!UNRESOLVED_REFERENCE!>length<!>
+            s.length
         }
     }
 }
 
 fun test_1(d: Derived, s: String?) {
     d.checkNotNull(s)
-    s<!UNSAFE_CALL!>.<!>length
+    s.length
 }
 
 fun test_2(d: Derived, s: Any?) {
@@ -78,6 +78,6 @@ fun test_3(d: Derived, s: Any?) {
 
 fun test_4(d: Derived, s: Any?) {
     if (d.checkIsReifiedR<String>(s)) {
-        s.<!UNRESOLVED_REFERENCE!>length<!>
+        s.length
     }
 }