mutable type in collection (#89)

kopykat-kt · Jul 5, 2023 · 0d53fd6 · 0d53fd6
1 parent 935db8c
commit 0d53fd6
Show file tree

Hide file tree

Showing 4 changed files with 115 additions and 37 deletions.
diff --git a/kopykat-ksp/src/main/kotlin/at/kopyk/utils/ClassNameExtensions.kt b/kopykat-ksp/src/main/kotlin/at/kopyk/utils/ClassNameExtensions.kt
@@ -2,6 +2,16 @@ package at.kopyk.utils
 
 import at.kopyk.poet.flattenWithSuffix
 import com.squareup.kotlinpoet.ClassName
+import com.squareup.kotlinpoet.ParameterizedTypeName
+import com.squareup.kotlinpoet.TypeName
 
 internal val ClassName.mutable: ClassName get() = flattenWithSuffix("Mutable")
+internal val ParameterizedTypeName.mutable: ParameterizedTypeName get() = flattenWithSuffix("Mutable")
 internal val ClassName.dslMarker: ClassName get() = flattenWithSuffix("DslMarker")
+
+internal val TypeName.mutable: TypeName?
+  get() = when (this) {
+    is ClassName -> this.mutable
+    is ParameterizedTypeName -> this.mutable
+    else -> null
+  }
diff --git a/kopykat-ksp/src/main/kotlin/at/kopyk/utils/TypeCompileScope.kt b/kopykat-ksp/src/main/kotlin/at/kopyk/utils/TypeCompileScope.kt
@@ -22,6 +22,7 @@ import com.squareup.kotlinpoet.FunSpec
 import com.squareup.kotlinpoet.KModifier
 import com.squareup.kotlinpoet.LIST
 import com.squareup.kotlinpoet.MAP
+import com.squareup.kotlinpoet.ParameterizedTypeName
 import com.squareup.kotlinpoet.SET
 import com.squareup.kotlinpoet.TypeName
 import com.squareup.kotlinpoet.TypeVariableName
@@ -201,56 +202,84 @@ internal class FileCompilerScope(
 internal fun KSType.toClassNameRespectingNullability(): ClassName =
   toClassName().copy(this.isMarkedNullable, emptyList(), emptyMap())
 
+internal fun KSType.toTypeNameRespectingNullability(typeParamResolver: TypeParameterResolver = TypeParameterResolver.EMPTY): TypeName =
+  toTypeName(typeParamResolver).copy(this.isMarkedNullable, emptyList(), emptyMap())
+
 internal fun TypeCompileScope.mutationInfo(ty: KSType): MutationInfo<TypeName> =
   when (ty.declaration) {
     is KSClassDeclaration -> {
-      val className: ClassName = ty.toClassNameRespectingNullability()
+      val typeName: TypeName = ty.toTypeNameRespectingNullability(typeParameterResolver)
+      val className = when (typeName) {
+        is ClassName -> typeName
+        is ParameterizedTypeName -> typeName.rawType
+        else -> ty.toClassNameRespectingNullability()
+      }
       val nullableLessClassName = className.copy(nullable = false)
-      infix fun String.dot(function: String) =
-        if (ty.isMarkedNullable) {
-          "$this?.$function()"
-        } else {
-          "$this.$function()"
-        }
-      val intermediate: MutationInfo<ClassName> = when {
-        nullableLessClassName == LIST ->
-          MutationInfo(
-            ClassName(className.packageName, "MutableList")
-              .copy(nullable = ty.isMarkedNullable, annotations = emptyList(), tags = emptyMap()),
-            { it dot "toMutableList" },
-            { it }
-          )
-        nullableLessClassName == MAP ->
-          MutationInfo(
-            ClassName(className.packageName, "MutableMap")
-              .copy(nullable = ty.isMarkedNullable, annotations = emptyList(), tags = emptyMap()),
-            { it dot "toMutableMap" },
-            { it }
-          )
-        nullableLessClassName == SET ->
-          MutationInfo(
-            ClassName(className.packageName, "MutableSet")
-              .copy(nullable = ty.isMarkedNullable, annotations = emptyList(), tags = emptyMap()),
-            { it dot "toMutableSet" },
-            { it }
-          )
+      infix fun String.dot(function: String) = dot(ty, function)
+      when {
+        nullableLessClassName == LIST -> mutationInfoOfCollection(ty, className, className.simpleName)
+        nullableLessClassName == MAP -> mutationInfoOfCollection(ty, className, className.simpleName)
+        nullableLessClassName == SET -> mutationInfoOfCollection(ty, className, className.simpleName)
         ty.hasMutableCopy() ->
           MutationInfo(
             className.mutable,
             { it dot "toMutable" },
             { it dot "freeze" }
           )
         else ->
-          MutationInfo(className, { it }, { it })
+          MutationInfo(
+            className.parameterizedWhenNotEmpty(
+              ty.arguments.map { it.toTypeName(typeParameterResolver) }
+            ),
+            { it },
+            { it }
+          )
       }
-      MutationInfo(
-        className = intermediate.className.parameterizedWhenNotEmpty(
-          ty.arguments.map { it.toTypeName(typeParameterResolver) }
-        ),
-        toMutable = intermediate.toMutable,
-        freeze = intermediate.freeze
-      )
     }
     else ->
       MutationInfo(ty.toTypeName(typeParameterResolver), { it }, { it })
   }
+
+internal fun TypeCompileScope.mutationInfoOfCollection(
+  ty: KSType,
+  className: ClassName,
+  collectionType: String
+): MutationInfo<TypeName> {
+  infix fun String.dot(function: String) = dot(ty, function)
+  infix fun String.dotMap(map: String) = dotMap(ty, map)
+  val type = ty.arguments[0].type?.resolve()
+  val isMutableCollection =
+    (ty.arguments.size == 1 && type?.hasMutableCopy() == true)
+  val transform: (String) -> String = if (isMutableCollection) {
+    { it dotMap "it.toMutable()" dot "toMutable$collectionType" }
+  } else {
+    { it dot "toMutable$collectionType" }
+  }
+  val freeze: (String) -> String = if (isMutableCollection) {
+    { it dotMap "it.freeze()" }
+  } else {
+    { it }
+  }
+  return MutationInfo(
+    ClassName(className.packageName, "Mutable$collectionType")
+      .copy(nullable = ty.isMarkedNullable, annotations = emptyList(), tags = emptyMap()).parameterizedWhenNotEmpty(
+        type?.takeIf { isMutableCollection }?.toTypeName(typeParameterResolver)?.mutable?.let(::listOf)
+          ?: ty.arguments.map { it.toTypeName(typeParameterResolver) }
+      ),
+    transform,
+    freeze
+  )
+}
+
+internal fun String.dot(ty: KSType, function: String) = if (ty.isMarkedNullable) {
+  "$this?.$function()"
+} else {
+  "$this.$function()"
+}
+
+internal fun String.dotMap(ty: KSType, map: String) =
+  if (ty.isMarkedNullable) {
+    "$this?.map { $map }"
+  } else {
+    "$this.map { $map }"
+  }
diff --git a/kopykat-ksp/src/test/kotlin/at/kopyk/MutableCollectionCopyTest.kt b/kopykat-ksp/src/test/kotlin/at/kopyk/MutableCollectionCopyTest.kt
@@ -0,0 +1,33 @@
+package at.kopyk
+
+import org.junit.jupiter.api.Test
+
+/**
+ * @author: xiaozhikang
+ * @create: 2023/7/3
+ */
+class MutableCollectionCopyTest {
+  @Test
+  fun `copy property in collection`() {
+    """
+      data class Group(val p: List<Person>)
+      data class Person(val age: Int)
+      
+      val g1 = Group(listOf(Person(1), Person(2)))
+      val g2 = g1.copy { p[1].age ++ }
+      val age = g2.p[1].age
+    """.trimIndent().evals("age" to 3)
+  }
+
+  @Test
+  fun `copy property in collection with generic type`() {
+    """
+      data class Group(val p: List<Person<String>>)
+      data class Person<T>(val mark: T)
+      
+      val g1 = Group(listOf(Person("old"), Person("old")))
+      val g2 = g1.copy { p[1].mark = "new" }
+      val mark = g2.p[1].mark
+    """.trimIndent().evals("mark" to "new")
+  }
+}
diff --git a/utils/kotlin-poet/src/main/kotlin/at/kopyk/poet/KotlinPoetExtensions.kt b/utils/kotlin-poet/src/main/kotlin/at/kopyk/poet/KotlinPoetExtensions.kt
@@ -10,6 +10,7 @@ import com.squareup.kotlinpoet.ClassName
 import com.squareup.kotlinpoet.FileSpec
 import com.squareup.kotlinpoet.KModifier
 import com.squareup.kotlinpoet.LambdaTypeName
+import com.squareup.kotlinpoet.ParameterizedTypeName
 import com.squareup.kotlinpoet.ParameterizedTypeName.Companion.parameterizedBy
 import com.squareup.kotlinpoet.TypeName
 import com.squareup.kotlinpoet.UNIT
@@ -70,6 +71,11 @@ public fun ClassName.flattenWithSuffix(suffix: String): ClassName {
   return ClassName(packageName, mutableSimpleName).copy(this.isNullable, emptyList(), emptyMap())
 }
 
+public fun ParameterizedTypeName.flattenWithSuffix(suffix: String): ParameterizedTypeName {
+  val mutableSimpleName = (rawType.simpleNames + suffix).joinToString(separator = "$")
+  return ClassName(rawType.packageName, mutableSimpleName).copy(this.isNullable, emptyList(), emptyMap()).parameterizedBy(this.typeArguments)
+}
+
 // https://kotlinlang.org/docs/reference/keyword-reference.html
 private val KEYWORDS = setOf(
   // Hard keywords