[FIR] Properly set source for inferred return type of function

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) {
+            val returnTypeRef = if (returnExpression?.typeRef is FirResolvedTypeRef) {
-                result.transformReturnTypeRef(
+                returnExpression.resultType.approximateDeclarationType(
-                    transformer,
+                    session,
-                    withExpectedType(
+                    simpleFunction?.visibilityForApproximation(),
-                        returnExpression.resultType.approximateDeclarationType(
+                    isLocal = simpleFunction?.isLocal == true,
-                            session,
+                    isInlineFunction = simpleFunction?.isInline == true
-                            simpleFunction?.visibilityForApproximation(),
+                ).copyWithNewSource(result.returnTypeRef.source)
-                            isLocal = simpleFunction?.isLocal == true,
-                            isInlineFunction = simpleFunction?.isInline == true
-                        )
-                    )
-                )
             } else {
-                result.transformReturnTypeRef(
+                buildErrorTypeRef {
-                    transformer,
+                    source = result.returnTypeRef.source
-                    withExpectedType(buildErrorTypeRef { diagnostic = ConeSimpleDiagnostic("empty body", DiagnosticKind.Other) })
+                    diagnostic = ConeSimpleDiagnostic("empty body", DiagnosticKind.Other)
-                )
+                }
             }
+            result.transformReturnTypeRef(transformer, withExpectedType(returnTypeRef))
         }
 
         return result