Add Resolver.getPackageAnnotations Resolver.getPackagesWithAnnotation…

… api.

These new APIs provides capability to inspect java package level annotations.

(cherry picked from commit 4f4fbe8)

api/api.base

@@ -164,6 +164,8 @@ package com.google.devtools.ksp.processing {
     method @NonNull public kotlin.sequences.Sequence<com.google.devtools.ksp.symbol.KSFile> getNewFiles();
     method @Nullable @com.google.devtools.ksp.KspExperimental public String getOwnerJvmClassName(@NonNull com.google.devtools.ksp.symbol.KSPropertyDeclaration declaration);
     method @Nullable @com.google.devtools.ksp.KspExperimental public String getOwnerJvmClassName(@NonNull com.google.devtools.ksp.symbol.KSFunctionDeclaration declaration);
+    method @NonNull @com.google.devtools.ksp.KspExperimental public kotlin.sequences.Sequence<com.google.devtools.ksp.symbol.KSAnnotation> getPackageAnnotations(@NonNull String packageName);
+    method @NonNull @com.google.devtools.ksp.KspExperimental public kotlin.sequences.Sequence<java.lang.String> getPackagesWithAnnotation(@NonNull String annotationName);
     method @Nullable public com.google.devtools.ksp.symbol.KSPropertyDeclaration getPropertyDeclarationByName(@NonNull com.google.devtools.ksp.symbol.KSName name, boolean includeTopLevel = false);
     method @NonNull public kotlin.sequences.Sequence<com.google.devtools.ksp.symbol.KSAnnotated> getSymbolsWithAnnotation(@NonNull String annotationName, boolean inDepth = false);
     method @NonNull public com.google.devtools.ksp.symbol.KSTypeArgument getTypeArgument(@NonNull com.google.devtools.ksp.symbol.KSTypeReference typeRef, @NonNull com.google.devtools.ksp.symbol.Variance variance);

api/src/main/kotlin/com/google/devtools/ksp/processing/Resolver.kt

@@ -300,4 +300,22 @@ interface Resolver {
      */
     @KspExperimental
     fun isJavaRawType(type: KSType): Boolean
+
+    /**
+     * Returns annotations applied in package-info.java (if applicable) for given package name.
+     *
+     * @param packageName package name to check.
+     * @return a sequence of KSAnnotations applied in corresponding package-info.java file.
+     */
+    @KspExperimental
+    fun getPackageAnnotations(packageName: String): Sequence<KSAnnotation>
+
+    @KspExperimental
+    /**
+     * Returns name of packages with given annotation.
+     *
+     * @param annotationName name of the annotation to be queried.
+     * @return a sequence of package names with corresponding annotation name.
+     */
+    fun getPackagesWithAnnotation(annotationName: String): Sequence<String>
 }

compiler-plugin/src/main/kotlin/com/google/devtools/ksp/processing/impl/ResolverImpl.kt

@@ -24,10 +24,9 @@ import com.google.devtools.ksp.processing.SymbolProcessor
 import com.google.devtools.ksp.symbol.*
 import com.google.devtools.ksp.symbol.ClassKind
 import com.google.devtools.ksp.symbol.Variance
+import com.google.devtools.ksp.symbol.impl.*
 import com.google.devtools.ksp.symbol.impl.binary.*
 import com.google.devtools.ksp.symbol.impl.declarationsInSourceOrder
-import com.google.devtools.ksp.symbol.impl.findParentAnnotated
-import com.google.devtools.ksp.symbol.impl.findPsi
 import com.google.devtools.ksp.symbol.impl.getInstanceForCurrentRound
 import com.google.devtools.ksp.symbol.impl.java.*
 import com.google.devtools.ksp.symbol.impl.jvmAccessFlag
@@ -132,6 +131,9 @@ class ResolverImpl(
     val psiDocumentManager = PsiDocumentManager.getInstance(project)
     private val nameToKSMap: MutableMap<KSName, KSClassDeclaration>
     private val javaTypeParameterMap: MutableMap<LazyJavaTypeParameterDescriptor, PsiTypeParameter> = mutableMapOf()
+    private val packageInfoFiles by lazy {
+        allKSFiles.filter { it.fileName == "package-info.java" }.asSequence().memoized()
+    }
 
     /**
      * Checking as member of is an expensive operation, hence we cache result values in this map.
@@ -1448,6 +1450,23 @@ class ResolverImpl(
         return type is KSTypeImpl && type.kotlinType.unwrap() is RawType
     }
 
+    @KspExperimental
+    override fun getPackageAnnotations(packageName: String): Sequence<KSAnnotation> {
+        return packageInfoFiles.singleOrNull { it.packageName.asString() == packageName }
+            ?.getPackageAnnotations()?.asSequence() ?: emptySequence()
+    }
+
+    @KspExperimental
+    override fun getPackagesWithAnnotation(annotationName: String): Sequence<String> {
+        return packageInfoFiles.filter {
+            it.getPackageAnnotations().any {
+                (it.annotationType.element as? KSClassifierReference)?.referencedName()
+                    ?.substringAfterLast(".") == annotationName.substringAfterLast(".") &&
+                    it.annotationType.resolve().declaration.qualifiedName?.asString() == annotationName
+            }
+        }.map { it.packageName.asString() }
+    }
+
     private val psiJavaFiles = allKSFiles.filterIsInstance<KSFileJavaImpl>().map {
         Pair(it.psi.virtualFile.path, it.psi)
     }.toMap()

compiler-plugin/src/main/kotlin/com/google/devtools/ksp/symbol/impl/utils.kt

@@ -557,3 +557,6 @@ fun DeclarationDescriptor.findPsi(): PsiElement? {
     val leaf = containingFile.findElementAt(psi.textOffset) ?: return null
     return leaf.parentsWithSelf.firstOrNull { psi.manager.areElementsEquivalent(it, psi) }
 }
+
+internal fun KSFile.getPackageAnnotations() = (this as? KSFileJavaImpl)?.psi?.packageStatement
+    ?.annotationList?.annotations?.map { KSAnnotationJavaImpl.getCached(it) } ?: emptyList<KSAnnotation>()

compiler-plugin/src/test/kotlin/com/google/devtools/ksp/test/KSPCompilerPluginTest.kt

@@ -399,6 +399,12 @@ class KSPCompilerPluginTest : AbstractKSPCompilerPluginTest() {
         runTest("../test-utils/testData/api/overridee.kt")
     }
 
+    @TestMetadata("packageAnnotations.kt")
+    @Test
+    fun testPackageAnnotation() {
+        runTest("../test-utils/testData/api/packageAnnotations.kt")
+    }
+
     @TestMetadata("parameterTypes.kt")
     @Test
     fun testParameterTypes() {

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

@@ -0,0 +1,34 @@
+package com.google.devtools.ksp.processor
+
+import com.google.devtools.ksp.KspExperimental
+import com.google.devtools.ksp.processing.Resolver
+import com.google.devtools.ksp.symbol.KSAnnotated
+import com.google.devtools.ksp.symbol.KSFile
+import com.google.devtools.ksp.symbol.KSVisitorVoid
+
+class PackageAnnotationProcessor : AbstractTestProcessor() {
+    val result = mutableListOf<String>()
+    override fun toResult(): List<String> {
+        return result
+    }
+
+    @OptIn(KspExperimental::class)
+    override fun process(resolver: Resolver): List<KSAnnotated> {
+        resolver.getAllFiles().forEach {
+            it.accept(
+                object : KSVisitorVoid() {
+                    override fun visitFile(file: KSFile, data: Unit) {
+                        result.add(
+                            "${file.fileName}:${resolver.getPackageAnnotations(file.packageName.asString())
+                                .joinToString { it.toString() }}"
+                        )
+                    }
+                },
+                Unit
+            )
+        }
+        result.sort()
+        result.add(resolver.getPackagesWithAnnotation("PackageAnnotation").joinToString())
+        return emptyList()
+    }
+}

test-utils/testData/api/packageAnnotations.kt

@@ -0,0 +1,40 @@
+/*
+ * Copyright 2023 Google LLC
+ * Copyright 2010-2023 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.
+ */
+
+// WITH_RUNTIME
+// TEST PROCESSOR: PackageAnnotationProcessor
+// EXPECTED:
+// Foo.java:@PackageAnnotation
+// annotations.kt:
+// package-info.java:@PackageAnnotation
+// example
+// END
+
+// FILE: annotations.kt
+@Retention(RetentionPolicy.CLASS)
+@Target(ElementType.PACKAGE)
+annotation class PackageAnnotation
+
+// FILE: example/package-info.java
+@PackageAnnotation
+package example;
+
+import PackageAnnotation;
+
+// FILE: example/Foo.java
+package example;
+public class Foo