handle receiver annotation use site target

Author: neetopia

kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSFunctionDeclarationImpl.kt

@@ -77,9 +77,22 @@ class KSFunctionDeclarationImpl private constructor(internal val ktFunctionSymbo
                 null
             } else {
                 (ktFunctionSymbol.psiIfSource() as? KtFunction)?.receiverTypeReference
-                    ?.let { KSTypeReferenceImpl.getCached(it, this@KSFunctionDeclarationImpl) }
+                    ?.let {
+                        // receivers are modeled as parameter in AA therefore annotations are stored in
+                        // the corresponding receiver parameter, need to pass it to the `KSTypeReferenceImpl`
+                        KSTypeReferenceImpl.getCached(
+                            it,
+                            this@KSFunctionDeclarationImpl,
+                            ktFunctionSymbol.receiverParameter?.annotations ?: emptyList()
+                        )
+                    }
                     ?: ktFunctionSymbol.receiverType?.let {
-                        KSTypeReferenceResolvedImpl.getCached(it, this@KSFunctionDeclarationImpl)
+                        KSTypeReferenceResolvedImpl.getCached(
+                            it,
+                            this@KSFunctionDeclarationImpl,
+                            -1,
+                            ktFunctionSymbol.receiverParameter?.annotations ?: emptyList()
+                        )
                     }
             }
         }

kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSPropertyDeclarationImpl.kt

@@ -102,9 +102,23 @@ class KSPropertyDeclarationImpl private constructor(internal val ktPropertySymbo
 
     override val extensionReceiver: KSTypeReference? by lazy {
         (ktPropertySymbol.psiIfSource() as? KtProperty)?.receiverTypeReference
-            ?.let { KSTypeReferenceImpl.getCached(it, this) }
-            ?: ktPropertySymbol.receiverType
-                ?.let { KSTypeReferenceResolvedImpl.getCached(it, this@KSPropertyDeclarationImpl) }
+            ?.let {
+                // receivers are modeled as parameter in AA therefore annotations are stored in
+                // the corresponding receiver parameter, need to pass it to the `KSTypeReferenceImpl`
+                KSTypeReferenceImpl.getCached(
+                    it,
+                    this,
+                    ktPropertySymbol.receiverParameter?.annotations ?: emptyList()
+                )
+            }
+            ?: ktPropertySymbol.receiverType?.let {
+                KSTypeReferenceResolvedImpl.getCached(
+                    it,
+                    this@KSPropertyDeclarationImpl,
+                    -1,
+                    ktPropertySymbol.receiverParameter?.annotations ?: emptyList()
+                )
+            }
     }
 
     override val type: KSTypeReference by lazy {

kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSTypeReferenceImpl.kt

@@ -29,6 +29,7 @@ import com.google.devtools.ksp.symbol.KSVisitor
 import com.google.devtools.ksp.symbol.Location
 import com.google.devtools.ksp.symbol.Modifier
 import com.google.devtools.ksp.symbol.Origin
+import org.jetbrains.kotlin.analysis.api.annotations.KaAnnotation
 import org.jetbrains.kotlin.analysis.api.types.KtType
 import org.jetbrains.kotlin.psi.KtAnnotationEntry
 import org.jetbrains.kotlin.psi.KtDynamicType
@@ -38,12 +39,17 @@ import org.jetbrains.kotlin.psi.KtUserType
 
 class KSTypeReferenceImpl(
     private val ktTypeReference: KtTypeReference,
-    override val parent: KSNode?
+    override val parent: KSNode?,
+    private val additionalAnnotations: List<KaAnnotation>
 ) : KSTypeReference {
     companion object : KSObjectCache<IdKeyPair<KtTypeReference, KSNode?>, KSTypeReference>() {
-        fun getCached(ktTypeReference: KtTypeReference, parent: KSNode? = null): KSTypeReference {
+        fun getCached(
+            ktTypeReference: KtTypeReference,
+            parent: KSNode? = null,
+            additionalAnnotations: List<KaAnnotation> = emptyList()
+        ): KSTypeReference {
             return cache.getOrPut(IdKeyPair(ktTypeReference, parent)) {
-                KSTypeReferenceImpl(ktTypeReference, parent)
+                KSTypeReferenceImpl(ktTypeReference, parent, additionalAnnotations)
             }
         }
     }
@@ -79,7 +85,7 @@ class KSTypeReferenceImpl(
         (ktTypeReference.annotationEntries.asSequence() + innerAnnotations.asSequence().flatten())
             .map { annotationEntry ->
                 KSAnnotationImpl.getCached(annotationEntry, this@KSTypeReferenceImpl) {
-                    ktType.annotations.single {
+                    (ktType.annotations + additionalAnnotations).single {
                         it.psi == annotationEntry
                     }
                 }

kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/resolved/KSTypeReferenceResolvedImpl.kt

@@ -34,18 +34,27 @@ import com.google.devtools.ksp.symbol.*
 import com.intellij.psi.PsiClass
 import com.intellij.psi.PsiTypeParameter
 import com.intellij.psi.impl.source.PsiClassReferenceType
+import org.jetbrains.kotlin.analysis.api.annotations.KaAnnotation
 import org.jetbrains.kotlin.analysis.api.types.*
 import org.jetbrains.kotlin.psi.KtClassOrObject
 import org.jetbrains.kotlin.psi.KtTypeParameter
 
 class KSTypeReferenceResolvedImpl private constructor(
     private val ktType: KtType,
     override val parent: KSNode?,
-    private val index: Int
+    private val index: Int,
+    private val additionalAnnotations: List<KaAnnotation>
 ) : KSTypeReference, Deferrable {
     companion object : KSObjectCache<IdKeyTriple<KtType, KSNode?, Int>, KSTypeReference>() {
-        fun getCached(type: KtType, parent: KSNode? = null, index: Int = -1): KSTypeReference =
-            cache.getOrPut(IdKeyTriple(type, parent, index)) { KSTypeReferenceResolvedImpl(type, parent, index) }
+        fun getCached(
+            type: KtType,
+            parent: KSNode? = null,
+            index: Int = -1,
+            additionalAnnotations: List<KaAnnotation> = emptyList()
+        ): KSTypeReference =
+            cache.getOrPut(IdKeyTriple(type, parent, index)) {
+                KSTypeReferenceResolvedImpl(type, parent, index, additionalAnnotations)
+            }
     }
 
     override val element: KSReferenceElement? by lazy {
@@ -62,7 +71,8 @@ class KSTypeReferenceResolvedImpl private constructor(
     }
 
     override val annotations: Sequence<KSAnnotation> by lazy {
-        ktType.annotations(this)
+        ktType.annotations(this) +
+            additionalAnnotations.asSequence().map { KSAnnotationResolvedImpl.getCached(it, this) }
     }
 
     override val origin: Origin = parent?.origin ?: Origin.SYNTHETIC

kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/util.kt

@@ -354,7 +354,7 @@ internal fun KtSymbol.toKSNode(): KSNode {
 internal fun ClassId.toKtClassSymbol(): KtClassOrObjectSymbol? {
     return analyze {
         if (this@toKtClassSymbol.isLocal) {
-            this@toKtClassSymbol.outerClassId?.toKtClassSymbol()?.getDeclaredMemberScope()?.getClassifierSymbols {
+            this@toKtClassSymbol.outerClassId?.toKtClassSymbol()?.getDeclaredMemberScope()?.classifiers {
                 it.asString() == this@toKtClassSymbol.shortClassName.asString()
             }?.singleOrNull() as? KtClassOrObjectSymbol
         } else {

kotlin-analysis-api/src/test/kotlin/com/google/devtools/ksp/test/KSPAATest.kt

@@ -76,6 +76,12 @@ class KSPAATest : AbstractKSPAATest() {
         runTest("../test-utils/testData/api/annotationOnConstructorParameter.kt")
     }
 
+    @TestMetadata("annotationOnReceiver.kt")
+    @Test
+    fun testAnnotationOnReceiver() {
+        runTest("../test-utils/testData/api/annotationOnReceiver.kt")
+    }
+
     @TestMetadata("annotationWithArbitraryClassValue.kt")
     @Test
     fun testAnnotationWithArbitraryClassValue() {

test-utils/src/main/kotlin/com/google/devtools/ksp/processor/AnnotationOnReceiverProcessor.kt

@@ -0,0 +1,45 @@
+package com.google.devtools.ksp.processor
+
+import com.google.devtools.ksp.getClassDeclarationByName
+import com.google.devtools.ksp.getDeclaredFunctions
+import com.google.devtools.ksp.getDeclaredProperties
+import com.google.devtools.ksp.processing.Resolver
+import com.google.devtools.ksp.symbol.KSAnnotated
+
+class AnnotationOnReceiverProcessor : AbstractTestProcessor() {
+    val results = mutableListOf<String>()
+
+    override fun toResult(): List<String> {
+        return results
+    }
+
+    override fun process(resolver: Resolver): List<KSAnnotated> {
+        listOf("Test", "TestLib").forEach {
+            resolver.getClassDeclarationByName(it)!!.let { cls ->
+                cls.getDeclaredFunctions().forEach { method ->
+                    method.extensionReceiver.let { receiver ->
+                        if (receiver != null) {
+                            results.add(
+                                receiver.annotations.map {
+                                    it.annotationType.toString() + it.arguments.toString()
+                                }.joinToString()
+                            )
+                        }
+                    }
+                }
+                cls.getDeclaredProperties().forEach { prop ->
+                    prop.extensionReceiver.let { receiver ->
+                        if (receiver != null) {
+                            results.add(
+                                receiver.annotations.map {
+                                    it.annotationType.toString() + it.arguments.toString()
+                                }.joinToString()
+                            )
+                        }
+                    }
+                }
+            }
+        }
+        return emptyList()
+    }
+}

test-utils/testData/api/annotationOnReceiver.kt

@@ -0,0 +1,40 @@
+/*
+ * Copyright 2024 Google LLC
+ * Copyright 2010-2024 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+// TEST PROCESSOR: AnnotationOnReceiverProcessor
+// EXPECTED:
+// MyAnnotation[]
+// MyAnnotation[]
+// MyAnnotation[]
+// MyAnnotation[]
+// END
+
+// MODULE: lib
+// FILE: TestLib.kt
+annotation class MyAnnotation
+class TestLib {
+    fun @receiver:MyAnnotation String.extF(): String = TODO()
+    val @receiver:MyAnnotation String.extP: String
+        get() = TODO()
+}
+// MODULE: main(lib)
+// FILE: Test.kt
+class Test {
+    fun @receiver:MyAnnotation String.extF(): String = TODO()
+    val @receiver:MyAnnotation String.extP: String
+        get() = TODO()
+}