FIR: Use return type calculator in FirTypeMismatchOnOverrideChecker

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/context/CheckerContext.kt

@@ -13,12 +13,14 @@ import org.jetbrains.kotlin.fir.resolve.ImplicitReceiverStack
 import org.jetbrains.kotlin.fir.resolve.PersistentImplicitReceiverStack
 import org.jetbrains.kotlin.fir.resolve.SessionHolder
 import org.jetbrains.kotlin.fir.resolve.calls.ImplicitReceiverValue
+import org.jetbrains.kotlin.fir.resolve.transformers.ReturnTypeCalculator
 import org.jetbrains.kotlin.name.Name
 
 abstract class CheckerContext {
     abstract val implicitReceiverStack: ImplicitReceiverStack
     abstract val containingDeclarations: List<FirDeclaration>
     abstract val sessionHolder: SessionHolder
+    abstract val returnTypeCalculator: ReturnTypeCalculator
 
     val session: FirSession
         get() = sessionHolder.session
@@ -39,19 +41,22 @@ abstract class CheckerContext {
 class PersistentCheckerContext(
     override val implicitReceiverStack: PersistentImplicitReceiverStack = PersistentImplicitReceiverStack(),
     override val containingDeclarations: PersistentList<FirDeclaration> = persistentListOf(),
-    override val sessionHolder: SessionHolder
+    override val sessionHolder: SessionHolder,
+    override val returnTypeCalculator: ReturnTypeCalculator,
 ) : CheckerContext() {
-    constructor(sessionHolder: SessionHolder) : this(
+    constructor(sessionHolder: SessionHolder, returnTypeCalculator: ReturnTypeCalculator) : this(
         PersistentImplicitReceiverStack(),
         persistentListOf(),
-        sessionHolder
+        sessionHolder,
+        returnTypeCalculator
     )
 
     fun addImplicitReceiver(name: Name?, value: ImplicitReceiverValue<*>): PersistentCheckerContext {
         return PersistentCheckerContext(
             implicitReceiverStack.add(name, value),
             containingDeclarations,
-            sessionHolder
+            sessionHolder,
+            returnTypeCalculator
         )
     }
 
@@ -59,7 +64,8 @@ class PersistentCheckerContext(
         return PersistentCheckerContext(
             implicitReceiverStack,
             containingDeclarations.add(declaration),
-            sessionHolder
+            sessionHolder,
+            returnTypeCalculator
         )
     }
 }

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/declaration/FirTypeMismatchOnOverrideChecker.kt

@@ -18,7 +18,10 @@ import org.jetbrains.kotlin.fir.symbols.impl.FirFunctionSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirTypeParameterSymbol
 import org.jetbrains.kotlin.fir.typeContext
-import org.jetbrains.kotlin.fir.types.*
+import org.jetbrains.kotlin.fir.types.ConeKotlinType
+import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
+import org.jetbrains.kotlin.fir.types.coneType
+import org.jetbrains.kotlin.fir.types.upperBoundIfFlexible
 import org.jetbrains.kotlin.types.AbstractTypeChecker
 import org.jetbrains.kotlin.types.AbstractTypeCheckerContext
 import org.jetbrains.kotlin.utils.addToStdlib.min
@@ -38,8 +41,8 @@ object FirTypeMismatchOnOverrideChecker : FirRegularClassChecker() {
 
         for (it in declaration.declarations) {
             when (it) {
-                is FirSimpleFunction -> checkFunction(it, reporter, typeCheckerContext, firTypeScope)
+                is FirSimpleFunction -> checkFunction(it, reporter, typeCheckerContext, firTypeScope, context)
-                is FirProperty -> checkProperty(it, reporter, typeCheckerContext, firTypeScope)
+                is FirProperty -> checkProperty(it, reporter, typeCheckerContext, firTypeScope, context)
             }
         }
     }
@@ -95,11 +98,12 @@ object FirTypeMismatchOnOverrideChecker : FirRegularClassChecker() {
     private fun FirCallableMemberDeclaration<*>.checkReturnType(
         overriddenSymbols: List<FirCallableSymbol<*>>,
         typeCheckerContext: AbstractTypeCheckerContext,
+        context: CheckerContext,
     ): FirMemberDeclaration? {
         val returnType = returnTypeRef.safeAs<FirResolvedTypeRef>()?.type
             ?: return null
 
-        val bounds = overriddenSymbols.map { it.fir.returnTypeRef.coneType.upperBoundIfFlexible() }
+        val bounds = overriddenSymbols.map { context.returnTypeCalculator.tryCalculateReturnType(it.fir).coneType.upperBoundIfFlexible() }
 
         for (it in bounds.indices) {
             val restriction = bounds[it]
@@ -117,7 +121,8 @@ object FirTypeMismatchOnOverrideChecker : FirRegularClassChecker() {
         function: FirSimpleFunction,
         reporter: DiagnosticReporter,
         typeCheckerContext: AbstractTypeCheckerContext,
-        firTypeScope: FirTypeScope
+        firTypeScope: FirTypeScope,
+        context: CheckerContext,
     ) {
         if (!function.isOverride) {
             return
@@ -131,7 +136,8 @@ object FirTypeMismatchOnOverrideChecker : FirRegularClassChecker() {
 
         val restriction = function.checkReturnType(
             overriddenSymbols = overriddenFunctionSymbols,
-            typeCheckerContext = typeCheckerContext
+            typeCheckerContext = typeCheckerContext,
+            context = context,
         )
 
         restriction?.let {
@@ -147,7 +153,8 @@ object FirTypeMismatchOnOverrideChecker : FirRegularClassChecker() {
         property: FirProperty,
         reporter: DiagnosticReporter,
         typeCheckerContext: AbstractTypeCheckerContext,
-        firTypeScope: FirTypeScope
+        firTypeScope: FirTypeScope,
+        context: CheckerContext,
     ) {
         if (!property.isOverride) {
             return
@@ -161,7 +168,8 @@ object FirTypeMismatchOnOverrideChecker : FirRegularClassChecker() {
 
         val restriction = property.checkReturnType(
             overriddenSymbols = overriddenPropertySymbols,
-            typeCheckerContext = typeCheckerContext
+            typeCheckerContext = typeCheckerContext,
+            context = context,
         )
 
         restriction?.let {

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/collectors/AbstractDiagnosticCollector.kt

@@ -19,6 +19,8 @@ import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.SessionHolder
 import org.jetbrains.kotlin.fir.resolve.collectImplicitReceivers
 import org.jetbrains.kotlin.fir.resolve.defaultType
+import org.jetbrains.kotlin.fir.resolve.transformers.ReturnTypeCalculator
+import org.jetbrains.kotlin.fir.resolve.transformers.ReturnTypeCalculatorForFullBodyResolve
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
 import org.jetbrains.kotlin.fir.types.FirTypeRef
 import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
@@ -28,7 +30,8 @@ import org.jetbrains.kotlin.name.Name
 
 abstract class AbstractDiagnosticCollector(
     override val session: FirSession,
-    override val scopeSession: ScopeSession = ScopeSession()
+    override val scopeSession: ScopeSession = ScopeSession(),
+    returnTypeCalculator: ReturnTypeCalculator = ReturnTypeCalculatorForFullBodyResolve()
 ) : SessionHolder {
     fun collectDiagnostics(firFile: FirFile): Iterable<FirDiagnostic<*>> {
         if (!componentsInitialized) {
@@ -48,7 +51,7 @@ abstract class AbstractDiagnosticCollector(
     private val visitor = Visitor()
 
     @Suppress("LeakingThis")
-    private var context = PersistentCheckerContext(this)
+    private var context = PersistentCheckerContext(this, returnTypeCalculator)
 
     fun initializeComponents(vararg components: AbstractDiagnosticCollectorComponent) {
         if (componentsInitialized) {

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/diagnostics/FirIdeDiagnosticsCollector.kt

@@ -14,13 +14,18 @@ import org.jetbrains.kotlin.fir.analysis.diagnostics.DiagnosticReporter
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirDiagnostic
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirPsiDiagnostic
 import org.jetbrains.kotlin.fir.declarations.FirFile
+import org.jetbrains.kotlin.fir.resolve.ScopeSession
+import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.createReturnTypeCalculatorForIDE
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.addValueFor
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.checkCanceled
 import org.jetbrains.kotlin.psi.KtElement
 
 internal class FirIdeDiagnosticsCollector private constructor(
     session: FirSession,
-) : AbstractDiagnosticCollector(session) {
+) : AbstractDiagnosticCollector(
+    session,
+    returnTypeCalculator = createReturnTypeCalculatorForIDE(session, ScopeSession())
+) {
     private val result = mutableMapOf<KtElement, MutableList<Diagnostic>>()
 
     init {