ksp: more accurately represent function types

When converting a KSTypeReference, represent a function type with LambdaTypeName. The better matches the source code and allows for receiver types (ex: String.() -> Unit) to be represented correctly.
square · Nov 27, 2023 · 09fd8ff · 09fd8ff
1 parent 945c182
commit 09fd8ff
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 7 deletions.
diff --git a/interop/ksp/src/main/kotlin/com/squareup/kotlinpoet/ksp/KsTypes.kt b/interop/ksp/src/main/kotlin/com/squareup/kotlinpoet/ksp/KsTypes.kt
@@ -15,6 +15,7 @@
  */
 package com.squareup.kotlinpoet.ksp
 
+import com.google.devtools.ksp.symbol.KSCallableReference
 import com.google.devtools.ksp.symbol.KSClassDeclaration
 import com.google.devtools.ksp.symbol.KSType
 import com.google.devtools.ksp.symbol.KSTypeAlias
@@ -27,6 +28,8 @@ import com.google.devtools.ksp.symbol.Variance.COVARIANT
 import com.google.devtools.ksp.symbol.Variance.INVARIANT
 import com.squareup.kotlinpoet.ClassName
 import com.squareup.kotlinpoet.KModifier
+import com.squareup.kotlinpoet.LambdaTypeName
+import com.squareup.kotlinpoet.ParameterSpec
 import com.squareup.kotlinpoet.STAR
 import com.squareup.kotlinpoet.TypeName
 import com.squareup.kotlinpoet.TypeVariableName
@@ -179,5 +182,14 @@ public fun KSTypeArgument.toTypeName(
 public fun KSTypeReference.toTypeName(
   typeParamResolver: TypeParameterResolver = TypeParameterResolver.EMPTY,
 ): TypeName {
-  return resolve().toTypeName(typeParamResolver, element?.typeArguments.orEmpty())
+  return when (val elem = element) {
+    is KSCallableReference -> {
+      LambdaTypeName.get(
+        receiver = elem.receiverType?.toTypeName(typeParamResolver),
+        parameters = elem.functionParameters.map { ParameterSpec.unnamed(it.type.toTypeName(typeParamResolver)) },
+        returnType = elem.returnType.toTypeName(typeParamResolver),
+      )
+    }
+    else -> resolve().toTypeName(typeParamResolver, element?.typeArguments.orEmpty())
+  }
 }
diff --git a/...est-processor/src/main/kotlin/com/squareup/kotlinpoet/ksp/test/processor/TestProcessor.kt b/...est-processor/src/main/kotlin/com/squareup/kotlinpoet/ksp/test/processor/TestProcessor.kt
@@ -171,7 +171,14 @@ class TestProcessor(private val env: SymbolProcessorEnvironment) : SymbolProcess
           )
           .addParameters(
             function.parameters.map { parameter ->
-              val parameterType = parameter.type.toValidatedTypeName(functionTypeParams).let {
+              // function references can't be obtained from a resolved KSType because it resolves to FunctionN<> which
+              // loses the necessary context, skip validation in these cases as we know they won't match.
+              val typeName = if (parameter.type.resolve().isFunctionType) {
+                parameter.type.toTypeName(functionTypeParams)
+              } else {
+                parameter.type.toValidatedTypeName(functionTypeParams)
+              }
+              val parameterType = typeName.let {
                 if (unwrapTypeAliases) {
                   it.unwrapTypeAlias()
                 } else {

diff --git a/...processor/src/test/kotlin/com/squareup/kotlinpoet/ksp/test/processor/TestProcessorTest.kt b/...processor/src/test/kotlin/com/squareup/kotlinpoet/ksp/test/processor/TestProcessorTest.kt
@@ -183,8 +183,6 @@ class TestProcessorTest {
       import kotlin.Boolean
       import kotlin.Enum
       import kotlin.Float
-      import kotlin.Function0
-      import kotlin.Function1
       import kotlin.Int
       import kotlin.IntArray
       import kotlin.String
@@ -261,9 +259,9 @@ class TestProcessorTest {
         }
 
         public suspend fun functionD(
-          param1: Function0<String>,
-          param2: Function1<String, String>,
-          param3: Function1<String, String>,
+          param1: () -> String,
+          param2: (String) -> String,
+          param3: String.() -> String,
         ) {
         }