Raw FIR: introduce lambda argument expressions

Without it we cannot distinguish lambda argument from just last argument

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrVisitor.kt

@@ -617,9 +617,9 @@ internal class Fir2IrVisitor(
         }
     }
 
-    override fun visitNamedArgumentExpression(namedArgumentExpression: FirNamedArgumentExpression, data: Any?): IrElement {
+    override fun visitWrappedArgumentExpression(wrappedArgumentExpression: FirWrappedArgumentExpression, data: Any?): IrElement {
         // TODO: change this temporary hack to something correct
-        return namedArgumentExpression.expression.toIrExpression()
+        return wrappedArgumentExpression.expression.toIrExpression()
     }
 
     private fun FirQualifiedAccess.toIrExpression(typeRef: FirTypeRef): IrExpression {

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

@@ -312,7 +312,11 @@ class RawFirBuilder(val session: FirSession, val stubMode: Boolean) {
 
         private fun KtCallElement.extractArgumentsTo(container: FirAbstractCall) {
             for (argument in this.valueArguments) {
-                container.arguments += argument.toFirExpression()
+                val argumentExpression = argument.toFirExpression()
+                container.arguments += when (argument) {
+                    is KtLambdaArgument -> FirLambdaArgumentExpressionImpl(session, argument, argumentExpression)
+                    else -> argumentExpression
+                }
             }
         }
 
@@ -1244,9 +1248,7 @@ class RawFirBuilder(val session: FirSession, val stubMode: Boolean) {
                 }
                 this.calleeReference = calleeReference
                 firFunctionCalls += this
-                for (argument in expression.valueArguments) {
+                expression.extractArgumentsTo(this)
-                    arguments += argument.toFirExpression()
-                }
                 for (typeArgument in expression.typeArguments) {
                     typeArguments += typeArgument.convert<FirTypeProjection>()
                 }

compiler/fir/psi2fir/testData/rawBuilder/expressions/lambda.txt

@@ -24,7 +24,7 @@ FILE: lambda.kt
 
         }
 
-        use#(use@fun <implicit>.<anonymous>(<destruct>: <implicit>): <implicit> {
+        use#(<L> = use@fun <implicit>.<anonymous>(<destruct>: <implicit>): <implicit> {
             ^  {
                 lval x: <implicit> = <destruct>#.component1()
                 lval y: <implicit> = <destruct>#.component2()
@@ -33,7 +33,7 @@ FILE: lambda.kt
 
         }
         )
-        ^foo use#(use@fun <implicit>.<anonymous>(): <implicit> {
+        ^foo use#(<L> = use@fun <implicit>.<anonymous>(): <implicit> {
             ^  {
                 when () {
                     ==(it#.x#, Int(0)) ->  {
@@ -50,7 +50,7 @@ FILE: lambda.kt
         )
     }
     public? final? fun bar(): Int {
-        ^bar use#(lambda@fun <implicit>.<anonymous>(): <implicit> {
+        ^bar use#(<L> = lambda@fun <implicit>.<anonymous>(): <implicit> {
             ^  {
                 when () {
                     ==(it#.x#, Int(0)) ->  {
@@ -68,7 +68,7 @@ FILE: lambda.kt
     }
     public? final? fun test(list: List<Int>): kotlin/Unit {
         lval map: <implicit> = mutableMapOf#<Int, String>()
-        list#.forEach#(forEach@fun <implicit>.<anonymous>(): <implicit> {
+        list#.forEach#(<L> = forEach@fun <implicit>.<anonymous>(): <implicit> {
             ^  {
                 lval <complex-set>: <implicit> = map#.getOrPut#(it#, getOrPut@fun <implicit>.<anonymous>(): <implicit> {
                     ^  {

compiler/fir/psi2fir/testData/rawBuilder/expressions/these.txt

@@ -26,7 +26,7 @@ FILE: these.kt
         ^extension this#.bar#()
     }
     public? final? fun test(some: Some): Int {
-        ^test with#(some#, with@fun <implicit>.<anonymous>(): <implicit> {
+        ^test with#(some#, <L> = with@fun <implicit>.<anonymous>(): <implicit> {
             ^  {
                 this#.foo#().plus#(this@with.extension#())
             }

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

@@ -630,6 +630,11 @@ class FirRenderer(builder: StringBuilder) : FirVisitorVoid() {
         namedArgumentExpression.expression.accept(this)
     }
 
+    override fun visitLambdaArgumentExpression(lambdaArgumentExpression: FirLambdaArgumentExpression) {
+        print("<L> = ")
+        lambdaArgumentExpression.expression.accept(this)
+    }
+
     override fun visitCall(call: FirCall) {
         print("(")
         call.arguments.renderSeparated()

compiler/fir/tree/src/org/jetbrains/kotlin/fir/expressions/FirLambdaArgumentExpression.kt

@@ -0,0 +1,13 @@
+/*
+ * Copyright 2010-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.fir.expressions
+
+import org.jetbrains.kotlin.fir.visitors.FirVisitor
+
+interface FirLambdaArgumentExpression : FirWrappedArgumentExpression {
+    override fun <R, D> accept(visitor: FirVisitor<R, D>, data: D): R =
+        visitor.visitLambdaArgumentExpression(this, data)
+}

compiler/fir/tree/src/org/jetbrains/kotlin/fir/expressions/FirNamedArgumentExpression.kt

@@ -8,16 +8,9 @@ package org.jetbrains.kotlin.fir.expressions
 import org.jetbrains.kotlin.fir.visitors.FirVisitor
 import org.jetbrains.kotlin.name.Name
 
-interface FirNamedArgumentExpression : FirExpression {
+interface FirNamedArgumentExpression : FirWrappedArgumentExpression {
     val name: Name
 
-    val expression: FirExpression
-
     override fun <R, D> accept(visitor: FirVisitor<R, D>, data: D): R =
         visitor.visitNamedArgumentExpression(this, data)
-
-    override fun <R, D> acceptChildren(visitor: FirVisitor<R, D>, data: D) {
-        expression.accept(visitor, data)
-        super.acceptChildren(visitor, data)
-    }
 }

compiler/fir/tree/src/org/jetbrains/kotlin/fir/expressions/FirWrappedArgumentExpression.kt

@@ -0,0 +1,24 @@
+/*
+ * Copyright 2010-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.fir.expressions
+
+import org.jetbrains.kotlin.fir.types.FirTypeRef
+import org.jetbrains.kotlin.fir.visitors.FirVisitor
+
+interface FirWrappedArgumentExpression : FirExpression {
+    val expression: FirExpression
+
+    override val typeRef: FirTypeRef
+        get() = expression.typeRef
+
+    override fun <R, D> accept(visitor: FirVisitor<R, D>, data: D): R =
+        visitor.visitWrappedArgumentExpression(this, data)
+
+    override fun <R, D> acceptChildren(visitor: FirVisitor<R, D>, data: D) {
+        expression.accept(visitor, data)
+        super.acceptChildren(visitor, data)
+    }
+}

compiler/fir/tree/src/org/jetbrains/kotlin/fir/expressions/impl/FirLambdaArgumentExpressionImpl.kt

@@ -0,0 +1,30 @@
+/*
+ * Copyright 2010-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.fir.expressions.impl
+
+import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.expressions.FirExpression
+import org.jetbrains.kotlin.fir.expressions.FirLambdaArgumentExpression
+import org.jetbrains.kotlin.fir.transformSingle
+import org.jetbrains.kotlin.fir.types.FirTypeRef
+import org.jetbrains.kotlin.fir.visitors.FirTransformer
+
+class FirLambdaArgumentExpressionImpl(
+    session: FirSession,
+    psi: PsiElement?,
+    override var expression: FirExpression
+) : FirLambdaArgumentExpression, FirAbstractExpression(session, psi) {
+    override var typeRef: FirTypeRef
+        get() = super<FirLambdaArgumentExpression>.typeRef
+        set(_) {}
+
+    override fun <D> transformChildren(transformer: FirTransformer<D>, data: D): FirElement {
+        expression = expression.transformSingle(transformer, data)
+        return super<FirAbstractExpression>.transformChildren(transformer, data)
+    }
+}

compiler/fir/tree/src/org/jetbrains/kotlin/fir/expressions/impl/FirNamedArgumentExpressionImpl.kt

@@ -9,8 +9,10 @@ import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.expressions.FirExpression
+import org.jetbrains.kotlin.fir.expressions.FirLambdaArgumentExpression
 import org.jetbrains.kotlin.fir.expressions.FirNamedArgumentExpression
 import org.jetbrains.kotlin.fir.transformSingle
+import org.jetbrains.kotlin.fir.types.FirTypeRef
 import org.jetbrains.kotlin.fir.visitors.FirTransformer
 import org.jetbrains.kotlin.name.Name
 
@@ -20,6 +22,10 @@ class FirNamedArgumentExpressionImpl(
     override val name: Name,
     override var expression: FirExpression
 ) : FirNamedArgumentExpression, FirAbstractExpression(session, psi) {
+    override var typeRef: FirTypeRef
+        get() = super<FirNamedArgumentExpression>.typeRef
+        set(_) {}
+
     override fun <D> transformChildren(transformer: FirTransformer<D>, data: D): FirElement {
         expression = expression.transformSingle(transformer, data)
         return super<FirAbstractExpression>.transformChildren(transformer, data)

compiler/fir/tree/visitors/org/jetbrains/kotlin/fir/visitors/FirTransformerGenerated.kt

@@ -260,10 +260,6 @@ abstract class FirTransformer<in D> : FirVisitor<CompositeTransformResult<FirEle
         return transformJump(returnExpression, data)
     }
 
-    open fun transformNamedArgumentExpression(namedArgumentExpression: FirNamedArgumentExpression, data: D): CompositeTransformResult<FirStatement> {
-        return transformExpression(namedArgumentExpression, data)
-    }
-
     open fun transformThrowExpression(throwExpression: FirThrowExpression, data: D): CompositeTransformResult<FirStatement> {
         return transformExpression(throwExpression, data)
     }
@@ -276,6 +272,18 @@ abstract class FirTransformer<in D> : FirVisitor<CompositeTransformResult<FirEle
         return transformExpression(whenExpression, data)
     }
 
+    open fun transformWrappedArgumentExpression(wrappedArgumentExpression: FirWrappedArgumentExpression, data: D): CompositeTransformResult<FirStatement> {
+        return transformExpression(wrappedArgumentExpression, data)
+    }
+
+    open fun transformLambdaArgumentExpression(lambdaArgumentExpression: FirLambdaArgumentExpression, data: D): CompositeTransformResult<FirStatement> {
+        return transformWrappedArgumentExpression(lambdaArgumentExpression, data)
+    }
+
+    open fun transformNamedArgumentExpression(namedArgumentExpression: FirNamedArgumentExpression, data: D): CompositeTransformResult<FirStatement> {
+        return transformWrappedArgumentExpression(namedArgumentExpression, data)
+    }
+
     open fun transformLoop(loop: FirLoop, data: D): CompositeTransformResult<FirStatement> {
         return transformStatement(loop, data)
     }
@@ -556,6 +564,10 @@ abstract class FirTransformer<in D> : FirVisitor<CompositeTransformResult<FirEle
         return transformLabeledElement(labeledElement, data)
     }
 
+    final override fun visitLambdaArgumentExpression(lambdaArgumentExpression: FirLambdaArgumentExpression, data: D): CompositeTransformResult<FirElement> {
+        return transformLambdaArgumentExpression(lambdaArgumentExpression, data)
+    }
+
     final override fun visitLoop(loop: FirLoop, data: D): CompositeTransformResult<FirElement> {
         return transformLoop(loop, data)
     }
@@ -736,6 +748,10 @@ abstract class FirTransformer<in D> : FirVisitor<CompositeTransformResult<FirEle
         return transformWhileLoop(whileLoop, data)
     }
 
+    final override fun visitWrappedArgumentExpression(wrappedArgumentExpression: FirWrappedArgumentExpression, data: D): CompositeTransformResult<FirElement> {
+        return transformWrappedArgumentExpression(wrappedArgumentExpression, data)
+    }
+
     final override fun visitElement(element: FirElement, data: D): CompositeTransformResult<FirElement> {
         return transformElement(element, data)
     }

compiler/fir/tree/visitors/org/jetbrains/kotlin/fir/visitors/FirVisitorGenerated.kt

@@ -260,10 +260,6 @@ abstract class FirVisitor<out R, in D> {
         return visitJump(returnExpression, data)
     }
 
-    open fun visitNamedArgumentExpression(namedArgumentExpression: FirNamedArgumentExpression, data: D): R {
-        return visitExpression(namedArgumentExpression, data)
-    }
-
     open fun visitThrowExpression(throwExpression: FirThrowExpression, data: D): R {
         return visitExpression(throwExpression, data)
     }
@@ -276,6 +272,18 @@ abstract class FirVisitor<out R, in D> {
         return visitExpression(whenExpression, data)
     }
 
+    open fun visitWrappedArgumentExpression(wrappedArgumentExpression: FirWrappedArgumentExpression, data: D): R {
+        return visitExpression(wrappedArgumentExpression, data)
+    }
+
+    open fun visitLambdaArgumentExpression(lambdaArgumentExpression: FirLambdaArgumentExpression, data: D): R {
+        return visitWrappedArgumentExpression(lambdaArgumentExpression, data)
+    }
+
+    open fun visitNamedArgumentExpression(namedArgumentExpression: FirNamedArgumentExpression, data: D): R {
+        return visitWrappedArgumentExpression(namedArgumentExpression, data)
+    }
+
     open fun visitLoop(loop: FirLoop, data: D): R {
         return visitStatement(loop, data)
     }

compiler/fir/tree/visitors/org/jetbrains/kotlin/fir/visitors/FirVisitorVoidGenerated.kt

@@ -260,10 +260,6 @@ abstract class FirVisitorVoid : FirVisitor<Unit, Nothing?>() {
         visitJump(returnExpression, null)
     }
 
-    open fun visitNamedArgumentExpression(namedArgumentExpression: FirNamedArgumentExpression) {
-        visitExpression(namedArgumentExpression, null)
-    }
-
     open fun visitThrowExpression(throwExpression: FirThrowExpression) {
         visitExpression(throwExpression, null)
     }
@@ -276,6 +272,18 @@ abstract class FirVisitorVoid : FirVisitor<Unit, Nothing?>() {
         visitExpression(whenExpression, null)
     }
 
+    open fun visitWrappedArgumentExpression(wrappedArgumentExpression: FirWrappedArgumentExpression) {
+        visitExpression(wrappedArgumentExpression, null)
+    }
+
+    open fun visitLambdaArgumentExpression(lambdaArgumentExpression: FirLambdaArgumentExpression) {
+        visitWrappedArgumentExpression(lambdaArgumentExpression, null)
+    }
+
+    open fun visitNamedArgumentExpression(namedArgumentExpression: FirNamedArgumentExpression) {
+        visitWrappedArgumentExpression(namedArgumentExpression, null)
+    }
+
     open fun visitLoop(loop: FirLoop) {
         visitStatement(loop, null)
     }
@@ -556,6 +564,10 @@ abstract class FirVisitorVoid : FirVisitor<Unit, Nothing?>() {
         visitLabeledElement(labeledElement)
     }
 
+    final override fun visitLambdaArgumentExpression(lambdaArgumentExpression: FirLambdaArgumentExpression, data: Nothing?) {
+        visitLambdaArgumentExpression(lambdaArgumentExpression)
+    }
+
     final override fun visitLoop(loop: FirLoop, data: Nothing?) {
         visitLoop(loop)
     }
@@ -736,6 +748,10 @@ abstract class FirVisitorVoid : FirVisitor<Unit, Nothing?>() {
         visitWhileLoop(whileLoop)
     }
 
+    final override fun visitWrappedArgumentExpression(wrappedArgumentExpression: FirWrappedArgumentExpression, data: Nothing?) {
+        visitWrappedArgumentExpression(wrappedArgumentExpression)
+    }
+
     final override fun visitElement(element: FirElement, data: Nothing?) {
         visitElement(element)
     }