Uast: handling @receiver annotations

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

@@ -3,6 +3,7 @@ package org.jetbrains.uast.kotlin
 import com.intellij.psi.PsiClass
 import org.jetbrains.kotlin.asJava.toLightAnnotation
 import org.jetbrains.kotlin.descriptors.annotations.AnnotationDescriptor
+import org.jetbrains.kotlin.descriptors.annotations.AnnotationUseSiteTarget
 import org.jetbrains.kotlin.psi.KtAnnotationEntry
 import org.jetbrains.kotlin.psi.KtParameter
 import org.jetbrains.kotlin.psi.ValueArgument
@@ -12,7 +13,9 @@ import org.jetbrains.kotlin.resolve.calls.model.ArgumentMatch
 import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
 import org.jetbrains.kotlin.resolve.descriptorUtil.annotationClass
 import org.jetbrains.kotlin.resolve.source.getPsi
+import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstance
 import org.jetbrains.uast.*
+import org.jetbrains.uast.kotlin.declarations.KotlinUMethod
 
 class KotlinUAnnotation(
         override val psi: KtAnnotationEntry,
@@ -78,5 +81,15 @@ class KotlinUAnnotation(
         val defaultValue = (parameter.source.getPsi() as? KtParameter)?.defaultValue ?: return null
         return getLanguagePlugin().convertWithParent(defaultValue)
     }
+
+    override fun convertParent(): UElement? {
+        val superParent = super.convertParent() ?: return null
+        if (psi.useSiteTarget?.getAnnotationUseSiteTarget() == AnnotationUseSiteTarget.RECEIVER) {
+            (superParent.uastParent as? KotlinUMethod)?.uastParameters?.firstIsInstance<KotlinReceiverUParameter>()?.let {
+                return it
+            }
+        }
+        return superParent
+    }
 }
 

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

@@ -27,6 +27,7 @@ import org.jetbrains.kotlin.asJava.elements.isSetter
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.containingClassOrObject
+import org.jetbrains.kotlin.utils.SmartList
 import org.jetbrains.uast.*
 import org.jetbrains.uast.java.internal.JavaUElementWithComments
 import org.jetbrains.uast.kotlin.*
@@ -59,8 +60,17 @@ open class KotlinUMethod(
                 .map { KotlinUAnnotation(it, this) }
     }
 
+    private val receiver by lz { (sourcePsi as? KtCallableDeclaration)?.receiverTypeReference }
+
     override val uastParameters by lz {
-        psi.parameterList.parameters.map { KotlinUParameter(it, (it as? KtLightElement<*, *>)?.kotlinOrigin, this) }
+        val lightParams = psi.parameterList.parameters
+        val receiver = receiver ?: return@lz lightParams.map {
+            KotlinUParameter(it, (it as? KtLightElement<*, *>)?.kotlinOrigin, 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) }
+        uParameters
     }
 
     override val uastAnchor: UElement

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

@@ -201,9 +201,27 @@ open class KotlinUParameter(
     }
 }
 
+class KotlinReceiverUParameter(
+        psi: PsiParameter,
+        private val receiver: KtTypeReference,
+        givenParent: UElement?
+) : KotlinUParameter(psi, receiver, givenParent) {
+
+    override val annotations: List<UAnnotation> by lz {
+        receiver.annotationEntries
+                .filter { it.useSiteTarget?.getAnnotationUseSiteTarget() == AnnotationUseSiteTarget.RECEIVER }
+                .map { KotlinUAnnotation(it, this) } +
+        super.annotations
+    }
+
+}
+
 class KotlinNullabilityUAnnotation(val annotatedElement: PsiElement, override val uastParent: UElement) : UAnnotation, JvmDeclarationUElement {
 
     private fun getTargetType(annotatedElement: PsiElement): KotlinType? {
+        if (annotatedElement is KtTypeReference) {
+            annotatedElement.getType()?.let { return it }
+        }
         if (annotatedElement is KtCallableDeclaration) {
             annotatedElement.typeReference?.getType()?.let { return it }
         }

plugins/uast-kotlin/testData/ReceiverFun.kt

@@ -0,0 +1,3 @@
+annotation class MyReceiverAnnotation
+
+fun @receiver:MyReceiverAnnotation String.foo() = this.length

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

@@ -0,0 +1,10 @@
+UFile (package = )
+    UClass (name = ReceiverFunKt)
+        UAnnotationMethod (name = foo)
+            UParameter (name = $receiver)
+                UAnnotation (fqName = MyReceiverAnnotation)
+                UAnnotation (fqName = org.jetbrains.annotations.NotNull)
+            UQualifiedReferenceExpression
+                UThisExpression (label = null)
+                USimpleNameReferenceExpression (identifier = length)
+    UClass (name = MyReceiverAnnotation)

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

@@ -0,0 +1,6 @@
+public final class ReceiverFunKt {
+    public static final fun foo(@MyReceiverAnnotation @org.jetbrains.annotations.NotNull $receiver: java.lang.String) : int = this.length
+}
+
+public abstract annotation MyReceiverAnnotation {
+}

plugins/uast-kotlin/tests/SimpleKotlinRenderLogTest.kt

@@ -64,4 +64,7 @@ class SimpleKotlinRenderLogTest : AbstractKotlinRenderLogTest() {
 
     @Test
     fun testClassAnnotation() = doTest("ClassAnnotation")
+
+    @Test
+    fun testReceiverFun() = doTest("ReceiverFun")
 }