From 76d360d7937e99b6031a8dc88aa463a9fe6f4dcf Mon Sep 17 00:00:00 2001
From: Dmitriy Novozhilov <dmitriy.novozhilov@jetbrains.com>
Date: Fri, 3 Feb 2023 10:05:19 +0200
Subject: [PATCH] [FIR] Properly set source for inferred return type of
 function

---
 .../FirDeclarationsResolveTransformer.kt      | 28 ++++++++-----------
 1 file changed, 12 insertions(+), 16 deletions(-)

diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirDeclarationsResolveTransformer.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirDeclarationsResolveTransformer.kt
index 04061e29b43..dd449ebfdf9 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirDeclarationsResolveTransformer.kt
+++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirDeclarationsResolveTransformer.kt
@@ -596,24 +596,20 @@ open class FirDeclarationsResolveTransformer(transformer: FirAbstractBodyResolve
         if (result.returnTypeRef is FirImplicitTypeRef) {
             val simpleFunction = function as? FirSimpleFunction
             val returnExpression = (body?.statements?.single() as? FirReturnExpression)?.result
-            if (returnExpression?.typeRef is FirResolvedTypeRef) {
-                result.transformReturnTypeRef(
-                    transformer,
-                    withExpectedType(
-                        returnExpression.resultType.approximateDeclarationType(
-                            session,
-                            simpleFunction?.visibilityForApproximation(),
-                            isLocal = simpleFunction?.isLocal == true,
-                            isInlineFunction = simpleFunction?.isInline == true
-                        )
-                    )
-                )
+            val returnTypeRef = if (returnExpression?.typeRef is FirResolvedTypeRef) {
+                returnExpression.resultType.approximateDeclarationType(
+                    session,
+                    simpleFunction?.visibilityForApproximation(),
+                    isLocal = simpleFunction?.isLocal == true,
+                    isInlineFunction = simpleFunction?.isInline == true
+                ).copyWithNewSource(result.returnTypeRef.source)
             } else {
-                result.transformReturnTypeRef(
-                    transformer,
-                    withExpectedType(buildErrorTypeRef { diagnostic = ConeSimpleDiagnostic("empty body", DiagnosticKind.Other) })
-                )
+                buildErrorTypeRef {
+                    source = result.returnTypeRef.source
+                    diagnostic = ConeSimpleDiagnostic("empty body", DiagnosticKind.Other)
+                }
             }
+            result.transformReturnTypeRef(transformer, withExpectedType(returnTypeRef))
         }
 
         return result