FIR: use real source element for return statement

fix fir

compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/RawFirBuilder.kt

@@ -1528,7 +1528,7 @@ class RawFirBuilder(
             val source = expression.toFirSourceElement(FirFakeSourceElementKind.ImplicitUnit)
             val result = expression.returnedExpression?.toFirExpression("Incorrect return expression")
                 ?: buildUnitExpression { this.source = source }
-            return result.toReturn(source, expression.getTargetLabel()?.getReferencedName())
+            return result.toReturn(source, expression.getTargetLabel()?.getReferencedName(), fromKtReturnExpression = true)
         }
 
         override fun visitTryExpression(expression: KtTryExpression, data: Unit): FirElement {

compiler/fir/raw-fir/raw-fir.common/src/org/jetbrains/kotlin/fir/builder/BaseFirBuilder.kt

@@ -148,7 +148,11 @@ abstract class BaseFirBuilder<T>(val baseSession: FirSession, val context: Conte
         val ANONYMOUS_OBJECT_NAME = Name.special("<anonymous>")
     }
 
-    fun FirExpression.toReturn(baseSource: FirSourceElement? = source, labelName: String? = null): FirReturnExpression {
+    fun FirExpression.toReturn(
+        baseSource: FirSourceElement? = source,
+        labelName: String? = null,
+        fromKtReturnExpression: Boolean = false
+    ): FirReturnExpression {
         return buildReturnExpression {
             fun FirFunctionTarget.bindToErrorFunction(message: String, kind: DiagnosticKind) {
                 bind(
@@ -162,7 +166,9 @@ abstract class BaseFirBuilder<T>(val baseSession: FirSession, val context: Conte
                 )
             }
 
-            source = baseSource?.fakeElement(FirFakeSourceElementKind.ImplicitReturn)
+            source =
+                if (fromKtReturnExpression) baseSource?.realElement()
+                else baseSource?.fakeElement(FirFakeSourceElementKind.ImplicitReturn)
             result = this@toReturn
             if (labelName == null) {
                 target = context.firFunctionTargets.lastOrNull { !it.isLambda } ?: FirFunctionTarget(labelName, isLambda = false).apply {

compiler/fir/tree/src/org/jetbrains/kotlin/fir/FirSourceElement.kt

@@ -264,6 +264,13 @@ fun FirSourceElement.fakeElement(newKind: FirFakeSourceElementKind): FirSourceEl
     }
 }
 
+fun FirSourceElement.realElement(): FirSourceElement = when (this) {
+    is FirRealPsiSourceElement<*> -> this
+    is FirLightSourceElement -> FirLightSourceElement(lighterASTNode, startOffset, endOffset, treeStructure, FirRealSourceElementKind)
+    is FirPsiSourceElement<*> -> FirRealPsiSourceElement(psi)
+}
+
+
 class FirLightSourceElement(
     override val lighterASTNode: LighterASTNode,
     override val startOffset: Int,