From c1754ad427ddf5dc457d08ccbefb9f807b3c0443 Mon Sep 17 00:00:00 2001
From: Tianyu Geng <tgeng@google.com>
Date: Fri, 24 Sep 2021 13:54:42 -0700
Subject: [PATCH] FIR: fix label on non-lambda anonymous function

---
 .../fir/lightTree/converter/DeclarationsConverter.kt   |  5 +++--
 .../org/jetbrains/kotlin/fir/builder/RawFirBuilder.kt  |  3 ++-
 .../tests/labels/labeledFunctionLiteral.fir.kt         | 10 +++++-----
 3 files changed, 10 insertions(+), 8 deletions(-)

diff --git a/compiler/fir/raw-fir/light-tree2fir/src/org/jetbrains/kotlin/fir/lightTree/converter/DeclarationsConverter.kt b/compiler/fir/raw-fir/light-tree2fir/src/org/jetbrains/kotlin/fir/lightTree/converter/DeclarationsConverter.kt
index 0c0498ca928..8985a949a7c 100644
--- a/compiler/fir/raw-fir/light-tree2fir/src/org/jetbrains/kotlin/fir/lightTree/converter/DeclarationsConverter.kt
+++ b/compiler/fir/raw-fir/light-tree2fir/src/org/jetbrains/kotlin/fir/lightTree/converter/DeclarationsConverter.kt
@@ -1507,14 +1507,15 @@ class DeclarationsConverter(
         val functionSymbol: FirFunctionSymbol<*>
         val isAnonymousFunction = identifier == null && isLocal
         val functionBuilder = if (isAnonymousFunction) {
-            val labelName = functionDeclaration.getLabelName() ?: context.calleeNamesForLambda.lastOrNull()?.identifier
-            target = FirFunctionTarget(labelName = labelName, isLambda = false)
             functionSymbol = FirAnonymousFunctionSymbol()
             FirAnonymousFunctionBuilder().apply {
                 source = functionSource
                 receiverTypeRef = receiverType
                 symbol = functionSymbol
                 isLambda = false
+                label = context.getLastLabel(functionDeclaration)
+                val labelName = label?.name ?: context.calleeNamesForLambda.lastOrNull()?.identifier
+                target = FirFunctionTarget(labelName = labelName, isLambda = false)
             }
         } else {
             val functionName = identifier.nameAsSafeName()
diff --git a/compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/RawFirBuilder.kt b/compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/RawFirBuilder.kt
index 2b9a671a4be..08decde77b4 100644
--- a/compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/RawFirBuilder.kt
+++ b/compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/RawFirBuilder.kt
@@ -1174,7 +1174,8 @@ open class RawFirBuilder(
                     receiverTypeRef = receiverType
                     symbol = FirAnonymousFunctionSymbol()
                     isLambda = false
-                    labelName = function.getLabelName() ?: context.calleeNamesForLambda.lastOrNull()?.identifier
+                    label = context.getLastLabel(function)
+                    labelName = label?.name ?: context.calleeNamesForLambda.lastOrNull()?.identifier
                 }
             } else {
                 FirSimpleFunctionBuilder().apply {
diff --git a/compiler/testData/diagnostics/tests/labels/labeledFunctionLiteral.fir.kt b/compiler/testData/diagnostics/tests/labels/labeledFunctionLiteral.fir.kt
index 8f3c8be28da..4857a8e49d9 100644
--- a/compiler/testData/diagnostics/tests/labels/labeledFunctionLiteral.fir.kt
+++ b/compiler/testData/diagnostics/tests/labels/labeledFunctionLiteral.fir.kt
@@ -1,17 +1,17 @@
 // !DIAGNOSTICS: -UNUSED_VARIABLE
 
 val funLit = lambda@ fun String.() {
-    val d1 = this<!UNRESOLVED_LABEL!>@lambda<!>
+    val d1 = this@lambda
 }
 
 fun test() {
-    val funLit = lambda@ fun String.(): <!UNRESOLVED_LABEL!>String<!> {
-        return this<!UNRESOLVED_LABEL!>@lambda<!>
+    val funLit = lambda@ fun String.(): String {
+        return this@lambda
     }
 }
 
 fun lambda() {
-    val funLit = lambda@ fun String.(): <!UNRESOLVED_LABEL!>String<!> {
-        return this<!UNRESOLVED_LABEL!>@lambda<!>
+    val funLit = lambda@ fun String.(): String {
+        return this@lambda
     }
 }