KSP2: Record implicit Java lookups

(cherry picked from commit 643f36052835b5e4385584b2090248d0647f21f2)

18 files changed, 236 insertions, 67 deletions

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 113e12c7..eea0ca5e 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
@@ -25,20 +25,32 @@ import com.google.devtools.ksp.LookupTrackerWrapper
 import com.google.devtools.ksp.containingFile
 import com.google.devtools.ksp.impl.symbol.kotlin.KSFileJavaImpl
 import com.google.devtools.ksp.impl.symbol.kotlin.KSFunctionDeclarationImpl
+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.KSTypeImpl
-import com.google.devtools.ksp.impl.symbol.kotlin.classifierSymbol
+import com.google.devtools.ksp.impl.symbol.kotlin.analyze
 import com.google.devtools.ksp.impl.symbol.kotlin.typeArguments
 import com.google.devtools.ksp.symbol.KSClassDeclaration
 import com.google.devtools.ksp.symbol.KSDeclaration
-import com.google.devtools.ksp.symbol.KSFile
 import com.google.devtools.ksp.symbol.KSNode
-import com.google.devtools.ksp.symbol.KSType
-import com.google.devtools.ksp.symbol.KSTypeReference
 import com.intellij.psi.PsiJavaFile
 import com.intellij.util.containers.MultiMap
 import com.intellij.util.io.DataExternalizer
 import com.intellij.util.io.IOUtil
+import org.jetbrains.kotlin.analysis.api.symbols.KtFunctionLikeSymbol
+import org.jetbrains.kotlin.analysis.api.symbols.KtJavaFieldSymbol
+import org.jetbrains.kotlin.analysis.api.symbols.KtPropertySymbol
+import org.jetbrains.kotlin.analysis.api.symbols.KtSymbol
+import org.jetbrains.kotlin.analysis.api.symbols.markers.KtSymbolWithMembers
+import org.jetbrains.kotlin.analysis.api.types.KtCapturedType
+import org.jetbrains.kotlin.analysis.api.types.KtDefinitelyNotNullType
+import org.jetbrains.kotlin.analysis.api.types.KtDynamicType
+import org.jetbrains.kotlin.analysis.api.types.KtErrorType
+import org.jetbrains.kotlin.analysis.api.types.KtFlexibleType
+import org.jetbrains.kotlin.analysis.api.types.KtIntegerLiteralType
+import org.jetbrains.kotlin.analysis.api.types.KtIntersectionType
+import org.jetbrains.kotlin.analysis.api.types.KtNonErrorClassType
+import org.jetbrains.kotlin.analysis.api.types.KtType
+import org.jetbrains.kotlin.analysis.api.types.KtTypeParameterType
 import org.jetbrains.kotlin.incremental.IncrementalCompilationContext
 import org.jetbrains.kotlin.incremental.LookupStorage
 import org.jetbrains.kotlin.incremental.LookupSymbol
@@ -51,17 +63,6 @@ import org.jetbrains.kotlin.incremental.update
 import java.io.DataInput
 import java.io.DataOutput
 import java.io.File
-import org.jetbrains.kotlin.analysis.api.symbols.KtClassLikeSymbol
-import org.jetbrains.kotlin.analysis.api.types.KtCapturedType
-import org.jetbrains.kotlin.analysis.api.types.KtClassType
-import org.jetbrains.kotlin.analysis.api.types.KtDefinitelyNotNullType
-import org.jetbrains.kotlin.analysis.api.types.KtDynamicType
-import org.jetbrains.kotlin.analysis.api.types.KtErrorType
-import org.jetbrains.kotlin.analysis.api.types.KtFlexibleType
-import org.jetbrains.kotlin.analysis.api.types.KtIntegerLiteralType
-import org.jetbrains.kotlin.analysis.api.types.KtIntersectionType
-import org.jetbrains.kotlin.analysis.api.types.KtType
-import org.jetbrains.kotlin.analysis.api.types.KtTypeParameterType
 
 class IncrementalContextAA(
     override val isIncremental: Boolean,
@@ -117,9 +118,8 @@ class IncrementalContextAA(
             it.type?.let { recordWithArgs(it, file) }
         }
         when (type) {
-            is KtClassType -> {
-                val classifier = type.classifierSymbol() as? KtClassLikeSymbol ?: return
-                val fqn = classifier.classIdIfNonLocal?.asFqNameString() ?: return
+            is KtNonErrorClassType -> {
+                val fqn = type.classId.asFqNameString()
                 recordLookup(file, fqn)
             }
             is KtFlexibleType -> {
@@ -143,15 +143,114 @@ class IncrementalContextAA(
         }
     }
 
-    fun recordLookup(type: KtType, parent: KSNode?) {
-        val file = (parent?.containingFile as? KSFileJavaImpl)?.psi ?: return
+    fun recordLookup(type: KtType, context: KSNode?) {
+        val file = (context?.containingFile as? KSFileJavaImpl)?.psi ?: return
 
         recordWithArgs(type, file)
     }
+
+    private fun recordLookupForDeclaration(symbol: KtSymbol, file: PsiJavaFile) {
+        when (symbol) {
+            is KtJavaFieldSymbol -> recordWithArgs(symbol.returnType, file)
+            is KtPropertySymbol -> recordWithArgs(symbol.returnType, file)
+            is KtFunctionLikeSymbol -> {
+                recordWithArgs(symbol.returnType, file)
+                symbol.valueParameters.forEach { recordWithArgs(it.returnType, file) }
+                symbol.typeParameters.forEach {
+                    it.upperBounds.forEach {
+                        recordWithArgs(it, file)
+                    }
+                }
+            }
+        }
+    }
+
+    fun recordLookupForPropertyOrMethod(declaration: KSDeclaration) {
+        val file = (declaration.containingFile as? KSFileJavaImpl)?.psi ?: return
+
+        val symbol = when (declaration) {
+            is KSPropertyDeclarationJavaImpl -> declaration.ktJavaFieldSymbol
+            is KSPropertyDeclarationImpl -> declaration.ktPropertySymbol
+            is KSFunctionDeclarationImpl -> declaration.ktFunctionSymbol
+            else -> return
+        }
+        recordLookupForDeclaration(symbol, file)
+    }
+
+    internal fun recordLookupForGetAll(supers: List<KtType>, predicate: (KtSymbol) -> Boolean) {
+        val visited: MutableSet<KtType> = mutableSetOf()
+        analyze {
+            supers.forEach {
+                recordLookupWithSupertypes(it, visited) { type, file ->
+                    if (type is KtNonErrorClassType) {
+                        (type.classSymbol as? KtSymbolWithMembers)?.let {
+                            val declared = it.getDeclaredMemberScope().getAllSymbols() +
+                                it.getStaticDeclaredMemberScope().getAllSymbols()
+                            declared.forEach {
+                                if (predicate(it))
+                                    recordLookupForDeclaration(it, file)
+                            }
+                        }
+                    }
+                }
+            }
+        }
+    }
+
+    private val KtSymbol.psiJavaFile: PsiJavaFile?
+        get() = psi?.containingFile as? PsiJavaFile
+
+    private val KtType.psiJavaFiles: List<PsiJavaFile>
+        get() {
+            return when (this) {
+                is KtNonErrorClassType -> classSymbol.psiJavaFile?.let { listOf(it) } ?: emptyList()
+                is KtFlexibleType -> lowerBound.psiJavaFiles + upperBound.psiJavaFiles
+                is KtIntersectionType -> conjuncts.flatMap { it.psiJavaFiles }
+                is KtCapturedType -> projection.type?.psiJavaFiles ?: emptyList()
+                is KtDefinitelyNotNullType -> original.psiJavaFiles
+                is KtErrorType, is KtIntegerLiteralType, is KtDynamicType, is KtTypeParameterType -> emptyList()
+            }
+        }
+
+    fun recordLookupWithSupertypes(ktType: KtType, visited: MutableSet<KtType>, extra: (KtType, PsiJavaFile) -> Unit) {
+        analyze {
+            fun record(type: KtType) {
+                if (type in visited)
+                    return
+                visited.add(type)
+                for (superType in type.getDirectSuperTypes()) {
+                    for (file in type.psiJavaFiles) {
+                        recordWithArgs(superType, file)
+                    }
+                    record(superType)
+                }
+                for (file in type.psiJavaFiles) {
+                    extra(type, file)
+                }
+            }
+            record(ktType)
+        }
+    }
 }
 
-internal fun recordLookup(ktType: KtType, parent: KSNode?) =
-    ResolverAAImpl.instance.incrementalContext.recordLookup(ktType, parent)
+internal fun recordLookup(ktType: KtType, context: KSNode?) =
+    ResolverAAImpl.instance.incrementalContext.recordLookup(ktType, context)
+
+internal fun recordLookupWithSupertypes(ktType: KtType, extra: (KtType, PsiJavaFile) -> Unit = { _, _ -> }) =
+    ResolverAAImpl.instance.incrementalContext.recordLookupWithSupertypes(ktType, mutableSetOf(), extra)
+
+internal fun recordLookupForPropertyOrMethod(declaration: KSDeclaration) =
+    ResolverAAImpl.instance.incrementalContext.recordLookupForPropertyOrMethod(declaration)
+
+internal fun recordLookupForGetAllProperties(supers: List<KtType>) =
+    ResolverAAImpl.instance.incrementalContext.recordLookupForGetAll(supers) {
+        it is KtPropertySymbol || it is KtJavaFieldSymbol
+    }
+
+internal fun recordLookupForGetAllFunctions(supers: List<KtType>) =
+    ResolverAAImpl.instance.incrementalContext.recordLookupForGetAll(supers) {
+        it is KtFunctionLikeSymbol
+    }
 
 internal fun recordGetSealedSubclasses(classDeclaration: KSClassDeclaration) =
     ResolverAAImpl.instance.incrementalContext.recordGetSealedSubclasses(classDeclaration)
diff --git a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/KotlinSymbolProcessing.kt b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/KotlinSymbolProcessing.kt
index 3fe8dc8f..98773227 100644
--- a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/KotlinSymbolProcessing.kt
+++ b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/KotlinSymbolProcessing.kt
@@ -524,7 +524,11 @@ class KotlinSymbolProcessing(
         }
 
         if (!logger.hasError) {
-            incrementalContext.updateCachesAndOutputs(initialDirtySet, codeGenerator.outputs, codeGenerator.sourceToOutputs)
+            incrementalContext.updateCachesAndOutputs(
+                initialDirtySet,
+                codeGenerator.outputs,
+                codeGenerator.sourceToOutputs
+            )
         }
 
         codeGenerator.closeFiles()
diff --git a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/ResolverAAImpl.kt b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/ResolverAAImpl.kt
index d1ba2726..22e4d7f2 100644
--- a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/ResolverAAImpl.kt
+++ b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/ResolverAAImpl.kt
@@ -530,6 +530,11 @@ class ResolverAAImpl(
             is KSPropertyDeclarationImpl -> overridee.ktPropertySymbol
             else -> return false
         }
+        overrider.closestClassDeclaration()?.asStarProjectedType()?.let {
+            recordLookupWithSupertypes((it as KSTypeImpl).type)
+        }
+        recordLookupForPropertyOrMethod(overrider)
+        recordLookupForPropertyOrMethod(overridee)
         return analyze {
             overriderSymbol.getAllOverriddenSymbols().contains(overrideeSymbol) ||
                 overriderSymbol.getIntersectionOverriddenSymbols().contains(overrideeSymbol)
@@ -541,8 +546,11 @@ class ResolverAAImpl(
         overridee: KSDeclaration,
         containingClass: KSClassDeclaration
     ): Boolean {
+        recordLookupForPropertyOrMethod(overrider)
+        recordLookupForPropertyOrMethod(overridee)
         return when (overrider) {
             is KSPropertyDeclaration -> containingClass.getAllProperties().singleOrNull {
+                recordLookupForPropertyOrMethod(it)
                 it.simpleName.asString() == overrider.simpleName.asString()
             }?.let { overrides(it, overridee) } ?: false
             is KSFunctionDeclaration -> {
@@ -561,7 +569,10 @@ class ResolverAAImpl(
                         }
                     ).any { overrides(it, overridee) }
                 } else {
-                    candidates.any { overrides(it, overridee) }
+                    candidates.any {
+                        recordLookupForPropertyOrMethod(it)
+                        overrides(it, overridee)
+                    }
                 }
             }
             else -> false
@@ -639,6 +650,8 @@ class ResolverAAImpl(
         return propertyAsMemberOfCache.getOrPut(key) {
             val resolved = property.type.resolve()
             if (containing is KSTypeImpl && resolved is KSTypeImpl) {
+                recordLookupWithSupertypes(containing.type)
+                recordLookupForPropertyOrMethod(property)
                 val isSubTypeOf = analyze {
                     (declaredIn.asStarProjectedType() as? KSTypeImpl)?.type?.let {
                         containing.type.isSubTypeOf(it)
@@ -676,6 +689,8 @@ class ResolverAAImpl(
         val key = function to containing
         return functionAsMemberOfCache.getOrPut(key) {
             if (containing is KSTypeImpl) {
+                recordLookupWithSupertypes(containing.type)
+                recordLookupForPropertyOrMethod(function)
                 val isSubTypeOf = analyze {
                     (propertyDeclaredIn.asStarProjectedType() as? KSTypeImpl)?.type?.let {
                         containing.type.isSubTypeOf(it)
diff --git a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSClassDeclarationEnumEntryImpl.kt b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSClassDeclarationEnumEntryImpl.kt
index eb946ad7..800aa9e1 100644
--- a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSClassDeclarationEnumEntryImpl.kt
+++ b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSClassDeclarationEnumEntryImpl.kt
@@ -1,6 +1,9 @@
 package com.google.devtools.ksp.impl.symbol.kotlin
 
 import com.google.devtools.ksp.KSObjectCache
+import com.google.devtools.ksp.impl.recordLookup
+import com.google.devtools.ksp.impl.recordLookupForGetAllFunctions
+import com.google.devtools.ksp.impl.recordLookupForGetAllProperties
 import com.google.devtools.ksp.processing.impl.KSNameImpl
 import com.google.devtools.ksp.symbol.ClassKind
 import com.google.devtools.ksp.symbol.KSAnnotation
@@ -47,11 +50,23 @@ class KSClassDeclarationEnumEntryImpl private constructor(private val ktEnumEntr
     }
 
     override fun getAllFunctions(): Sequence<KSFunctionDeclaration> {
+        analyze {
+            ktEnumEntrySymbol.returnType.getDirectSuperTypes().forEach {
+                recordLookup(it, this@KSClassDeclarationEnumEntryImpl)
+            }
+            recordLookupForGetAllFunctions(ktEnumEntrySymbol.returnType.getDirectSuperTypes())
+        }
         return ktEnumEntrySymbol.enumEntryInitializer?.declarations()?.filterIsInstance<KSFunctionDeclaration>()
             ?: emptySequence()
     }
 
     override fun getAllProperties(): Sequence<KSPropertyDeclaration> {
+        analyze {
+            ktEnumEntrySymbol.returnType.getDirectSuperTypes().forEach {
+                recordLookup(it, this@KSClassDeclarationEnumEntryImpl)
+            }
+            recordLookupForGetAllProperties(ktEnumEntrySymbol.returnType.getDirectSuperTypes())
+        }
         return ktEnumEntrySymbol.enumEntryInitializer?.declarations()?.filterIsInstance<KSPropertyDeclaration>()
             ?: emptySequence()
     }
diff --git a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSClassDeclarationImpl.kt b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSClassDeclarationImpl.kt
index 783affcf..ca935e3b 100644
--- a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSClassDeclarationImpl.kt
+++ b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSClassDeclarationImpl.kt
@@ -20,6 +20,9 @@ package com.google.devtools.ksp.impl.symbol.kotlin
 import com.google.devtools.ksp.KSObjectCache
 import com.google.devtools.ksp.impl.ResolverAAImpl
 import com.google.devtools.ksp.impl.recordGetSealedSubclasses
+import com.google.devtools.ksp.impl.recordLookup
+import com.google.devtools.ksp.impl.recordLookupForGetAllFunctions
+import com.google.devtools.ksp.impl.recordLookupForGetAllProperties
 import com.google.devtools.ksp.impl.symbol.kotlin.resolved.KSTypeReferenceResolvedImpl
 import com.google.devtools.ksp.processing.impl.KSNameImpl
 import com.google.devtools.ksp.processing.impl.KSTypeReferenceSyntheticImpl
@@ -109,10 +112,14 @@ class KSClassDeclarationImpl private constructor(internal val ktClassOrObjectSym
     }
 
     override fun getAllFunctions(): Sequence<KSFunctionDeclaration> {
+        ktClassOrObjectSymbol.superTypes.forEach { recordLookup(it, this) }
+        recordLookupForGetAllFunctions(ktClassOrObjectSymbol.superTypes)
         return ktClassOrObjectSymbol.getAllFunctions()
     }
 
     override fun getAllProperties(): Sequence<KSPropertyDeclaration> {
+        ktClassOrObjectSymbol.superTypes.forEach { recordLookup(it, this) }
+        recordLookupForGetAllProperties(ktClassOrObjectSymbol.superTypes)
         return ktClassOrObjectSymbol.getAllProperties()
     }
 
diff --git a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSFunctionDeclarationImpl.kt b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSFunctionDeclarationImpl.kt
index fe55cfac..93bc6390 100644
--- a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSFunctionDeclarationImpl.kt
+++ b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSFunctionDeclarationImpl.kt
@@ -19,6 +19,8 @@ package com.google.devtools.ksp.impl.symbol.kotlin
 
 import com.google.devtools.ksp.*
 import com.google.devtools.ksp.impl.ResolverAAImpl
+import com.google.devtools.ksp.impl.recordLookupForPropertyOrMethod
+import com.google.devtools.ksp.impl.recordLookupWithSupertypes
 import com.google.devtools.ksp.impl.symbol.kotlin.resolved.KSTypeReferenceResolvedImpl
 import com.google.devtools.ksp.processing.impl.KSNameImpl
 import com.google.devtools.ksp.symbol.*
@@ -99,9 +101,13 @@ class KSFunctionDeclarationImpl private constructor(internal val ktFunctionSymbo
     }
 
     override fun findOverridee(): KSDeclaration? {
+        closestClassDeclaration()?.asStarProjectedType()?.let {
+            recordLookupWithSupertypes((it as KSTypeImpl).type)
+        }
+        recordLookupForPropertyOrMethod(this)
         return analyze {
             ktFunctionSymbol.getDirectlyOverriddenSymbols().firstOrNull()?.unwrapFakeOverrides?.toKSDeclaration()
-        }
+        }?.also { recordLookupForPropertyOrMethod(it) }
     }
 
     override fun asMemberOf(containing: KSType): KSFunction {
diff --git a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSPropertyDeclarationImpl.kt b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSPropertyDeclarationImpl.kt
index 88e6f5e0..139fddb5 100644
--- a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSPropertyDeclarationImpl.kt
+++ b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSPropertyDeclarationImpl.kt
@@ -18,7 +18,10 @@
 package com.google.devtools.ksp.impl.symbol.kotlin
 
 import com.google.devtools.ksp.KSObjectCache
+import com.google.devtools.ksp.closestClassDeclaration
 import com.google.devtools.ksp.impl.ResolverAAImpl
+import com.google.devtools.ksp.impl.recordLookupForPropertyOrMethod
+import com.google.devtools.ksp.impl.recordLookupWithSupertypes
 import com.google.devtools.ksp.impl.symbol.kotlin.resolved.KSTypeReferenceResolvedImpl
 import com.google.devtools.ksp.impl.symbol.util.BinaryClassInfoCache
 import com.google.devtools.ksp.processing.impl.KSNameImpl
@@ -121,10 +124,14 @@ class KSPropertyDeclarationImpl private constructor(internal val ktPropertySymbo
     }
 
     override fun findOverridee(): KSPropertyDeclaration? {
+        closestClassDeclaration()?.asStarProjectedType()?.let {
+            recordLookupWithSupertypes((it as KSTypeImpl).type)
+        }
+        recordLookupForPropertyOrMethod(this)
         return analyze {
             ktPropertySymbol.getDirectlyOverriddenSymbols().firstOrNull()
                 ?.unwrapFakeOverrides?.toKSDeclaration() as? KSPropertyDeclaration
-        }
+        }?.also { recordLookupForPropertyOrMethod(it) }
     }
 
     override fun asMemberOf(containing: KSType): KSType {
diff --git a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSPropertyDeclarationJavaImpl.kt b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSPropertyDeclarationJavaImpl.kt
index 3b5456d9..4cfe8ec5 100644
--- a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSPropertyDeclarationJavaImpl.kt
+++ b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSPropertyDeclarationJavaImpl.kt
@@ -1,7 +1,10 @@
 package com.google.devtools.ksp.impl.symbol.kotlin
 
 import com.google.devtools.ksp.KSObjectCache
+import com.google.devtools.ksp.closestClassDeclaration
 import com.google.devtools.ksp.impl.ResolverAAImpl
+import com.google.devtools.ksp.impl.recordLookupForPropertyOrMethod
+import com.google.devtools.ksp.impl.recordLookupWithSupertypes
 import com.google.devtools.ksp.impl.symbol.kotlin.resolved.KSTypeReferenceResolvedImpl
 import com.google.devtools.ksp.processing.impl.KSNameImpl
 import com.google.devtools.ksp.symbol.*
@@ -9,7 +12,7 @@ import org.jetbrains.kotlin.analysis.api.symbols.KtJavaFieldSymbol
 import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.descriptors.java.JavaVisibilities
 
-class KSPropertyDeclarationJavaImpl private constructor(private val ktJavaFieldSymbol: KtJavaFieldSymbol) :
+class KSPropertyDeclarationJavaImpl private constructor(val ktJavaFieldSymbol: KtJavaFieldSymbol) :
     KSPropertyDeclaration,
     AbstractKSDeclarationImpl(ktJavaFieldSymbol),
     KSExpectActual by KSExpectActualImpl(ktJavaFieldSymbol) {
@@ -41,7 +44,12 @@ class KSPropertyDeclarationJavaImpl private constructor(private val ktJavaFieldS
     }
 
     override fun findOverridee(): KSPropertyDeclaration? {
+        closestClassDeclaration()?.asStarProjectedType()?.let {
+            recordLookupWithSupertypes((it as KSTypeImpl).type)
+        }
+        recordLookupForPropertyOrMethod(this)
         TODO("Not yet implemented")
+        // ?.also { recordLookupForPropertyOrMethod(it) }
     }
 
     override fun asMemberOf(containing: KSType): KSType {
diff --git a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSTypeArgumentImpl.kt b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSTypeArgumentImpl.kt
index 4c06a983..1a96cba9 100644
--- a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSTypeArgumentImpl.kt
+++ b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSTypeArgumentImpl.kt
@@ -17,7 +17,6 @@
 package com.google.devtools.ksp.impl.symbol.kotlin
 
 import com.google.devtools.ksp.KSObjectCache
-import com.google.devtools.ksp.findParentOfType
 import com.google.devtools.ksp.symbol.KSAnnotation
 import com.google.devtools.ksp.symbol.KSNode
 import com.google.devtools.ksp.symbol.KSTypeArgument
@@ -28,7 +27,6 @@ import com.google.devtools.ksp.symbol.Origin
 import com.google.devtools.ksp.symbol.Variance
 import org.jetbrains.kotlin.psi.KtProjectionKind
 import org.jetbrains.kotlin.psi.KtTypeProjection
-import org.jetbrains.kotlin.psi.KtUserType
 
 class KSTypeArgumentImpl(private val ktTypeArgument: KtTypeProjection, override val parent: KSNode) : KSTypeArgument {
     companion object : KSObjectCache<KtTypeProjection, KSTypeArgument>() {
diff --git a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSTypeImpl.kt b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSTypeImpl.kt
index 9620f614..50788410 100644
--- a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSTypeImpl.kt
+++ b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSTypeImpl.kt
@@ -20,6 +20,7 @@ package com.google.devtools.ksp.impl.symbol.kotlin
 import com.google.devtools.ksp.IdKeyPair
 import com.google.devtools.ksp.KSObjectCache
 import com.google.devtools.ksp.impl.ResolverAAImpl
+import com.google.devtools.ksp.impl.recordLookupWithSupertypes
 import com.google.devtools.ksp.impl.symbol.kotlin.resolved.KSTypeArgumentResolvedImpl
 import com.google.devtools.ksp.impl.symbol.kotlin.synthetic.getExtensionFunctionTypeAnnotation
 import com.google.devtools.ksp.symbol.KSAnnotation
@@ -97,6 +98,7 @@ class KSTypeImpl private constructor(internal val type: KtType) : KSType {
         if (that.isError || this.isError) {
             return false
         }
+        recordLookupWithSupertypes((that as KSTypeImpl).type)
         return type.isAssignableFrom((that as KSTypeImpl).type)
     }
 
diff --git a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSTypeReferenceImpl.kt b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSTypeReferenceImpl.kt
index 7eaccf23..e4e764ea 100644
--- a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSTypeReferenceImpl.kt
+++ b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSTypeReferenceImpl.kt
@@ -19,7 +19,6 @@ package com.google.devtools.ksp.impl.symbol.kotlin
 import com.google.devtools.ksp.ExceptionMessage
 import com.google.devtools.ksp.IdKeyPair
 import com.google.devtools.ksp.KSObjectCache
-import com.google.devtools.ksp.impl.ResolverAAImpl
 import com.google.devtools.ksp.impl.recordLookup
 import com.google.devtools.ksp.impl.symbol.util.toKSModifiers
 import com.google.devtools.ksp.symbol.KSAnnotation
diff --git a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/resolved/KSTypeReferenceResolvedImpl.kt b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/resolved/KSTypeReferenceResolvedImpl.kt
index f947d048..68d9d9c0 100644
--- a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/resolved/KSTypeReferenceResolvedImpl.kt
+++ b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/resolved/KSTypeReferenceResolvedImpl.kt
@@ -19,7 +19,6 @@ package com.google.devtools.ksp.impl.symbol.kotlin.resolved
 
 import com.google.devtools.ksp.IdKeyTriple
 import com.google.devtools.ksp.KSObjectCache
-import com.google.devtools.ksp.impl.ResolverAAImpl
 import com.google.devtools.ksp.impl.recordLookup
 import com.google.devtools.ksp.impl.symbol.kotlin.Deferrable
 import com.google.devtools.ksp.impl.symbol.kotlin.KSClassDeclarationImpl
diff --git a/test-utils/src/main/kotlin/com/google/devtools/ksp/processor/RecordJavaAsMemberOfProcessor.kt b/test-utils/src/main/kotlin/com/google/devtools/ksp/processor/RecordJavaAsMemberOfProcessor.kt
index 8da0f726..9a34cd19 100644
--- a/test-utils/src/main/kotlin/com/google/devtools/ksp/processor/RecordJavaAsMemberOfProcessor.kt
+++ b/test-utils/src/main/kotlin/com/google/devtools/ksp/processor/RecordJavaAsMemberOfProcessor.kt
@@ -17,6 +17,7 @@
 
 package com.google.devtools.ksp.processor
 
+import com.google.devtools.ksp.impl.ResolverAAImpl
 import com.google.devtools.ksp.processing.Resolver
 import com.google.devtools.ksp.processing.impl.ResolverImpl
 import com.google.devtools.ksp.symbol.*
@@ -53,12 +54,14 @@ class RecordJavaAsMemberOfProcessor : AbstractTestProcessor() {
 
         function!!.asMemberOf(type!!)
 
-        if (resolver is ResolverImpl) {
-            val m = resolver.incrementalContext.dumpLookupRecords()
-            m.toSortedMap().forEach { symbol, files ->
-                files.filter { it.endsWith(".java") }.sorted().forEach {
-                    results.add("$symbol: $it")
-                }
+        val m = when (resolver) {
+            is ResolverImpl -> resolver.incrementalContext.dumpLookupRecords().toSortedMap()
+            is ResolverAAImpl -> resolver.incrementalContext.dumpLookupRecords().toSortedMap()
+            else -> throw IllegalStateException("Unknown Resolver: $resolver")
+        }
+        m.forEach { symbol, files ->
+            files.filter { it.endsWith(".java") }.sorted().forEach {
+                results.add("$symbol: $it")
             }
         }
         return emptyList()
diff --git a/test-utils/src/main/kotlin/com/google/devtools/ksp/processor/RecordJavaGetAllMembersProcessor.kt b/test-utils/src/main/kotlin/com/google/devtools/ksp/processor/RecordJavaGetAllMembersProcessor.kt
index e172e0ee..8369fd92 100644
--- a/test-utils/src/main/kotlin/com/google/devtools/ksp/processor/RecordJavaGetAllMembersProcessor.kt
+++ b/test-utils/src/main/kotlin/com/google/devtools/ksp/processor/RecordJavaGetAllMembersProcessor.kt
@@ -17,6 +17,7 @@
 
 package com.google.devtools.ksp.processor
 
+import com.google.devtools.ksp.impl.ResolverAAImpl
 import com.google.devtools.ksp.processing.Resolver
 import com.google.devtools.ksp.processing.impl.ResolverImpl
 import com.google.devtools.ksp.symbol.*
@@ -41,12 +42,14 @@ class RecordJavaGetAllMembersProcessor : AbstractTestProcessor() {
             }
         }
 
-        if (resolver is ResolverImpl) {
-            val m = resolver.incrementalContext.dumpLookupRecords().toSortedMap()
-            m.toSortedMap().forEach { symbol, files ->
-                files.filter { it.endsWith(".java") }.sorted().forEach {
-                    results.add("$symbol: $it")
-                }
+        val m = when (resolver) {
+            is ResolverImpl -> resolver.incrementalContext.dumpLookupRecords().toSortedMap()
+            is ResolverAAImpl -> resolver.incrementalContext.dumpLookupRecords().toSortedMap()
+            else -> throw IllegalStateException("Unknown Resolver: $resolver")
+        }
+        m.forEach { symbol, files ->
+            files.filter { it.endsWith(".java") }.sorted().forEach {
+                results.add("$symbol: $it")
             }
         }
         return emptyList()
diff --git a/test-utils/src/main/kotlin/com/google/devtools/ksp/processor/RecordJavaOverridesProcessor.kt b/test-utils/src/main/kotlin/com/google/devtools/ksp/processor/RecordJavaOverridesProcessor.kt
index a1872e85..50c9dd19 100644
--- a/test-utils/src/main/kotlin/com/google/devtools/ksp/processor/RecordJavaOverridesProcessor.kt
+++ b/test-utils/src/main/kotlin/com/google/devtools/ksp/processor/RecordJavaOverridesProcessor.kt
@@ -17,6 +17,7 @@
 
 package com.google.devtools.ksp.processor
 
+import com.google.devtools.ksp.impl.ResolverAAImpl
 import com.google.devtools.ksp.processing.Resolver
 import com.google.devtools.ksp.processing.impl.ResolverImpl
 import com.google.devtools.ksp.symbol.*
@@ -58,12 +59,14 @@ class RecordJavaOverridesProcessor : AbstractTestProcessor() {
         resolver.overrides(A_f1!!, C_f1!!)
         A_f2!!.findOverridee()
 
-        if (resolver is ResolverImpl) {
-            val m = resolver.incrementalContext.dumpLookupRecords().toSortedMap()
-            m.toSortedMap().forEach { symbol, files ->
-                files.filter { it.endsWith(".java") }.sorted().forEach {
-                    results.add("$symbol: $it")
-                }
+        val m = when (resolver) {
+            is ResolverImpl -> resolver.incrementalContext.dumpLookupRecords().toSortedMap()
+            is ResolverAAImpl -> resolver.incrementalContext.dumpLookupRecords().toSortedMap()
+            else -> throw IllegalStateException("Unknown Resolver: $resolver")
+        }
+        m.forEach { symbol, files ->
+            files.filter { it.endsWith(".java") }.sorted().forEach {
+                results.add("$symbol: $it")
             }
         }
         return emptyList()
diff --git a/test-utils/src/main/kotlin/com/google/devtools/ksp/processor/RecordJavaProcessor.kt b/test-utils/src/main/kotlin/com/google/devtools/ksp/processor/RecordJavaProcessor.kt
index 15f06bf2..8e975e93 100644
--- a/test-utils/src/main/kotlin/com/google/devtools/ksp/processor/RecordJavaProcessor.kt
+++ b/test-utils/src/main/kotlin/com/google/devtools/ksp/processor/RecordJavaProcessor.kt
@@ -17,6 +17,7 @@
 
 package com.google.devtools.ksp.processor
 
+import com.google.devtools.ksp.impl.ResolverAAImpl
 import com.google.devtools.ksp.processing.Resolver
 import com.google.devtools.ksp.processing.impl.ResolverImpl
 import com.google.devtools.ksp.symbol.*
@@ -37,12 +38,14 @@ class RecordJavaProcessor : AbstractTestProcessor() {
                 it.validate()
             }
         }
-        if (resolver is ResolverImpl) {
-            val m = resolver.incrementalContext.dumpLookupRecords().toSortedMap()
-            m.forEach { symbol, files ->
-                files.filter { it.endsWith(".java") }.sorted().forEach {
-                    results.add("$symbol: $it")
-                }
+        val m = when (resolver) {
+            is ResolverImpl -> resolver.incrementalContext.dumpLookupRecords().toSortedMap()
+            is ResolverAAImpl -> resolver.incrementalContext.dumpLookupRecords().toSortedMap()
+            else -> throw IllegalStateException("Unknown Resolver: $resolver")
+        }
+        m.forEach { symbol, files ->
+            files.filter { it.endsWith(".java") }.sorted().forEach {
+                results.add("$symbol: $it")
             }
         }
         return emptyList()
diff --git a/test-utils/src/main/kotlin/com/google/devtools/ksp/processor/RecordJavaSupertypesProcessor.kt b/test-utils/src/main/kotlin/com/google/devtools/ksp/processor/RecordJavaSupertypesProcessor.kt
index a3f5528d..a9e9103f 100644
--- a/test-utils/src/main/kotlin/com/google/devtools/ksp/processor/RecordJavaSupertypesProcessor.kt
+++ b/test-utils/src/main/kotlin/com/google/devtools/ksp/processor/RecordJavaSupertypesProcessor.kt
@@ -17,6 +17,7 @@
 
 package com.google.devtools.ksp.processor
 
+import com.google.devtools.ksp.impl.ResolverAAImpl
 import com.google.devtools.ksp.processing.Resolver
 import com.google.devtools.ksp.processing.impl.ResolverImpl
 import com.google.devtools.ksp.symbol.*
@@ -38,12 +39,14 @@ class RecordJavaSupertypesProcessor : AbstractTestProcessor() {
         types.forEach {
             resolver.builtIns.anyType.isAssignableFrom(it)
         }
-        if (resolver is ResolverImpl) {
-            val m = resolver.incrementalContext.dumpLookupRecords().toSortedMap()
-            m.forEach { symbol, files ->
-                files.filter { it.endsWith(".java") }.sorted().forEach {
-                    results.add("$symbol: $it")
-                }
+        val m = when (resolver) {
+            is ResolverImpl -> resolver.incrementalContext.dumpLookupRecords().toSortedMap()
+            is ResolverAAImpl -> resolver.incrementalContext.dumpLookupRecords().toSortedMap()
+            else -> throw IllegalStateException("Unknown Resolver: $resolver")
+        }
+        m.forEach { symbol, files ->
+            files.filter { it.endsWith(".java") }.sorted().forEach {
+                results.add("$symbol: $it")
             }
         }
         return emptyList()
diff --git a/test-utils/src/test/kotlin/com/google/devtools/ksp/test/KSPAATest.kt b/test-utils/src/test/kotlin/com/google/devtools/ksp/test/KSPAATest.kt
index 2a121349..40cb116a 100644
--- a/test-utils/src/test/kotlin/com/google/devtools/ksp/test/KSPAATest.kt
+++ b/test-utils/src/test/kotlin/com/google/devtools/ksp/test/KSPAATest.kt
@@ -490,35 +490,30 @@ class KSPAATest : AbstractKSPAATest() {
         runTest("../test-utils/testData/api/rawTypes.kt")
     }
 
-    @Disabled
     @TestMetadata("recordJavaAnnotationTypes.kt")
     @Test
     fun testRecordJavaAnnotationTypes() {
         runTest("../test-utils/testData/api/recordJavaAnnotationTypes.kt")
     }
 
-    @Disabled
     @TestMetadata("recordJavaAsMemberOf.kt")
     @Test
     fun testRecordJavaAsMemberOf() {
         runTest("../test-utils/testData/api/recordJavaAsMemberOf.kt")
     }
 
-    @Disabled
     @TestMetadata("recordJavaGetAllMembers.kt")
     @Test
     fun testRecordJavaGetAllMembers() {
         runTest("../test-utils/testData/api/recordJavaGetAllMembers.kt")
     }
 
-    @Disabled
     @TestMetadata("recordJavaOverrides.kt")
     @Test
     fun testRecordJavaOverrides() {
         runTest("../test-utils/testData/api/recordJavaOverrides.kt")
     }
 
-    @Disabled
     @TestMetadata("recordJavaSupertypes.kt")
     @Test
     fun testRecordJavaSupertypes() {