FIR: Fix label/receiver for lambda within infix calls

While this fix only changes label name, it's important for proper
lambda resolution because receiver is set for a lambda
only in presence of a label

This commit fixes a lot of checkType calls in diagnostic tests

compiler/fir/analysis-tests/testData/resolve/labelAndReceiverForInfix.kt

@@ -0,0 +1,14 @@
+class A {
+    fun bar() {}
+}
+
+infix fun (() -> Unit).foo(x: A.() -> Unit) {}
+
+fun main() {
+    {
+        return@foo
+    } foo {
+        bar()
+        return@foo
+    }
+}

compiler/fir/analysis-tests/testData/resolve/labelAndReceiverForInfix.txt

@@ -0,0 +1,22 @@
+FILE: labelAndReceiverForInfix.kt
+    public final class A : R|kotlin/Any| {
+        public constructor(): R|A| {
+            super<R|kotlin/Any|>()
+        }
+
+        public final fun bar(): R|kotlin/Unit| {
+        }
+
+    }
+    public final infix fun R|() -> kotlin/Unit|.foo(x: R|A.() -> kotlin/Unit|): R|kotlin/Unit| {
+    }
+    public final fun main(): R|kotlin/Unit| {
+        foo@fun <anonymous>(): R|kotlin/Unit| {
+            ^@foo Unit
+        }
+        .R|/foo|(foo@fun R|A|.<anonymous>(): R|kotlin/Unit| {
+            this@R|special/anonymous|.R|/A.bar|()
+            ^@foo Unit
+        }
+        )
+    }

compiler/fir/analysis-tests/tests/org/jetbrains/kotlin/fir/FirDiagnosticsTestGenerated.java

@@ -228,6 +228,11 @@ public class FirDiagnosticsTestGenerated extends AbstractFirDiagnosticsTest {
         runTest("compiler/fir/analysis-tests/testData/resolve/javaStaticScopeInheritance.kt");
     }
 
+    @TestMetadata("labelAndReceiverForInfix.kt")
+    public void testLabelAndReceiverForInfix() throws Exception {
+        runTest("compiler/fir/analysis-tests/testData/resolve/labelAndReceiverForInfix.kt");
+    }
+
     @TestMetadata("lambdaArgInScopeFunction.kt")
     public void testLambdaArgInScopeFunction() throws Exception {
         runTest("compiler/fir/analysis-tests/testData/resolve/lambdaArgInScopeFunction.kt");

compiler/fir/analysis-tests/tests/org/jetbrains/kotlin/fir/FirDiagnosticsWithLightTreeTestGenerated.java

@@ -228,6 +228,11 @@ public class FirDiagnosticsWithLightTreeTestGenerated extends AbstractFirDiagnos
         runTest("compiler/fir/analysis-tests/testData/resolve/javaStaticScopeInheritance.kt");
     }
 
+    @TestMetadata("labelAndReceiverForInfix.kt")
+    public void testLabelAndReceiverForInfix() throws Exception {
+        runTest("compiler/fir/analysis-tests/testData/resolve/labelAndReceiverForInfix.kt");
+    }
+
     @TestMetadata("lambdaArgInScopeFunction.kt")
     public void testLambdaArgInScopeFunction() throws Exception {
         runTest("compiler/fir/analysis-tests/testData/resolve/lambdaArgInScopeFunction.kt");

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

@@ -10,10 +10,10 @@ import org.jetbrains.kotlin.fir.FirFunctionTarget
 import org.jetbrains.kotlin.fir.FirLabel
 import org.jetbrains.kotlin.fir.FirLoopTarget
 import org.jetbrains.kotlin.fir.expressions.FirExpression
-import org.jetbrains.kotlin.fir.expressions.builder.FirFunctionCallBuilder
 import org.jetbrains.kotlin.fir.symbols.impl.FirTypeParameterSymbol
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.name.Name
 
 class Context<T> {
     lateinit var packageFqName: FqName
@@ -21,9 +21,9 @@ class Context<T> {
     val currentClassId get() = ClassId(packageFqName, className, firFunctionTargets.isNotEmpty())
 
     val firFunctionTargets = mutableListOf<FirFunctionTarget>()
-    val firFunctionCalls = mutableListOf<FirFunctionCallBuilder>()
+    val calleeNamesForLambda = mutableListOf<Name>()
     val firLabels = mutableListOf<FirLabel>()
     val firLoopTargets = mutableListOf<FirLoopTarget>()
     var capturedTypeParameters = persistentListOf<FirTypeParameterSymbol>()
     val arraySetArgument = mutableMapOf<T, FirExpression>()
-}
+}

compiler/fir/raw-fir/light-tree2fir/src/org/jetbrains/kotlin/fir/lightTree/converter/ExpressionsConverter.kt

@@ -137,7 +137,7 @@ class ExpressionsConverter(
             receiverTypeRef = implicitType
             symbol = FirAnonymousFunctionSymbol()
             isLambda = true
-            label = context.firLabels.pop() ?: context.firFunctionCalls.lastOrNull()?.calleeReference?.name?.let {
+            label = context.firLabels.pop() ?: context.calleeNamesForLambda.lastOrNull()?.let {
                 buildLabel { name = it.asString() }
             }
             target = FirFunctionTarget(labelName = label?.name, isLambda = true)
@@ -203,7 +203,6 @@ class ExpressionsConverter(
         var isLeftArgument = true
         lateinit var operationTokenName: String
         var leftArgNode: LighterASTNode? = null
-        var rightArgAsFir: FirExpression = buildErrorExpression(null, ConeSimpleDiagnostic("No right operand", DiagnosticKind.Syntax))
         var rightArg: LighterASTNode? = null
         var operationReferenceSource: FirLightSourceElement? = null
         binaryExpression.forEachChildren {
@@ -217,7 +216,6 @@ class ExpressionsConverter(
                     if (isLeftArgument) {
                         leftArgNode = it
                     } else {
-                        rightArgAsFir = getAsFirExpression(it, "No right operand")
                         rightArg = it
                     }
                 }
@@ -226,7 +224,23 @@ class ExpressionsConverter(
 
         val baseSource = binaryExpression.toFirSourceElement()
         val operationToken = operationTokenName.getOperationSymbol()
+        if (operationToken == IDENTIFIER) {
+            context.calleeNamesForLambda += operationTokenName.nameAsSafeName()
+        }
+
+        val rightArgAsFir =
+            if (rightArg != null)
+                getAsFirExpression<FirExpression>(rightArg, "No right operand")
+            else
+                buildErrorExpression(null, ConeSimpleDiagnostic("No right operand", DiagnosticKind.Syntax))
+
         val leftArgAsFir = getAsFirExpression<FirExpression>(leftArgNode, "No left operand")
+
+        if (operationToken == IDENTIFIER) {
+            // No need for the callee name since arguments are already generated
+            context.calleeNamesForLambda.removeLast()
+        }
+
         when (operationToken) {
             ELVIS ->
                 return leftArgAsFir.generateNotNullOrOther(baseSession, rightArgAsFir, "elvis", baseSource)
@@ -544,9 +558,9 @@ class ExpressionsConverter(
                 this.source = source
                 this.calleeReference = calleeReference
 
-                context.firFunctionCalls += this
+                context.calleeNamesForLambda += calleeReference.name
                 this.extractArgumentsFrom(valueArguments.flatMap { convertValueArguments(it) }, stubMode)
-                context.firFunctionCalls.removeLast()
+                context.calleeNamesForLambda.removeLast()
             }
         } else {
             FirQualifiedAccessExpressionBuilder().apply {

compiler/fir/raw-fir/light-tree2fir/tests/org/jetbrains/kotlin/fir/lightTree/LightTree2FirConverterTestCaseGenerated.java

@@ -248,6 +248,11 @@ public class LightTree2FirConverterTestCaseGenerated extends AbstractLightTree2F
             runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/init.kt");
         }
 
+        @TestMetadata("labelForInfix.kt")
+        public void testLabelForInfix() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/labelForInfix.kt");
+        }
+
         @TestMetadata("lambda.kt")
         public void testLambda() throws Exception {
             runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/lambda.kt");

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

@@ -911,7 +911,7 @@ class RawFirBuilder(
                     }
                 }
                 val expressionSource = expression.toFirSourceElement()
-                label = context.firLabels.pop() ?: context.firFunctionCalls.lastOrNull()?.calleeReference?.name?.let {
+                label = context.firLabels.pop() ?: context.calleeNamesForLambda.lastOrNull()?.let {
                     buildLabel {
                         source = expressionSource
                         name = it.asString()
@@ -1472,9 +1472,21 @@ class RawFirBuilder(
 
         override fun visitBinaryExpression(expression: KtBinaryExpression, data: Unit): FirElement {
             val operationToken = expression.operationToken
+
+            if (operationToken == IDENTIFIER) {
+                context.calleeNamesForLambda += expression.operationReference.getReferencedNameAsName()
+            }
+
             val leftArgument = expression.left.toFirExpression("No left operand")
             val rightArgument = expression.right.toFirExpression("No right operand")
+
+            if (operationToken == IDENTIFIER) {
+                // No need for the callee name since arguments are already generated
+                context.calleeNamesForLambda.removeLast()
+            }
+
             val source = expression.toFirSourceElement()
+
             when (operationToken) {
                 ELVIS ->
                     return leftArgument.generateNotNullOrOther(baseSession, rightArgument, "elvis", source)
@@ -1617,9 +1629,9 @@ class RawFirBuilder(
                 FirFunctionCallBuilder().apply {
                     this.source = source
                     this.calleeReference = calleeReference
-                    context.firFunctionCalls += this
+                    context.calleeNamesForLambda += calleeReference.name
                     expression.extractArgumentsTo(this)
-                    context.firFunctionCalls.removeLast()
+                    context.calleeNamesForLambda.removeLast()
                 }
             }
 

compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/labelForInfix.kt

@@ -0,0 +1,3 @@
+fun main() {
+    { length } foo { bar() }
+}

compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/labelForInfix.txt

@@ -0,0 +1,10 @@
+FILE: labelForInfix.kt
+    public? final? fun main(): R|kotlin/Unit| {
+        foo@fun <implicit>.<anonymous>(): <implicit> {
+            length#
+        }
+        .foo#(foo@fun <implicit>.<anonymous>(): <implicit> {
+            bar#()
+        }
+        )
+    }

compiler/fir/raw-fir/psi2fir/tests/org/jetbrains/kotlin/fir/builder/RawFirBuilderTestCaseGenerated.java

@@ -248,6 +248,11 @@ public class RawFirBuilderTestCaseGenerated extends AbstractRawFirBuilderTestCas
             runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/init.kt");
         }
 
+        @TestMetadata("labelForInfix.kt")
+        public void testLabelForInfix() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/labelForInfix.kt");
+        }
+
         @TestMetadata("lambda.kt")
         public void testLambda() throws Exception {
             runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/lambda.kt");

compiler/testData/codegen/box/labels/infixCallLabelling.kt

@@ -1,4 +1,3 @@
-// IGNORE_BACKEND_FIR: JVM_IR
 fun test(x: Int): Int {
     x myMap {
         return@myMap
@@ -22,4 +21,4 @@ fun box(): String {
     myMap(0)
 
     return "OK"
-}
+}

compiler/visualizer/tests/org/jetbrains/kotlin/visualizer/fir/FirVisualizerForRawFirDataGenerated.java

@@ -248,6 +248,11 @@ public class FirVisualizerForRawFirDataGenerated extends AbstractFirVisualizer {
             runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/init.kt");
         }
 
+        @TestMetadata("labelForInfix.kt")
+        public void testLabelForInfix() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/labelForInfix.kt");
+        }
+
         @TestMetadata("lambda.kt")
         public void testLambda() throws Exception {
             runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/lambda.kt");

compiler/visualizer/tests/org/jetbrains/kotlin/visualizer/psi/PsiVisualizerForRawFirDataGenerated.java

@@ -248,6 +248,11 @@ public class PsiVisualizerForRawFirDataGenerated extends AbstractPsiVisualizer {
             runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/init.kt");
         }
 
+        @TestMetadata("labelForInfix.kt")
+        public void testLabelForInfix() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/labelForInfix.kt");
+        }
+
         @TestMetadata("lambda.kt")
         public void testLambda() throws Exception {
             runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/lambda.kt");