refactor type replace logic to handle more types

Author: neetopia

kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSTypeImpl.kt

@@ -31,7 +31,6 @@ import com.google.devtools.ksp.symbol.Nullability
 import org.jetbrains.kotlin.analysis.api.KtStarTypeProjection
 import org.jetbrains.kotlin.analysis.api.annotations.KtAnnotationsList
 import org.jetbrains.kotlin.analysis.api.annotations.annotations
-import org.jetbrains.kotlin.analysis.api.components.buildClassType
 import org.jetbrains.kotlin.analysis.api.symbols.*
 import org.jetbrains.kotlin.analysis.api.types.*
 
@@ -111,24 +110,12 @@ class KSTypeImpl private constructor(internal val type: KtType) : KSType {
     }
 
     override fun replace(arguments: List<KSTypeArgument>): KSType {
-        if (arguments.isNotEmpty() && arguments.size != type.typeArguments().size) {
-            return KSErrorType
-        }
-        return analyze {
-            analysisSession.buildClassType((type as KtNonErrorClassType).classSymbol) {
-                arguments.forEach { arg -> argument(arg.toKtTypeProjection()) }
-            }.let { getCached(it) }
-        }
+        return type.replace(arguments.map { it.toKtTypeProjection() })?.let { getCached(it) } ?: KSErrorType
     }
 
     override fun starProjection(): KSType {
-        return analyze {
-            analysisSession.buildClassType((type as KtNonErrorClassType).classSymbol) {
-                type.typeArguments().forEach {
-                    argument(KtStarTypeProjection(type.token))
-                }
-            }.let { getCached(it) }
-        }
+        return type.replace(List(type.typeArguments().size) { KtStarTypeProjection(type.token) })
+            ?.let { getCached(it) } ?: KSErrorType
     }
 
     override fun makeNullable(): KSType {

kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/util.kt

@@ -37,6 +37,7 @@ import org.jetbrains.kotlin.analysis.api.analyze
 import org.jetbrains.kotlin.analysis.api.annotations.*
 import org.jetbrains.kotlin.analysis.api.components.KtSubstitutorBuilder
 import org.jetbrains.kotlin.analysis.api.components.buildClassType
+import org.jetbrains.kotlin.analysis.api.components.buildTypeParameterType
 import org.jetbrains.kotlin.analysis.api.fir.evaluate.FirAnnotationValueConverter
 import org.jetbrains.kotlin.analysis.api.fir.symbols.KtFirValueParameterSymbol
 import org.jetbrains.kotlin.analysis.api.lifetime.KtAlwaysAccessibleLifetimeToken
@@ -549,3 +550,19 @@ internal fun KtType.isAssignableFrom(that: KtType): Boolean {
         }
     }
 }
+
+// TODO: fix flexible type creation once upstream available.
+internal fun KtType.replace(newArgs: List<KtTypeProjection>): KtType? {
+    if (newArgs.isNotEmpty() && newArgs.size != this.typeArguments().size) {
+        return null
+    }
+    return analyze {
+        when (val symbol = classifierSymbol()) {
+            is KtClassLikeSymbol -> analysisSession.buildClassType(symbol) {
+                newArgs.forEach { arg -> argument(arg) }
+            }
+            is KtTypeParameterSymbol -> analysisSession.buildTypeParameterType(symbol)
+            else -> throw IllegalStateException("Unexpected type $this")
+        }
+    }
+}

kotlin-analysis-api/testData/replaceWithErrorTypeArgs.kt

@@ -98,6 +98,8 @@
 // KLE.E.asType([INVARIANT Int, INVARIANT String]): <ERROR TYPE>
 // KLE.E.asType([INVARIANT NotExist1, INVARIANT NotExist2]): <ERROR TYPE>
 // KLE.E.asType(emptyList()): KLE
+// default type:A
+// flexible type star:T
 // END
 
 // MODULE: lib
@@ -118,12 +120,15 @@ enum class KLE {
 // MODULE: main(lib)
 // FILE: JS.java
 class JS<T1, T2> {}
-class JS1<T> {}
+class JS1<T> {
+    T p;
+}
 enum JSE {
     E
 }
 
 // FILE: KS.kt
+fun <A> f(a: A) = TODO()
 class KS<T1, T2>
 class KS1<T>
 enum class KSE {

test-utils/src/main/kotlin/com/google/devtools/ksp/processor/ReplaceWithErrorTypeArgsProcessor.kt

@@ -18,6 +18,7 @@
 package com.google.devtools.ksp.processor
 
 import com.google.devtools.ksp.getClassDeclarationByName
+import com.google.devtools.ksp.getDeclaredProperties
 import com.google.devtools.ksp.processing.Resolver
 import com.google.devtools.ksp.symbol.*
 
@@ -58,6 +59,11 @@ open class ReplaceWithErrorTypeArgsProcessor : AbstractTestProcessor() {
             results.add("$declName.asType($yargs): ${decl.asType(yargs)}")
             results.add("$declName.asType(emptyList()): ${decl.asType(emptyList())}")
         }
+        val function = resolver.getFunctionDeclarationsByName(resolver.getKSNameFromString("f"), true).single()
+        results.add("default type:${function.parameters.single().type.resolve().replace(emptyList())}")
+        // TODO: fix flexible type creation once upstream available.
+        val js1 = resolver.getClassDeclarationByName("JS1")!!
+        results.add("flexible type star:${js1.getDeclaredProperties().single().type.resolve().starProjection()}")
         return emptyList()
     }
 

test-utils/testData/api/replaceWithErrorTypeArgs.kt

@@ -98,6 +98,8 @@
 // KLE.E.asType([INVARIANT Int, INVARIANT String]): <ERROR TYPE>
 // KLE.E.asType([INVARIANT NotExist1, INVARIANT NotExist2]): <ERROR TYPE>
 // KLE.E.asType(emptyList()): KLE.E
+// default type:A
+// flexible type star:(T..T?)
 // END
 
 // MODULE: lib
@@ -118,12 +120,15 @@ enum class KLE {
 // MODULE: main(lib)
 // FILE: JS.java
 class JS<T1, T2> {}
-class JS1<T> {}
+class JS1<T> {
+    T p;
+}
 enum JSE {
     E
 }
 
 // FILE: KS.kt
+fun <A> f(a: A) = TODO()
 class KS<T1, T2>
 class KS1<T>
 enum class KSE {