Uast: implicit receiver for function variables calls (KT-25524)

plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/expressions/KotlinUFunctionCallExpression.kt

@@ -16,7 +16,9 @@
 
 package org.jetbrains.uast.kotlin
 
+import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiMethod
+import com.intellij.psi.PsiNamedElement
 import com.intellij.psi.PsiType
 import org.jetbrains.kotlin.descriptors.CallableDescriptor
 import org.jetbrains.kotlin.descriptors.ConstructorDescriptor
@@ -25,6 +27,7 @@ import org.jetbrains.kotlin.psi.psiUtil.parents
 import org.jetbrains.kotlin.resolve.CompileTimeConstantUtils
 import org.jetbrains.kotlin.resolve.calls.callUtil.getResolvedCall
 import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
+import org.jetbrains.kotlin.resolve.calls.model.VariableAsFunctionResolvedCall
 import org.jetbrains.uast.*
 import org.jetbrains.uast.internal.acceptList
 import org.jetbrains.uast.kotlin.declarations.KotlinUIdentifier
@@ -101,7 +104,28 @@ class KotlinUFunctionCallExpression(
     }
 
     override val receiver: UExpression?
-        get() = (uastParent as? UQualifiedReferenceExpression)?.takeIf { it.selector == this }?.receiver
+        get() {
+            (uastParent as? UQualifiedReferenceExpression)?.let {
+                if (it.selector == this) return it.receiver
+            }
+
+            val ktNameReferenceExpression = psi.calleeExpression as? KtNameReferenceExpression ?: return null
+            val variableAsFunctionResolvedCall = resolvedCall as? VariableAsFunctionResolvedCall ?: return null
+
+            // an implicit receiver for variables calls (KT-25524)
+            return object : KotlinAbstractUExpression(this), UReferenceExpression {
+
+                private val resolvedDeclaration = variableAsFunctionResolvedCall.variableCall.resultingDescriptor.toSource()
+
+                override val psi: KtNameReferenceExpression get() = ktNameReferenceExpression
+
+                override val resolvedName: String? get() = (resolvedDeclaration as? PsiNamedElement)?.name
+
+                override fun resolve(): PsiElement? = resolvedDeclaration
+
+            }
+
+        }
 
     override fun resolve(): PsiMethod? {
         val descriptor = resolvedCall?.resultingDescriptor ?: return null

plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/expressions/KotlinUFunctionCallExpression.kt.173

@@ -31,6 +31,7 @@ import org.jetbrains.kotlin.psi.psiUtil.parents
 import org.jetbrains.kotlin.resolve.CompileTimeConstantUtils
 import org.jetbrains.kotlin.resolve.calls.callUtil.getResolvedCall
 import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
+import org.jetbrains.kotlin.resolve.calls.model.VariableAsFunctionResolvedCall
 import org.jetbrains.uast.*
 import org.jetbrains.uast.internal.acceptList
 import org.jetbrains.uast.visitor.UastVisitor
@@ -87,7 +88,28 @@ class KotlinUFunctionCallExpression(
     }
 
     override val receiver: UExpression?
-        get() = (uastParent as? UQualifiedReferenceExpression)?.takeIf { it.selector == this }?.receiver
+        get() {
+            (uastParent as? UQualifiedReferenceExpression)?.let {
+                if (it.selector == this) return it.receiver
+            }
+
+            val ktNameReferenceExpression = psi.calleeExpression as? KtNameReferenceExpression ?: return null
+            val variableAsFunctionResolvedCall = resolvedCall as? VariableAsFunctionResolvedCall ?: return null
+
+            // an implicit receiver for variables calls (KT-25524)
+            return object : KotlinAbstractUExpression(this), UReferenceExpression {
+
+                private val resolvedDeclaration = variableAsFunctionResolvedCall.variableCall.resultingDescriptor.toSource()
+
+                override val psi: KtNameReferenceExpression get() = ktNameReferenceExpression
+
+                override val resolvedName: String? get() = (resolvedDeclaration as? PsiNamedElement)?.name
+
+                override fun resolve(): PsiElement? = resolvedDeclaration
+
+            }
+
+        }
 
     override fun resolve(): PsiMethod? {
         val descriptor = resolvedCall?.resultingDescriptor ?: return null

plugins/uast-kotlin/testData/Lambdas.kt

@@ -2,4 +2,8 @@ import java.util.stream.Stream
 
 fun foo() {
     Stream.empty<String>().filter { it.isEmpty() }
+}
+
+fun doSelectItem(selectItemFunction: () -> Unit) {
+    selectItemFunction()
 }

plugins/uast-kotlin/testData/Lambdas.log.txt

@@ -0,0 +1,28 @@
+UFile (package = )
+    UImportStatement (isOnDemand = false)
+    UClass (name = LambdasKt)
+        UAnnotationMethod (name = foo)
+            UBlockExpression
+                UQualifiedReferenceExpression
+                    UQualifiedReferenceExpression
+                        USimpleNameReferenceExpression (identifier = Stream)
+                        UCallExpression (kind = UastCallKind(name='method_call'), argCount = 0))
+                            UIdentifier (Identifier (empty))
+                            USimpleNameReferenceExpression (identifier = empty)
+                    UCallExpression (kind = UastCallKind(name='method_call'), argCount = 1))
+                        UIdentifier (Identifier (filter))
+                        USimpleNameReferenceExpression (identifier = filter)
+                        ULambdaExpression
+                            UBlockExpression
+                                UQualifiedReferenceExpression
+                                    USimpleNameReferenceExpression (identifier = it)
+                                    UCallExpression (kind = UastCallKind(name='method_call'), argCount = 0))
+                                        UIdentifier (Identifier (isEmpty))
+                                        USimpleNameReferenceExpression (identifier = isEmpty)
+        UAnnotationMethod (name = doSelectItem)
+            UParameter (name = selectItemFunction)
+                UAnnotation (fqName = org.jetbrains.annotations.NotNull)
+            UBlockExpression
+                UCallExpression (kind = UastCallKind(name='method_call'), argCount = 0))
+                    UIdentifier (Identifier (selectItemFunction))
+                    USimpleNameReferenceExpression (identifier = invoke)

plugins/uast-kotlin/testData/Lambdas.render.txt

@@ -0,0 +1,12 @@
+import java.util.stream.Stream
+
+public final class LambdasKt {
+    public static final fun foo() : void {
+        Stream.empty().filter({ 
+            it.isEmpty()
+        })
+    }
+    public static final fun doSelectItem(@org.jetbrains.annotations.NotNull selectItemFunction: kotlin.jvm.functions.Function0<kotlin.Unit>) : void {
+        invoke()
+    }
+}

plugins/uast-kotlin/tests/KotlinUastApiTest.kt

@@ -384,6 +384,29 @@ class KotlinUastApiTest : AbstractKotlinUastTest() {
         }
     }
 
+    @Test
+    fun testLambdaParamCall() {
+        doTest("Lambdas") { _, file ->
+            val lambdaCall = file.findElementByTextFromPsi<UCallExpression>("selectItemFunction()")
+            assertEquals(
+                "UIdentifier (Identifier (selectItemFunction))",
+                lambdaCall.methodIdentifier?.asLogString()
+            )
+            assertEquals(
+                "selectItemFunction",
+                lambdaCall.methodIdentifier?.name
+            )
+            assertEquals(
+                "invoke",
+                lambdaCall.methodName
+            )
+            val receiver = lambdaCall.receiver ?: kotlin.test.fail("receiver expected")
+            assertEquals("UReferenceExpression", receiver.asLogString())
+            val uParameter = (receiver as UReferenceExpression).resolve().toUElement() ?: kotlin.test.fail("uelement expected")
+            assertEquals("UParameter (name = selectItemFunction)", uParameter.asLogString())
+        }
+    }
+
 }
 
 fun <T, R> Iterable<T>.assertedFind(value: R, transform: (T) -> R): T =

plugins/uast-kotlin/tests/KotlinUastApiTest.kt.172

@@ -227,6 +227,29 @@ class KotlinUastApiTest : AbstractKotlinUastTest() {
             val uCallExpression = lambda.uastParent.assertedCast<UCallExpression> { "UCallExpression expected" }
             assertTrue(uCallExpression.valueArguments.contains(lambda))
         }
-
     }
+
+    @Test
+    fun testLambdaParamCall() {
+        doTest("Lambdas") { _, file ->
+            val lambdaCall = file.findElementByTextFromPsi<UCallExpression>("selectItemFunction()")
+            assertEquals(
+                "UIdentifier (Identifier (selectItemFunction))",
+                lambdaCall.methodIdentifier?.asLogString()
+            )
+            assertEquals(
+                "selectItemFunction",
+                lambdaCall.methodIdentifier?.name
+            )
+            assertEquals(
+                "invoke",
+                lambdaCall.methodName
+            )
+            val receiver = lambdaCall.receiver ?: kotlin.test.fail("receiver expected")
+            assertEquals("UReferenceExpression", receiver.asLogString())
+            val uParameter = (receiver as UReferenceExpression).resolve().toUElement() ?: kotlin.test.fail("uelement expected")
+            assertEquals("UParameter (name = selectItemFunction)", uParameter.asLogString())
+        }
+    }
+
 }

plugins/uast-kotlin/tests/KotlinUastApiTest.kt.173

@@ -348,6 +348,29 @@ class KotlinUastApiTest : AbstractKotlinUastTest() {
         }
     }
 
+    @Test
+    fun testLambdaParamCall() {
+        doTest("Lambdas") { _, file ->
+            val lambdaCall = file.findElementByTextFromPsi<UCallExpression>("selectItemFunction()")
+            assertEquals(
+                "UIdentifier (Identifier (selectItemFunction))",
+                lambdaCall.methodIdentifier?.asLogString()
+            )
+            assertEquals(
+                "selectItemFunction",
+                lambdaCall.methodIdentifier?.name
+            )
+            assertEquals(
+                "invoke",
+                lambdaCall.methodName
+            )
+            val receiver = lambdaCall.receiver ?: kotlin.test.fail("receiver expected")
+            assertEquals("UReferenceExpression", receiver.asLogString())
+            val uParameter = (receiver as UReferenceExpression).resolve().toUElement() ?: kotlin.test.fail("uelement expected")
+            assertEquals("UParameter (name = selectItemFunction)", uParameter.asLogString())
+        }
+    }
+
 }
 
 fun <T, R> Iterable<T>.assertedFind(value: R, transform: (T) -> R): T = find { transform(it) == value } ?: throw AssertionError("'$value' not found, only ${this.joinToString { transform(it).toString() }}")

plugins/uast-kotlin/tests/SimpleKotlinRenderLogTest.kt

@@ -81,4 +81,7 @@ class SimpleKotlinRenderLogTest : AbstractKotlinRenderLogTest() {
         // disabled due to inconsistent parents for 2-receivers call (KT-22344)
         check(testName, file, false)
     }
+
+    @Test
+    fun testLambdas() = doTest("Lambdas")
 }

plugins/uast-kotlin/tests/SimpleKotlinRenderLogTest.kt.172

@@ -59,4 +59,7 @@ class SimpleKotlinRenderLogTest : AbstractKotlinRenderLogTest() {
 
     @Test
     fun testConstructors() = doTest("Constructors")
+
+    @Test
+    fun testLambdas() = doTest("Lambdas")
 }