Unify TypeAliasTag API between metadata and KSP (#1152)

* Extract TypeAliasTag for reuse and simplify API in KSP

* Update TestProcessor implementation for new alias API

This is based on the implementation for metadata in Moshi, and possibly worth promoting as a util to KotlinPoet itself

* I am once again asking IntelliJ to actually rename imports

* Spotless

interop/kotlinx-metadata/src/main/kotlin/com/squareup/kotlinpoet/metadata/specs/kmTypes.kt

@@ -30,6 +30,7 @@ import com.squareup.kotlinpoet.metadata.isNullable
 import com.squareup.kotlinpoet.metadata.isPrimary
 import com.squareup.kotlinpoet.metadata.isReified
 import com.squareup.kotlinpoet.metadata.isSuspend
+import com.squareup.kotlinpoet.tags.TypeAliasTag
 import kotlinx.metadata.KmClass
 import kotlinx.metadata.KmClassifier
 import kotlinx.metadata.KmClassifier.Class
@@ -158,7 +159,7 @@ internal fun KmType.toTypeName(
     // type in tags for reference.
     val abbreviatedTypeName = it.toTypeName(typeParamResolver)
     abbreviatedTypeName.copy(
-      tags = mapOf(TypeNameAliasTag::class to TypeNameAliasTag(finalType))
+      tags = mapOf(TypeAliasTag::class to TypeAliasTag(finalType))
     )
   } ?: finalType
 }

interop/kotlinx-metadata/src/test/kotlin/com/squareup/kotlinpoet/metadata/specs/KotlinPoetMetadataSpecsTest.kt

@@ -25,6 +25,7 @@ import com.squareup.kotlinpoet.metadata.KotlinPoetMetadataPreview
 import com.squareup.kotlinpoet.metadata.specs.MultiClassInspectorTest.ClassInspectorType.ELEMENTS
 import com.squareup.kotlinpoet.metadata.specs.MultiClassInspectorTest.ClassInspectorType.REFLECTIVE
 import com.squareup.kotlinpoet.tag
+import com.squareup.kotlinpoet.tags.TypeAliasTag
 import kotlinx.metadata.KmClass
 import kotlinx.metadata.KmConstructor
 import kotlinx.metadata.KmFunction
@@ -200,14 +201,14 @@ class KotlinPoetMetadataSpecsTest : MultiClassInspectorTest() {
     )
 
     val fooPropertyType = typeSpec.propertySpecs.first { it.name == "foo" }.type
-    val fooAliasData = fooPropertyType.tag<TypeNameAliasTag>()
+    val fooAliasData = fooPropertyType.tag<TypeAliasTag>()
     checkNotNull(fooAliasData)
-    assertThat(fooAliasData.type).isEqualTo(STRING)
+    assertThat(fooAliasData.abbreviatedType).isEqualTo(STRING)
 
     val barPropertyType = typeSpec.propertySpecs.first { it.name == "bar" }.type
-    val barAliasData = barPropertyType.tag<TypeNameAliasTag>()
+    val barAliasData = barPropertyType.tag<TypeAliasTag>()
     checkNotNull(barAliasData)
-    assertThat(barAliasData.type).isEqualTo(LIST.parameterizedBy(STRING))
+    assertThat(barAliasData.abbreviatedType).isEqualTo(LIST.parameterizedBy(STRING))
   }
 
   class TypeAliases(val foo: TypeAliasName, val bar: GenericTypeAlias)

interop/ksp/src/main/kotlin/com/squareup/kotlinpoet/ksp/ksTypes.kt

@@ -31,6 +31,7 @@ import com.squareup.kotlinpoet.STAR
 import com.squareup.kotlinpoet.TypeName
 import com.squareup.kotlinpoet.TypeVariableName
 import com.squareup.kotlinpoet.WildcardTypeName
+import com.squareup.kotlinpoet.tags.TypeAliasTag
 
 /** Returns the [ClassName] representation of this [KSType] IFF it's a [KSClassDeclaration]. */
 @KotlinPoetKspPreview
@@ -49,13 +50,10 @@ public fun KSType.toClassName(): ClassName {
  * @param typeParamResolver an optional resolver for enclosing declarations' type parameters. Parent
  *                          declarations can be anything with generics that child nodes declare as
  *                          defined by [KSType.arguments].
- * @param unwrapTypeAliases optionally controls whether typealiases should be unwrapped to their
- *                          aliased type. Useful in cases where only the aliased type matters
  */
 @KotlinPoetKspPreview
 public fun KSType.toTypeName(
-  typeParamResolver: TypeParameterResolver = TypeParameterResolver.EMPTY,
-  unwrapTypeAliases: Boolean = false
+  typeParamResolver: TypeParameterResolver = TypeParameterResolver.EMPTY
 ): TypeName {
   val type = when (val decl = declaration) {
     is KSClassDeclaration -> {
@@ -68,16 +66,17 @@ public fun KSType.toTypeName(
       } else {
         decl.typeParameters.toTypeParameterResolver(typeParamResolver)
       }
-      val firstPass = if (unwrapTypeAliases) {
-        decl.type.resolve()
-          .toTypeName(extraResolver, unwrapTypeAliases)
-          .copy(nullable = isMarkedNullable)
-          .rawType()
-      } else {
-        decl.toClassNameInternal()
-      }
-      firstPass
-        .withTypeArguments(arguments.map { it.toTypeName(typeParamResolver) })
+      val mappedArgs = arguments.map { it.toTypeName(typeParamResolver) }
+
+      val abbreviatedType = decl.type.resolve()
+        .toTypeName(extraResolver)
+        .copy(nullable = isMarkedNullable)
+        .rawType()
+        .withTypeArguments(mappedArgs)
+
+      decl.toClassNameInternal()
+        .withTypeArguments(mappedArgs)
+        .copy(tags = mapOf(TypeAliasTag::class to TypeAliasTag(abbreviatedType)))
     }
     else -> error("Unsupported type: $declaration")
   }
@@ -92,16 +91,13 @@ public fun KSType.toTypeName(
  * @param typeParamResolver an optional resolver for enclosing declarations' type parameters. Parent
  *                          declarations can be anything with generics that child nodes declare as
  *                          defined by [KSType.arguments].
- * @param unwrapTypeAliases optionally controls whether typealiases should be unwrapped to their
- *                          aliased type. Useful in cases where only the aliased type matters
  */
 @KotlinPoetKspPreview
 public fun KSTypeParameter.toTypeVariableName(
-  typeParamResolver: TypeParameterResolver = TypeParameterResolver.EMPTY,
-  unwrapTypeAliases: Boolean = false
+  typeParamResolver: TypeParameterResolver = TypeParameterResolver.EMPTY
 ): TypeVariableName {
   val typeVarName = name.getShortName()
-  val typeVarBounds = bounds.map { it.toTypeName(typeParamResolver, unwrapTypeAliases) }.toList()
+  val typeVarBounds = bounds.map { it.toTypeName(typeParamResolver) }.toList()
   val typeVarVariance = when (variance) {
     COVARIANT -> KModifier.OUT
     CONTRAVARIANT -> KModifier.IN
@@ -117,15 +113,12 @@ public fun KSTypeParameter.toTypeVariableName(
  * @param typeParamResolver an optional resolver for enclosing declarations' type parameters. Parent
  *                          declarations can be anything with generics that child nodes declare as
  *                          defined by [KSType.arguments].
- * @param unwrapTypeAliases optionally controls whether typealiases should be unwrapped to their
- *                          aliased type. Useful in cases where only the aliased type matters
  */
 @KotlinPoetKspPreview
 public fun KSTypeArgument.toTypeName(
-  typeParamResolver: TypeParameterResolver,
-  unwrapTypeAliases: Boolean = false
+  typeParamResolver: TypeParameterResolver
 ): TypeName {
-  val typeName = type?.resolve()?.toTypeName(typeParamResolver, unwrapTypeAliases) ?: return STAR
+  val typeName = type?.resolve()?.toTypeName(typeParamResolver) ?: return STAR
   return when (variance) {
     COVARIANT -> WildcardTypeName.producerOf(typeName)
     CONTRAVARIANT -> WildcardTypeName.consumerOf(typeName)
@@ -141,14 +134,10 @@ public fun KSTypeArgument.toTypeName(
  * @param typeParamResolver an optional resolver for enclosing declarations' type parameters. Parent
  *                          declarations can be anything with generics that child nodes declare as
  *                          defined by [KSType.arguments].
- * @param unwrapTypeAliases optionally controls whether typealiases should be unwrapped to their
- *                          aliased type. Useful in cases where only the aliased type matters
  */
 @KotlinPoetKspPreview
 public fun KSTypeReference.toTypeName(
-  typeParamResolver: TypeParameterResolver = TypeParameterResolver.EMPTY,
-  unwrapTypeAliases: Boolean = false
+  typeParamResolver: TypeParameterResolver = TypeParameterResolver.EMPTY
 ): TypeName {
-  val type = resolve()
-  return type.toTypeName(typeParamResolver, unwrapTypeAliases)
+  return resolve().toTypeName(typeParamResolver)
 }

interop/ksp/test-processor/src/main/kotlin/com/squareup/kotlinpoet/ksp/test/processor/TestProcessor.kt

@@ -64,14 +64,30 @@ class TestProcessor(private val env: SymbolProcessorEnvironment) : SymbolProcess
         addModifiers(decl.modifiers.mapNotNull { it.toKModifier() })
       }
     val classTypeParams = decl.typeParameters.toTypeParameterResolver()
-    classBuilder.addTypeVariables(decl.typeParameters.map { it.toTypeVariableName(classTypeParams, unwrapTypeAliases) })
+    classBuilder.addTypeVariables(
+      decl.typeParameters.map { typeParam ->
+        typeParam.toTypeVariableName(classTypeParams).let {
+          if (unwrapTypeAliases) {
+            it.unwrapTypeAlias()
+          } else {
+            it
+          }
+        }
+      }
+    )
 
     // Add properties
     for (property in decl.getDeclaredProperties()) {
       classBuilder.addProperty(
         PropertySpec.builder(
           property.simpleName.getShortName(),
-          property.type.toTypeName(classTypeParams, unwrapTypeAliases)
+          property.type.toTypeName(classTypeParams).let {
+            if (unwrapTypeAliases) {
+              it.unwrapTypeAlias()
+            } else {
+              it
+            }
+          }
         )
           .addOriginatingKSFile(decl.containingFile!!)
           .mutable(property.isMutable)
@@ -93,16 +109,40 @@ class TestProcessor(private val env: SymbolProcessorEnvironment) : SymbolProcess
             function.getVisibility().toKModifier()?.let { addModifiers(it) }
             addModifiers(function.modifiers.mapNotNull { it.toKModifier() })
           }
-          .addTypeVariables(function.typeParameters.map { it.toTypeVariableName(functionTypeParams, unwrapTypeAliases) })
+          .addTypeVariables(
+            function.typeParameters.map { typeParam ->
+              typeParam.toTypeVariableName(functionTypeParams).let {
+                if (unwrapTypeAliases) {
+                  it.unwrapTypeAlias()
+                } else {
+                  it
+                }
+              }
+            }
+          )
           .addParameters(
             function.parameters.map { parameter ->
-              val parameterType = parameter.type.toTypeName(functionTypeParams, unwrapTypeAliases)
+              val parameterType = parameter.type.toTypeName(functionTypeParams).let {
+                if (unwrapTypeAliases) {
+                  it.unwrapTypeAlias()
+                } else {
+                  it
+                }
+              }
               parameter.name?.let {
                 ParameterSpec.builder(it.getShortName(), parameterType).build()
               } ?: ParameterSpec.unnamed(parameterType)
             }
           )
-          .returns(function.returnType!!.toTypeName(functionTypeParams, unwrapTypeAliases))
+          .returns(
+            function.returnType!!.toTypeName(functionTypeParams).let {
+              if (unwrapTypeAliases) {
+                it.unwrapTypeAlias()
+              } else {
+                it
+              }
+            }
+          )
           .build()
       )
     }

interop/ksp/test-processor/src/main/kotlin/com/squareup/kotlinpoet/ksp/test/processor/typeAliasUnwrapping.kt

@@ -0,0 +1,53 @@
+package com.squareup.kotlinpoet.ksp.test.processor
+
+import com.squareup.kotlinpoet.AnnotationSpec
+import com.squareup.kotlinpoet.ClassName
+import com.squareup.kotlinpoet.LambdaTypeName
+import com.squareup.kotlinpoet.ParameterizedTypeName
+import com.squareup.kotlinpoet.TypeName
+import com.squareup.kotlinpoet.TypeVariableName
+import com.squareup.kotlinpoet.WildcardTypeName
+import com.squareup.kotlinpoet.tag
+import com.squareup.kotlinpoet.tags.TypeAliasTag
+import java.util.TreeSet
+
+/*
+ * Example implementation of how to unwrap a typealias from TypeNameAliasTag
+ */
+
+internal fun TypeName.unwrapTypeAliasReal(): TypeName {
+  return tag<TypeAliasTag>()?.abbreviatedType?.let { unwrappedType ->
+    // If any type is nullable, then the whole thing is nullable
+    var isAnyNullable = isNullable
+    // Keep track of all annotations across type levels. Sort them too for consistency.
+    val runningAnnotations = TreeSet<AnnotationSpec>(compareBy { it.toString() }).apply {
+      addAll(annotations)
+    }
+    val nestedUnwrappedType = unwrappedType.unwrapTypeAlias()
+    runningAnnotations.addAll(nestedUnwrappedType.annotations)
+    isAnyNullable = isAnyNullable || nestedUnwrappedType.isNullable
+    nestedUnwrappedType.copy(nullable = isAnyNullable, annotations = runningAnnotations.toList())
+  } ?: this
+}
+
+// TypeVariableName gets a special overload because these usually need to be kept in a type-safe
+// manner.
+internal fun TypeVariableName.unwrapTypeAlias(): TypeVariableName {
+  return TypeVariableName(
+    name = name,
+    bounds = bounds.map { it.unwrapTypeAlias() },
+    variance = variance
+  )
+    .copy(nullable = isNullable, annotations = annotations, tags = tags)
+}
+
+internal fun TypeName.unwrapTypeAlias(): TypeName {
+  return when (this) {
+    is ClassName -> unwrapTypeAliasReal()
+    is ParameterizedTypeName -> unwrapTypeAliasReal()
+    is TypeVariableName -> unwrapTypeAlias()
+    is WildcardTypeName -> unwrapTypeAliasReal()
+    is LambdaTypeName -> unwrapTypeAliasReal()
+    else -> throw UnsupportedOperationException("Type '${javaClass.simpleName}' is illegal. Only classes, parameterized types, wildcard types, or type variables are allowed.")
+  }
+}

interop/ksp/test-processor/src/test/kotlin/com/squareup/kotlinpoet/ksp/test/processor/TestProcessorTest.kt

@@ -208,12 +208,14 @@ class TestProcessorTest {
 
            typealias TypeAliasName = String
            typealias GenericTypeAlias = List<String>
+           typealias GenericMapTypeAlias<V, K> = Map<K, V>
 
            @ExampleAnnotation
            class Example {
              fun aliases(
                aliasedName: TypeAliasName,
-               genericAlias: GenericTypeAlias
+               genericAlias: GenericTypeAlias,
+               genericMapAlias: GenericMapTypeAlias<String, Int>,
              ) {
              }
            }
@@ -229,12 +231,18 @@ class TestProcessorTest {
       """
       package test
 
+      import kotlin.Int
       import kotlin.String
       import kotlin.Unit
       import kotlin.collections.List
+      import kotlin.collections.Map
 
       public class Example {
-        public fun aliases(aliasedName: String, genericAlias: List<String>): Unit {
+        public fun aliases(
+          aliasedName: String,
+          genericAlias: List<String>,
+          genericMapAlias: Map<Int, String>
+        ): Unit {
         }
       }
 

kotlinpoet/src/main/java/com/squareup/kotlinpoet/tags/TypeAliasTag.kt

@@ -0,0 +1,10 @@
+package com.squareup.kotlinpoet.tags
+
+import com.squareup.kotlinpoet.TypeName
+
+/**
+ * This tag indicates that this [TypeName] represents a `typealias` type.
+ *
+ * @property [abbreviatedType] the underlying type for this alias.
+ */
+public class TypeAliasTag(public val abbreviatedType: TypeName)