[FIR] Use proper ReturnTypeCalculator inside FakeOverrideTypeCalculator

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

@@ -6,10 +6,13 @@
 package org.jetbrains.kotlin.fir.resolve.transformers
 
 import org.jetbrains.kotlin.fir.declarations.FirCallableDeclaration
+import org.jetbrains.kotlin.fir.scopes.FakeOverrideTypeCalculator
 import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
 import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
 
 abstract class ReturnTypeCalculator {
+    abstract val fakeOverrideTypeCalculator: FakeOverrideTypeCalculator
+
     abstract fun tryCalculateReturnTypeOrNull(declaration: FirCallableDeclaration): FirResolvedTypeRef?
 
     fun tryCalculateReturnType(declaration: FirCallableDeclaration): FirResolvedTypeRef {

compiler/fir/providers/src/org/jetbrains/kotlin/fir/resolve/transformers/ReturnTypeCalculatorForFullBodyResolve.kt

@@ -14,6 +14,9 @@ import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
 import org.jetbrains.kotlin.fir.types.builder.buildErrorTypeRef
 
 object ReturnTypeCalculatorForFullBodyResolve : ReturnTypeCalculator() {
+    override val fakeOverrideTypeCalculator: FakeOverrideTypeCalculator
+        get() = FakeOverrideTypeCalculator.Forced
+
     override fun tryCalculateReturnTypeOrNull(declaration: FirCallableDeclaration): FirResolvedTypeRef? {
         val returnTypeRef = declaration.returnTypeRef
         if (returnTypeRef is FirResolvedTypeRef) return returnTypeRef

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

@@ -24,10 +24,10 @@ abstract class FakeOverrideTypeCalculator {
         }
     }
 
-    object Forced : FakeOverrideTypeCalculator() {
+    abstract class AbstractFakeOverrideTypeCalculator : FakeOverrideTypeCalculator() {
         override fun computeReturnType(declaration: FirCallableDeclaration): FirResolvedTypeRef? {
             val fakeOverrideSubstitution = declaration.attributes.fakeOverrideSubstitution
-                ?: return declaration.returnTypeRef as? FirResolvedTypeRef
+                ?: return declaration.getResolvedTypeRef()
             synchronized(fakeOverrideSubstitution) {
                 if (declaration.attributes.fakeOverrideSubstitution == null) {
                     return declaration.returnTypeRef as FirResolvedTypeRef
@@ -42,6 +42,15 @@ abstract class FakeOverrideTypeCalculator {
                 return returnType
             }
         }
+
+        protected abstract fun FirCallableDeclaration.getResolvedTypeRef(): FirResolvedTypeRef?
+    }
+
+
+    object Forced : AbstractFakeOverrideTypeCalculator() {
+        override fun FirCallableDeclaration.getResolvedTypeRef(): FirResolvedTypeRef? {
+            return returnTypeRef as? FirResolvedTypeRef
+        }
     }
 }
 

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

@@ -20,6 +20,7 @@ import org.jetbrains.kotlin.fir.resolve.transformers.FirTransformerBasedResolveP
 import org.jetbrains.kotlin.fir.resolve.transformers.ReturnTypeCalculator
 import org.jetbrains.kotlin.fir.resolve.transformers.TransformImplicitType
 import org.jetbrains.kotlin.fir.resolve.transformers.contracts.runContractResolveForLocalClass
+import org.jetbrains.kotlin.fir.scopes.FakeOverrideTypeCalculator
 import org.jetbrains.kotlin.fir.scopes.fakeOverrideSubstitution
 import org.jetbrains.kotlin.fir.symbols.impl.FirSyntheticPropertySymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
@@ -205,6 +206,7 @@ private class ReturnTypeCalculatorWithJump(
         outerBodyResolveContext: BodyResolveContext?
     ) -> FirDesignatedBodyResolveTransformerForReturnTypeCalculator = ::FirDesignatedBodyResolveTransformerForReturnTypeCalculator,
 ) : ReturnTypeCalculator() {
+    override val fakeOverrideTypeCalculator: FakeOverrideTypeCalculator = FakeOverrideTypeCalculatorWithJump()
 
     @OptIn(PrivateForInline::class)
     var outerBodyResolveContext: BodyResolveContext? = null
@@ -307,6 +309,12 @@ private class ReturnTypeCalculatorWithJump(
         }
         return newReturnTypeRef
     }
+
+    private inner class FakeOverrideTypeCalculatorWithJump : FakeOverrideTypeCalculator.AbstractFakeOverrideTypeCalculator() {
+        override fun FirCallableDeclaration.getResolvedTypeRef(): FirResolvedTypeRef? {
+            return this@ReturnTypeCalculatorWithJump.computeReturnTypeRef(this)
+        }
+    }
 }
 
 open class FirDesignatedBodyResolveTransformerForReturnTypeCalculator(