FIR: check DELEGATE_SPECIAL_FUNCTION_RETURN_TYPE_MISMATCH

compiler/fir/checkers/checkers-component-generator/src/org/jetbrains/kotlin/fir/checkers/generator/diagnostics/FirDiagnosticsList.kt

@@ -706,6 +706,12 @@ object DIAGNOSTICS_LIST : DiagnosticList() {
             parameter<String>("expectedFunctionSignature")
             parameter<Collection<AbstractFirBasedSymbol<*>>>("candidates")
         }
+        val DELEGATE_SPECIAL_FUNCTION_RETURN_TYPE_MISMATCH by error<KtExpression> {
+            parameter<String>("delegateFunction")
+            parameter<ConeKotlinType>("expected")
+            parameter<ConeKotlinType>("actual")
+        }
+
         val UNDERSCORE_IS_RESERVED by error<KtExpression>(PositioningStrategy.RESERVED_UNDERSCORE)
         val UNDERSCORE_USAGE_WITHOUT_BACKTICKS by error<KtExpression>(PositioningStrategy.RESERVED_UNDERSCORE)
     }

compiler/fir/checkers/gen/org/jetbrains/kotlin/fir/analysis/diagnostics/FirErrors.kt

@@ -415,6 +415,7 @@ object FirErrors {
     val DELEGATE_SPECIAL_FUNCTION_MISSING by error3<KtExpression, String, ConeKotlinType, String>()
     val DELEGATE_SPECIAL_FUNCTION_AMBIGUITY by error2<KtExpression, String, Collection<AbstractFirBasedSymbol<*>>>()
     val DELEGATE_SPECIAL_FUNCTION_NONE_APPLICABLE by error2<KtExpression, String, Collection<AbstractFirBasedSymbol<*>>>()
+    val DELEGATE_SPECIAL_FUNCTION_RETURN_TYPE_MISMATCH by error3<KtExpression, String, ConeKotlinType, ConeKotlinType>()
     val UNDERSCORE_IS_RESERVED by error0<KtExpression>(SourceElementPositioningStrategies.RESERVED_UNDERSCORE)
     val UNDERSCORE_USAGE_WITHOUT_BACKTICKS by error0<KtExpression>(SourceElementPositioningStrategies.RESERVED_UNDERSCORE)
 

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

@@ -21,18 +21,23 @@ import org.jetbrains.kotlin.fir.resolve.diagnostics.ConeAmbiguityError
 import org.jetbrains.kotlin.fir.resolve.diagnostics.ConeInapplicableCandidateError
 import org.jetbrains.kotlin.fir.resolve.diagnostics.ConeUnresolvedNameError
 import org.jetbrains.kotlin.fir.symbols.AbstractFirBasedSymbol
+import org.jetbrains.kotlin.fir.typeContext
 import org.jetbrains.kotlin.fir.types.ConeKotlinErrorType
 import org.jetbrains.kotlin.fir.types.coneType
 import org.jetbrains.kotlin.fir.types.render
 import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
 import org.jetbrains.kotlin.resolve.calls.tower.CandidateApplicability
 import org.jetbrains.kotlin.resolve.calls.tower.isSuccess
+import org.jetbrains.kotlin.types.AbstractTypeChecker
 
 object FirDelegatedPropertyChecker : FirPropertyChecker() {
     override fun check(declaration: FirProperty, context: CheckerContext, reporter: DiagnosticReporter) {
         val delegate = declaration.delegate ?: return
         val delegateType = delegate.typeRef.coneType
 
+        // TODO: Also suppress delegate issue if type inference failed. For example, in
+        //  compiler/testData/diagnostics/tests/delegatedProperty/inference/differentDelegatedExpressions.fir.kt, no delegate issues are
+        //  reported due to the inference issue.
         if (delegateType is ConeKotlinErrorType) {
             val delegateSource = delegate.source
             // Implicit recursion type is not reported since the type ref does not have a real source.
@@ -47,8 +52,13 @@ object FirDelegatedPropertyChecker : FirPropertyChecker() {
             override fun visitElement(element: FirElement) = element.acceptChildren(this)
 
             override fun visitFunctionCall(functionCall: FirFunctionCall) {
-                val errorNamedReference = functionCall.calleeReference as? FirErrorNamedReference ?: return
+                val hasReferenceError = hasFunctionReferenceErrors(functionCall)
-                if (errorNamedReference.source?.kind != FirFakeSourceElementKind.DelegatedPropertyAccessor) return
+                if (isGet && !hasReferenceError) checkReturnType(functionCall)
+            }
+
+            private fun hasFunctionReferenceErrors(functionCall: FirFunctionCall): Boolean {
+                val errorNamedReference = functionCall.calleeReference as? FirErrorNamedReference ?: return false
+                if (errorNamedReference.source?.kind != FirFakeSourceElementKind.DelegatedPropertyAccessor) return false
                 val expectedFunctionSignature =
                     (if (isGet) "getValue" else "setValue") + "(${functionCall.arguments.joinToString(", ") { it.typeRef.coneType.render() }})"
                 val delegateDescription = if (isGet) "delegate" else "delegate for var (read-write property)"
@@ -77,7 +87,7 @@ object FirDelegatedPropertyChecker : FirPropertyChecker() {
                     }
                 }
 
-                when (val diagnostic = errorNamedReference.diagnostic) {
+                return when (val diagnostic = errorNamedReference.diagnostic) {
                     is ConeUnresolvedNameError -> {
                         reporter.reportOn(
                             errorNamedReference.source,
@@ -87,6 +97,7 @@ object FirDelegatedPropertyChecker : FirPropertyChecker() {
                             delegateDescription,
                             context
                         )
+                        true
                     }
                     is ConeAmbiguityError -> {
                         if (diagnostic.applicability.isSuccess) {
@@ -101,10 +112,28 @@ object FirDelegatedPropertyChecker : FirPropertyChecker() {
                         } else {
                             reportInapplicableDiagnostics(diagnostic.applicability, diagnostic.candidates.map { it.symbol })
                         }
+                        true
                     }
                     is ConeInapplicableCandidateError -> {
                         reportInapplicableDiagnostics(diagnostic.applicability, listOf(diagnostic.candidate.symbol))
+                        true
                     }
+                    else -> false
+                }
+            }
+
+            private fun checkReturnType(functionCall: FirFunctionCall) {
+                val returnType = functionCall.typeRef.coneType
+                val propertyType = declaration.returnTypeRef.coneType
+                if (!AbstractTypeChecker.isSubtypeOf(context.session.typeContext, returnType, propertyType)) {
+                    reporter.reportOn(
+                        delegate.source,
+                        FirErrors.DELEGATE_SPECIAL_FUNCTION_RETURN_TYPE_MISMATCH,
+                        "getValue",
+                        propertyType,
+                        returnType,
+                        context
+                    )
                 }
             }
         }

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/diagnostics/FirDefaultErrorMessages.kt

@@ -82,6 +82,7 @@ import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.DELEGATED_PROPERT
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.DELEGATE_SPECIAL_FUNCTION_AMBIGUITY
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.DELEGATE_SPECIAL_FUNCTION_MISSING
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.DELEGATE_SPECIAL_FUNCTION_NONE_APPLICABLE
+import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.DELEGATE_SPECIAL_FUNCTION_RETURN_TYPE_MISMATCH
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.DELEGATION_IN_INTERFACE
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.DELEGATION_SUPER_CALL_IN_ENUM_CONSTRUCTOR
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.DEPRECATED_MODIFIER_PAIR
@@ -913,6 +914,13 @@ class FirDefaultErrorMessages : DefaultErrorMessages.Extension {
                 TO_STRING,
                 SYMBOLS
             )
+            map.put(
+                DELEGATE_SPECIAL_FUNCTION_RETURN_TYPE_MISMATCH,
+                "The ''{0}'' function of property delegate is expected to return ''{1}'', but returns ''{2}''",
+                TO_STRING,
+                RENDER_TYPE,
+                RENDER_TYPE,
+            )
 
             // Type alias
             map.put(TOPLEVEL_TYPEALIASES_ONLY, "Nested and local type aliases are not supported")

compiler/testData/diagnostics/tests/delegatedProperty/inference/differentDelegatedExpressions.fir.kt

@@ -12,8 +12,8 @@ class A(outer: Outer) {
 
     var b: String by  foo(getMyProperty())
     var r: String by  foo(outer.getContainer().getMyProperty())
-    var e: String by  <!DELEGATE_SPECIAL_FUNCTION_NONE_APPLICABLE!>+ foo(getMyProperty())<!>
+    var e: String by  <!DELEGATE_SPECIAL_FUNCTION_NONE_APPLICABLE, DELEGATE_SPECIAL_FUNCTION_RETURN_TYPE_MISMATCH!>+ foo(getMyProperty())<!>
-    var f: String by  <!DELEGATE_SPECIAL_FUNCTION_NONE_APPLICABLE!>foo(getMyProperty()) - 1<!>
+    var f: String by  <!DELEGATE_SPECIAL_FUNCTION_NONE_APPLICABLE, DELEGATE_SPECIAL_FUNCTION_RETURN_TYPE_MISMATCH!>foo(getMyProperty()) - 1<!>
 }
 
 fun <A, B> foo(a: Any?) = MyProperty<A, B>()

compiler/testData/diagnostics/tests/delegatedProperty/inference/genericMethodInGenericClass.fir.kt

@@ -25,4 +25,4 @@ class C<R>() {
 }
 
 var c1: Int by C()
-var c2: Int by <!DELEGATE_SPECIAL_FUNCTION_NONE_APPLICABLE!>C<Number>()<!>
+var c2: Int by <!DELEGATE_SPECIAL_FUNCTION_NONE_APPLICABLE, DELEGATE_SPECIAL_FUNCTION_RETURN_TYPE_MISMATCH!>C<Number>()<!>

compiler/testData/diagnostics/tests/delegatedProperty/typeMismatchForGetReturnType.fir.kt

@@ -3,7 +3,7 @@
 
 import kotlin.reflect.KProperty
 
-val c: Int by Delegate()
+val c: Int by <!DELEGATE_SPECIAL_FUNCTION_RETURN_TYPE_MISMATCH!>Delegate()<!>
 
 class Delegate {
   operator fun getValue(t: Any?, p: KProperty<*>): String {

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/diagnostics/KtFirDataClassConverters.kt

@@ -1985,6 +1985,15 @@ internal val KT_DIAGNOSTIC_CONVERTER = KtDiagnosticConverterBuilder.buildConvert
             token,
         )
     }
+    add(FirErrors.DELEGATE_SPECIAL_FUNCTION_RETURN_TYPE_MISMATCH) { firDiagnostic ->
+        DelegateSpecialFunctionReturnTypeMismatchImpl(
+            firDiagnostic.a,
+            firSymbolBuilder.typeBuilder.buildKtType(firDiagnostic.b),
+            firSymbolBuilder.typeBuilder.buildKtType(firDiagnostic.c),
+            firDiagnostic as FirPsiDiagnostic<*>,
+            token,
+        )
+    }
     add(FirErrors.UNDERSCORE_IS_RESERVED) { firDiagnostic ->
         UnderscoreIsReservedImpl(
             firDiagnostic as FirPsiDiagnostic<*>,

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/diagnostics/KtFirDiagnostics.kt

@@ -1388,6 +1388,13 @@ sealed class KtFirDiagnostic<PSI: PsiElement> : KtDiagnosticWithPsi<PSI> {
         abstract val candidates: List<KtSymbol>
     }
 
+    abstract class DelegateSpecialFunctionReturnTypeMismatch : KtFirDiagnostic<KtExpression>() {
+        override val diagnosticClass get() = DelegateSpecialFunctionReturnTypeMismatch::class
+        abstract val delegateFunction: String
+        abstract val expected: KtType
+        abstract val actual: KtType
+    }
+
     abstract class UnderscoreIsReserved : KtFirDiagnostic<KtExpression>() {
         override val diagnosticClass get() = UnderscoreIsReserved::class
     }

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/diagnostics/KtFirDiagnosticsImpl.kt

@@ -2253,6 +2253,16 @@ internal class DelegateSpecialFunctionNoneApplicableImpl(
     override val firDiagnostic: FirPsiDiagnostic<*> by weakRef(firDiagnostic)
 }
 
+internal class DelegateSpecialFunctionReturnTypeMismatchImpl(
+    override val delegateFunction: String,
+    override val expected: KtType,
+    override val actual: KtType,
+    firDiagnostic: FirPsiDiagnostic<*>,
+    override val token: ValidityToken,
+) : KtFirDiagnostic.DelegateSpecialFunctionReturnTypeMismatch(), KtAbstractFirDiagnostic<KtExpression> {
+    override val firDiagnostic: FirPsiDiagnostic<*> by weakRef(firDiagnostic)
+}
+
 internal class UnderscoreIsReservedImpl(
     firDiagnostic: FirPsiDiagnostic<*>,
     override val token: ValidityToken,