From 5a9ff3471a411cead670f898f9bad4806b751df1 Mon Sep 17 00:00:00 2001
From: Ilya Kirillov <ilya.kirillov@jetbrains.com>
Date: Mon, 14 Dec 2020 21:26:45 +0100
Subject: [PATCH] FIR IDE: fix function targets on context element copy

---
 .../api/LowLevelFirApiFacadeForCompletion.kt  | 25 +++++++++++++------
 1 file changed, 18 insertions(+), 7 deletions(-)

diff --git a/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/api/LowLevelFirApiFacadeForCompletion.kt b/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/api/LowLevelFirApiFacadeForCompletion.kt
index 61c1ea3cea8..ec529402f4b 100644
--- a/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/api/LowLevelFirApiFacadeForCompletion.kt
+++ b/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/api/LowLevelFirApiFacadeForCompletion.kt
@@ -6,16 +6,16 @@
 package org.jetbrains.kotlin.idea.fir.low.level.api.api
 
 import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.declarations.FirFile
-import org.jetbrains.kotlin.fir.declarations.FirProperty
-import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
-import org.jetbrains.kotlin.fir.declarations.FirSimpleFunction
+import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.builder.buildPropertyAccessorCopy
 import org.jetbrains.kotlin.fir.declarations.builder.buildPropertyCopy
 import org.jetbrains.kotlin.fir.declarations.builder.buildSimpleFunctionCopy
+import org.jetbrains.kotlin.fir.expressions.FirReturnExpression
 import org.jetbrains.kotlin.fir.resolve.FirTowerDataContext
 import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
+import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
 import org.jetbrains.kotlin.idea.fir.low.level.api.FirModuleResolveStateForCompletion
 import org.jetbrains.kotlin.idea.fir.low.level.api.FirModuleResolveStateImpl
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirTowerDataContextCollector
@@ -106,10 +106,9 @@ object LowLevelFirApiFacadeForCompletion {
             resolvePhase = minOf(originalFunction.resolvePhase, FirResolvePhase.DECLARATIONS)
             source = builtFunction.source
             session = state.rootModuleSession
-        }
+        }.apply { reassignAllReturnTargets (builtFunction) }
     }
 
-
     private fun buildPropertyCopyForCompletion(
         firIdeProvider: FirIdeProvider,
         element: KtProperty,
@@ -129,7 +128,7 @@ object LowLevelFirApiFacadeForCompletion {
                 resolvePhase = minOf(builtSetter.resolvePhase, FirResolvePhase.DECLARATIONS)
                 source = builtSetter.source
                 session = state.rootModuleSession
-            }
+            }.apply { reassignAllReturnTargets(builtSetter) }
         } else {
             builtSetter
         }
@@ -146,4 +145,16 @@ object LowLevelFirApiFacadeForCompletion {
             session = state.rootModuleSession
         }
     }
+
+
+    private fun FirFunction<*>.reassignAllReturnTargets(from: FirFunction<*>) {
+        this.accept(object : FirVisitorVoid() {
+            override fun visitElement(element: FirElement) {
+                if (element is FirReturnExpression && element.target.labeledElement == from) {
+                    element.target.bind(this@reassignAllReturnTargets)
+                }
+                element.acceptChildren(this)
+            }
+        })
+    }
 }