use KtType for hashCode and equals checking for KSType`

* handles type alias for getting declarations for KSType.
* unmutes typeAlias test.

5 files changed, 60 insertions, 45 deletions

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 6096c18c..9d881ec8 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
@@ -25,14 +25,8 @@ import com.google.devtools.ksp.symbol.KSTypeArgument
 import com.google.devtools.ksp.symbol.Nullability
 import org.jetbrains.kotlin.analysis.api.KtStarProjectionTypeArgument
 import org.jetbrains.kotlin.analysis.api.components.buildClassType
-import org.jetbrains.kotlin.analysis.api.symbols.KtClassOrObjectSymbol
-import org.jetbrains.kotlin.analysis.api.symbols.KtTypeAliasSymbol
-import org.jetbrains.kotlin.analysis.api.types.KtClassErrorType
-import org.jetbrains.kotlin.analysis.api.types.KtFlexibleType
-import org.jetbrains.kotlin.analysis.api.types.KtFunctionalType
-import org.jetbrains.kotlin.analysis.api.types.KtNonErrorClassType
-import org.jetbrains.kotlin.analysis.api.types.KtType
-import org.jetbrains.kotlin.analysis.api.types.KtTypeNullability
+import org.jetbrains.kotlin.analysis.api.symbols.*
+import org.jetbrains.kotlin.analysis.api.types.*
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
 class KSTypeImpl private constructor(internal val type: KtType) : KSType {
@@ -45,10 +39,11 @@ class KSTypeImpl private constructor(internal val type: KtType) : KSType {
             when (this@toDeclaration) {
                 is KtNonErrorClassType -> {
                     when (val symbol = this@toDeclaration.classSymbol) {
-                        is KtTypeAliasSymbol -> symbol.expandedType.toDeclaration()
+                        is KtTypeAliasSymbol -> KSTypeAliasImpl.getCached(symbol)
                         is KtClassOrObjectSymbol -> KSClassDeclarationImpl.getCached(symbol)
                     }
                 }
+                is KtTypeParameterType -> KSTypeParameterImpl.getCached(symbol)
                 is KtClassErrorType -> KSErrorTypeClassDeclaration
                 is KtFlexibleType ->
                     type.lowerBoundIfFlexible().toDeclaration()
@@ -135,7 +130,35 @@ class KSTypeImpl private constructor(internal val type: KtType) : KSType {
     override val isSuspendFunctionType: Boolean
         get() = type is KtFunctionalType && type.isSuspend
 
+    override fun hashCode(): Int {
+        return type.toAbbreviatedType().hashCode()
+    }
+
+    override fun equals(other: Any?): Boolean {
+        fun KtTypeAliasSymbol.toAbbreviatedType(): KtType {
+            val result = this.expandedType
+            if (result.classifierSymbol() is KtTypeAliasSymbol) {
+                return (result.classifierSymbol() as KtTypeAliasSymbol).toAbbreviatedType()
+            }
+            return result
+        }
+
+        if (other !is KSTypeImpl)
+            return false
+        val thisType = type.toAbbreviatedType()
+        val otherType = other.type.toAbbreviatedType()
+        return thisType.equals(otherType)
+    }
+
     override fun toString(): String {
         return type.render()
     }
+
+    private fun KtType.toAbbreviatedType(): KtType {
+        val symbol = this.classifierSymbol()
+        return when (symbol) {
+            is KtTypeAliasSymbol -> symbol.expandedType.toAbbreviatedType()
+            else -> this
+        }
+    }
 }
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 46386207..79991385 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
@@ -27,10 +27,7 @@ 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 org.jetbrains.kotlin.analysis.api.types.KtClassErrorType
-import org.jetbrains.kotlin.analysis.api.types.KtFunctionalType
-import org.jetbrains.kotlin.analysis.api.types.KtNonErrorClassType
-import org.jetbrains.kotlin.analysis.api.types.KtType
+import org.jetbrains.kotlin.analysis.api.types.*
 
 class KSTypeReferenceImpl(private val ktType: KtType) : KSTypeReference {
     companion object : KSObjectCache<KtType, KSTypeReference>() {
@@ -43,10 +40,7 @@ class KSTypeReferenceImpl(private val ktType: KtType) : KSTypeReference {
         // TODO: non exist type returns KtNonErrorClassType, check upstream for KtClassErrorType usage.
         return if (
             analyze {
-                ktType is KtClassErrorType || (
-                    (ktType is KtNonErrorClassType) &&
-                        ktType.classId.toKtClassSymbol() == null
-                    )
+                ktType is KtClassErrorType || (ktType.classifierSymbol() == null)
             }
         ) {
             KSErrorType
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 77ce8b49..3aa38639 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
@@ -28,36 +28,16 @@ import org.jetbrains.kotlin.analysis.api.KtStarProjectionTypeArgument
 import org.jetbrains.kotlin.analysis.api.KtTypeArgument
 import org.jetbrains.kotlin.analysis.api.KtTypeArgumentWithVariance
 import org.jetbrains.kotlin.analysis.api.analyze
-import org.jetbrains.kotlin.analysis.api.annotations.KtAnnotated
-import org.jetbrains.kotlin.analysis.api.annotations.annotations
+import org.jetbrains.kotlin.analysis.api.annotations.*
 import org.jetbrains.kotlin.analysis.api.lifetime.KtAlwaysAccessibleLifetimeToken
 import org.jetbrains.kotlin.analysis.api.lifetime.KtAlwaysAccessibleLifetimeTokenFactory
-import org.jetbrains.kotlin.analysis.api.symbols.KtClassOrObjectSymbol
-import org.jetbrains.kotlin.analysis.api.symbols.KtEnumEntrySymbol
-import org.jetbrains.kotlin.analysis.api.symbols.KtFileSymbol
-import org.jetbrains.kotlin.analysis.api.symbols.KtFunctionLikeSymbol
-import org.jetbrains.kotlin.analysis.api.symbols.KtJavaFieldSymbol
-import org.jetbrains.kotlin.analysis.api.symbols.KtNamedClassOrObjectSymbol
-import org.jetbrains.kotlin.analysis.api.symbols.KtPropertySymbol
-import org.jetbrains.kotlin.analysis.api.symbols.KtSymbol
-import org.jetbrains.kotlin.analysis.api.symbols.KtSymbolOrigin
-import org.jetbrains.kotlin.analysis.api.symbols.KtTypeAliasSymbol
-import org.jetbrains.kotlin.analysis.api.symbols.KtTypeParameterSymbol
+import org.jetbrains.kotlin.analysis.api.symbols.*
 import org.jetbrains.kotlin.analysis.api.symbols.markers.KtSymbolWithMembers
-import org.jetbrains.kotlin.analysis.api.types.KtCapturedType
-import org.jetbrains.kotlin.analysis.api.types.KtClassErrorType
-import org.jetbrains.kotlin.analysis.api.types.KtDefinitelyNotNullType
-import org.jetbrains.kotlin.analysis.api.types.KtDynamicType
-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.KtTypeNullability
-import org.jetbrains.kotlin.analysis.api.types.KtTypeParameterType
+import org.jetbrains.kotlin.analysis.api.types.*
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.types.Variance
+import org.jetbrains.kotlin.utils.KotlinExceptionWithAttachments
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
 internal val ktSymbolOriginToOrigin = mapOf(
@@ -234,3 +214,22 @@ internal fun ClassId.toKtClassSymbol(): KtClassOrObjectSymbol? {
         }
     }
 }
+
+internal fun KtType.classifierSymbol(): KtClassifierSymbol? {
+    return try {
+        when (this) {
+            is KtTypeParameterType -> this.symbol
+            is KtCapturedType -> TODO("fix in upstream")
+            is KtClassErrorType -> null
+            is KtFunctionalType -> classSymbol
+            is KtUsualClassType -> classSymbol
+            is KtDefinitelyNotNullType -> original.classifierSymbol()
+            is KtDynamicType -> null
+            is KtFlexibleType -> lowerBound.classifierSymbol()
+            is KtIntegerLiteralType -> null
+            is KtIntersectionType -> null
+        }
+    } catch (e: KotlinExceptionWithAttachments) {
+        null
+    }
+}
diff --git a/kotlin-analysis-api/src/test/kotlin/com/google/devtools/ksp/impl/test/KSPAATest.kt b/kotlin-analysis-api/src/test/kotlin/com/google/devtools/ksp/impl/test/KSPAATest.kt
index 7383a9b4..718a66fc 100644
--- a/kotlin-analysis-api/src/test/kotlin/com/google/devtools/ksp/impl/test/KSPAATest.kt
+++ b/kotlin-analysis-api/src/test/kotlin/com/google/devtools/ksp/impl/test/KSPAATest.kt
@@ -547,7 +547,6 @@ class KSPAATest : AbstractKSPAATest() {
         runTest("../test-utils/testData/api/topLevelMembers.kt")
     }
 
-    @Disabled
     @TestMetadata("typeAlias.kt")
     @Test
     fun testTypeAlias() {
diff --git a/test-utils/testData/api/typeAlias.kt b/test-utils/testData/api/typeAlias.kt
index 22bfe456..3d1e2777 100644
--- a/test-utils/testData/api/typeAlias.kt
+++ b/test-utils/testData/api/typeAlias.kt
@@ -20,7 +20,7 @@
 // EXPECTED:
 // a : A = String
 // b : B = String
-// c : C = A = String
+// c : CC = A = String
 // d : String
 // listOfInt : ListOfInt = List<Int>
 // listOfInt_B : ListOfInt_B = ListOfInt = List<Int>
@@ -33,7 +33,7 @@
 
 typealias A = String
 typealias B = String
-typealias C = A
+typealias CC = A
 typealias ListOfInt = List<Int>
 typealias ListOfInt_B = ListOfInt
 typealias ListOfInt_C = ListOfInt_B
@@ -44,7 +44,7 @@ typealias MyList_B_String = MyList_B<String>
 
 val a: A = ""
 val b: B = ""
-val c: C = ""
+val c: CC = ""
 val d: String = ""
 val listOfInt: ListOfInt = TODO()
 val listOfInt_B: ListOfInt_B = TODO()