Uast: handling reified methods receivers (KT-37613)

plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/KotlinUastLanguagePlugin.kt

@@ -54,6 +54,7 @@ import org.jetbrains.uast.kotlin.declarations.KotlinUMethodWithFakeLightDelegate
 import org.jetbrains.uast.kotlin.expressions.*
 import org.jetbrains.uast.kotlin.psi.UastFakeLightMethod
 import org.jetbrains.uast.kotlin.psi.UastKotlinPsiParameter
+import org.jetbrains.uast.kotlin.psi.UastKotlinPsiParameterBase
 import org.jetbrains.uast.kotlin.psi.UastKotlinPsiVariable
 
 interface KotlinUastResolveProviderService {
@@ -251,14 +252,18 @@ internal object KotlinConverter {
             is KtStringTemplateEntry -> element.expression?.let { convertExpression(it, givenParent, expectedTypes) } ?: expr<UExpression> { UastEmptyExpression(givenParent) }
             is KtWhenEntry -> el<USwitchClauseExpressionWithBody>(build(::KotlinUSwitchEntry))
             is KtWhenCondition -> convertWhenCondition(element, givenParent, expectedTypes)
-            is KtTypeReference -> el<UTypeReferenceExpression> { LazyKotlinUTypeReferenceExpression(element, givenParent) }
+            is KtTypeReference ->
+                expectedTypes.accommodate(
+                    alternative { LazyKotlinUTypeReferenceExpression(element, givenParent) },
+                    alternative { convertReceiverParameter(element) }
+                ).firstOrNull()
             is KtConstructorDelegationCall ->
                 el<UCallExpression> { KotlinUFunctionCallExpression(element, givenParent) }
             is KtSuperTypeCallEntry ->
                 el<UExpression> {
                     (element.getParentOfType<KtClassOrObject>(true)?.parent as? KtObjectLiteralExpression)
-                            ?.toUElementOfType<UExpression>()
+                        ?.toUElementOfType<UExpression>()
-                    ?: KotlinUFunctionCallExpression(element, givenParent)
+                        ?: KotlinUFunctionCallExpression(element, givenParent)
                 }
             is KtImportDirective -> el<UImportStatement>(build(::KotlinUImportStatement))
             else -> {
@@ -285,6 +290,11 @@ internal object KotlinConverter {
         }}
     }
 
+    internal fun convertReceiverParameter(receiver: KtTypeReference): UParameter? {
+        val call = (receiver.parent as? KtCallableDeclaration) ?: return null
+        if (call.receiverTypeReference != receiver) return null
+        return call.toUElementOfType<UMethod>()?.uastParameters?.firstOrNull()
+    }
 
     internal fun convertEntry(entry: KtStringTemplateEntry,
                               givenParent: UElement?,
@@ -504,6 +514,9 @@ internal object KotlinConverter {
                         el<UField>(buildKtOpt(original.kotlinOrigin, ::KotlinUField))
                 is KtLightParameter -> el<UParameter>(buildKtOpt(original.kotlinOrigin, ::KotlinUParameter))
                 is UastKotlinPsiParameter -> el<UParameter>(buildKt(original.ktParameter, ::KotlinUParameter))
+                is UastKotlinPsiParameterBase<*> -> el<UParameter> {
+                    original.ktOrigin.safeAs<KtTypeReference>()?.let { convertReceiverParameter(it) }
+                }
                 is UastKotlinPsiVariable -> el<UVariable>(buildKt(original.ktElement, ::KotlinUVariable))
 
                 is KtEnumEntry -> el<UEnumConstant> {

plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/declarations/KotlinUMethod.kt

@@ -63,19 +63,18 @@ open class KotlinUMethod(
     private val receiver by lz { (sourcePsi as? KtCallableDeclaration)?.receiverTypeReference }
 
     override val uastParameters by lz {
-        val lightParams = psi.parameterList.parameters
+
-        val receiver = receiver ?: return@lz lightParams.map {
+        fun parameterOrigin(psiParameter: PsiParameter?): KtElement? = when (psiParameter) {
-            KotlinUParameter(
+            is KtLightElement<*, *> -> psiParameter.kotlinOrigin
-                it, when (it) {
+            is UastKotlinPsiParameter -> psiParameter.ktParameter
-                    is KtLightElement<*, *> -> it.kotlinOrigin
+            else -> null
-                    is UastKotlinPsiParameter -> it.ktParameter
-                    else -> null
-                }, this
-            )
         }
+
+        val lightParams = psi.parameterList.parameters
+        val receiver = receiver ?: return@lz lightParams.map { KotlinUParameter(it, parameterOrigin(it), this) }
         val receiverLight = lightParams.firstOrNull() ?: return@lz emptyList<UParameter>()
         val uParameters = SmartList<UParameter>(KotlinReceiverUParameter(receiverLight, receiver, this))
-        lightParams.drop(1).mapTo(uParameters) { KotlinUParameter(it, (it as? KtLightElement<*, *>)?.kotlinOrigin, this) }
+        lightParams.drop(1).mapTo(uParameters) { KotlinUParameter(it, parameterOrigin(it), this) }
         uParameters
     }
 

plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/declarations/KotlinUMethod.kt.191

@@ -64,19 +64,18 @@ open class KotlinUMethod(
     private val receiver by lz { (sourcePsi as? KtCallableDeclaration)?.receiverTypeReference }
 
     override val uastParameters by lz {
-        val lightParams = psi.parameterList.parameters
+
-        val receiver = receiver ?: return@lz lightParams.map {
+        fun parameterOrigin(psiParameter: PsiParameter?): KtElement? = when (psiParameter) {
-            KotlinUParameter(
+            is KtLightElement<*, *> -> psiParameter.kotlinOrigin
-                it, when (it) {
+            is UastKotlinPsiParameter -> psiParameter.ktParameter
-                    is KtLightElement<*, *> -> it.kotlinOrigin
+            else -> null
-                    is UastKotlinPsiParameter -> it.ktParameter
-                    else -> null
-                }, this
-            )
         }
+
+        val lightParams = psi.parameterList.parameters
+        val receiver = receiver ?: return@lz lightParams.map { KotlinUParameter(it, parameterOrigin(it), this) }
         val receiverLight = lightParams.firstOrNull() ?: return@lz emptyList<UParameter>()
         val uParameters = SmartList<UParameter>(KotlinReceiverUParameter(receiverLight, receiver, this))
-        lightParams.drop(1).mapTo(uParameters) { KotlinUParameter(it, (it as? KtLightElement<*, *>)?.kotlinOrigin, this) }
+        lightParams.drop(1).mapTo(uParameters) { KotlinUParameter(it, parameterOrigin(it), this) }
         uParameters
     }
 

plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/psi/UastFakeLightMethod.kt

@@ -13,6 +13,7 @@ import com.intellij.psi.impl.light.LightTypeParameterBuilder
 import org.jetbrains.kotlin.asJava.elements.KotlinLightTypeParameterListBuilder
 import org.jetbrains.kotlin.descriptors.CallableDescriptor
 import org.jetbrains.kotlin.psi.KtFunction
+import org.jetbrains.kotlin.psi.KtTypeReference
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 import org.jetbrains.uast.UastErrorType
@@ -53,6 +54,19 @@ internal class UastFakeLightMethod(internal val original: KtFunction, containing
 
             init {
                 val parameterList = this
+
+                original.receiverTypeReference?.let { receiver ->
+                    this.addParameter(
+                        UastKotlinPsiParameterBase(
+                            "\$this\$${original.name}",
+                            receiver.getType()
+                                ?.toPsiType(this@UastFakeLightMethod, original, false)
+                                ?: UastErrorType,
+                            parameterList, receiver
+                        )
+                    )
+                }
+
                 for ((i, p) in original.valueParameters.withIndex()) {
                     this.addParameter(
                         UastKotlinPsiParameter(

plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/psi/UastKotlinPsiParameter.kt

@@ -8,6 +8,7 @@ import com.intellij.psi.PsiType
 import com.intellij.psi.impl.light.LightParameter
 import org.jetbrains.kotlin.descriptors.VariableDescriptor
 import org.jetbrains.kotlin.idea.KotlinLanguage
+import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtExpression
 import org.jetbrains.kotlin.psi.KtParameter
 import org.jetbrains.kotlin.resolve.BindingContext
@@ -21,12 +22,12 @@ import org.jetbrains.uast.kotlin.toPsiType
 class UastKotlinPsiParameter(
     name: String,
     type: PsiType,
-    private val parent: PsiElement,
+    parent: PsiElement,
     language: Language,
     isVarArgs: Boolean,
-    val ktDefaultValue: KtExpression?,
+    ktDefaultValue: KtExpression?,
-    val ktParameter: KtParameter
+    ktParameter: KtParameter
-) : LightParameter(name, type, parent, language, isVarArgs) {
+) : UastKotlinPsiParameterBase<KtParameter>(name, type, parent, ktParameter, language, isVarArgs, ktDefaultValue) {
     companion object {
         fun create(parameter: KtParameter, parent: PsiElement, containingElement: UElement, index: Int): PsiParameter {
             val psiParent = containingElement.getParentOfType<UDeclaration>()?.javaPsi ?: parent
@@ -42,15 +43,29 @@ class UastKotlinPsiParameter(
         }
     }
 
+    val ktParameter: KtParameter get() = ktOrigin
+
+}
+
+open class UastKotlinPsiParameterBase<T : KtElement>(
+    name: String,
+    type: PsiType,
+    private val parent: PsiElement,
+    val ktOrigin: T,
+    language: Language = ktOrigin.language,
+    isVarArgs: Boolean = false,
+    val ktDefaultValue: KtExpression? = null,
+) : LightParameter(name, type, parent, language, isVarArgs) {
+
     override fun getParent(): PsiElement = parent
 
-    override fun getContainingFile(): PsiFile? = ktParameter.containingFile
+    override fun getContainingFile(): PsiFile? = ktOrigin.containingFile
 
     override fun equals(other: Any?): Boolean {
         if (this === other) return true
         if (other == null || other::class.java != this::class.java) return false
-        return ktParameter == (other as? UastKotlinPsiParameter)?.ktParameter
+        return ktOrigin == (other as? UastKotlinPsiParameterBase<*>)?.ktOrigin
     }
 
-    override fun hashCode() = ktParameter.hashCode()
+    override fun hashCode(): Int = ktOrigin.hashCode()
 }

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

@@ -12,6 +12,8 @@ UFile (package = test.pkg)
                             UClassLiteralExpression
                             USimpleNameReferenceExpression (identifier = java)
         UMethod (name = systemService1)
+            UParameter (name = $this$systemService1)
+                UAnnotation (fqName = org.jetbrains.annotations.NotNull)
             UBlockExpression
                 UReturnExpression
                     UCallExpression (kind = UastCallKind(name='method_call'), argCount = 1))

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

@@ -4,7 +4,7 @@ public final class ReifiedKt {
     public static final fun systemService2(@org.jetbrains.annotations.NotNull $this$systemService2: test.pkg.Context) : error.NonExistentClass {
         return <anonymous class>(java.lang.String.java)
     }
-    fun systemService1() : <ErrorType> {
+    fun systemService1(@org.jetbrains.annotations.NotNull $this$systemService1: test.pkg.Context) : <ErrorType> {
         return <anonymous class>(T.java)
     }
 }

plugins/uast-kotlin/testData/ReifiedParameters.kt

@@ -2,6 +2,10 @@ inline fun <reified T> functionWithLambda(t: T, process: (T) -> Int): Int = proc
 
 inline fun <reified T> functionWithVararg(i: Int?, vararg t: T): T = t[0]
 
+inline fun <reified T> T.functionWithReceiver(i: Int?): T = this
+
+inline fun <reified T> T.`name with spaces`(i: Int?): T = this
+
 inline fun <reified T> functionWithParamAnnotation(@Suppress("s") t: T): T = t
 
 inline fun <reified T> functionUnresolved(@Suppress("s") t: Unresolved<T>): T = t

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

@@ -21,6 +21,22 @@ UFile (package = )
                     UArrayAccessExpression
                         USimpleNameReferenceExpression (identifier = t)
                         ULiteralExpression (value = 0)
+        UMethod (name = functionWithReceiver)
+            UParameter (name = $this$functionWithReceiver)
+                UAnnotation (fqName = org.jetbrains.annotations.Nullable)
+            UParameter (name = i)
+                UAnnotation (fqName = org.jetbrains.annotations.Nullable)
+            UBlockExpression
+                UReturnExpression
+                    UThisExpression (label = null)
+        UMethod (name = name with spaces)
+            UParameter (name = $this$name with spaces)
+                UAnnotation (fqName = org.jetbrains.annotations.Nullable)
+            UParameter (name = i)
+                UAnnotation (fqName = org.jetbrains.annotations.Nullable)
+            UBlockExpression
+                UReturnExpression
+                    UThisExpression (label = null)
         UMethod (name = functionWithParamAnnotation)
             UParameter (name = t)
                 UAnnotation (fqName = org.jetbrains.annotations.Nullable)

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

@@ -5,6 +5,12 @@ public final class ReifiedParametersKt {
     fun functionWithVararg(@org.jetbrains.annotations.Nullable i: int, @org.jetbrains.annotations.Nullable t: T) : T {
         return t[0]
     }
+    fun functionWithReceiver(@org.jetbrains.annotations.Nullable $this$functionWithReceiver: T, @org.jetbrains.annotations.Nullable i: int) : T {
+        return this
+    }
+    fun name with spaces(@org.jetbrains.annotations.Nullable $this$name with spaces: T, @org.jetbrains.annotations.Nullable i: int) : T {
+        return this
+    }
     fun functionWithParamAnnotation(@org.jetbrains.annotations.Nullable @kotlin.Suppress(names = "s") t: T) : T {
         return t
     }

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

@@ -47,10 +47,14 @@ UFile (package = )
                 UReturnExpression
                     ULiteralExpression (value = 42)
         UMethod (name = function6)
+            UParameter (name = $this$function6)
+                UAnnotation (fqName = org.jetbrains.annotations.Nullable)
             UParameter (name = t)
                 UAnnotation (fqName = org.jetbrains.annotations.Nullable)
             UParameter (name = i)
+                UAnnotation (fqName = org.jetbrains.annotations.NotNull)
             UParameter (name = s)
+                UAnnotation (fqName = org.jetbrains.annotations.NotNull)
             UBlockExpression
                 UReturnExpression
                     USimpleNameReferenceExpression (identifier = t)
@@ -95,10 +99,14 @@ UFile (package = )
                 UReturnExpression
                     USimpleNameReferenceExpression (identifier = t)
         UMethod (name = function11)
+            UParameter (name = $this$function11)
+                UAnnotation (fqName = org.jetbrains.annotations.Nullable)
             UParameter (name = t)
                 UAnnotation (fqName = org.jetbrains.annotations.Nullable)
             UParameter (name = i)
+                UAnnotation (fqName = org.jetbrains.annotations.NotNull)
             UParameter (name = s)
+                UAnnotation (fqName = org.jetbrains.annotations.NotNull)
             UBlockExpression
                 UReturnExpression
                     USimpleNameReferenceExpression (identifier = t)

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

@@ -12,7 +12,7 @@ public final class ReifiedReturnTypeKt {
     fun function5(@org.jetbrains.annotations.Nullable t: T, @org.jetbrains.annotations.NotNull i: int, @org.jetbrains.annotations.NotNull s: java.lang.String) : int {
         return 42
     }
-    fun function6(@org.jetbrains.annotations.Nullable t: T, i: int, s: java.lang.String) : T {
+    fun function6(@org.jetbrains.annotations.Nullable $this$function6: T, @org.jetbrains.annotations.Nullable t: T, @org.jetbrains.annotations.NotNull i: int, @org.jetbrains.annotations.NotNull s: java.lang.String) : T {
         return t
     }
     fun function7(@org.jetbrains.annotations.Nullable t: T, @org.jetbrains.annotations.NotNull i: int, @org.jetbrains.annotations.NotNull s: java.lang.String) : T {
@@ -27,7 +27,7 @@ public final class ReifiedReturnTypeKt {
     fun function10(@org.jetbrains.annotations.Nullable t: T, @org.jetbrains.annotations.NotNull i: int, @org.jetbrains.annotations.NotNull s: java.lang.String) : T {
         return t
     }
-    fun function11(@org.jetbrains.annotations.Nullable t: T, i: int, s: java.lang.String) : T {
+    fun function11(@org.jetbrains.annotations.Nullable $this$function11: T, @org.jetbrains.annotations.Nullable t: T, @org.jetbrains.annotations.NotNull i: int, @org.jetbrains.annotations.NotNull s: java.lang.String) : T {
         return t
     }
 }

plugins/uast-kotlin/tests/KotlinUastApiTest.kt

@@ -659,7 +659,7 @@ class KotlinUastApiTest : AbstractKotlinUastTest() {
                     assertEquals(
                         "parameter ${parameter.name} of method ${method.name} should be equals to converted from sourcePsi",
                         parameter,
-                        parameter.sourcePsi.toUElement()
+                        parameter.sourcePsi.toUElementOfType<UParameter>()
                     )
                     assertEquals(
                         "parameter ${parameter.name} of method ${method.name} should be equals to converted from javaPsi",