improve classcasts

ksp/classcasts/generator/src/main/kotlin/ClassCastsFiller.kt

@@ -32,14 +32,9 @@ private fun FileSpec.Builder.createTypeDefinition(ksClassDeclaration: KSClassDec
 
 internal fun FileSpec.Builder.fill(
     sourceKSClassDeclaration: KSClassDeclaration,
-    subtypesMap: Map<KSClassDeclaration, Set<KSClassDeclaration>>,
+    subtypes: Set<KSClassDeclaration>
-    targetClassDeclaration: KSClassDeclaration = sourceKSClassDeclaration
 ) {
-    if (sourceKSClassDeclaration == targetClassDeclaration) {
+    subtypes.forEach { targetClassDeclaration ->
-        subtypesMap[sourceKSClassDeclaration] ?.forEach {
-            fill(sourceKSClassDeclaration, subtypesMap, it)
-        }
-    } else {
         val sourceClassName = sourceKSClassDeclaration.toClassName()
         val targetClassClassName = targetClassDeclaration.toClassName()
         val targetClassTypeDefinition = createTypeDefinition(targetClassDeclaration)
@@ -91,16 +86,5 @@ internal fun FileSpec.Builder.fill(
                 addModifiers(KModifier.INLINE)
             }.build()
         )
-
-        subtypesMap[targetClassDeclaration] ?.let {
-            if (it.count { it.classKind == ClassKind.CLASS } > 1) {
-                it
-            } else {
-                it.filter { it.classKind != ClassKind.CLASS }
-            }
-        } ?.forEach {
-            fill(sourceKSClassDeclaration, subtypesMap, it)
-            fill(targetClassDeclaration, subtypesMap, it)
-        }
     }
 }

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

@@ -9,67 +9,72 @@ import com.google.devtools.ksp.processing.Resolver
 import com.google.devtools.ksp.processing.SymbolProcessor
 import com.google.devtools.ksp.symbol.*
 import com.squareup.kotlinpoet.*
+import com.squareup.kotlinpoet.ksp.toClassName
 import dev.inmo.micro_ksp.generator.writeFile
 import dev.inmo.micro_utils.ksp.classcasts.ClassCastsExcluded
 import dev.inmo.micro_utils.ksp.classcasts.ClassCastsIncluded
+import java.io.File
 
 class Processor(
     private val codeGenerator: CodeGenerator
 ) : SymbolProcessor {
-    private val classCastsIncludedClassName = ClassCastsIncluded::class.asClassName()
-
     @OptIn(KspExperimental::class)
     private fun FileSpec.Builder.generateClassCasts(
         ksClassDeclaration: KSClassDeclaration,
         resolver: Resolver
     ) {
-        val classes = resolver.getSymbolsWithAnnotation(classCastsIncludedClassName.canonicalName).filterIsInstance<KSClassDeclaration>()
+        val rootAnnotation = ksClassDeclaration.getAnnotationsByType(ClassCastsIncluded::class).first()
-        val classesRegexes: Map<KSClassDeclaration, Pair<Regex?, Regex?>> = classes.mapNotNull {
+        val (includeRegex: Regex?, excludeRegex: Regex?) = rootAnnotation.let {
-            it to (it.getAnnotationsByType(ClassCastsIncluded::class).firstNotNullOfOrNull {
+            it.typesRegex.takeIf { it.isNotEmpty() } ?.let(::Regex) to it.excludeRegex.takeIf { it.isNotEmpty() } ?.let(::Regex)
-                it.typesRegex.takeIf { it.isNotEmpty() } ?.let(::Regex) to it.excludeRegex.takeIf { it.isNotEmpty() } ?.let(::Regex)
+        }
-            } ?: return@mapNotNull null)
-        }.toMap()
         val classesSubtypes = mutableMapOf<KSClassDeclaration, MutableSet<KSClassDeclaration>>()
 
-        resolver.getAllFiles().forEach {
+        fun KSClassDeclaration.checkSupertypeLevel(levelsAllowed: Int?): Boolean {
-            it.declarations.forEach { potentialSubtype ->
+            val supertypes by lazy {
-                if (
+                superTypes.map { it.resolve().declaration }
-                    potentialSubtype is KSClassDeclaration
+            }
-                    && potentialSubtype.isAnnotationPresent(ClassCastsExcluded::class).not()
+            return when {
-                ) {
+                levelsAllowed == null -> true
-                    val allSupertypes = potentialSubtype.getAllSuperTypes().map { it.declaration }
+                levelsAllowed <= 0 -> false
-
+                supertypes.any { it == ksClassDeclaration } -> true
-                    for (currentClass in classes) {
+                else -> supertypes.any {
-                        val regexes = classesRegexes[currentClass]
+                    (it as? KSClassDeclaration) ?.checkSupertypeLevel(levelsAllowed - 1) == true
-                        val simpleName = potentialSubtype.simpleName.getShortName()
-                        when {
-                            currentClass !in allSupertypes
-                                    || regexes ?.first ?.matches(simpleName) == false
-                                    || regexes ?.second ?.matches(simpleName) == true -> continue
-                            else -> {
-                                classesSubtypes.getOrPut(currentClass) { mutableSetOf() }.add(potentialSubtype)
-                            }
-                        }
-                    }
                 }
             }
         }
-        fun fillWithSealeds(source: KSClassDeclaration, current: KSClassDeclaration = source) {
+
-            val regexes = classesRegexes[source]
+        fun handleDeclaration(ksDeclarationContainer: KSDeclarationContainer) {
+            ksDeclarationContainer.declarations.forEach { potentialSubtype ->
+                val simpleName = potentialSubtype.simpleName.getShortName()
+                when {
+                    potentialSubtype === ksClassDeclaration -> {}
+                    potentialSubtype.isAnnotationPresent(ClassCastsExcluded::class) -> return@forEach
+                    potentialSubtype !is KSClassDeclaration || !potentialSubtype.checkSupertypeLevel(rootAnnotation.levelsToInclude.takeIf { it >= 0 }) -> return@forEach
+                    excludeRegex ?.matches(simpleName) == true -> return@forEach
+                    includeRegex ?.matches(simpleName) == false -> {}
+                    else -> classesSubtypes.getOrPut(ksClassDeclaration) { mutableSetOf() }.add(potentialSubtype)
+                }
+                handleDeclaration(potentialSubtype as? KSDeclarationContainer ?: return@forEach)
+            }
+        }
+        resolver.getAllFiles().forEach {
+            handleDeclaration(it)
+        }
+        fun fillWithSealeds(current: KSClassDeclaration) {
             current.getSealedSubclasses().forEach {
                 val simpleName = it.simpleName.getShortName()
                 if (
-                    regexes ?.first ?.matches(simpleName) == false
+                    includeRegex ?.matches(simpleName) == false
-                    || regexes ?.second ?.matches(simpleName) == true
+                    || excludeRegex ?.matches(simpleName) == true
                     || it.isAnnotationPresent(ClassCastsExcluded::class)
                 ) {
                     return@forEach
                 }
-                classesSubtypes.getOrPut(source) { mutableSetOf() }.add(it)
+                classesSubtypes.getOrPut(ksClassDeclaration) { mutableSetOf() }.add(it)
-                fillWithSealeds(source, it)
+                fillWithSealeds(it)
             }
         }
-        classes.forEach { fillWithSealeds(it) }
+        fillWithSealeds(ksClassDeclaration)
 
         addAnnotation(
             AnnotationSpec.builder(Suppress::class).apply {
@@ -82,12 +87,10 @@ class Processor(
                 useSiteTarget(AnnotationSpec.UseSiteTarget.FILE)
             }.build()
         )
-        classes.forEach {
+        fill(
-            fill(
+            ksClassDeclaration,
-                it,
+            classesSubtypes.values.flatten().toSet()
-                classesSubtypes.toMap()
+        )
-            )
-        }
     }
 
     @OptIn(KspExperimental::class)

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

@@ -3,10 +3,13 @@ package dev.inmo.micro_utils.ksp.classcasts.generator.test
 import dev.inmo.micro_utils.ksp.classcasts.ClassCastsExcluded
 import dev.inmo.micro_utils.ksp.classcasts.ClassCastsIncluded
 
-@ClassCastsIncluded
+@ClassCastsIncluded(levelsToInclude = 1)
 sealed interface Test {
     object A : Test
     @ClassCastsExcluded
-    object B : Test
+    object B : Test // Will not be included in class casts due to annotation ClassCastsExcluded
     object C : Test
+    interface D : Test {
+        object DD : D // Will not be included in class casts due to levelsToInclude
+    }
 }

ksp/classcasts/generator/test/src/commonMain/kotlin/TestClassCasts.kt

@@ -30,3 +30,11 @@ public inline fun Test.cOrThrow(): Test.C = this as
     dev.inmo.micro_utils.ksp.classcasts.generator.test.Test.C
 
 public inline fun <T> Test.ifC(block: (Test.C) -> T): T? = cOrNull() ?.let(block)
+
+public inline fun Test.dOrNull(): Test.D? = this as?
+    dev.inmo.micro_utils.ksp.classcasts.generator.test.Test.D
+
+public inline fun Test.dOrThrow(): Test.D = this as
+    dev.inmo.micro_utils.ksp.classcasts.generator.test.Test.D
+
+public inline fun <T> Test.ifD(block: (Test.D) -> T): T? = dOrNull() ?.let(block)

ksp/classcasts/src/commonMain/kotlin/ClassCastsIncluded.kt

@@ -2,4 +2,9 @@ package dev.inmo.micro_utils.ksp.classcasts
 
 @Target(AnnotationTarget.CLASS)
 @Retention(AnnotationRetention.SOURCE)
-annotation class ClassCastsIncluded(val typesRegex: String = "", val excludeRegex: String = "", val outputFilePrefix: String = "")
+annotation class ClassCastsIncluded(
+    val typesRegex: String = "",
+    val excludeRegex: String = "",
+    val outputFilePrefix: String = "",
+    val levelsToInclude: Int = -1
+)