Record references to classes in annotation arguments
diff --git a/integration-tests/src/test/kotlin/com/google/devtools/ksp/test/secondary/IncrementalAnnotationArgumentClassReferences.kt b/integration-tests/src/test/kotlin/com/google/devtools/ksp/test/secondary/IncrementalAnnotationArgumentClassReferences.kt
index 8d4f91b..2932590 100644
--- a/integration-tests/src/test/kotlin/com/google/devtools/ksp/test/secondary/IncrementalAnnotationArgumentClassReferences.kt
+++ b/integration-tests/src/test/kotlin/com/google/devtools/ksp/test/secondary/IncrementalAnnotationArgumentClassReferences.kt
@@ -84,7 +84,7 @@
         gradleRunner.withArguments(ASSEMBLE).build().let { result ->
             Assert.assertEquals(TaskOutcome.SUCCESS, result.task(KSP_KOTLIN)?.outcome)
             val actual = result.output.lines().filter { it.startsWith(PROCESSOR_LABEL) }
-            Assert.assertNotEquals(
+            Assert.assertEquals(
                 "\n${expected.joinToString("\n")}\n[SEPARATOR]\n${actual.joinToString("\n")}",
                 expected,
                 actual
diff --git a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/common/IncrementalContextBase.kt b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/common/IncrementalContextBase.kt
index e95f23f..182ffa5 100644
--- a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/common/IncrementalContextBase.kt
+++ b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/common/IncrementalContextBase.kt
@@ -18,6 +18,7 @@
 package com.google.devtools.ksp.common
 
 import com.google.devtools.ksp.IncrementalContextLoggingOptions
+import com.google.devtools.ksp.impl.symbol.kotlin.separateQualifierAndName
 import com.google.devtools.ksp.symbol.KSClassDeclaration
 import com.google.devtools.ksp.symbol.KSDeclaration
 import com.google.devtools.ksp.symbol.KSDeclarationContainer
@@ -253,8 +254,7 @@
 
         // Calculate dirty files by dirty classes in CP.
         val dirtyFilesByCP = changedClasses.flatMap { fqn ->
-            val name = fqn.substringAfterLast('.')
-            val scope = fqn.substringBeforeLast('.', "<anonymous>")
+            val (scope, name) = separateQualifierAndName(fqn)
             classLookupCache[LookupSymbolWrapper(name, scope)].map { it.toRelativeFile() } +
                 symbolLookupCache[LookupSymbolWrapper(name, scope)].map { it.toRelativeFile() }
         }.toSet()
@@ -521,8 +521,7 @@
             return
 
         val path = psiFile.virtualFile.path
-        val name = fqn.substringAfterLast('.')
-        val scope = fqn.substringBeforeLast('.', "<anonymous>")
+        val (scope, name) = separateQualifierAndName(fqn)
 
         // Java types are classes. Therefore lookups only happen in packages.
         fun record(scope: String, name: String) =
@@ -615,8 +614,14 @@
         visitedFiles.add(file)
 
         // Propagate by dependencies
-        symbolsMap[file]?.forEach {
-            visit(it)
+        if (NoSourceFile.isSyntheticFile(file.name)) {
+            val fqn = stripSyntheticFileNameModifiers(file.name)
+            val (scope, name) = separateQualifierAndName(fqn)
+            visit(LookupSymbolWrapper(name, scope))
+        } else {
+            symbolsMap[file]?.forEach {
+                visit(it)
+            }
         }
 
         // Propagate by input-output relations
@@ -635,4 +640,17 @@
         initialSet.forEach { visit(it) }
         return visitedFiles
     }
+
+    private val noSourceFilePrefix = "<NoSourceFile for "
+    private val noSourceFileSuffix = " is a virtual file; DO NOT USE.>"
+
+    /**
+     * Returns `true` if the file represents is a class literal reference in an annotation argument.
+     */
+    private fun isClassLiteralReferenceInAnnotationArgument(file: File): Boolean {
+        return file != anyChangesWildcard &&
+            file != removedOutputsKey &&
+            file.path.startsWith(noSourceFilePrefix) &&
+            file.path.endsWith(noSourceFileSuffix)
+    }
 }
diff --git a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/DualLookupTracker.kt b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/DualLookupTracker.kt
index 51e35ec..006bae2 100644
--- a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/DualLookupTracker.kt
+++ b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/DualLookupTracker.kt
@@ -17,6 +17,7 @@
 
 package com.google.devtools.ksp.impl
 
+import com.google.devtools.ksp.impl.symbol.kotlin.separateQualifierAndName
 import org.jetbrains.kotlin.incremental.LookupTrackerImpl
 import org.jetbrains.kotlin.incremental.components.LookupTracker
 import org.jetbrains.kotlin.incremental.components.Position
@@ -32,8 +33,7 @@
     override fun record(filePath: String, position: Position, scopeFqName: String, scopeKind: ScopeKind, name: String) {
         symbolTracker.record(filePath, position, scopeFqName, scopeKind, name)
         if (scopeKind == ScopeKind.CLASSIFIER) {
-            val className = scopeFqName.substringAfterLast('.')
-            val outerScope = scopeFqName.substringBeforeLast('.', "<anonymous>")
+            val (outerScope, className) = separateQualifierAndName(scopeFqName)
             // DO NOT USE: ScopeKind is meaningless
             classTracker.record(filePath, position, outerScope, scopeKind, className)
         }
diff --git a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/IncrementalContextAA.kt b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/IncrementalContextAA.kt
index 4cccab3..d6eb59f 100644
--- a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/IncrementalContextAA.kt
+++ b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/IncrementalContextAA.kt
@@ -18,12 +18,12 @@
 package com.google.devtools.ksp.impl
 
 import com.google.devtools.ksp.IncrementalContextLoggingOptions
+import com.google.devtools.ksp.InternalKSPException
 import com.google.devtools.ksp.common.IncrementalContextBase
 import com.google.devtools.ksp.common.LookupStorageWrapper
 import com.google.devtools.ksp.common.LookupSymbolWrapper
 import com.google.devtools.ksp.common.LookupTrackerWrapper
 import com.google.devtools.ksp.common.NoSourceFile
-import com.google.devtools.ksp.common.isSyntheticFileName
 import com.google.devtools.ksp.common.stripSyntheticFileNameModifiers
 import com.google.devtools.ksp.containingFile
 import com.google.devtools.ksp.impl.symbol.kotlin.KSFileJavaImpl
@@ -31,6 +31,8 @@
 import com.google.devtools.ksp.impl.symbol.kotlin.KSPropertyDeclarationImpl
 import com.google.devtools.ksp.impl.symbol.kotlin.KSPropertyDeclarationJavaImpl
 import com.google.devtools.ksp.impl.symbol.kotlin.analyze
+import com.google.devtools.ksp.impl.symbol.kotlin.getFqn
+import com.google.devtools.ksp.impl.symbol.kotlin.separateQualifierAndName
 import com.google.devtools.ksp.impl.symbol.kotlin.typeArguments
 import com.google.devtools.ksp.symbol.KSClassDeclaration
 import com.google.devtools.ksp.symbol.KSDeclaration
@@ -54,6 +56,7 @@
 import org.jetbrains.kotlin.analysis.api.types.KaIntersectionType
 import org.jetbrains.kotlin.analysis.api.types.KaType
 import org.jetbrains.kotlin.analysis.api.types.KaTypeParameterType
+import org.jetbrains.kotlin.analysis.api.types.symbol
 import org.jetbrains.kotlin.incremental.IncrementalCompilationContext
 import org.jetbrains.kotlin.incremental.LookupStorage
 import org.jetbrains.kotlin.incremental.LookupSymbol
@@ -182,6 +185,75 @@
         recordWithArgs(type, file)
     }
 
+    /**
+     * Records a reference to `MyClass::class`.
+     * 
+     * @param type the referenced type, e.g., `MyClass`
+     * @param context the parent of [type], i.e., where the reference occurs.
+     */
+    fun recordClassReferenceLookup(type: KaType, context: KSNode) {
+        if (!isIncremental) {
+            return
+        }
+
+        val fqn = type.symbol?.classId?.asFqNameString()
+            ?: type.getFqn()
+            ?: return
+
+        recordClassReferenceLookup(fqn, context)
+    }
+
+    /**
+     * Records a reference to `MyClass::class`.
+     * 
+     * @param fqn the fully qualified name of the referenced type, e.g., the fully qualified name of `MyClass`.
+     * @param context the parent of [fqn], i.e., where the reference occurs.
+     */
+    fun recordClassReferenceLookup(fqn: String, context: KSNode) {
+        if (!isIncremental) {
+            return
+        }
+
+        val (scope, name) = separateQualifierAndName(fqn)
+
+        symbolLookupTracker.record(filePathFor(context), scope, name)
+    }
+
+    /**
+     * Returns a string representing a file path for [context].
+     *
+     * If the filepath is already available, it is returned. Otherwise, a synthetic filepath is generated.
+     */
+    private fun filePathFor(context: KSNode): String {
+        // Try directly getting the filepath
+        val maybeAvailableFilePath = context.containingFile?.filePath
+        if (maybeAvailableFilePath != null) {
+            return maybeAvailableFilePath
+        }
+
+        // Construct synthetic filepath
+        val closestParentDeclaration = context.findFirstParentDeclaration() ?: throw InternalKSPException(
+            "Unexpected missing parent declaration for KSNode '$context'",
+            context.location,
+            context.javaClass
+        )
+
+        val parentDeclarationName = closestParentDeclaration.qualifiedName?.asString()
+            ?: (
+                closestParentDeclaration.packageName.asString()
+                    + "."
+                    + closestParentDeclaration.simpleName.asString()
+                )
+
+        return NoSourceFile(baseDir, parentDeclarationName).filePath
+    }
+
+    /** Returns the nearest parent declaration if it exists */
+    private fun KSNode.findFirstParentDeclaration(): KSDeclaration? = when (this) {
+        is KSDeclaration -> this
+        else -> parent?.findFirstParentDeclaration()
+    }
+
     @OptIn(KaExperimentalApi::class)
     private fun recordLookupForDeclaration(symbol: KaSymbol, file: PsiJavaFile) {
         when (symbol) {
@@ -379,6 +451,24 @@
 internal fun recordLookup(ktType: KaType, context: KSNode?) =
     ResolverAAImpl.instance.incrementalContext.recordLookup(ktType, context)
 
+/**
+ * Records a reference to `MyClass::class`.
+ * 
+ * @param ktType the referenced type, e.g., `MyClass`
+ * @param context the parent of [ktType], i.e., where the reference occurs.
+ */
+internal fun recordClassReferenceLookup(ktType: KaType, context: KSNode) =
+    ResolverAAImpl.instance.incrementalContext.recordClassReferenceLookup(ktType, context)
+
+/**
+ * Records a reference to `MyClass::class`.
+ * 
+ * @param fqn the fully qualified name of the referenced type, e.g., the fully qualified name of `MyClass`.
+ * @param context the parent of [fqn], i.e., where the reference occurs.
+ */
+internal fun recordClassReferenceLookup(fqn: String, context: KSNode) =
+    ResolverAAImpl.instance.incrementalContext.recordClassReferenceLookup(fqn, context)
+
 internal fun recordLookupWithSupertypes(ktType: KaType, extra: (KaType, PsiJavaFile) -> Unit = { _, _ -> }) =
     ResolverAAImpl.instance.incrementalContext.recordLookupWithSupertypes(ktType, mutableSetOf(), extra)
 
diff --git a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/java/KSAnnotationJavaImpl.kt b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/java/KSAnnotationJavaImpl.kt
index f3defe0..28ae42b 100644
--- a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/java/KSAnnotationJavaImpl.kt
+++ b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/java/KSAnnotationJavaImpl.kt
@@ -4,6 +4,7 @@
 import com.google.devtools.ksp.common.impl.KSNameImpl
 import com.google.devtools.ksp.getClassDeclarationByName
 import com.google.devtools.ksp.impl.ResolverAAImpl
+import com.google.devtools.ksp.impl.recordClassReferenceLookup
 import com.google.devtools.ksp.impl.symbol.kotlin.KSClassDeclarationEnumEntryImpl
 import com.google.devtools.ksp.impl.symbol.kotlin.KSErrorType
 import com.google.devtools.ksp.impl.symbol.kotlin.KSValueArgumentImpl
@@ -11,6 +12,7 @@
 import com.google.devtools.ksp.impl.symbol.kotlin.classifierSymbol
 import com.google.devtools.ksp.impl.symbol.kotlin.getDefaultValue
 import com.google.devtools.ksp.impl.symbol.kotlin.resolved.KSTypeReferenceResolvedImpl
+import com.google.devtools.ksp.impl.symbol.kotlin.toKaType
 import com.google.devtools.ksp.impl.symbol.kotlin.toLocation
 import com.google.devtools.ksp.symbol.AnnotationUseSiteTarget
 import com.google.devtools.ksp.symbol.ClassKind
@@ -39,6 +41,7 @@
 import com.intellij.psi.PsiType
 import com.intellij.psi.impl.compiled.ClsClassImpl
 import org.jetbrains.kotlin.analysis.api.KaImplementationDetail
+import org.jetbrains.kotlin.analysis.api.annotations.KaAnnotationValue
 import org.jetbrains.kotlin.analysis.api.impl.base.annotations.KaBaseNamedAnnotationValue
 import org.jetbrains.kotlin.analysis.api.symbols.KaClassSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KaSymbolOrigin
@@ -82,10 +85,10 @@
                 val value = it.value
                 val calculatedValue: Any? = if (value is PsiArrayInitializerMemberValue) {
                     value.initializers.map {
-                        calcValue(it)
+                        calcValue(it, this@KSAnnotationJavaImpl)
                     }
                 } else {
-                    calcValue(it.value)
+                    calcValue(it.value, this@KSAnnotationJavaImpl)
                 }
                 KSValueArgumentLiteImpl(
                     name?.let { KSNameImpl.getCached(it) },
@@ -120,10 +123,10 @@
                                 annoMethod.defaultValue?.let { value ->
                                     val calculatedValue: Any? = if (value is PsiArrayInitializerMemberValue) {
                                         value.initializers.map {
-                                            calcValue(it)
+                                            calcValue(it, this@KSAnnotationJavaImpl)
                                         }
                                     } else {
-                                        calcValue(value)
+                                        calcValue(value, this@KSAnnotationJavaImpl)
                                     }
                                     KSValueArgumentLiteImpl(
                                         KSNameImpl.getCached(annoMethod.name),
@@ -137,6 +140,9 @@
                     } else {
                         symbol.valueParameters.mapNotNull { valueParameterSymbol ->
                             val constantValue = valueParameterSymbol.getDefaultValue() ?: return@mapNotNull null
+                            if (constantValue is KaAnnotationValue.ClassLiteralValue) {
+                                recordClassReferenceLookup(constantValue.type, this@KSAnnotationJavaImpl)
+                            }
                             KSValueArgumentImpl.getCached(
                                 KaBaseNamedAnnotationValue(
                                     valueParameterSymbol.name,
@@ -175,8 +181,13 @@
     }
 }
 
-fun calcValue(value: PsiAnnotationMemberValue?): Any? {
+fun calcValue(value: PsiAnnotationMemberValue?, parent: KSNode? = null): Any? {
     if (value is PsiAnnotation) {
+        parent?.let { ctx ->
+            value.qualifiedName?.let { fqn ->
+                recordClassReferenceLookup(fqn, ctx)
+            }
+        }
         return KSAnnotationJavaImpl.getCached(value, null)
     }
     val result = when (value) {
@@ -207,6 +218,13 @@
                 else -> {
                     ResolverAAImpl.instance
                         .getClassDeclarationByName(component.canonicalText)?.asStarProjectedType()
+                        ?.also { type ->
+                            parent?.let { ctx ->
+                                type.toKaType()?.let { tpe ->
+                                    recordClassReferenceLookup(tpe, ctx)
+                                }
+                            }
+                        }
                         ?: KSErrorType(component.canonicalText)
                 }
             }
@@ -219,6 +237,13 @@
         is PsiType -> {
             ResolverAAImpl.instance
                 .getClassDeclarationByName(result.canonicalText)?.asStarProjectedType()
+                ?.also { type ->
+                    parent?.let { ctx ->
+                        type.toKaType()?.let { tpe ->
+                            recordClassReferenceLookup(tpe, ctx)
+                        }
+                    }
+                }
                 ?: KSErrorType(result.canonicalText)
         }
 
diff --git a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/resolved/KSAnnotationResolvedImpl.kt b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/resolved/KSAnnotationResolvedImpl.kt
index be613b9..37a1924 100644
--- a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/resolved/KSAnnotationResolvedImpl.kt
+++ b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/resolved/KSAnnotationResolvedImpl.kt
@@ -3,6 +3,7 @@
 import com.google.devtools.ksp.common.IdKeyPair
 import com.google.devtools.ksp.common.KSObjectCache
 import com.google.devtools.ksp.common.impl.KSNameImpl
+import com.google.devtools.ksp.impl.recordClassReferenceLookup
 import com.google.devtools.ksp.impl.symbol.java.KSValueArgumentLiteImpl
 import com.google.devtools.ksp.impl.symbol.java.calcValue
 import com.google.devtools.ksp.impl.symbol.kotlin.*
@@ -25,6 +26,7 @@
 import com.intellij.psi.impl.compiled.ClsClassImpl
 import org.jetbrains.kotlin.analysis.api.KaImplementationDetail
 import org.jetbrains.kotlin.analysis.api.annotations.KaAnnotation
+import org.jetbrains.kotlin.analysis.api.annotations.KaAnnotationValue
 import org.jetbrains.kotlin.analysis.api.impl.base.annotations.KaBaseNamedAnnotationValue
 import org.jetbrains.kotlin.analysis.api.symbols.KaSymbolOrigin
 import org.jetbrains.kotlin.descriptors.annotations.AnnotationUseSiteTarget.*
@@ -88,10 +90,10 @@
                             annoMethod.defaultValue?.let { value ->
                                 val calculatedValue: Any? = if (value is PsiArrayInitializerMemberValue) {
                                     value.initializers.map {
-                                        calcValue(it)
+                                        calcValue(it, this@KSAnnotationResolvedImpl)
                                     }
                                 } else {
-                                    calcValue(value)
+                                    calcValue(value, this@KSAnnotationResolvedImpl)
                                 }
                                 KSValueArgumentLiteImpl(
                                     KSNameImpl.getCached(annoMethod.name),
@@ -106,6 +108,9 @@
                     symbol.memberScope.constructors.singleOrNull()?.let {
                         it.valueParameters.mapNotNull { valueParameterSymbol ->
                             val constantValue = valueParameterSymbol.getDefaultValue() ?: return@mapNotNull null
+                            if (constantValue is KaAnnotationValue.ClassLiteralValue) {
+                                recordClassReferenceLookup(constantValue.type, this@KSAnnotationResolvedImpl)
+                            }
                             KSValueArgumentImpl.getCached(
                                 KaBaseNamedAnnotationValue(
                                     valueParameterSymbol.name,
diff --git a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/util.kt b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/util.kt
index 4fc828c..5220151 100644
--- a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/util.kt
+++ b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/util.kt
@@ -24,6 +24,7 @@
 import com.google.devtools.ksp.containingFile
 import com.google.devtools.ksp.impl.KSPCoreEnvironment
 import com.google.devtools.ksp.impl.ResolverAAImpl
+import com.google.devtools.ksp.impl.recordClassReferenceLookup
 import com.google.devtools.ksp.impl.symbol.kotlin.resolved.KSAnnotationResolvedImpl
 import com.google.devtools.ksp.impl.symbol.kotlin.resolved.KSClassifierParameterImpl
 import com.google.devtools.ksp.impl.symbol.kotlin.resolved.KSClassifierReferenceResolvedImpl
@@ -746,7 +747,10 @@
     } ?: KSErrorType
 
     is KaAnnotationValue.ClassLiteralValue -> {
-        KSTypeImpl.getCached(this@toValue.type)
+        parent?.let { ctx ->
+            recordClassReferenceLookup(type, ctx)
+        }
+        KSTypeImpl.getCached(type)
     }
 
     is KaAnnotationValue.ConstantValue -> this.value.value
@@ -1336,3 +1340,14 @@
 // Annotations on deeply synthesized members like getter of Java annotation arguments can be defined in src.
 internal val KSNode.definitionOrigin: Origin
     get() = containingFile?.origin ?: origin
+
+/**
+ * Returns the qualifier and name of [fqn].
+ *
+ * E.g., if `fqn = a.b.c.f`, then it returns `a.b.c, f`.
+ */
+internal fun separateQualifierAndName(fqn: String): Pair<String, String> {
+    val scope = fqn.substringBeforeLast('.', "<anonymous>")
+    val name = fqn.substringAfterLast('.')
+    return scope to name
+}