[FIR] Force calculation of return type of fake overrides in return type calculator

compiler/fir/analysis-tests/testData/resolve/implicitTypeInFakeOverride.kt

@@ -3,3 +3,7 @@ fun <K> extract(x: Out<K>) = x.get()
 class Out<out T>(val x: T) {
     fun get() = x
 }
+
+fun test(out: Out<String>) {
+    <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.String")!>extract(out)<!>
+}

compiler/fir/analysis-tests/testData/resolve/implicitTypeInFakeOverride.txt

@@ -1,6 +1,6 @@
 FILE: implicitTypeInFakeOverride.kt
-    public final fun <K> extract(x: R|Out<K>|): R|T| {
+    public final fun <K> extract(x: R|Out<K>|): R|K| {
-        ^extract R|<local>/x|.R|/Out.get|()
+        ^extract R|<local>/x|.R|FakeOverride</Out.get: R|K|>|()
     }
     public final class Out<out T> : R|kotlin/Any| {
         public constructor<out T>(x: R|T|): R|Out<T>| {
@@ -15,3 +15,6 @@ FILE: implicitTypeInFakeOverride.kt
         }
 
     }
+    public final fun test(out: R|Out<kotlin/String>|): R|kotlin/Unit| {
+        R|/extract|<R|kotlin/String|>(R|<local>/out|)
+    }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/ReturnTypeCalculator.kt

@@ -9,6 +9,8 @@ import org.jetbrains.kotlin.fir.declarations.FirTypedDeclaration
 import org.jetbrains.kotlin.fir.diagnostics.DiagnosticKind
 import org.jetbrains.kotlin.fir.diagnostics.ConeSimpleDiagnostic
 import org.jetbrains.kotlin.fir.render
+import org.jetbrains.kotlin.fir.scopes.FakeOverrideTypeCalculator
+import org.jetbrains.kotlin.fir.scopes.ForcedFakeOverrideTypeCalculator
 import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
 import org.jetbrains.kotlin.fir.types.builder.buildErrorTypeRef
 

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirImplicitBodyResolve.kt

@@ -17,6 +17,8 @@ import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.firProvider
 import org.jetbrains.kotlin.fir.resolve.transformers.*
 import org.jetbrains.kotlin.fir.resolve.transformers.contracts.runContractResolveForLocalClass
+import org.jetbrains.kotlin.fir.scopes.ForcedFakeOverrideTypeCalculator
+import org.jetbrains.kotlin.fir.symbols.PossiblyFirFakeOverrideSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirAccessorSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
 import org.jetbrains.kotlin.fir.types.FirImplicitTypeRef
@@ -25,6 +27,7 @@ import org.jetbrains.kotlin.fir.types.builder.buildErrorTypeRef
 import org.jetbrains.kotlin.fir.visitors.CompositeTransformResult
 import org.jetbrains.kotlin.fir.visitors.FirTransformer
 import org.jetbrains.kotlin.fir.visitors.compose
+import org.jetbrains.kotlin.utils.addToStdlib.runIf
 
 @OptIn(AdapterForResolveProcessor::class)
 class FirImplicitTypeBodyResolveProcessor(
@@ -216,6 +219,13 @@ private class ReturnTypeCalculatorWithJump(
             return result
         }
 
+        runIf(declaration.origin == FirDeclarationOrigin.FakeOverride) {
+            val possiblyFirFakeOverrideSymbol = declaration.symbol as PossiblyFirFakeOverrideSymbol<*, *>
+            val overriddenDeclaration = possiblyFirFakeOverrideSymbol.overriddenSymbol?.fir as FirTypedDeclaration? ?: return@runIf
+            tryCalculateReturnType(overriddenDeclaration)
+            return ForcedFakeOverrideTypeCalculator.computeReturnType(declaration)
+        }
+
         return when (val status = implicitBodyResolveComputationSession.getStatus(declaration.symbol)) {
             is ImplicitBodyResolveComputationStatus.Computed -> status.resolvedTypeRef
             is ImplicitBodyResolveComputationStatus.Computing ->

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/scopes/FakeOverrideTypeCalculator.kt

@@ -10,28 +10,32 @@ import org.jetbrains.kotlin.fir.declarations.FirDeclarationDataKey
 import org.jetbrains.kotlin.fir.declarations.FirDeclarationDataRegistry
 import org.jetbrains.kotlin.fir.declarations.FirTypedDeclaration
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
+import org.jetbrains.kotlin.fir.resolvedTypeFromPrototype
 import org.jetbrains.kotlin.fir.symbols.AbstractFirBasedSymbol
-import org.jetbrains.kotlin.fir.types.ConeKotlinType
+import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
-import org.jetbrains.kotlin.fir.types.coneTypeSafe
+import org.jetbrains.kotlin.fir.types.FirTypeRef
 
 abstract class FakeOverrideTypeCalculator {
-    abstract fun computeReturnType(declaration: FirTypedDeclaration): ConeKotlinType?
+    abstract fun computeReturnType(declaration: FirTypedDeclaration): FirTypeRef
 
-    class DoNothing : FakeOverrideTypeCalculator() {
+    object DoNothing : FakeOverrideTypeCalculator() {
-        override fun computeReturnType(declaration: FirTypedDeclaration): ConeKotlinType? {
+        override fun computeReturnType(declaration: FirTypedDeclaration): FirTypeRef {
-            return declaration.returnTypeRef.coneTypeSafe()
+            return declaration.returnTypeRef
         }
     }
 }
 
-class ForcedFakeOverrideTypeCalculator : FakeOverrideTypeCalculator() {
+object ForcedFakeOverrideTypeCalculator : FakeOverrideTypeCalculator() {
-    override fun computeReturnType(declaration: FirTypedDeclaration): ConeKotlinType {
+    override fun computeReturnType(declaration: FirTypedDeclaration): FirResolvedTypeRef {
-        declaration.returnTypeRef.coneTypeSafe<ConeKotlinType>()?.let { return it }
+        (declaration.returnTypeRef as? FirResolvedTypeRef)?.let { return it }
         val (substitutor, baseSymbol) = declaration.attributes.fakeOverrideSubstitution ?: error("")
         val baseDeclaration = baseSymbol.fir as FirTypedDeclaration
-        val returnType = computeReturnType(baseDeclaration)
+        val baseReturnType = computeReturnType(baseDeclaration).type
         declaration.attributes.fakeOverrideSubstitution = null
-        return substitutor.substituteOrSelf(returnType)
+        val coneType = substitutor.substituteOrSelf(baseReturnType)
+        val returnType = declaration.returnTypeRef.resolvedTypeFromPrototype(coneType)
+        declaration.replaceReturnTypeRef(returnType)
+        return returnType
     }
 }