AA: use object cache for KSTypeReference.

* extract IdKey to common-util module.
* add additional information for creating type reference to support location info.
* implement location for KSTypeReference.

17 files changed, 86 insertions, 53 deletions

diff --git a/common-util/src/main/kotlin/com/google/devtools/ksp/KSPUtils.kt b/common-util/src/main/kotlin/com/google/devtools/ksp/KSPUtils.kt
new file mode 100644
index 00000000..007a3eca
--- /dev/null
+++ b/common-util/src/main/kotlin/com/google/devtools/ksp/KSPUtils.kt
@@ -0,0 +1,18 @@
+package com.google.devtools.ksp
+
+class IdKey<T>(private val k: T) {
+    override fun equals(other: Any?): Boolean = if (other is IdKey<*>) k === other.k else false
+    override fun hashCode(): Int = k.hashCode()
+}
+
+class IdKeyPair<T, P>(private val k1: T, private val k2: P) {
+    override fun equals(other: Any?): Boolean = if (other is IdKeyPair<*, *>) k1 === other.k1 &&
+        k2 === other.k2 else false
+    override fun hashCode(): Int = k1.hashCode() * 31 + k2.hashCode()
+}
+
+class IdKeyTriple<T, P, Q>(private val k1: T, private val k2: P, private val k3: Q) {
+    override fun equals(other: Any?): Boolean = if (other is IdKeyTriple<*, *, *>) k1 === other.k1 &&
+        k2 === other.k2 && k3 === other.k3 else false
+    override fun hashCode(): Int = k1.hashCode() * 31 * 31 + k2.hashCode() * 31 + k3.hashCode()
+}
diff --git a/compiler-plugin/src/main/kotlin/com/google/devtools/ksp/symbol/impl/binary/KSTypeArgumentDescriptorImpl.kt b/compiler-plugin/src/main/kotlin/com/google/devtools/ksp/symbol/impl/binary/KSTypeArgumentDescriptorImpl.kt
index 3f50b60b..5b76e27a 100644
--- a/compiler-plugin/src/main/kotlin/com/google/devtools/ksp/symbol/impl/binary/KSTypeArgumentDescriptorImpl.kt
+++ b/compiler-plugin/src/main/kotlin/com/google/devtools/ksp/symbol/impl/binary/KSTypeArgumentDescriptorImpl.kt
@@ -17,6 +17,7 @@
 
 package com.google.devtools.ksp.symbol.impl.binary
 
+import com.google.devtools.ksp.IdKeyTriple
 import com.google.devtools.ksp.KSObjectCache
 import com.google.devtools.ksp.symbol.KSAnnotation
 import com.google.devtools.ksp.symbol.KSNode
@@ -25,7 +26,6 @@ import com.google.devtools.ksp.symbol.Location
 import com.google.devtools.ksp.symbol.NonExistLocation
 import com.google.devtools.ksp.symbol.Origin
 import com.google.devtools.ksp.symbol.Variance
-import com.google.devtools.ksp.symbol.impl.kotlin.IdKeyTriple
 import com.google.devtools.ksp.symbol.impl.kotlin.KSTypeArgumentImpl
 import org.jetbrains.kotlin.types.TypeProjection
 
diff --git a/compiler-plugin/src/main/kotlin/com/google/devtools/ksp/symbol/impl/binary/KSTypeReferenceDescriptorImpl.kt b/compiler-plugin/src/main/kotlin/com/google/devtools/ksp/symbol/impl/binary/KSTypeReferenceDescriptorImpl.kt
index b2260cf9..91f68fa4 100644
--- a/compiler-plugin/src/main/kotlin/com/google/devtools/ksp/symbol/impl/binary/KSTypeReferenceDescriptorImpl.kt
+++ b/compiler-plugin/src/main/kotlin/com/google/devtools/ksp/symbol/impl/binary/KSTypeReferenceDescriptorImpl.kt
@@ -17,6 +17,7 @@
 
 package com.google.devtools.ksp.symbol.impl.binary
 
+import com.google.devtools.ksp.IdKeyTriple
 import com.google.devtools.ksp.KSObjectCache
 import com.google.devtools.ksp.memoized
 import com.google.devtools.ksp.symbol.KSAnnotation
@@ -29,7 +30,6 @@ import com.google.devtools.ksp.symbol.Location
 import com.google.devtools.ksp.symbol.Modifier
 import com.google.devtools.ksp.symbol.NonExistLocation
 import com.google.devtools.ksp.symbol.Origin
-import com.google.devtools.ksp.symbol.impl.kotlin.IdKeyTriple
 import com.google.devtools.ksp.symbol.impl.kotlin.getKSTypeCached
 import org.jetbrains.kotlin.builtins.isSuspendFunctionTypeOrSubtype
 import org.jetbrains.kotlin.types.AbbreviatedType
diff --git a/compiler-plugin/src/main/kotlin/com/google/devtools/ksp/symbol/impl/kotlin/KSTypeImpl.kt b/compiler-plugin/src/main/kotlin/com/google/devtools/ksp/symbol/impl/kotlin/KSTypeImpl.kt
index 9704f229..9e8929a8 100644
--- a/compiler-plugin/src/main/kotlin/com/google/devtools/ksp/symbol/impl/kotlin/KSTypeImpl.kt
+++ b/compiler-plugin/src/main/kotlin/com/google/devtools/ksp/symbol/impl/kotlin/KSTypeImpl.kt
@@ -17,6 +17,7 @@
 
 package com.google.devtools.ksp.symbol.impl.kotlin
 
+import com.google.devtools.ksp.IdKey
 import com.google.devtools.ksp.KSObjectCache
 import com.google.devtools.ksp.processing.impl.ResolverImpl
 import com.google.devtools.ksp.symbol.KSAnnotation
@@ -131,23 +132,6 @@ class KSTypeImpl private constructor(
         get() = kotlinType.isSuspendFunctionType || kotlinType.isKSuspendFunctionType
 }
 
-class IdKey<T>(private val k: T) {
-    override fun equals(other: Any?): Boolean = if (other is IdKey<*>) k === other.k else false
-    override fun hashCode(): Int = k.hashCode()
-}
-
-class IdKeyPair<T, P>(private val k1: T, private val k2: P) {
-    override fun equals(other: Any?): Boolean = if (other is IdKeyPair<*, *>) k1 === other.k1 &&
-        k2 === other.k2 else false
-    override fun hashCode(): Int = k1.hashCode() * 31 + k2.hashCode()
-}
-
-class IdKeyTriple<T, P, Q>(private val k1: T, private val k2: P, private val k3: Q) {
-    override fun equals(other: Any?): Boolean = if (other is IdKeyTriple<*, *, *>) k1 === other.k1 &&
-        k2 === other.k2 && k3 === other.k3 else false
-    override fun hashCode(): Int = k1.hashCode() * 31 * 31 + k2.hashCode() * 31 + k3.hashCode()
-}
-
 fun getKSTypeCached(
     kotlinType: KotlinType,
     ksTypeArguments: List<KSTypeArgument>? = null,
diff --git a/compiler-plugin/src/main/kotlin/com/google/devtools/ksp/symbol/impl/synthetic/KSTypeReferenceSyntheticImpl.kt b/compiler-plugin/src/main/kotlin/com/google/devtools/ksp/symbol/impl/synthetic/KSTypeReferenceSyntheticImpl.kt
index fd7a6f6f..aa2631e4 100644
--- a/compiler-plugin/src/main/kotlin/com/google/devtools/ksp/symbol/impl/synthetic/KSTypeReferenceSyntheticImpl.kt
+++ b/compiler-plugin/src/main/kotlin/com/google/devtools/ksp/symbol/impl/synthetic/KSTypeReferenceSyntheticImpl.kt
@@ -1,5 +1,6 @@
 package com.google.devtools.ksp.symbol.impl.synthetic
 
+import com.google.devtools.ksp.IdKey
 import com.google.devtools.ksp.KSObjectCache
 import com.google.devtools.ksp.symbol.KSAnnotation
 import com.google.devtools.ksp.symbol.KSNode
@@ -11,7 +12,6 @@ import com.google.devtools.ksp.symbol.Location
 import com.google.devtools.ksp.symbol.Modifier
 import com.google.devtools.ksp.symbol.NonExistLocation
 import com.google.devtools.ksp.symbol.Origin
-import com.google.devtools.ksp.symbol.impl.kotlin.IdKey
 
 class KSTypeReferenceSyntheticImpl(val ksType: KSType, override val parent: KSNode?) : KSTypeReference {
     companion object : KSObjectCache<Pair<IdKey<KSType>, KSNode?>, KSTypeReferenceSyntheticImpl>() {
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 a64a571a..9f97e4cb 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
@@ -23,7 +23,7 @@ import org.jetbrains.kotlin.analysis.api.components.buildClassType
 import org.jetbrains.kotlin.analysis.api.symbols.*
 import org.jetbrains.kotlin.psi.KtObjectDeclaration
 
-class KSClassDeclarationImpl private constructor(private val ktClassOrObjectSymbol: KtClassOrObjectSymbol) :
+class KSClassDeclarationImpl private constructor(internal val ktClassOrObjectSymbol: KtClassOrObjectSymbol) :
     KSClassDeclaration,
     AbstractKSDeclarationImpl(ktClassOrObjectSymbol),
     KSExpectActual by KSExpectActualImpl(ktClassOrObjectSymbol) {
@@ -56,7 +56,9 @@ class KSClassDeclarationImpl private constructor(private val ktClassOrObjectSymb
 
     override val superTypes: Sequence<KSTypeReference> by lazy {
         analyze {
-            ktClassOrObjectSymbol.superTypes.map { KSTypeReferenceImpl(it) }.asSequence()
+            ktClassOrObjectSymbol.superTypes.mapIndexed { index, type ->
+                KSTypeReferenceImpl.getCached(type, this@KSClassDeclarationImpl, index)
+            }.asSequence()
         }
     }
 
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 8bbade9d..1000bd48 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
@@ -52,14 +52,16 @@ class KSFunctionDeclarationImpl private constructor(private val ktFunctionSymbol
             if (!ktFunctionSymbol.isExtension) {
                 null
             } else {
-                ktFunctionSymbol.receiverType?.let { KSTypeReferenceImpl(it) }
+                ktFunctionSymbol.receiverType?.let {
+                    KSTypeReferenceImpl.getCached(it, this@KSFunctionDeclarationImpl)
+                }
             }
         }
     }
 
     override val returnType: KSTypeReference? by lazy {
         analyze {
-            KSTypeReferenceImpl(ktFunctionSymbol.returnType)
+            KSTypeReferenceImpl.getCached(ktFunctionSymbol.returnType, this@KSFunctionDeclarationImpl)
         }
     }
 
diff --git a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSPropertyAccessorImpl.kt b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSPropertyAccessorImpl.kt
index d4944a8c..76a38715 100644
--- a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSPropertyAccessorImpl.kt
+++ b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSPropertyAccessorImpl.kt
@@ -52,7 +52,7 @@ abstract class KSPropertyAccessorImpl(
     }
 
     override val parent: KSNode?
-        get() = TODO("Not yet implemented")
+        get() = ktPropertyAccessorSymbol.getContainingKSSymbol()
 }
 
 class KSPropertySetterImpl private constructor(
@@ -87,7 +87,7 @@ class KSPropertyGetterImpl private constructor(
     }
 
     override val returnType: KSTypeReference? by lazy {
-        KSTypeReferenceImpl(getter.returnType)
+        KSTypeReferenceImpl.getCached(getter.returnType, this@KSPropertyGetterImpl)
     }
 
     override fun <D, R> accept(visitor: KSVisitor<D, R>, data: D): R {
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 9291ee7d..896a70f0 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
@@ -39,11 +39,11 @@ class KSPropertyDeclarationImpl private constructor(private val ktPropertySymbol
     }
 
     override val extensionReceiver: KSTypeReference? by lazy {
-        ktPropertySymbol.receiverType?.let { KSTypeReferenceImpl(it) }
+        ktPropertySymbol.receiverType?.let { KSTypeReferenceImpl.getCached(it, this@KSPropertyDeclarationImpl) }
     }
 
     override val type: KSTypeReference by lazy {
-        KSTypeReferenceImpl(ktPropertySymbol.returnType)
+        KSTypeReferenceImpl.getCached(ktPropertySymbol.returnType, this@KSPropertyDeclarationImpl)
     }
 
     override val isMutable: Boolean by lazy {
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 baf9a4df..316d6c58 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
@@ -31,7 +31,7 @@ class KSPropertyDeclarationJavaImpl private constructor(private val ktJavaFieldS
         get() = null
 
     override val type: KSTypeReference by lazy {
-        KSTypeReferenceImpl.getCached(ktJavaFieldSymbol.returnType)
+        KSTypeReferenceImpl.getCached(ktJavaFieldSymbol.returnType, this@KSPropertyDeclarationJavaImpl)
     }
 
     override val isMutable: Boolean
diff --git a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSTypeAliasImpl.kt b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSTypeAliasImpl.kt
index 44777c40..e5dbbee3 100644
--- a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSTypeAliasImpl.kt
+++ b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSTypeAliasImpl.kt
@@ -39,7 +39,7 @@ class KSTypeAliasImpl private constructor(private val ktTypeAliasSymbol: KtTypeA
     }
 
     override val type: KSTypeReference by lazy {
-        KSTypeReferenceImpl.getCached(ktTypeAliasSymbol.expandedType)
+        KSTypeReferenceImpl.getCached(ktTypeAliasSymbol.expandedType, this)
     }
 
     override val simpleName: KSName
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 1270fa2c..e9ab22b4 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
@@ -51,7 +51,7 @@ class KSTypeArgumentImpl private constructor(private val ktTypeArgument: KtTypeA
     }
 
     override val type: KSTypeReference? by lazy {
-        ktTypeArgument.type?.let { KSTypeReferenceImpl(it) }
+        ktTypeArgument.type?.let { KSTypeReferenceImpl.getCached(it, this@KSTypeArgumentImpl) }
     }
 
     override val annotations: Sequence<KSAnnotation> by lazy {
diff --git a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSTypeArgumentLiteImpl.kt b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSTypeArgumentLiteImpl.kt
index f19285cb..5eb239e1 100644
--- a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSTypeArgumentLiteImpl.kt
+++ b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSTypeArgumentLiteImpl.kt
@@ -11,7 +11,8 @@ import com.google.devtools.ksp.symbol.NonExistLocation
 import com.google.devtools.ksp.symbol.Origin
 import com.google.devtools.ksp.symbol.Variance
 
-class KSTypeArgumentLiteImpl(override val type: KSTypeReference, override val variance: Variance) : KSTypeArgument {
+class KSTypeArgumentLiteImpl private constructor(override val type: KSTypeReference, override val variance: Variance) :
+    KSTypeArgument {
     companion object : KSObjectCache<Pair<KSTypeReference, Variance>, KSTypeArgument>() {
         fun getCached(type: KSTypeReference, variance: Variance) = cache.getOrPut(Pair(type, variance)) {
             KSTypeArgumentLiteImpl(type, variance)
diff --git a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSTypeParameterImpl.kt b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSTypeParameterImpl.kt
index 5f5dda9d..73e5470e 100644
--- a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSTypeParameterImpl.kt
+++ b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSTypeParameterImpl.kt
@@ -21,7 +21,7 @@ import com.google.devtools.ksp.KSObjectCache
 import com.google.devtools.ksp.symbol.*
 import org.jetbrains.kotlin.analysis.api.symbols.KtTypeParameterSymbol
 
-class KSTypeParameterImpl private constructor(private val ktTypeParameterSymbol: KtTypeParameterSymbol) :
+class KSTypeParameterImpl private constructor(internal val ktTypeParameterSymbol: KtTypeParameterSymbol) :
     KSTypeParameter,
     AbstractKSDeclarationImpl(ktTypeParameterSymbol),
     KSExpectActual by KSExpectActualImpl(ktTypeParameterSymbol) {
@@ -45,7 +45,9 @@ class KSTypeParameterImpl private constructor(private val ktTypeParameterSymbol:
     override val isReified: Boolean = ktTypeParameterSymbol.isReified
 
     override val bounds: Sequence<KSTypeReference> by lazy {
-        ktTypeParameterSymbol.upperBounds.asSequence().map { KSTypeReferenceImpl(it) }
+        ktTypeParameterSymbol.upperBounds.asSequence().mapIndexed { index, type ->
+            KSTypeReferenceImpl.getCached(type, this@KSTypeParameterImpl, index)
+        }
     }
 
     override val typeParameters: List<KSTypeParameter> = emptyList()
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 79991385..4485f87d 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
@@ -17,22 +17,26 @@
 
 package com.google.devtools.ksp.impl.symbol.kotlin
 
+import com.google.devtools.ksp.IdKeyTriple
 import com.google.devtools.ksp.KSObjectCache
-import com.google.devtools.ksp.symbol.KSAnnotation
-import com.google.devtools.ksp.symbol.KSNode
-import com.google.devtools.ksp.symbol.KSReferenceElement
-import com.google.devtools.ksp.symbol.KSType
-import com.google.devtools.ksp.symbol.KSTypeReference
-import com.google.devtools.ksp.symbol.KSVisitor
-import com.google.devtools.ksp.symbol.Location
-import com.google.devtools.ksp.symbol.Modifier
-import com.google.devtools.ksp.symbol.Origin
+import com.google.devtools.ksp.symbol.*
+import com.intellij.psi.PsiClass
+import com.intellij.psi.PsiTypeParameter
+import com.intellij.psi.impl.source.PsiClassReferenceType
 import org.jetbrains.kotlin.analysis.api.types.*
+import org.jetbrains.kotlin.psi.KtClassOrObject
+import org.jetbrains.kotlin.psi.KtTypeParameter
 
-class KSTypeReferenceImpl(private val ktType: KtType) : KSTypeReference {
-    companion object : KSObjectCache<KtType, KSTypeReference>() {
-        fun getCached(type: KtType): KSTypeReference = cache.getOrPut(type) { KSTypeReferenceImpl(type) }
+class KSTypeReferenceImpl private constructor(
+    private val ktType: KtType,
+    override val parent: KSNode?,
+    private val index: Int
+) : KSTypeReference {
+    companion object : KSObjectCache<IdKeyTriple<KtType, KSNode?, Int>, KSTypeReference>() {
+        fun getCached(type: KtType, parent: KSNode? = null, index: Int = -1): KSTypeReference =
+            cache.getOrPut(IdKeyTriple(type, parent, index)) { KSTypeReferenceImpl(type, parent, index) }
     }
+
     // FIXME: return correct reference element.
     override val element: KSReferenceElement? = null
 
@@ -53,13 +57,32 @@ class KSTypeReferenceImpl(private val ktType: KtType) : KSTypeReference {
         ktType.annotations()
     }
 
-    override val origin: Origin = Origin.KOTLIN
-
-    override val location: Location
-        get() = TODO("Not yet implemented")
+    override val origin: Origin = parent?.origin ?: Origin.SYNTHETIC
 
-    override val parent: KSNode?
-        get() = TODO("Not yet implemented")
+    override val location: Location by lazy {
+        if (index != -1) {
+            parent?.location ?: NonExistLocation
+        } else {
+            when (parent) {
+                is KSClassDeclarationImpl -> {
+                    when (val psi = parent.ktClassOrObjectSymbol.psi) {
+                        is KtClassOrObject -> psi.superTypeListEntries.get(index).toLocation()
+                        is PsiClass -> (psi as? PsiClassReferenceType)?.reference?.toLocation() ?: NonExistLocation
+                        else -> NonExistLocation
+                    }
+                }
+                is KSTypeParameterImpl -> {
+                    when (val psi = parent.ktTypeParameterSymbol.psi) {
+                        is KtTypeParameter -> parent.location
+                        is PsiTypeParameter -> (psi.extendsListTypes[index] as? PsiClassReferenceType)
+                            ?.reference?.toLocation() ?: NonExistLocation
+                        else -> NonExistLocation
+                    }
+                }
+                else -> NonExistLocation
+            }
+        }
+    }
 
     override fun <D, R> accept(visitor: KSVisitor<D, R>, data: D): R {
         return visitor.visitTypeReference(this, data)
diff --git a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSValueParameterImpl.kt b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSValueParameterImpl.kt
index 3a6e12ce..75c277ca 100644
--- a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSValueParameterImpl.kt
+++ b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSValueParameterImpl.kt
@@ -34,7 +34,7 @@ class KSValueParameterImpl private constructor(
     }
 
     override val type: KSTypeReference by lazy {
-        KSTypeReferenceImpl(ktValueParameterSymbol.returnType)
+        KSTypeReferenceImpl.getCached(ktValueParameterSymbol.returnType, this@KSValueParameterImpl)
     }
 
     override val isVararg: Boolean by lazy {
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 3aa38639..8e39b566 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
@@ -182,6 +182,7 @@ internal fun KtSymbol.getContainingKSSymbol(): KSDeclaration? {
         when (val containingSymbol = this@getContainingKSSymbol.getContainingSymbol()) {
             is KtNamedClassOrObjectSymbol -> KSClassDeclarationImpl.getCached(containingSymbol)
             is KtFunctionLikeSymbol -> KSFunctionDeclarationImpl.getCached(containingSymbol)
+            is KtPropertySymbol -> KSPropertyDeclarationImpl.getCached(containingSymbol)
             else -> null
         }
     }