[FIR] Bind return expression to function symbol instead of function itself

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

@@ -8,9 +8,12 @@ package org.jetbrains.kotlin.fir
 abstract class FirAbstractTarget<E : FirTargetElement>(
     override val labelName: String?
 ) : FirTarget<E> {
-    override lateinit var labeledElement: E
+    protected abstract var _labeledElement: E
+
+    final override val labeledElement: E
+        get() = _labeledElement
 
     override fun bind(element: E) {
-        labeledElement = element
+        _labeledElement = element
     }
-}
+}

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

@@ -6,8 +6,17 @@
 package org.jetbrains.kotlin.fir
 
 import org.jetbrains.kotlin.fir.declarations.FirFunction
+import org.jetbrains.kotlin.fir.symbols.impl.FirFunctionSymbol
 
 class FirFunctionTarget(
     labelName: String?,
     val isLambda: Boolean
-) : FirAbstractTarget<FirFunction<*>>(labelName)
+) : FirAbstractTarget<FirFunction<*>>(labelName) {
+    private lateinit var targetSymbol: FirFunctionSymbol<*>
+
+    override var _labeledElement: FirFunction<*>
+        get() = targetSymbol.fir
+        set(value) {
+            targetSymbol = value.symbol
+        }
+}

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

@@ -9,4 +9,6 @@ import org.jetbrains.kotlin.fir.expressions.FirLoop
 
 class FirLoopTarget(
     labelName: String?
-) : FirAbstractTarget<FirLoop>(labelName)
+) : FirAbstractTarget<FirLoop>(labelName) {
+    override lateinit var _labeledElement: FirLoop
+}