support multiple depth search for annotated symbols

Author: neetopia

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

@@ -34,10 +34,11 @@ interface Resolver {
     /**
      * Get all symbols with specified annotation.
      *
-     * @param annotationName is the full qualified name of the annotation; using '.' as separator.
+     * @param annotationName is the fully qualified name of the annotation; using '.' as separator.
+     * @param inDepth whether to check symbols in depth, i.e. check symbols from local declarations. Operation can be expensive if true.
      * @return Elements annotated with the specified annotation.
      */
-    fun getSymbolsWithAnnotation(annotationName: String): List<KSAnnotated>
+    fun getSymbolsWithAnnotation(annotationName: String, inDepth: Boolean = false): List<KSAnnotated>
 
     /**
      * Find a class in the compilation classpath for the given name.

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

@@ -176,10 +176,10 @@ class ResolverImpl(
             }
     }
 
-    override fun getSymbolsWithAnnotation(annotationName: String): List<KSAnnotated> {
+    override fun getSymbolsWithAnnotation(annotationName: String, inDepth: Boolean): List<KSAnnotated> {
         val ksName = KSNameImpl.getCached(annotationName)
 
-        val visitor = object : BaseVisitor() {
+        val visitor = object : KSVisitorVoid() {
             val symbols = mutableSetOf<KSAnnotated>()
             override fun visitAnnotated(annotated: KSAnnotated, data: Unit) {
                 if (annotated.annotations.any {
@@ -192,31 +192,46 @@ class ResolverImpl(
                 }
             }
 
-            override fun visitClassDeclaration(type: KSClassDeclaration, data: Unit) {
-                visitAnnotated(type, data)
-                super.visitClassDeclaration(type, data)
+            override fun visitFile(file: KSFile, data: Unit) {
+                visitAnnotated(file, data)
+                file.declarations.map { it.accept(this, data) }
+            }
+
+            override fun visitClassDeclaration(classDeclaration: KSClassDeclaration, data: Unit) {
+                visitAnnotated(classDeclaration, data)
+                classDeclaration.typeParameters.map { it.accept(this, data) }
+                classDeclaration.declarations.map { it.accept(this, data) }
+                classDeclaration.primaryConstructor?.let { it.accept(this, data) }
+            }
+
+            override fun visitPropertyGetter(getter: KSPropertyGetter, data: Unit) {
+                visitAnnotated(getter, data)
+            }
+
+            override fun visitPropertySetter(setter: KSPropertySetter, data: Unit) {
+                visitAnnotated(setter, data)
             }
 
             override fun visitFunctionDeclaration(function: KSFunctionDeclaration, data: Unit) {
                 visitAnnotated(function, data)
-                super.visitFunctionDeclaration(function, data)
+                function.typeParameters.map { it.accept(this, data) }
+                function.parameters.map { it.accept(this, data) }
+                if (inDepth) {
+                    function.declarations.map { it.accept(this, data) }
+                }
             }
 
             override fun visitPropertyDeclaration(property: KSPropertyDeclaration, data: Unit) {
                 visitAnnotated(property, data)
-                super.visitPropertyDeclaration(property, data)
+                property.typeParameters.map { it.accept(this, data) }
+                property.getter?.let { it.accept(this, data) }
+                property.setter?.let { it.accept(this, data) }
             }
 
             override fun visitTypeParameter(typeParameter: KSTypeParameter, data: Unit) {
                 visitAnnotated(typeParameter, data)
                 super.visitTypeParameter(typeParameter, data)
             }
-
-            override fun visitTypeReference(typeReference: KSTypeReference, data: Unit) {
-                visitAnnotated(typeReference, data)
-                super.visitTypeReference(typeReference, data)
-            }
-
         }
 
         for (file in ksFiles) {

compiler-plugin/src/main/kotlin/com/google/devtools/ksp/processor/AnnotationArgumentProcessor.kt

@@ -29,7 +29,7 @@ class AnnotationArgumentProcessor : AbstractTestProcessor() {
     val visitor = ArgumentVisitor()
 
     override fun process(resolver: Resolver) {
-        val symbol = resolver.getSymbolsWithAnnotation("Bar").single()
+        val symbol = resolver.getSymbolsWithAnnotation("Bar", true).single()
         val annotation = (symbol as KSClassDeclaration).annotations.single()
         annotation.arguments.map { it.accept(visitor, Unit) }
         val C = resolver.getClassDeclarationByName("C")!!

compiler-plugin/src/main/kotlin/com/google/devtools/ksp/processor/GetSymbolsFromAnnotationProcessor.kt

@@ -0,0 +1,19 @@
+package com.google.devtools.ksp.processor
+
+import com.google.devtools.ksp.processing.Resolver
+
+class GetSymbolsFromAnnotationProcessor : AbstractTestProcessor() {
+    val result = mutableListOf<String>()
+    override fun toResult(): List<String> = result
+
+    override fun process(resolver: Resolver) {
+        result.add("==== Anno superficial====")
+        resolver.getSymbolsWithAnnotation("Anno").map { result.add(it.toString()) }
+        result.add("==== Anno in depth ====")
+        resolver.getSymbolsWithAnnotation("Anno", true).map { result.add(it.toString()) }
+        result.add("==== Bnno superficial====")
+        resolver.getSymbolsWithAnnotation("Bnno").map { result.add(it.toString()) }
+        result.add("==== Bnno in depth ====")
+        resolver.getSymbolsWithAnnotation("Bnno", true).map { result.add(it.toString()) }
+    }
+}

compiler-plugin/src/test/java/com/google/devtools/ksp/test/KotlinKSPTestGenerated.java

@@ -97,6 +97,11 @@ public void testFunctionTypeAlias() throws Exception {
         runTest("testData/api/functionTypeAlias.kt");
     }
 
+    @TestMetadata("getSymbolsFromAnnotation.kt")
+    public void testGetSymbolsFromAnnotation() throws Exception {
+        runTest("testData/api/getSymbolsFromAnnotation.kt");
+    }
+
     @TestMetadata("hello.kt")
     public void testHello() throws Exception {
         runTest("testData/api/hello.kt");

compiler-plugin/testData/api/getSymbolsFromAnnotation.kt

@@ -0,0 +1,46 @@
+// TEST PROCESSOR: GetSymbolsFromAnnotationProcessor
+// EXPECTED:
+// ==== Anno superficial====
+// Foo
+// propertyFoo
+// functionFoo
+// constructorParameterFoo
+// Foo
+// ==== Anno in depth ====
+// Foo
+// propertyFoo
+// functionFoo
+// local
+// constructorParameterFoo
+// Foo
+// ==== Bnno superficial====
+// File: Foo.kt
+// Foo
+// propertyFoo.getter()
+// ==== Bnno in depth ====
+// File: Foo.kt
+// Foo
+// propertyFoo.getter()
+// END
+//FILE: annotations.kt
+annotation class Anno
+annotation class Bnno
+
+//FILE: Foo.kt
+@file:Bnno
+
+@Anno
+class Foo @Anno constructor(@Anno val constructorParameterFoo: Int){
+    @Bnno constructor() {
+
+    }
+
+    @Anno
+    val propertyFoo: String
+    @Bnno get() = TODO()
+
+    @Anno
+    fun functionFoo(@Anno p1: Int, @Bnno p2: Int) {
+        @Anno val local = 1
+    }
+}