Simplify checks of descriptor and psi correspondance.

Author: ting-yuan

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

@@ -536,15 +536,15 @@ class ResolverImpl(
                         kindFilter = DescriptorKindFilter.VARIABLES,
                         name = propName
                     ) {
-                        synthesizedPropPrefix[prefix]!!(it as? PropertyDescriptor)?.findPsi() == psi
+                        synthesizedPropPrefix[prefix]!!(it as? PropertyDescriptor)?.correspondsTo(psi) == true
                     }
                 }
                 property ?: moduleClassResolver
                     .resolveContainingClass(psi)?.let { containingClass ->
                         val filter = if (psi is SyntheticElement) {
                             { declaration: DeclarationDescriptor -> declaration.name.asString() == psi.name }
                         } else {
-                            { declaration: DeclarationDescriptor -> declaration.findPsi() == psi }
+                            { declaration: DeclarationDescriptor -> declaration.correspondsTo(psi) }
                         }
                         containingClass.findEnclosedDescriptor(
                             kindFilter = DescriptorKindFilter.FUNCTIONS,
@@ -559,7 +559,7 @@ class ResolverImpl(
                     ?.findEnclosedDescriptor(
                         kindFilter = DescriptorKindFilter.VARIABLES,
                         name = psi.name,
-                        filter = { it.findPsi() == psi }
+                        filter = { it.correspondsTo(psi) }
                     )
             }
             else -> throw IllegalStateException("unhandled psi element kind: ${psi.javaClass}")
@@ -756,7 +756,7 @@ class ResolverImpl(
                             ?.findEnclosedDescriptor(
                                 kindFilter = DescriptorKindFilter.FUNCTIONS,
                                 name = owner.name,
-                                filter = { it.findPsi() == owner }
+                                filter = { it.correspondsTo(owner) }
                             ) as FunctionDescriptor
                     } as DeclarationDescriptor
                     val typeParameterDescriptor = LazyJavaTypeParameterDescriptor(

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

@@ -559,6 +559,11 @@ fun DeclarationDescriptor.findPsi(): PsiElement? {
     return leaf.parentsWithSelf.firstOrNull { psi.manager.areElementsEquivalent(it, psi) }
 }
 
+fun DeclarationDescriptor.correspondsTo(candidate: PsiElement): Boolean {
+    val psi = (this as? DeclarationDescriptorWithSource)?.source?.getPsi() ?: return false
+    return psi.manager.areElementsEquivalent(psi, candidate)
+}
+
 internal fun KSFile.getPackageAnnotations() = (this as? KSFileJavaImpl)?.psi?.packageStatement
     ?.annotationList?.annotations?.map { KSAnnotationJavaImpl.getCached(it) } ?: emptyList<KSAnnotation>()