KSP2: return type aliases if possible

(cherry picked from commit e5503dd)
google · Feb 26, 2025 · 6f2ea1e · 6f2ea1e
1 parent 1afb280
commit 6f2ea1e
Show file tree

Hide file tree

Showing 10 changed files with 37 additions and 10 deletions.
diff --git a/...api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSCallableReferenceImpl.kt b/...api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSCallableReferenceImpl.kt
@@ -7,6 +7,7 @@ import com.google.devtools.ksp.impl.symbol.kotlin.resolved.KSTypeReferenceResolv
 import com.google.devtools.ksp.symbol.*
 import org.jetbrains.kotlin.analysis.api.types.KaFunctionType
 import org.jetbrains.kotlin.analysis.api.types.KaType
+import org.jetbrains.kotlin.analysis.api.types.abbreviationOrSelf
 
 // TODO: implement a psi based version, rename this class to resolved Impl.
 class KSCallableReferenceImpl private constructor(
@@ -18,15 +19,15 @@ class KSCallableReferenceImpl private constructor(
             cache.getOrPut(IdKeyPair(ktFunctionalType, parent)) { KSCallableReferenceImpl(ktFunctionalType, parent) }
     }
     override val receiverType: KSTypeReference?
-        get() = ktFunctionalType.receiverType?.let { KSTypeReferenceResolvedImpl.getCached(it) }
+        get() = ktFunctionalType.receiverType?.abbreviationOrSelf?.let { KSTypeReferenceResolvedImpl.getCached(it) }
 
     override val functionParameters: List<KSValueParameter>
         get() = ktFunctionalType.parameterTypes.map {
             KSValueParameterLiteImpl.getCached(it, this@KSCallableReferenceImpl)
         }
 
     override val returnType: KSTypeReference
-        get() = KSTypeReferenceResolvedImpl.getCached(ktFunctionalType.returnType)
+        get() = KSTypeReferenceResolvedImpl.getCached(ktFunctionalType.returnType.abbreviationOrSelf)
 
     override val typeArguments: List<KSTypeArgument>
         get() = ktFunctionalType.typeArguments().map { KSTypeArgumentResolvedImpl.getCached(it, this) }

diff --git a/...-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSClassDeclarationImpl.kt b/...-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSClassDeclarationImpl.kt
@@ -32,6 +32,7 @@ import org.jetbrains.kotlin.analysis.api.KaExperimentalApi
 import org.jetbrains.kotlin.analysis.api.KaImplementationDetail
 import org.jetbrains.kotlin.analysis.api.impl.base.types.KaBaseStarTypeProjection
 import org.jetbrains.kotlin.analysis.api.symbols.*
+import org.jetbrains.kotlin.analysis.api.types.abbreviationOrSelf
 import org.jetbrains.kotlin.psi.KtClassOrObject
 
 class KSClassDeclarationImpl private constructor(internal val ktClassOrObjectSymbol: KaClassSymbol) :
@@ -84,7 +85,7 @@ class KSClassDeclarationImpl private constructor(internal val ktClassOrObjectSym
             }
         } ?: analyze {
             val supers = ktClassOrObjectSymbol.superTypes.mapIndexed { index, type ->
-                KSTypeReferenceResolvedImpl.getCached(type, this@KSClassDeclarationImpl, index)
+                KSTypeReferenceResolvedImpl.getCached(type.abbreviationOrSelf, this@KSClassDeclarationImpl, index)
             }
             // AA is returning additional kotlin.Any for java classes, explicitly extending kotlin.Any will result in
             // compile error, therefore filtering by name should work.

diff --git a/...i/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSFunctionDeclarationImpl.kt b/...i/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSFunctionDeclarationImpl.kt
@@ -28,6 +28,7 @@ import com.google.devtools.ksp.symbol.*
 import com.intellij.psi.PsiClass
 import com.intellij.psi.PsiMethod
 import org.jetbrains.kotlin.analysis.api.symbols.*
+import org.jetbrains.kotlin.analysis.api.types.abbreviationOrSelf
 import org.jetbrains.kotlin.psi.KtDeclaration
 import org.jetbrains.kotlin.psi.KtFunction
 
@@ -82,7 +83,7 @@ class KSFunctionDeclarationImpl private constructor(internal val ktFunctionSymbo
                             ktFunctionSymbol.receiverParameter?.annotations ?: emptyList()
                         )
                     }
-                    ?: ktFunctionSymbol.receiverType?.let {
+                    ?: ktFunctionSymbol.receiverType?.abbreviationOrSelf?.let {
                         KSTypeReferenceResolvedImpl.getCached(
                             it,
                             this@KSFunctionDeclarationImpl,
@@ -101,7 +102,7 @@ class KSFunctionDeclarationImpl private constructor(internal val ktFunctionSymbo
                 if (ktFunctionSymbol is KaConstructorSymbol) {
                     ((parentDeclaration as KSClassDeclaration).asStarProjectedType() as KSTypeImpl).type
                 } else {
-                    ktFunctionSymbol.returnType
+                    ktFunctionSymbol.returnType.abbreviationOrSelf
                 }.let { KSTypeReferenceResolvedImpl.getCached(it, this@KSFunctionDeclarationImpl) }
             }
     }

diff --git a/...-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSPropertyAccessorImpl.kt b/...-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSPropertyAccessorImpl.kt
@@ -25,6 +25,7 @@ import com.google.devtools.ksp.symbol.*
 import org.jetbrains.kotlin.analysis.api.symbols.KaPropertyAccessorSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KaPropertyGetterSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KaPropertySetterSymbol
+import org.jetbrains.kotlin.analysis.api.types.abbreviationOrSelf
 import org.jetbrains.kotlin.descriptors.annotations.AnnotationUseSiteTarget
 import org.jetbrains.kotlin.psi.KtDeclaration
 import org.jetbrains.kotlin.psi.KtModifierListOwner
@@ -139,7 +140,7 @@ class KSPropertyGetterImpl private constructor(
     override val returnType: KSTypeReference? by lazy {
         ((owner as? KSPropertyDeclarationImpl)?.ktPropertySymbol?.psiIfSource() as? KtProperty)?.typeReference
             ?.let { KSTypeReferenceImpl.getCached(it, this) }
-            ?: KSTypeReferenceResolvedImpl.getCached(getter.returnType, this@KSPropertyGetterImpl)
+            ?: KSTypeReferenceResolvedImpl.getCached(getter.returnType.abbreviationOrSelf, this@KSPropertyGetterImpl)
     }
 
     override fun <D, R> accept(visitor: KSVisitor<D, R>, data: D): R {

diff --git a/...i/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSPropertyDeclarationImpl.kt b/...i/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSPropertyDeclarationImpl.kt
@@ -39,6 +39,7 @@ import org.jetbrains.kotlin.analysis.api.symbols.KaPropertySymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KaSymbolModality
 import org.jetbrains.kotlin.analysis.api.symbols.KaSymbolVisibility
 import org.jetbrains.kotlin.analysis.api.symbols.receiverType
+import org.jetbrains.kotlin.analysis.api.types.abbreviationOrSelf
 import org.jetbrains.kotlin.descriptors.annotations.AnnotationUseSiteTarget
 import org.jetbrains.kotlin.load.kotlin.JvmPackagePartSource
 import org.jetbrains.kotlin.load.kotlin.KotlinJvmBinarySourceElement
@@ -111,7 +112,7 @@ class KSPropertyDeclarationImpl private constructor(internal val ktPropertySymbo
                     ktPropertySymbol.receiverParameter?.annotations ?: emptyList()
                 )
             }
-            ?: ktPropertySymbol.receiverType?.let {
+            ?: ktPropertySymbol.receiverType?.abbreviationOrSelf?.let {
                 KSTypeReferenceResolvedImpl.getCached(
                     it,
                     this@KSPropertyDeclarationImpl,
@@ -123,7 +124,10 @@ class KSPropertyDeclarationImpl private constructor(internal val ktPropertySymbo
 
     override val type: KSTypeReference by lazy {
         (ktPropertySymbol.psiIfSource() as? KtProperty)?.typeReference?.let { KSTypeReferenceImpl.getCached(it, this) }
-            ?: KSTypeReferenceResolvedImpl.getCached(ktPropertySymbol.returnType, this@KSPropertyDeclarationImpl)
+            ?: KSTypeReferenceResolvedImpl.getCached(
+                ktPropertySymbol.returnType.abbreviationOrSelf,
+                this@KSPropertyDeclarationImpl
+            )
     }
 
     override val isMutable: Boolean by lazy {

diff --git a/...tlin/com/google/devtools/ksp/impl/symbol/kotlin/KSPropertyDeclarationLocalVariableImpl.kt b/...tlin/com/google/devtools/ksp/impl/symbol/kotlin/KSPropertyDeclarationLocalVariableImpl.kt
@@ -11,6 +11,7 @@ import com.google.devtools.ksp.symbol.KSType
 import com.google.devtools.ksp.symbol.KSTypeReference
 import com.google.devtools.ksp.symbol.KSVisitor
 import org.jetbrains.kotlin.analysis.api.symbols.KaLocalVariableSymbol
+import org.jetbrains.kotlin.analysis.api.types.abbreviationOrSelf
 import org.jetbrains.kotlin.psi.KtProperty
 
 class KSPropertyDeclarationLocalVariableImpl private constructor(
@@ -32,7 +33,7 @@ class KSPropertyDeclarationLocalVariableImpl private constructor(
     override val type: KSTypeReference by lazy {
         (ktLocalVariableSymbol.psiIfSource() as? KtProperty)?.typeReference
             ?.let { KSTypeReferenceImpl.getCached(it, this) }
-            ?: KSTypeReferenceResolvedImpl.getCached(ktLocalVariableSymbol.returnType, this)
+            ?: KSTypeReferenceResolvedImpl.getCached(ktLocalVariableSymbol.returnType.abbreviationOrSelf, this)
     }
 
     override val isMutable: Boolean = !ktLocalVariableSymbol.isVal

diff --git a/...is-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSValueParameterImpl.kt b/...is-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSValueParameterImpl.kt
@@ -24,6 +24,7 @@ import com.google.devtools.ksp.impl.symbol.kotlin.resolved.KSTypeReferenceResolv
 import com.google.devtools.ksp.symbol.*
 import org.jetbrains.kotlin.analysis.api.fir.symbols.KaFirValueParameterSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KaValueParameterSymbol
+import org.jetbrains.kotlin.analysis.api.types.abbreviationOrSelf
 import org.jetbrains.kotlin.fir.java.JavaTypeParameterStack
 import org.jetbrains.kotlin.fir.java.declarations.FirJavaValueParameter
 import org.jetbrains.kotlin.fir.java.resolveIfJavaType
@@ -59,7 +60,10 @@ class KSValueParameterImpl private constructor(
         }
         (ktValueParameterSymbol.psiIfSource() as? KtParameter)?.typeReference
             ?.let { KSTypeReferenceImpl.getCached(it, this) }
-            ?: KSTypeReferenceResolvedImpl.getCached(ktValueParameterSymbol.returnType, this@KSValueParameterImpl)
+            ?: KSTypeReferenceResolvedImpl.getCached(
+                ktValueParameterSymbol.returnType.abbreviationOrSelf,
+                this@KSValueParameterImpl
+            )
     }
 
     override val isVararg: Boolean by lazy {

diff --git a/kotlin-analysis-api/testData/typeAlias.kt b/kotlin-analysis-api/testData/typeAlias.kt
@@ -36,6 +36,7 @@
 // nested2 : List<ListOfInt> = (expanded) List<List<Int>>
 // param w.o. asMemberOf: MyAlias<String> = Foo<Bar<T>, Baz<T>> = (expanded) Foo<Bar<String>, Baz<String>>
 // param with asMemberOf: MyAlias<String> = Foo<Bar<T>, Baz<T>> = (expanded) Foo<Bar<String>, Baz<String>>
+// param: MyAlias: MyAlias<String> = Foo<Bar<T>, Baz<T>> = (expanded) Foo<Bar<String>, Baz<String>>
 // END
 
 // MODULE: module1

diff --git a/test-utils/src/main/kotlin/com/google/devtools/ksp/processor/TypeAliasProcessor.kt b/test-utils/src/main/kotlin/com/google/devtools/ksp/processor/TypeAliasProcessor.kt
@@ -18,6 +18,7 @@
 package com.google.devtools.ksp.processor
 
 import com.google.devtools.ksp.getConstructors
+import com.google.devtools.ksp.getDeclaredProperties
 import com.google.devtools.ksp.processing.Resolver
 import com.google.devtools.ksp.symbol.*
 
@@ -71,6 +72,17 @@ open class TypeAliasProcessor : AbstractTestProcessor() {
 
         results.add("param w.o. asMemberOf: $type1Signatures = (expanded) $type1Expanded")
         results.add("param with asMemberOf: $type2Signatures = (expanded) $type2Expanded")
+
+        val property = subject.getDeclaredProperties().single()
+        val propertyType = property.type.resolve()
+        val propertyTypeSignatures = propertyType.typeAliasSignatures().joinToString(" = ")
+        val propertyTypeExpanded = resolver.expandType(propertyType).toSignature()
+        results.add(
+            "${property.simpleName.asString()}: " +
+                "${propertyType.declaration.qualifiedName?.asString()}: " +
+                "$propertyTypeSignatures = (expanded) $propertyTypeExpanded"
+        )
+
         return emptyList()
     }
 

diff --git a/test-utils/testData/api/typeAlias.kt b/test-utils/testData/api/typeAlias.kt
@@ -36,6 +36,7 @@
 // nested2 : List<ListOfInt> = (expanded) List<List<Int>>
 // param w.o. asMemberOf: MyAlias<String> = Foo<Bar<T>, Baz<T>> = (expanded) Foo<Bar<String>, Baz<String>>
 // param with asMemberOf: MyAlias<String> = Foo<Bar<T>, Baz<T>> = (expanded) Foo<Bar<String>, Baz<String>>
+// param: MyAlias: MyAlias<String> = Foo<Bar<T>, Baz<T>> = (expanded) Foo<Bar<String>, Baz<String>>
 // END
 
 // MODULE: module1