add opportunity to generate types lists for sealed subtypes

CHANGELOG.md

@@ -2,6 +2,10 @@
 
 ## 0.24.2
 
+* `Ksp`:
+  * `Sealed`:
+    * Add annotation `GenerateSealedTypesWorkaround` which allow to generate `subtypes` lists
+
 ## 0.24.1
 
 * `Versions`:

ksp/sealed/generator/src/main/kotlin/KSClassDeclarationGetGenerateSealedWorkaroundAnnotation.kt

@@ -3,9 +3,15 @@ package dev.inmo.micro_utils.ksp.sealed.generator
 import com.google.devtools.ksp.KspExperimental
 import com.google.devtools.ksp.getAnnotationsByType
 import com.google.devtools.ksp.symbol.KSClassDeclaration
+import dev.inmo.micro_utils.ksp.sealed.GenerateSealedTypesWorkaround
 import dev.inmo.micro_utils.ksp.sealed.GenerateSealedWorkaround
 import dev.inmo.microutils.kps.sealed.GenerateSealedWorkaround as OldGenerateSealedWorkaround
 
 @OptIn(KspExperimental::class)
 val KSClassDeclaration.getGenerateSealedWorkaroundAnnotation
     get() = (getAnnotationsByType(GenerateSealedWorkaround::class).firstOrNull() ?: getAnnotationsByType(OldGenerateSealedWorkaround::class).firstOrNull())
+
+
+@OptIn(KspExperimental::class)
+val KSClassDeclaration.getGenerateSealedTypesWorkaroundAnnotation
+    get() = getAnnotationsByType(GenerateSealedTypesWorkaround::class).firstOrNull()

ksp/sealed/generator/src/main/kotlin/Processor.kt

@@ -6,21 +6,17 @@ import com.google.devtools.ksp.processing.CodeGenerator
 import com.google.devtools.ksp.processing.Resolver
 import com.google.devtools.ksp.processing.SymbolProcessor
 import com.google.devtools.ksp.symbol.*
-import com.squareup.kotlinpoet.ClassName
-import com.squareup.kotlinpoet.CodeBlock
-import com.squareup.kotlinpoet.FileSpec
-import com.squareup.kotlinpoet.FunSpec
-import com.squareup.kotlinpoet.KModifier
+import com.squareup.kotlinpoet.*
 import com.squareup.kotlinpoet.ParameterizedTypeName.Companion.parameterizedBy
-import com.squareup.kotlinpoet.PropertySpec
-import com.squareup.kotlinpoet.asTypeName
 import com.squareup.kotlinpoet.ksp.toClassName
 import dev.inmo.micro_ksp.generator.buildSubFileName
 import dev.inmo.micro_ksp.generator.companion
 import dev.inmo.micro_ksp.generator.findSubClasses
 import dev.inmo.micro_ksp.generator.writeFile
+import dev.inmo.micro_utils.ksp.sealed.GenerateSealedTypesWorkaround
 import dev.inmo.micro_utils.ksp.sealed.GenerateSealedWorkaround
 import java.io.File
+import kotlin.reflect.KClass
 
 class Processor(
     private val codeGenerator: CodeGenerator
@@ -109,6 +105,62 @@ class Processor(
         )
     }
 
+    @OptIn(KspExperimental::class)
+    private fun FileSpec.Builder.generateSealedTypesWorkaround(
+        ksClassDeclaration: KSClassDeclaration,
+        resolver: Resolver
+    ) {
+        val annotation = ksClassDeclaration.getGenerateSealedTypesWorkaroundAnnotation
+        val subClasses = ksClassDeclaration.resolveSubclasses(
+            searchIn = resolver.getAllFiles(),
+            allowNonSealed = annotation ?.includeNonSealedSubTypes ?: false
+        ).distinct()
+        val subClassesNames = subClasses.filter {
+            it.getAnnotationsByType(GenerateSealedWorkaround.Exclude::class).count() == 0
+        }.sortedBy {
+            (it.getAnnotationsByType(GenerateSealedWorkaround.Order::class).firstOrNull()) ?.order ?: 0
+        }.map {
+            it.toClassName()
+        }.toList()
+        val className = ksClassDeclaration.toClassName()
+        val setType = Set::class.asTypeName().parameterizedBy(
+            KClass::class.asTypeName().parameterizedBy(
+                TypeVariableName(
+                    "out ${ksClassDeclaration.asStarProjectedType().toClassName().simpleNames.joinToString(".")}",
+                )
+            )
+        )
+        addProperty(
+            PropertySpec.builder(
+                "subtypes",
+                setType
+            ).apply {
+                modifiers.add(
+                    KModifier.PRIVATE
+                )
+                initializer(
+                    CodeBlock.of(
+                        """setOf(${subClassesNames.joinToString(",\n") { it.simpleNames.joinToString(".") + "::class" }})"""
+                    )
+                )
+            }.build()
+        )
+        addFunction(
+            FunSpec.builder("subtypes").apply {
+                val companion = ksClassDeclaration.takeIf { it.isCompanionObject } ?.toClassName()
+                    ?: ksClassDeclaration.companion ?.toClassName()
+                    ?: ClassName(className.packageName, *className.simpleNames.toTypedArray(), "Companion")
+                receiver(companion)
+                returns(setType)
+                addCode(
+                    CodeBlock.of(
+                        """return subtypes"""
+                    )
+                )
+            }.build()
+        )
+    }
+
     @OptIn(KspExperimental::class)
     override fun process(resolver: Resolver): List<KSAnnotated> {
         (resolver.getSymbolsWithAnnotation(GenerateSealedWorkaround::class.qualifiedName!!)).filterIsInstance<KSClassDeclaration>().forEach {
@@ -131,6 +183,26 @@ class Processor(
                 }.build()
             }
         }
+        (resolver.getSymbolsWithAnnotation(GenerateSealedTypesWorkaround::class.qualifiedName!!)).filterIsInstance<KSClassDeclaration>().forEach {
+            val prefix = (it.getGenerateSealedTypesWorkaroundAnnotation) ?.prefix ?.takeIf {
+                it.isNotEmpty()
+            } ?: it.buildSubFileName.replaceFirst(it.simpleName.asString(), "")
+            it.writeFile(prefix = prefix, suffix = "SealedTypesWorkaround") {
+                FileSpec.builder(
+                    it.packageName.asString(),
+                    "${it.simpleName.getShortName()}SealedTypesWorkaround"
+                ).apply {
+                    addFileComment(
+                        """
+                        THIS CODE HAVE BEEN GENERATED AUTOMATICALLY
+                        TO REGENERATE IT JUST DELETE FILE
+                        ORIGINAL FILE: ${it.containingFile ?.fileName}
+                        """.trimIndent()
+                    )
+                    generateSealedTypesWorkaround(it, resolver)
+                }.build()
+            }
+        }
 
         return emptyList()
     }

ksp/sealed/generator/test/src/commonMain/kotlin/Test.kt

@@ -1,14 +1,19 @@
 package dev.inmo.micro_utils.ksp.sealed.generator.test
 
+import dev.inmo.micro_utils.ksp.sealed.GenerateSealedTypesWorkaround
 import dev.inmo.micro_utils.ksp.sealed.GenerateSealedWorkaround
 
 @GenerateSealedWorkaround
+@GenerateSealedTypesWorkaround
 sealed interface Test {
     @GenerateSealedWorkaround.Order(2)
+    @GenerateSealedTypesWorkaround.Order(2)
     object A : Test
     @GenerateSealedWorkaround.Exclude
+    @GenerateSealedTypesWorkaround.Exclude
     object B : Test
     @GenerateSealedWorkaround.Order(0)
+    @GenerateSealedTypesWorkaround.Order(0)
     object C : Test
 
     // Required for successful sealed workaround generation

ksp/sealed/generator/test/src/commonMain/kotlin/TestSealedTypesWorkaround.kt

@@ -0,0 +1,12 @@
+// THIS CODE HAVE BEEN GENERATED AUTOMATICALLY
+// TO REGENERATE IT JUST DELETE FILE
+// ORIGINAL FILE: Test.kt
+package dev.inmo.micro_utils.ksp.`sealed`.generator.test
+
+import kotlin.collections.Set
+import kotlin.reflect.KClass
+
+private val subtypes: Set<KClass<out Test>> = setOf(Test.C::class,
+    Test.A::class)
+
+public fun Test.Companion.subtypes(): Set<KClass<out Test>> = subtypes

ksp/sealed/src/commonMain/kotlin/GenerateSealedTypesWorkaround.kt

@@ -0,0 +1,15 @@
+package dev.inmo.micro_utils.ksp.sealed
+
+@Retention(AnnotationRetention.BINARY)
+@Target(AnnotationTarget.CLASS)
+annotation class GenerateSealedTypesWorkaround(
+    val prefix: String = "",
+    val includeNonSealedSubTypes: Boolean = false,
+) {
+    @Retention(AnnotationRetention.BINARY)
+    @Target(AnnotationTarget.CLASS)
+    annotation class Order(val order: Int)
+    @Retention(AnnotationRetention.BINARY)
+    @Target(AnnotationTarget.CLASS)
+    annotation class Exclude
+}

ksp/sealed/src/commonMain/kotlin/GenerateSealedWorkaround.kt

@@ -4,7 +4,7 @@ package dev.inmo.micro_utils.ksp.sealed
 @Target(AnnotationTarget.CLASS)
 annotation class GenerateSealedWorkaround(
     val prefix: String = "",
-    val includeNonSealedSubTypes: Boolean = false
+    val includeNonSealedSubTypes: Boolean = false,
 ) {
     @Retention(AnnotationRetention.BINARY)
     @Target(AnnotationTarget.CLASS)