KT UAST: build implicit (this|lambda) receiver

^KTIJ-31794 fixed

GitOrigin-RevId: 2f854f9407d737f30272a436f566e3bb6151d71f

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

@@ -262,10 +262,11 @@ class KotlinUFunctionCallExpression(
         val ktNameReferenceExpression = callee as? KtNameReferenceExpression ?: return null
         val callableDeclaration = baseResolveProviderService.resolveToDeclaration(ktNameReferenceExpression) ?: return null
 
+        val isStatic = (callableDeclaration as? PsiMethod)?.hasModifier(JvmModifier.STATIC) == true
+
         val variable = when (callableDeclaration) {
             is PsiVariable -> callableDeclaration
             is PsiMethod -> {
-                val isStatic = callableDeclaration.hasModifier(JvmModifier.STATIC)
                 callableDeclaration.containingClass?.let { containingClass ->
                     PropertyUtilBase.getPropertyName(callableDeclaration.name)?.let { propertyName ->
                         PropertyUtilBase.findPropertyField(containingClass, propertyName, isStatic)
@@ -274,16 +275,98 @@ class KotlinUFunctionCallExpression(
             }
 
             else -> null
-        } ?: return null
-
-        // an implicit receiver for variables calls (KT-25524)
-        return object : KotlinAbstractUExpression(this), UReferenceExpression {
-            override val sourcePsi: KtNameReferenceExpression get() = ktNameReferenceExpression
-
-            override val resolvedName: String? get() = variable.name
-
-            override fun resolve(): PsiElement = variable
         }
+
+        if (variable != null) {
+            // an implicit receiver for variables calls (KT-25524)
+            return object : KotlinAbstractUExpression(this), UReferenceExpression {
+                override val sourcePsi: KtNameReferenceExpression get() = ktNameReferenceExpression
+
+                override val resolvedName: String? get() = variable.name
+
+                override fun resolve(): PsiElement = variable
+            }
+        }
+
+        // Bail out for non-instance functions
+        if (isStatic || callableDeclaration !is PsiMethod) return null
+
+        val implicitReceiver =
+            if (baseResolveProviderService.isResolvedToExtension(sourcePsi)) {
+                // Extension receiver of PsiParameter
+                callableDeclaration.parameterList.parameters.firstOrNull()
+            } else {
+                callableDeclaration.containingClass
+            }
+                ?: return null
+        val containingUClass = getContainingUClass()
+        if (implicitReceiver is PsiClass && implicitReceiver == containingUClass?.javaPsi) {
+            // Implicit `this`
+            return KotlinUImplicitThis(implicitReceiver, receiverType, this)
+        }
+        // Lambda receiver
+        var enclosingLambda = this.getParentOfType<ULambdaExpression>()
+        while (enclosingLambda != null) {
+            val lambdaReceiver = enclosingLambda.parameters.firstOrNull()
+            val lambdaReceiverType = lambdaReceiver?.type as? PsiClassReferenceType
+            val receiverTypeMatched = if (implicitReceiver is PsiParameter) {
+                // implicit receiver is an extension receiver, e.g., $this$bar
+                lambdaReceiverType == implicitReceiver.type
+            } else {
+                lambdaReceiverType?.resolve() == implicitReceiver
+            }
+            if (receiverTypeMatched) {
+                // Implicit <this>
+                return KotlinUImplicitLambdaReceiver(lambdaReceiver, receiverType, this)
+            }
+            enclosingLambda = enclosingLambda.getParentOfType<ULambdaExpression>()
+        }
+
+        return null
+    }
+
+    private class KotlinUImplicitThis(
+        val implicitReceiver: PsiClass,
+        val receiverType: PsiType?,
+        uasParent: UElement,
+    ) : KotlinAbstractUExpression(uasParent), UThisExpression {
+        override val label: String?
+            get() = null
+
+        override val labelIdentifier: UIdentifier?
+            get() = null
+
+        override fun getExpressionType() = receiverType
+
+        override fun resolve(): PsiElement? = implicitReceiver
+
+        override val sourcePsi: PsiElement?
+            get() = null
+
+        override val javaPsi: PsiElement?
+            get() = implicitReceiver
+    }
+
+    private class KotlinUImplicitLambdaReceiver(
+        val lambdaReceiver: UParameter?,
+        val receiverType: PsiType?,
+        uasParent: UElement,
+    ) : KotlinAbstractUExpression(uasParent), UThisExpression {
+        override val label: String?
+            get() = null
+
+        override val labelIdentifier: UIdentifier?
+            get() = null
+
+        override fun getExpressionType() = receiverType
+
+        override fun resolve(): PsiElement? = lambdaReceiver?.javaPsi
+
+        override val sourcePsi: PsiElement?
+            get() = null
+
+        override val javaPsi: PsiElement?
+            get() = lambdaReceiver?.javaPsi
     }
 
     private fun getMultiResolved(): Iterable<TypedResolveResult<PsiMethod>> {

plugins/kotlin/uast/uast-kotlin-base/tests/test/org/jetbrains/uast/test/common/kotlin/UastApiFixtureTestBase.kt

@@ -4,6 +4,7 @@ package org.jetbrains.uast.test.common.kotlin
 import com.intellij.openapi.application.ApplicationManager
 import com.intellij.platform.uast.testFramework.env.findElementByTextFromPsi
 import com.intellij.psi.PsiArrayInitializerMemberValue
+import com.intellij.psi.PsiClass
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiEnumConstant
 import com.intellij.psi.PsiField
@@ -1130,6 +1131,83 @@ interface UastApiFixtureTestBase {
         TestCase.assertEquals(unusedLambda!!.asRecursiveLogString(), invokedLambda!!.asRecursiveLogString())
     }
 
+    fun checkImplicitReceiver(myFixture: JavaCodeInsightTestFixture) {
+        myFixture.configureByText(
+            "main.kt", """
+                interface BiggerIntent {
+                  fun overall()
+                }
+
+                interface MyIntent {
+                  fun foo()
+                }
+                
+                class Test {
+                  fun MyIntent.bar() {}
+                  
+                  fun baz() {}
+
+                  fun test(bigger: BiggerIntent, intent: MyIntent) {
+                    bigger.apply { // <this>: BiggerIntent
+                      overall() // BiggerIntent#overall
+                      intent.apply { // <this>: MyIntent
+                        overall() // BiggerIntent#overall
+                        foo() // MyIntent#foo
+                        bar() // MyIntent#bar
+                        baz() // Test#baz
+                      }
+                    }
+                    baz() // Test#baz
+                  }
+                }
+            """.trimIndent()
+        )
+        val names = listOf("overall", "foo", "bar", "baz")
+        val uFile = myFixture.file.toUElement()!!
+        var count = 0
+        uFile.accept(object : AbstractUastVisitor() {
+            override fun visitCallExpression(node: UCallExpression): Boolean {
+                if (node.methodName !in names)
+                    return super.visitCallExpression(node)
+
+                val rcvType = node.receiverType
+                TestCase.assertEquals(
+                    node.sourcePsi?.text,
+                    when (node.methodName) {
+                        "overall" -> "BiggerIntent"
+                        "baz" -> "Test"
+                        else -> "MyIntent"
+                    },
+                    rcvType?.canonicalText
+                )
+
+                val rcv = node.receiver
+                TestCase.assertNotNull(node.sourcePsi?.text, rcv)
+
+                TestCase.assertEquals(
+                    node.sourcePsi?.text,
+                    rcvType?.canonicalText,
+                    rcv?.getExpressionType()?.canonicalText
+                )
+
+                val resolvedRcv = (rcv as? UResolvable)?.resolve()
+                TestCase.assertNotNull(node.sourcePsi?.text, resolvedRcv)
+                TestCase.assertTrue(
+                    node.sourcePsi?.text,
+                    if (node.methodName == "baz") {
+                        resolvedRcv is PsiClass
+                    } else {
+                        resolvedRcv is PsiParameter
+                    }
+                )
+
+                count++
+                return super.visitCallExpression(node)
+            }
+        })
+        TestCase.assertEquals(6, count)
+    }
+
     fun checkLambdaImplicitParameters(myFixture: JavaCodeInsightTestFixture) {
         myFixture.configureByText(
             "main.kt", """

plugins/kotlin/uast/uast-kotlin-fir/tests/test/org/jetbrains/fir/uast/test/FirUastApiFixtureTest.kt

@@ -138,6 +138,10 @@ class FirUastApiFixtureTest : KotlinLightCodeInsightFixtureTestCase(), UastApiFi
         checkInvokedLambdaBody(myFixture)
     }
 
+    fun testImplicitReceiver() {
+        checkImplicitReceiver(myFixture)
+    }
+
     fun testLambdaImplicitParameters() {
         checkLambdaImplicitParameters(myFixture)
     }

plugins/kotlin/uast/uast-kotlin/tests/test/org/jetbrains/uast/test/kotlin/comparison/FE1UastApiFixtureTest.kt

@@ -138,6 +138,10 @@ class FE1UastApiFixtureTest : KotlinLightCodeInsightFixtureTestCase(), UastApiFi
         checkInvokedLambdaBody(myFixture)
     }
 
+    fun testImplicitReceiver() {
+        checkImplicitReceiver(myFixture)
+    }
+
     fun testLambdaImplicitParameters() {
         checkLambdaImplicitParameters(myFixture)
     }

uast/uast-common-ide/src/org/jetbrains/uast/analysis/DependencyGraphBuilder.kt

@@ -844,7 +844,8 @@ private const val UAST_KT_ELVIS_NAME = "elvis"
 private const val TEMP_VAR_NAME = "@$,()"
 
 private fun hasImplicitReceiver(callExpression: UCallExpression): Boolean =
-  callExpression.receiver == null && callExpression.receiverType != null
+  (callExpression.receiver == null || callExpression.receiver?.sourcePsi == null) &&
+    callExpression.receiverType != null
 
 private val UExpression?.referenceOrThisIdentifier: String?
   get() = when (this) {